[Feature] WS1: Stable Component Boundaries for External Loops
@coder-jayp is already working on this.
Since Aug 11, 2026.
- 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
TensorDictkeys 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.