[Doc] add tutorial on collector trajectory assembly internals by coder-jayp · Pull Request #3600 · pytorch/rl

coder-jayp · 2026-04-07T17:41:28Z

Description

This PR adds a new tutorial: "TorchRL Collectors Deep Dive: Trajectory IDs, Partial Chunks, and Emission".

The tutorial explains the internal mechanics of how TorchRL collectors handle trajectories, which is especially important for users training recurrent or sequence-based policies.

What it covers

Why collectors return fixed-size chunks instead of complete episodes
How split_trajectories() reassembles, pads, and masks trajectories
Internal buffering with _traj_ingest() and _traj_emit()
The meaning and usage of ("collector", "traj_ids") and ("collector", "mask")
done/terminated boundary handling
Padded tensors vs as_nested=True (with performance comparison)
One explicit edge case where a trajectory spans multiple collector batches
Troubleshooting section

The tutorial includes fully runnable code examples and follows the existing TorchRL documentation style.

Motivation and Context

Users training recurrent or sequence-based policies (RNNs, GRUs, LSTMs, Transformers, etc.) frequently need to understand how TorchRL collectors turn raw batches into clean trajectories.

The current documentation shows how to use the collectors, but does not explain the internal dataflow — specifically why trajectories are split across batches, how partial trajectories are buffered and reassembled, and the exact meaning of ("collector", "traj_ids") and ("collector", "mask").

This tutorial fills that gap with clear explanations, runnable examples, and an edge case where a trajectory spans multiple collector iterations.

Types of changes

Documentation (update in the documentation)

Checklist

I have read the CONTRIBUTION guide
My change requires a change to the documentation.
I have updated the documentation accordingly.

Closes pytorch#3588 - Explains split_trajectories(), _traj_ingest(), _traj_emit() - Includes edge case of trajectory spanning multiple batches - Covers mask semantics, padded vs nested outputs, and troubleshooting - Follows TorchRL documentation style

pytorch-bot · 2026-04-07T17:41:33Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

Rolling out OSDC (ARC) runners on pull workflow for PyTorch trunk commits

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

github-actions · 2026-04-07T17:41:44Z

⚠️ PR Title Label Error

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

Current title: docs: add tutorial on collector trajectory assembly internals

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]).

coder-jayp · 2026-04-08T03:07:40Z

Hi @vmoens, it would be great to get some feedback. Thanks.

vmoens

Thanks for this!
I think the goal was not only to do this but also cover recent new features such as

collectors that yield fully built trajectories
populating replay buffers with these within the buffer.

Another item that needs to be thoroughly documented here is the done/truncated handling in split_trajectories(). Eg, if a trajectory is incomplete, what do we do with its last truncated value?

