pytorch / pytorch/rl

[Feature] WS1: Stable Component Boundaries for External Loops

Open
#4,092 2 comments 0 reactions 2 assignees View on GitHub

@coder-jayp is already working on this.

Since Aug 11, 2026.

enhancement
Dominant language
Python
Stars
3.6k
Forks
487
Avg merge
1d 1h
Merged PRs (30d)
207

Description

Motivation

Spin-off from the Post-training efforts RFC (#3948) specifically for Workstream 1 (Stable component boundaries).

While working on the TRL adapters (WS2), it became clear that TorchRL components currently rely on implicit contracts. For example, knowing exactly what shape and keys a ReplayBuffer will yield, or exactly what keys a Loss module expects in a TensorDict (like "tokens" vs "inputs"), requires reading the implementation source code.

For external users to confidently pull a single TorchRL component into their own training loop (like TRL or NeMo-RL), these interfaces need to be explicit, versioned, and documented public contracts so that adapters stay thin and don't break unexpectedly on version bumps.

Solution

I recommend formalizing these boundaries through documentation and Python types. Specifically, I propose:

1. Architectural Documentation
Create a dedicated Sphinx guide (e.g., "Integrating TorchRL Components into External Loops") that explicitly documents the expected I/O contracts for:

  • Buffers: Extend/sample expectations over token or trajectory data.
  • Rollouts: Collector/inference-server protocol (including weight sync logic).
  • Losses: Explicit definitions of the TensorDict keys expected as input and the loss keys yielded as output.

2. Formal Protocols (Optional but recommended)
Introduce Python typing.Protocol or lightweight base classes for these interfaces. This allows external adapters (and internal checks) to use static type checking (e.g., mypy) to verify they are satisfying the TorchRL contract before runtime.

Alternatives

An alternative is to leave these contracts implicit and rely on users reading the source code of recipes like grpo_utils.py to figure out how to wire up a collector or buffer. However, this creates a high barrier to entry and leads to brittle external code that might break when internal TorchRL keys or shapes change.

Additional context

Having just built the TRL interop adapters, I have a fresh perspective on exactly which contracts an external loop needs to consume. I would like to build these documentations and protocols if the approach looks good to you.

Checklist

  • I have checked that there is no similar issue in the repo (required)

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.