[BugFix] GAE shifted=True: tolerate missing next obs, non-canonical strides, docs

[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

Three small improvements to shifted-GAE in _call_value_nets:

• Add _fill_missing_next_inputs(next_data, root_data, in_keys). When
  compact_obs=True drops ("next", obs) from the rollout, the
  shifted-GAE bootstrap path was dereferencing a None key and writing
  NaNs into the assembled batch. Fall back to the root key value at the
  done position for any missing next-input key.

• Replace data_copy.view(-1) with data_copy.reshape(-1).
  Replay-buffer reads and memmap-backed tensors can expose
  non-canonical strides that view rejects; reshape accepts both.
  Drop the with ... as data_copy_view context that was paired with
  the in-place view — flattened tensors don't need write-back.

• Hoist ndim resolution out of the inner _call_value_net closure
  so it's available to follow-up cleanups.

• Docstring update on shifted=True across TD0 / TD1 / TDLambda /
  GAE / VTrace clarifying that the path expects long, contiguous
  trajectory windows with valid boundary next states.

Tests cover: shifted-bootstrap parity between a compact rollout (no
("next", obs)) and a NaN-filled reference for TD0 / TD1 / TDLambda /
GAE, non-canonical strides via transpose(0, 1) through shifted-GAE,
and a vectorized-axis sweep on the recurrent shifted/non-shifted parity
test.

[pytorch-bot]

🔗 Helpful Links

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

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

❌ 4 New Failures, 1 Cancelled Job, 2 Unrelated Failures

As of commit 8d89292 with merge base 0a01ee8 ():

NEW FAILURES - The following jobs have failed:

• Build Windows Wheels / pytorch/rl / build-wheel-py3_10-cpu (gh)
  Process completed with exit code 1.
• Libs Tests on Linux / unittests-gym (3.10, 12.8) / linux-job (gh)
  test/libs/test_gym.py::TestGym::test_set_seed_and_reset_works
• Unit-tests on Linux / tests-olddeps (3.10, 11.8) / linux-job (gh)
  test/transforms/test_env_transforms.py::TestTargetReturn::test_parallel_trans_env_check[device0-constant]
• Unit-tests on Linux / tests-stable-gpu (3.12, 13.0, 2) / linux-job (gh)
  test/test_collectors.py::TestUpdateParams::test_param_sync_mixed_device[SharedMemWeightSyncScheme-cpu-env_device0-False-MultiAsyncCollector]

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

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

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• Build Windows Wheels / pytorch/rl / upload / upload-wheel-py3_10-cpu (gh) (detected as infra flaky with no log or failing log classifier)
• Unit-tests on Linux / tests-optdeps (3.12, 13.0) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

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