[Hang-get]

Problem

split_trajectories runs on each collector batch when split_trajs=True. It first found the consecutive trajectory IDs, then rescanned the full flattened ID tensor once per trajectory to compute each length. For N steps and T trajectories, that count was O(N * T) and called .item() T times.

Change

Use torch.unique_consecutive(..., return_counts=True) to obtain all trajectory lengths in the existing pass, followed by one conversion to a Python list. The output layout and padding behavior are unchanged. This also adds an irregular-length equivalence test and a collector benchmark.

Benchmark

Intel Xeon Platinum 8352Y, Linux, Python 3.11.18, PyTorch 2.15.0+cpu, one PyTorch thread. Input: 32,768 steps arranged as 1,024 adjacent trajectories of length 32, with a four-value observation. The benchmark calls the real split_trajectories function.

Implementation	Median
main	82.532 ms
This PR	35.231 ms

This is a 2.34x speedup. Both measurements used pytest-benchmark with at least five rounds; the same output assertions passed before and after the change.

Duplicate check

I searched open and closed issues and PRs for split_trajectories, trajectory splitting, collector performance, and slow collectors. I found no open implementation of this change.

Tests

• python -m pytest -q test/test_postprocs.py::TestSplits — 397 passed
• python -m pytest -q test/test_postprocs.py::TestSplits::test_consecutive_trajectory_lengths — passed on the baseline and this change
• python -m pytest -q benchmarks/test_collectors_benchmark.py::test_split_trajectories_many_short --benchmark-min-rounds=5 — passed
• pre-commit run --files torchrl/collectors/utils.py test/test_postprocs.py benchmarks/test_collectors_benchmark.py — passed

I did not run the full test suite or CUDA benchmarks.

AI assistance (Cursor) was used for this change.

Made with Cursor

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⚠️ 16 Awaiting Approval

As of commit a1814f8 with merge base a711ef9 ():

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

• Build Aarch64 Linux Wheels (gh)
• Build Linux Wheels (gh)
• Build M1 Wheels (gh)
• Build Windows Wheels (gh)
• Continuous Benchmark (PR) (gh)
• Examples Tests on Linux (gh)
• Generate documentation (gh)
• Habitat Tests on Linux (gh)
• Libs Tests on Linux (gh)
• Lint (gh)
• LLM Tests on Linux (gh)
• Push Binary Nightly (gh)
• SOTA Tests on Linux (gh)
• Tutorials Tests on Linux (gh)
• Unit-tests on Linux (gh)
• Validate Test Partitioning (gh)

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

[meta-cla]

Hi @Hang-get!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!