[Refactor] ActionChunkTransform as a CatFrames recipe

[vmoens]

Motivation

ActionChunkTransform re-implemented the moving-window gather that CatFrames already owns: forward action chunks are exactly backward frame windows up to time reversal (the chunk [t, ..., t+H-1] is the reversed window at T-1-t read backwards, and padding="same" start-padding of the reversed sequence is the chunk's repeat-last end padding). This PR keeps the two user-facing classes (ergonomics and discoverability) but leaves a single windowing implementation underneath, following the R3MTransform(Compose) recipe pattern.

Changes

CatFrames (offline path) gains two options

• future=True: forward-looking windows [t, ..., t+N-1], implemented as time-reverse -> existing unfold/pad core -> reverse back + flip the window axis, so the padding="same" machinery is reused verbatim. Offline-only: the env step path raises (an online buffer cannot see future frames). The done entry is optional in this mode (absent done = one contiguous trajectory per row).
• mask_key=...: exposes the per-window validity mask that unfold_done already computed and then discarded (True = fabricated/padded slot, the action_is_pad convention). Also available for history windows.

ActionChunkTransform becomes a recipe

ActionChunkTransform(Compose) =
    UnsqueezeTransform(dim=-2, in_keys=[action_key], out_keys=[chunk_key])
    + CatFrames(N=chunk_size, dim=-2, padding="same", future=True, mask_key=pad_key)

The constructor, keys and outputs are unchanged; an equivalence test pins the recipe byte-identical to the previous gather. Env attachment remains a documented no-op (the env-facing paths are overridden, and clone() rebuilds the recipe so the overrides survive re-wrapping).

Behavioral upside: boundary-aware chunks. When the sampled window carries ("next", done), chunks no longer leak actions across episode boundaries inside a window: the steps past a done are padded with the last in-trajectory action and flagged in action_is_pad (new done_key kwarg; None opts out; absent done keeps the previous behavior). ActionChunkTransform shipped in the VLA stack and has not been released, so this needs no deprecation cycle.

Fixes and cleanups picked up along the way

• _apply_same_padding is now a vectorized gather instead of a Python loop over batch*time samples, which also makes the offline path compile-friendly (covered by a compile smoke test).
• unfold_done builds its leading no-reset block explicitly instead of slicing the done entry, fixing offline windows longer than the trajectory (N > T used to crash).
• A missing done entry in (history-mode) unfolding now raises an informative KeyError.
• CatFramesConfig brought to kwarg parity (padding, padding_value, as_inverse, reset_key, done_key, future, mask_key).

Tests and validation

• New: byte-identical equivalence vs the previous gather (multiple batch shapes and chunk sizes), done-aware chunking against hand-computed tables, done_key=None opt-out, done-shape validation, clone-keeps-recipe, recipe under torch.compile; CatFrames future/mask_key value tests (with and without dones, both paddings), env-path raises, ("next", key) overlap raises, N > T, missing-done KeyError.
• Existing: full TestCatFrames (287) and TestActionChunkTransform suites pass unchanged, including the functional-equivalence tests that pin the vectorized padding byte-identical to the loop it replaces; test/objectives/test_vla.py, test/objectives/test_bc.py, config tests, doctests and the VLA tutorial run clean.
• Benchmark: offline windowing path added to benchmarks/test_replaybuffer_benchmark.py.

Generated with Claude Code

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 3 New Failures, 9 Pending, 3 Unrelated Failures

As of commit 3cb1fce with merge base 45aa2e2 ():

NEW FAILURES - The following jobs have failed:

• Build Windows Wheels / pytorch/rl / build-wheel-py3_10-cpu (gh)
  Process completed with exit code 1.
• Unit-tests on Linux / tests-olddeps (3.10, 11.8) / linux-job (gh)
  test/transforms/test_action_transforms.py::TestActionChunkTransform::test_compile
• Unit-tests on Windows / unittests-cpu (3.10, windows.4xlarge, cpu) / windows-job (gh)
  ImportError: Can't connect to HTTPS URL because the SSL module is not available.

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• Unit-tests on Linux / tests-cpu (3.12) / linux-job (gh) (similar failure)
  test/test_collectors.py::TestCollectorGeneric::test_multi_sync_data_collector_ordering[stack-True-8-5]

BROKEN TRUNK - The following jobs failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Build Windows Wheels / pytorch/rl / upload / upload-wheel-py3_10-cpu (gh) (trunk failure)
  Unable to download artifact(s): Artifact not found for name: pytorch_rl__3.10_cpu_x64
• Unit-tests on Linux / tests-cpu (3.10) / linux-job (gh) (trunk failure)
  test/test_collectors.py::TestCollectorGeneric::test_multi_sync_data_collector_ordering[stack-True-8-5]

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