[Feature] sequence rl composability by theap06 · Pull Request #3695 · pytorch/rl

theap06 · 2026-04-30T00:04:04Z

Fixes #3678

Motivation and Fix

This PR addresses three composability gaps that made it unnecessarily hard to train RNN/sequence-based policies (LSTM, GRU, ACT, Decision Transformer) with TorchRL's existing building blocks.

pytorch-bot · 2026-04-30T00:04:07Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3695

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 17 Awaiting Approval

As of commit 3cefe03 with merge base 4da311b ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

github-actions · 2026-04-30T00:04:14Z

⚠️ PR Title Label Error

PR title must start with a label prefix in brackets (e.g., [BugFix]).

Current title: Feat/sequence rl

Supported Prefixes (case-sensitive)

Your PR title must start with exactly one of these prefixes:

Prefix	Label Applied	Example
`[BugFix]`	BugFix	`[BugFix] Fix memory leak in collector`
`[Feature]`	Feature	`[Feature] Add new optimizer`
`[Doc]` or `[Docs]`	Documentation	`[Doc] Update installation guide`
`[Refactor]`	Refactoring	`[Refactor] Clean up module imports`
`[CI]`	CI	`[CI] Fix workflow permissions`
`[Test]` or `[Tests]`	Tests	`[Tests] Add unit tests for buffer`
`[Environment]` or `[Environments]`	Environments	`[Environments] Add Gymnasium support`
`[Data]`	Data	`[Data] Fix replay buffer sampling`
`[Performance]` or `[Perf]`	Performance	`[Performance] Optimize tensor ops`
`[BC-Breaking]`	bc breaking	`[BC-Breaking] Remove deprecated API`
`[Deprecation]`	Deprecation	`[Deprecation] Mark old function`
`[Quality]`	Quality	`[Quality] Fix typos and add codespell`

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

…n, padded output with mask, and get_env_transforms_from_module

github-actions · 2026-04-30T01:06:48Z

⚠️ PR Title Label Error

PR title must start with a label prefix in brackets (e.g., [BugFix]).

Current title: Feat/sequence rl

Supported Prefixes (case-sensitive)

Your PR title must start with exactly one of these prefixes:

Prefix	Label Applied	Example
`[BugFix]`	BugFix	`[BugFix] Fix memory leak in collector`
`[Feature]`	Feature	`[Feature] Add new optimizer`
`[Doc]` or `[Docs]`	Documentation	`[Doc] Update installation guide`
`[Refactor]`	Refactoring	`[Refactor] Clean up module imports`
`[CI]`	CI	`[CI] Fix workflow permissions`
`[Test]` or `[Tests]`	Tests	`[Tests] Add unit tests for buffer`
`[Environment]` or `[Environments]`	Environments	`[Environments] Add Gymnasium support`
`[Data]`	Data	`[Data] Fix replay buffer sampling`
`[Performance]` or `[Perf]`	Performance	`[Performance] Optimize tensor ops`
`[BC-Breaking]`	bc breaking	`[BC-Breaking] Remove deprecated API`
`[Deprecation]`	Deprecation	`[Deprecation] Mark old function`
`[Quality]`	Quality	`[Quality] Fix typos and add codespell`

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

- SliceSampler: emit FutureWarning when auto-detect picks ('collector', 'traj_ids') over 'episode' in storages that contain both, to flag the BC change vs <0.13. - SliceSampler: probe storage schema via _storage.keys() instead of fetching a row, falling back to storage[0:1] only when the cheap path is unavailable. Narrowed broad except Exception to specific exception types. - SliceSampler: mask shape changed from [B*T, 1] to [B*T] so it broadcasts against the sample's batch dim (matches the trainer's existing usage of ('collector', 'mask')). - SliceSampler: raise ValueError when pad_output=True is combined with strict_length=True. - SliceSampler.__repr__ includes pad_output (getattr default for old pickles). - get_env_transforms_from_module re-exported from torchrl.modules and added to docs/source/reference/modules_utils.rst. - New example examples/replay-buffers/recurrent_slice_sampler_pipeline.py showing collector.start() + replay_buffer=rb + auto-detected traj_key + pad_output + recurrent GRU end-to-end. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

