[Feature] Collector.fake_tensordict() / MultiCollector.fake_tensordict()

[vmoens]

Stack from ghstack (oldest at bottom):

• [Doc] Migrate shifted=True callers to legacy/compact + docstring polish #3776
• [Example] Expose compact GAE cat dimension #3775
• [Example] Add Isaac RNN PPO rollout mode flags #3769
• [Example] Isaac Lab RNN PPO with compact memory + knowledge-base notes #3763
• [BugFix] Fix GAE compact path bias on recurrent value nets at internal truncations #3771
• [Refactor] Simplify LSTM/GRUModule recurrent-mode shape normalization #3768
• [Refactor] Keep [B, T] dim in value estimators #3767
• -> [Feature] Collector.fake_tensordict() / MultiCollector.fake_tensordict() #3764
• [Feature] MultiCollector: track policy version per fresh continue command #3759
• [BugFix] GAE shifted=True: tolerate missing next obs, non-canonical strides, docs #3757
• [Feature] timeit.mark_start/mark_end for non-context-manager timing #3762
• [BugFix] PolicyVersion: int64 dtype + preserve tensordict device #3755
• [Performance] LSTM/GRU scan: canonical strides + cuDNN flat-storage clones + thread-local recurrent mode #3754
• [BugFix] Recurrent policy auto-register with policy_factory #3753

Public method that returns a zero-filled tensordict shaped exactly like
one batch yielded by the collector, useful for storage initialization
and torch.compile / cudagraph warmup without having to step the env
first.

Collector.fake_tensordict() (single-process):

• Reuses the existing _final_rollout template; builds it lazily via
  _maybe_make_final_rollout(make_rollout=True) even when
  use_buffers=False so the public API is consistent.
• Mirrors the rollout post-pipeline: _maybe_attach_final_obs,
  _maybe_set_truncated, then _postproc (which runs
  split_trajectories, the user postproc, and private-key
  exclusion).
• Result: env keys + policy out-keys + ("collector", "traj_ids"),
  compact_obs exclusions and final_obs UnbatchedTensor
  leaves applied, last dim named "time".

MultiCollector.fake_tensordict() raises NotImplementedError.
Honoring the contract on the parent process would either require
creating an env there (which defeats the purpose of a multi-process
collector — Isaac Lab / mujoco-mjx etc. can only run in workers) or
routing a request to a worker over the pipe (which requires live
workers and adds protocol surface). Neither is in scope here; users
who need a fake tensordict can call it on a single-process
:class:~torchrl.collectors.Collector.

Tests pin: shape / names / keys / zero-fill parity between
fake_tensordict() and next(iter(collector)) (with and without
buffers); compact_obs drops ("next", obs) and final_obs
attaches ("final", obs) as UnbatchedTensor; and that
MultiCollector.fake_tensordict() raises NotImplementedError.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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:

• Run pull request jobs on OSDC runners in shadow mode

❌ 3 New Failures, 1 Cancelled Job

As of commit 7ef328e with merge base 0a01ee8 ():

NEW FAILURES - The following jobs have failed:

• Build Windows Wheels / pytorch/rl / build-wheel-py3_10-cpu (gh)
• Build Windows Wheels / pytorch/rl / upload / upload-wheel-py3_10-cpu (gh)
  Unable to download artifact(s): Artifact not found for name: pytorch_rl__3.10_cpu_x64
• Unit-tests on Linux / tests-cpu (3.14) / linux-job (gh)
  test/objectives/test_dqn.py::TestQMixer::test_dqn_prioritized_weights

CANCELLED JOB - The following job was cancelled. Please retry:

• Unit-tests on Windows / unittests-cpu (3.10, windows.4xlarge, cpu) / windows-job (gh)

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