[Doc] Document CatFrames with collectors and image replay buffers

[YeonwooSung]

Description

Adds a focused Sphinx recipe for using CatFrames with a data collector and a replay buffer on images, which was missing from the reference docs.

Two legitimate placements are documented and cross-linked:

• Env-side (docs/source/reference/envs_transforms.rst): stateful stacking for the policy, reset / InitTracker behavior, and dim=-3 for CHW images.
• Buffer-side (docs/source/reference/data_replaybuffers.rst): store unstacked raw pixels, rebuild the stack in rb.sample(), and wire a Collector through extend / sample. Shows both CatFrames.make_rb_transform_and_sampler and the explicit SliceSampler form.

The recipe also covers why raw uint8 frames are what you store, how the collector's ("collector", "traj_ids") interact with the slice sampler, and the common pitfalls of stacking twice (env and buffer on the same key) or using the vector default dim=-1 on pixels.

Copy-paste snippets use # doctest: +SKIP where gym / pixel rendering is required. The CatFrames class docstring now points at both sections.

Code example

Rebuild four-frame CHW stacks when sampling, while storing individual frames. Synthetic numbered grayscale pixels make sequence continuity visible without rendering dependencies.

import torch
from tensordict import TensorDict
from torchrl.data import LazyTensorStorage, ReplayBuffer
from torchrl.envs import CatFrames

pixels = torch.arange(100, dtype=torch.float32).reshape(100, 1, 1, 1)
data = TensorDict({
    "pixels": pixels,
    ("next", "pixels"): pixels + 1,
    ("collector", "traj_ids"): torch.zeros(100, dtype=torch.long),
    ("next", "done"): torch.zeros(100, 1, dtype=torch.bool),
    ("next", "terminated"): torch.zeros(100, 1, dtype=torch.bool),
}, batch_size=[100])
catframes = CatFrames(
    N=4, dim=-3, in_keys=["pixels"], out_keys=["pixels_stack"], reset_key="_reset",
)
transform, sampler = catframes.make_rb_transform_and_sampler(
    batch_size=8, traj_key=("collector", "traj_ids"),
)
rb = ReplayBuffer(
    storage=LazyTensorStorage(100), batch_size=8, sampler=sampler, transform=transform,
)
rb.extend(data)  # In training, extend with each collector batch.
batch = rb.sample()
assert batch["pixels_stack"].shape == (8, 4, 1, 1)
frames = batch["pixels_stack"][..., 0, 0]
assert (frames[:, 1:] - frames[:, :-1] == 1).all()

Motivation and Context

Using CatFrames for inference was already documented; reconstructing a frame stack when sampling from a replay buffer was not, especially for images (stack dim -3, not the vector default). Visual RL relies on this pattern, and the collector + extend + sample path is easy to get wrong (double stacking, wrong dim, no time axis).

close #2618

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

• Documentation (update in the documentation)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 f5ef733 with merge base 1d3de3d ():

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.

[vmoens]

The helper-based collector/replay example runs and returns [32, 4, 84, 84] stacks. The explicit alternative is missing the reshape/window-selection steps and mixes unrelated samples.

[vmoens]

[P2] Reshape sampled windows before CatFrames and retain their last step

SliceSampler returns the B windows flattened as [B*N]; it does not make a [B, N] TensorDict for this transform. Consequently CatFrames treats the entire sample as one time sequence, so the first N-1 rows of later windows incorporate frames from an unrelated preceding window. A synthetic increasing-frame replay reproduces stacks such as [52, 53, 54, 66]; 12 of 16 rows were wrong with N=4. Match the helper by reshaping to [-1, N] before CatFrames and selecting [:, -1] afterward, or remove this purported equivalent recipe.

[YeonwooSung]

Fixed in f5ef733. The explicit recipe now matches make_rb_transform_and_sampler: reshape(-1, N) before CatFrames and [:, -1] after it. A synthetic increasing-frame buffer reproduced the bleed ([54, 55, 56, 61] and only 16/64 consecutive rows with N=4); after the reshape/last-step pair every window is consecutive and the sample shape is [B, N] like the helper. The surrounding prose and the pitfalls list call out that SliceSampler returns [B * N], not [B, N].

[YeonwooSung]

Addressed the CHANGES_REQUESTED review in f5ef733.

The helper-based collector/replay example was already correct. The explicit SliceSampler alternative was missing the reshape / last-step pair, so CatFrames treated the flattened [B * N] sample as one time sequence and mixed frames across windows. That recipe now does reshape(-1, N) then [:, -1], matching make_rb_transform_and_sampler.

[vmoens]

Re-reviewed f5ef733. The manual replay recipe now restores the [batch, time] slice layout before CatFrames and selects the last item of each slice afterward. Sampling batch_size * frame_stack raw transitions matches the helper recipe.

I ran both rendered CartPole replay examples with a smaller collection/storage size; both return 32 samples with pixels_trsf shape [32, 4, 84, 84]. No remaining actionables.