…t=True Recurrent modules (LSTMModule, GRUModule) under set_recurrent_mode("recurrent") already split a flat sequence on is_init via _get_num_per_traj_init and use each split's stored hidden state as the initial state. By writing is_init=True at every slice start (OR-ed with the storage's existing is_init so real episode resets within a slice are preserved), the sampler lets a recurrent policy consume the flat [B*T] sample directly — no manual reshape, no manual splitting in the training loop. The marker is only emitted when the storage already contains an is_init field (i.e. the env had an InitTracker), so we never introduce a key that wasn't there. The example is updated to drop the sample.reshape(B, T) line and pass the flat sample straight to the policy. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

The training loop should never have to touch ('collector', 'mask') itself — mask-aware loss modules consume it internally. The example now just calls policy(sample) under set_recurrent_mode and trusts downstream code to handle masking when relevant. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

…r example The training loop no longer reshapes anything and the mask is invisible to user code, so the example shouldn't keep mentioning either. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

…de too The mask was over-emphasized. The natural recurrent-RL flow does not need padding: trajectories sit end-to-end in the buffer, the sampler returns concatenated variable-length slices, and the RNN splits on is_init. Make SliceSampler write is_init=True at every slice start regardless of pad_output (still OR-ed with whatever InitTracker wrote, still skipped when the storage carries no is_init). The example now uses the default flow (pad_output=False, no mask) and passes the variable-length sample straight to the recurrent policy. The pad_output docstring is reframed: it's a niche convenience for users who specifically want uniform [B*T] output, not the headline feature. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

TorchRL's primitives (recurrent modules in recurrent_mode, split_trajectories, loss modules) all consume concatenated variable-length slices directly via the is_init / truncated markers SliceSampler emits. Padding is an escape hatch for downstream code that genuinely needs a fixed B*T product, not a recommended path. Reframe the docstring accordingly and clarify that the sample shape stays 1D either way — pad_output just makes B*T a fixed product, it does not reshape to [B, T]. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

The auto-wrap helper used by both EnvBase post-init and SyncDataCollector now decides what to add by inspecting the env's full_observation_spec and full_state_spec rather than walking its transform stack. This is correct in the presence of BatchedEnv / SerialEnv (which hides child transforms from the top level) and also handles multi-agent setups where init_keys appear nested under per-agent paths — the leaf-level match catches them. Drop the require_primer flag: the env hook always added an InitTracker but the collector hook didn't (when no primer was found), so users who passed a bare env to the collector with a non-recurrent policy got different behavior depending on whether they passed policy= at env construction. Both paths now have identical semantics — always add InitTracker if absent from the spec, append primers for any recurrent submodules. The helper documents two known limitations: renamed init_keys (we won't recognise them) and policy_factory inputs (we can't inspect a Callable without instantiating it). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

… ParallelEnv New TestEnvTransformAutoWrap class with four cases for the recurrent env auto-wrap path: - bare GymEnv + recurrent policy: collector appends InitTracker + primer - GymEnv(policy=...) then collector wrap: idempotent, single InitTracker - SerialEnv wrapping TransformedEnv(env, InitTracker): no double-wrap - ParallelEnv wrapping TransformedEnv(env, InitTracker): no double-wrap Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

The policy= keyword on EnvBase is metaclass-injected and therefore does not appear in subclass __init__ signatures (GymEnv, DMControlEnv, custom subclasses), so users won't discover it from those docstrings. Add a dedicated section to the Environment API reference describing the behavior, idempotency, the symmetric collector-side path, and the documented limitations (renamed init_keys, policy factories). Mirror the existing metaclass note pattern used for spec_locked / auto_reset on the policy parameter in EnvBase's docstring, and cross-reference the new doc section. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Add an `auto_register_policy_transforms` kwarg to Collector / MultiCollector (forwarded through `_main_async_collector` to per-worker inner collectors). Three states: - True -> spec-based, idempotent append of InitTracker + TensorDictPrimer. - False -> silent no-op. - None (default through v0.14) -> no-op, but emit a FutureWarning when the helper *would* have wrapped, naming the v0.15 flip and the opt-in/out flag values. Default flips to True in v0.15. This preserves pre-0.13 behavior for users who never asked for auto-wrap, gives them two minor releases of warning, and lets early adopters opt in explicitly today. Helper refactor: split `_maybe_append_env_transforms_from_module` into a pure-read `_compute_missing_env_transforms` (used to drive the warning) and the existing apply-side function. Example updated with a comment block laying out the two equivalent paths (`EnvBase(policy=...)` vs `SyncDataCollector(..., auto_register_policy_ transforms=True)`) and the deprecation timeline. Tests cover: opt-in append on a bare env, idempotency through env hook + collector hook, no-double-wrap inside SerialEnv / ParallelEnv, FutureWarning fires under the default, and `=False` is fully silent. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Resolve conflicts in torchrl/collectors/_runner.py, _multi_base.py and _single.py: this branch added an `auto_register_policy_transforms` kwarg to Collector / MultiCollector / `_main_async_collector`, while main added `pre_collect_hook` / `post_collect_hook`. All three are independent additions — keep them all in the constructor signatures, the kwargs forwarded to the worker, and the inner_collector(...) call site. No logic overlaps. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