There are also some efficiency considerations (eg, using nested-tensors instead of padding, there's a kwarg in split_trajectories for that).

More generally, the tutorial should be more written in plain english, with intro ("in this tutorial you will learn ...") and conclusion ("we have covered ..." , "useful resources: X, Y, Z in tutorials and documentation).

Internal buffering (private features like _traj_ingest and such) should not be part of the public doc.

Closes pytorch#3588 - Explains split_trajectories(), trajs_per_batch, replay buffer usage - Covers mask semantics, done vs truncated, padded vs nested outputs - Follows TorchRL documentation style

coder-jayp · 2026-04-09T17:08:27Z

Hi @vmoens,

Thank you very much for the detailed feedback!

I have updated the tutorial based on your comments:

Removed all references to private internals (_traj_ingest, _traj_emit)
Added sections on trajs_per_batch and how to populate replay buffers with complete trajectories
Added clear explanation of done vs truncated handling
Kept the note about as_nested=True for efficiency
Rewrote the tutorial in plainer English with a proper introduction and conclusion

The tutorial is now focused only on public APIs and practical usage.

Looking forward to your feedback.

Thanks again

…ajectory-assembly

vmoens

Let's not use print() statements but plain sphinx gallery sections.
I find the explanation pretty shallow. A tutorial should be a simple, realistic example (training, eval, ...) where we show how we're using a specific feature to solve a given problem.
The synthetic data here is a bit hard to understand. Ideally we'd like data that was generated from a real env, show what the data looks like, then move on to show what happens if we don't assemble trajectories and what we can do if we assemble them, then perhaps show a concrete example of what this enables.

Last point: we should register this in the list of tutorials in the index otherwise the tutorial is invisible in the docs

Closes pytorch#3588 - Explains split_trajectories(), trajs_per_batch, and replay buffer usage - Covers mask semantics, done vs truncated handling, padded vs nested outputs - Uses real environment example and plain Sphinx gallery sections - Follows TorchRL documentation style

vmoens · 2026-04-10T17:36:28Z

I have now removed the old file coding_dqn.py completely.
Only the clean collector_trajectory_assembly.py remains.

I'm not sure I'm following, why do we need to delete anything?

Sorry for the confusion.

I saw that coding_dqn.py contained my earlier code with private internals (_traj_ingest, _traj_emit) and many print() statements.
So I removed it to keep only the clean version collector_trajectory_assembly.py.

Closes pytorch#3588 - Explains split_trajectories(), trajs_per_batch, and replay buffer usage - Covers mask semantics, done vs truncated handling, padded vs nested outputs - Uses real environment example and plain Sphinx gallery sections - Follows TorchRL documentation style

…n.py - Rewrite tutorial to use real env data (GymEnv + SyncDataCollector) instead of synthetic TensorDict, replace print headers with Sphinx Gallery RST sections, and register in docs index - git rm _version.py (auto-generated by setuptools_scm), add to .gitignore, and add pre-commit hook to prevent future commits Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…ng_dqn.py Merge remote branch, resolve conflicts: - Keep _version.py deleted (auto-generated, must not be tracked) - Keep our rewritten tutorial (real env data, no synthetic TensorDict) - Restore coding_dqn.py that was accidentally deleted by remote Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Show how passing a replay buffer to the collector and calling start() enables fully asynchronous background collection, including with trajs_per_batch for complete trajectory batches. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

The replay buffer in the collector.start() example now uses a SliceSampler(slice_len=16, end_key=("next", "done")) to sample contiguous sub-sequences that respect episode boundaries, showcasing the natural fit with complete trajectory storage. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Add seealso in dqn_with_rnn.py pointing to the collector trajectory assembly tutorial for details on split_trajectories, trajs_per_batch, SliceSampler, and async collection. Also fix broken RNN tutorial cross-ref in collector_trajectory_assembly.py. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

The "Storing trajectories in a replay buffer" section now uses a SliceSampler so sampled batches are contiguous sub-sequences rather than isolated transitions. Added seealso linking to the replay buffer tutorial's trajectory storage section for more detail. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

vmoens · 2026-04-10T19:54:29Z

@coder-jayp I restored the dqn tutorial and edited the trajectory one. Can you give it a look?

SliceSampler operates on flat transition storage and uses end_key to locate episode boundaries. Extending the buffer with pre-assembled (trajs, max_len) tensors causes a shape error ("Expected the end-of-trajectory signal to be 1-dimensional"). Fix both the replay buffer section and the async collector.start() section to extend with flat collector batches instead of trajs_per_batch output. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

coder-jayp · 2026-04-11T16:38:18Z

Thanks @vmoens. Really appreciate it.

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 7, 2026

github-actions Bot added the tutorials/ label Apr 7, 2026

coder-jayp changed the title ~~docs: add tutorial on collector trajectory assembly internals~~ [Doc] add tutorial on collector trajectory assembly internals Apr 7, 2026

github-actions Bot added the Documentation Improvements or additions to documentation label Apr 7, 2026

vmoens reviewed Apr 8, 2026

View reviewed changes

coder-jayp added 2 commits April 9, 2026 22:23

docs: add tutorial on collector trajectory assembly internals

ad53c81

Closes pytorch#3588 - Explains split_trajectories(), trajs_per_batch, replay buffer usage - Covers mask semantics, done vs truncated, padded vs nested outputs - Follows TorchRL documentation style

docs: improve tutorial introduction and conclusion based on feedback

9d3ab10

Merge remote-tracking branch 'origin/main' into tutorial-collector-tr…

4cd2579

…ajectory-assembly

vmoens reviewed Apr 9, 2026

View reviewed changes

coder-jayp added 2 commits April 10, 2026 23:03

vmoens reviewed Apr 10, 2026

View reviewed changes

coder-jayp and others added 7 commits April 10, 2026 23:14

vmoens and others added 2 commits April 10, 2026 14:08

[Doc] Show trajectory count in SliceSampler sample

f28eceb

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

vmoens merged commit 4b1f0f4 into pytorch:main Apr 10, 2026
51 of 78 checks passed

Navigation Menu

Uh oh!

[Doc] add tutorial on collector trajectory assembly internals - #3600

[Doc] add tutorial on collector trajectory assembly internals#3600
vmoens merged 15 commits into
pytorch:mainfrom
coder-jayp:tutorial-collector-trajectory-assembly

coder-jayp commented Apr 7, 2026

Uh oh!

pytorch-bot Bot commented Apr 7, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Apr 7, 2026

Uh oh!

coder-jayp commented Apr 8, 2026

Uh oh!

vmoens left a comment

Uh oh!

coder-jayp commented Apr 9, 2026

Uh oh!

vmoens left a comment

Uh oh!

vmoens Apr 10, 2026

Uh oh!

coder-jayp Apr 10, 2026

Uh oh!

vmoens Apr 10, 2026

Uh oh!

coder-jayp Apr 10, 2026

Uh oh!

vmoens commented Apr 10, 2026

Uh oh!

Uh oh!

coder-jayp commented Apr 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

coder-jayp commented Apr 7, 2026

Description

What it covers

Motivation and Context

Types of changes

Checklist

Uh oh!

pytorch-bot Bot commented Apr 7, 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/3600

❗ 1 Active SEVs

Uh oh!

github-actions Bot commented Apr 7, 2026

⚠️ PR Title Label Error

Supported Prefixes (case-sensitive)

Uh oh!

coder-jayp commented Apr 8, 2026

Uh oh!

vmoens left a comment

Choose a reason for hiding this comment

Uh oh!

coder-jayp commented Apr 9, 2026

Uh oh!

vmoens left a comment

Choose a reason for hiding this comment

Uh oh!

vmoens Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

coder-jayp Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

vmoens Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

coder-jayp Apr 10, 2026

Choose a reason for hiding this comment

Uh oh!

vmoens commented Apr 10, 2026

Uh oh!

Uh oh!

coder-jayp commented Apr 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pytorch-bot Bot commented Apr 7, 2026 •

edited

Loading