…plicit shape `_compute_missing_env_transforms` walks the policy via `get_primers_from_module`, which calls every submodule's `make_tensordict_primer()`. Some implementations (e.g. `ConsistentDropoutModule` without `input_shape`) raise when they can't build a primer without further user input. Pre-PR the collector never walked for primers, so this was never exercised; the new dry-run path triggered six TestConsistentDropout failures on CI. Catch RuntimeError/ValueError from `get_primers_from_module` in the dry-run helper and fall back to "no auto-wired primers" — InitTracker detection still runs as normal. Users whose primer-bearing submodules need explicit configuration get the same behaviour they had before this PR (no auto-wrap for that primer); the collector itself doesn't fail. Also opt the two recurrent-collector tests in `test_tensordictmodules.py` into the new flag (`auto_register_policy_transforms=True`); they were relying on the old unconditional auto-wrap behavior. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

vmoens · 2026-05-01T15:17:04Z

Heads-up on the breadth of the auto-wrap change (raised after a CI failure I want to flag explicitly):

The collector-side auto-wrap in this PR is wired up via _compute_missing_env_transforms, which walks the whole policy (module.apply(...)) and calls make_tensordict_primer() on every submodule that exposes one. That ended up being more intrusive than I anticipated:

It's not limited to LSTMModule / GRUModule. Anything that exposes make_tensordict_primer() is invoked.
Some primer-bearing modules need explicit configuration to build their primer. ConsistentDropoutModule is the concrete example we hit: without input_shape it raises RuntimeError. Pre-PR the collector never ran this code path, so policies containing ConsistentDropoutModule worked fine through the collector. With this PR, even the dry-run (used to decide whether to emit the FutureWarning) called make_tensordict_primer() and failed.

Mitigations already in this branch:

The default for auto_register_policy_transforms is None through v0.14: the helper only runs as a dry-run to decide whether to emit a FutureWarning. The default flips to True in v0.15. Users who want today's behavior immediately pass True; users who want to opt out forever pass False.
The dry-run wraps get_primers_from_module in try/except (RuntimeError, ValueError) and falls back to "no auto-wired primers" for that policy. InitTracker detection still runs. So a policy with a ConsistentDropoutModule no longer breaks the collector — it just doesn't get an auto-wired dropout primer (which is the same behavior it had before this PR).

Worth a closer look during review:

Are there other primer-bearing modules in the wild whose make_tensordict_primer() has lazy/conditional init like ConsistentDropoutModule? If yes, those will silently not be auto-wired with the current fallback. Acceptable IMO (the user can wire them manually), but worth documenting.
Is module.apply(...) the right granularity? An alternative would be to look only at top-level submodules of TensorDictSequential, which would miss nested-but-attached recurrent submodules but also avoid surprises like this one.
The two-release deprecation window mitigates the BC risk but doesn't eliminate it. Anyone passing a recurrent policy + a bare env to a collector today is currently broken silently — they get a FutureWarning and a downstream KeyError: 'is_init' rather than a working collector. That's the intentional state through v0.14, but the warning is the only thing telling them.

Happy to tighten any of these if it feels too aggressive.

theap06 · 2026-05-01T22:31:16Z

Heads-up on the breadth of the auto-wrap change (raised after a CI failure I want to flag explicitly):

The collector-side auto-wrap in this PR is wired up via _compute_missing_env_transforms, which walks the whole policy (module.apply(...)) and calls make_tensordict_primer() on every submodule that exposes one. That ended up being more intrusive than I anticipated:

It's not limited to LSTMModule / GRUModule. Anything that exposes make_tensordict_primer() is invoked.

Some primer-bearing modules need explicit configuration to build their primer. ConsistentDropoutModule is the concrete example we hit: without input_shape it raises RuntimeError. Pre-PR the collector never ran this code path, so policies containing ConsistentDropoutModule worked fine through the collector. With this PR, even the dry-run (used to decide whether to emit the FutureWarning) called make_tensordict_primer() and failed.

Mitigations already in this branch:

The default for auto_register_policy_transforms is None through v0.14: the helper only runs as a dry-run to decide whether to emit a FutureWarning. The default flips to True in v0.15. Users who want today's behavior immediately pass True; users who want to opt out forever pass False.

The dry-run wraps get_primers_from_module in try/except (RuntimeError, ValueError) and falls back to "no auto-wired primers" for that policy. InitTracker detection still runs. So a policy with a ConsistentDropoutModule no longer breaks the collector — it just doesn't get an auto-wired dropout primer (which is the same behavior it had before this PR).

Worth a closer look during review:

Are there other primer-bearing modules in the wild whose make_tensordict_primer() has lazy/conditional init like ConsistentDropoutModule? If yes, those will silently not be auto-wired with the current fallback. Acceptable IMO (the user can wire them manually), but worth documenting.

Is module.apply(...) the right granularity? An alternative would be to look only at top-level submodules of TensorDictSequential, which would miss nested-but-attached recurrent submodules but also avoid surprises like this one.

The two-release deprecation window mitigates the BC risk but doesn't eliminate it. Anyone passing a recurrent policy + a bare env to a collector today is currently broken silently — they get a FutureWarning and a downstream KeyError: 'is_init' rather than a working collector. That's the intentional state through v0.14, but the warning is the only thing telling them.

Happy to tighten any of these if it feels too aggressive.

I'll add more exceptions per submodule and write up test cases to determine the issues when considering the whole walk. Thanks for the feedback! Scope makes sense

theap06 · 2026-05-02T00:09:00Z

@vmoens lmk what you think! Thanks for the detailed feedback.

These files (A3C Atari training, ACT demo + image, DreamerV3 image) are unrelated to the sequence-RL composability scope of this PR and were introduced in ab130bc alongside the legitimate "added the try and catch" fix. They were what was actually breaking lint: - sota-implementations/act/act_demo.py:151:20 F541 (f-string no placeholder) - ufmt-formatting drift on the same file Removing them cleans up CI without touching any of theap06's substantive work in `get_primers_from_module(strict=False)` or the new `test_get_primers_from_module_partial_failure` test. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Adds docs/source/reference/data_layout.rst as the cross-cutting reference on the contiguous 1-D trajectory layout: boundary keys (is_init / done / terminated / truncated / traj_ids), why ndim>=2 storages cannot work with multi-process collectors that share a storage, the replay_buffer= + trajs_per_batch handoff that fixes it, SliceSampler defaults, recurrent policy auto-wrapping, and the deprecation story for split_trajectories. Linked from collectors_replay.rst and envs_api.rst. Adds a UserWarning when split_trajs=True is passed explicitly to Collector or MultiCollector, redirecting to SliceSampler and noting it will become a DeprecationWarning in a future release. Default split_trajs=None remains silent. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

- ndim section: soften "does not work" → "runs into trouble"; document set_truncated as a real (if lossy) mitigation; add ReplayBufferEnsemble per-worker pattern with example. - buffer-to-collector handoff: clarify that workers extend the shared buffer directly in both sync and async modes, and that trajs_per_batch acts on the *contents* of each extend, not on parent-side aggregation. - split_trajectories: add the artificial-batch-boundary defect (cuts trajectories at frames_per_batch, baking collector hyperparams into the data shape). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Adds tutorials/sphinx-tutorials/recurrent_sequence_training.py — the multi-step / sequence-training complement to dqn_with_rnn.py (which covers single-step recurrent DQN at collection time). Walks through the post-pytorch#3695 recurrent contract end-to-end: - Collector auto-wiring of InitTracker + the recurrent-state primer via auto_register_policy_transforms=True - Trajectory-aware sampling with SliceSampler - Multi-step LSTM forward under set_recurrent_mode(True) - Boundary safety: hand-built two-trajectory packed batch + isolation check that proves hidden state does not leak across is_init markers - A tiny end-to-end training loop closing the BC-style sequence path Runs in ~3s on CPU. Cross-references the recurrent state lifecycle guide (pytorch#3792), collector internals page (pytorch#3796), and the glossary. Toctree entry added to docs/source/index.rst next to dqn_with_rnn.

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 30, 2026

github-actions Bot added sota-implementations/ ReplayBuffers Modules Trainers Integrations/torch_geometric Integrations labels Apr 30, 2026

theap06 marked this pull request as draft April 30, 2026 00:04

[Feature] Sequence-based RL composability: SliceSampler auto-detectio…

0864d3b

…n, padded output with mask, and get_env_transforms_from_module

theap06 force-pushed the feat/sequenceRL branch from c8bb82b to 0864d3b Compare April 30, 2026 01:06

theap06 marked this pull request as ready for review April 30, 2026 01:06

theap06 changed the title ~~Feat/sequence rl~~ [Feature] Feat/sequence rl Apr 30, 2026

github-actions Bot added the Feature New feature label Apr 30, 2026

vmoens changed the title ~~[Feature] Feat/sequence rl~~ [Feature] sequence rl composability Apr 30, 2026

github-actions Bot added Documentation Improvements or additions to documentation Examples labels Apr 30, 2026

vmoens and others added 3 commits April 30, 2026 17:18

[Doc] Strip stale reshape/mask references from recurrent slice sample…

55b5f4f

…r example The training loop no longer reshapes anything and the mask is invisible to user code, so the example shouldn't keep mentioning either. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

vmoens reviewed Apr 30, 2026

View reviewed changes

Comment thread examples/replay-buffers/recurrent_slice_sampler_pipeline.py Outdated

Comment thread examples/replay-buffers/recurrent_slice_sampler_pipeline.py

[Feature] Auto-wire recurrent env transforms

1352477

github-actions Bot added the Collectors label Apr 30, 2026

theap06 requested a review from vmoens May 1, 2026 00:12

vmoens and others added 3 commits May 1, 2026 10:18

vmoens and others added 5 commits May 1, 2026 13:51

fixed the ci failure and added the try and catch

ab130bc

fixed the linting issue - ready for running ci again

c0b4e4d

vmoens mentioned this pull request May 4, 2026

[RFC]Making sequence-based RL first-class in TorchRL #3678

Closed

vmoens and others added 2 commits May 5, 2026 18:27

vmoens reviewed May 6, 2026

View reviewed changes

Comment thread docs/source/reference/data_layout.rst Outdated

Comment thread docs/source/reference/data_layout.rst

Comment thread docs/source/reference/data_layout.rst Outdated

Comment thread docs/source/reference/data_layout.rst

vmoens merged commit a4164f0 into pytorch:main May 6, 2026
66 of 84 checks passed

This was referenced Jun 11, 2026

[Feature] Mask-aware BCLoss; LossModule._reduce_loss honors ("collect… #3850

Closed

[Feature] Mask-aware BCLoss; LossModule._reduce_loss honors ("collect… #3857

Closed

theap06 mentioned this pull request Jun 13, 2026

[RFC] Mask-aware reduction across LossModule subclasses #3866

Closed

coder-jayp mentioned this pull request Jun 21, 2026

[Refactor] Migrate remaining LossModules to mask-aware reduction (#3866) #3888

Merged

4 tasks

Navigation Menu

Uh oh!

[Feature] sequence rl composability - #3695

[Feature] sequence rl composability#3695
vmoens merged 19 commits into
pytorch:mainfrom
theap06:feat/sequenceRL

theap06 commented Apr 30, 2026

Uh oh!

pytorch-bot Bot commented Apr 30, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Apr 30, 2026

Uh oh!

github-actions Bot commented Apr 30, 2026

Uh oh!

Uh oh!

Uh oh!

vmoens commented May 1, 2026

Uh oh!

theap06 commented May 1, 2026

Uh oh!

theap06 commented May 2, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

theap06 commented Apr 30, 2026

Uh oh!

pytorch-bot Bot commented Apr 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3695

⚠️ 17 Awaiting Approval

Uh oh!

github-actions Bot commented Apr 30, 2026

⚠️ PR Title Label Error

Supported Prefixes (case-sensitive)

Uh oh!

github-actions Bot commented Apr 30, 2026

⚠️ PR Title Label Error

Supported Prefixes (case-sensitive)

Uh oh!

Uh oh!

Uh oh!

vmoens commented May 1, 2026

Uh oh!

theap06 commented May 1, 2026

Uh oh!

theap06 commented May 2, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pytorch-bot Bot commented Apr 30, 2026 •

edited

Loading