diff --git a/docs/source/reference/data_replaybuffers.rst b/docs/source/reference/data_replaybuffers.rst index 1f1425ba642..2d890591772 100644 --- a/docs/source/reference/data_replaybuffers.rst +++ b/docs/source/reference/data_replaybuffers.rst @@ -506,6 +506,9 @@ trajectory-edge padding, while ``rebin``'s stack keeps a separate frame axis:: ... ), ... ) # doctest: +SKIP +The same ``dim=-3`` convention, and the collector / ``extend`` / ``sample`` +wiring, is spelled out in :ref:`catframes-collector-replay`. + **Multiple files.** A clip is often split across many small files (one per episode) rather than one large mp4. :meth:`VideoClipRef.from_files` addresses a list of files as a single logical sequence, so slicing, :meth:`rebin` and decoding work across @@ -529,3 +532,246 @@ When camera and control loops run at different rates, prefer VideoClipRef clear_video_decoder_cache set_video_decoder_cache_size + +.. _catframes-collector-replay: + +Frame stacking images with CatFrames +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Visual off-policy training almost always stacks the last ``N`` frames so a +CNN can see motion. :class:`~torchrl.envs.transforms.CatFrames` can do that +in two places; pick **one** of them for the tensor that is stored. + +- **Env-side** (policy / inference): a stateful rolling buffer on the + :class:`~torchrl.envs.TransformedEnv`. Documented in + :ref:`CatFrames for visual RL `. +- **Buffer-side** (this section): store **unstacked** raw pixels and + rebuild the stack in ``rb.sample()``. A stack of ``N`` float32 frames + occupies ``N`` times the RAM of one uint8 image; putting + :class:`~torchrl.envs.transforms.CatFrames` on the sample path is how + you avoid paying that cost in the storage. + +The two placements share ``N`` and ``dim`` so the policy and the loss see +the same layout. They must not both write the stacked key: env stacking +**and** buffer stacking of the already-stacked tensor produces +``[N * N * C, H, W]``. + +For CHW images the stack dimension is ``dim=-3`` (channel), not the +vector default ``dim=-1``. After :class:`~torchrl.envs.transforms.ToTensorImage` +(and optional :class:`~torchrl.envs.transforms.GrayScale`) a frame is +``[C, H, W]``; concatenating along ``-3`` yields ``[N * C, H, W]``. + +Why the raw frame is what you store +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Keep the env transform and the stored tensor on **different keys**. +:class:`~torchrl.envs.transforms.ToTensorImage` should write a processed +copy (``out_keys=["pixels_trsf"]``) and leave the env's ``"pixels"`` +untouched. :class:`~torchrl.envs.transforms.CatFrames` then stacks +``pixels_trsf`` for the policy. On ``rb.extend(...)`` you drop +``pixels_trsf`` (and its ``("next", ...)`` counterpart) so the storage +holds only the uint8 frame. The buffer transform recreates +``pixels_trsf`` from ``"pixels"`` at sample time. + +:meth:`~torchrl.envs.transforms.CatFrames.make_rb_transform_and_sampler` +builds the sample-time half of that pipeline: a +:class:`~torchrl.data.replay_buffers.SliceSampler` with ``slice_len=N``, +and a transform that reshapes each sampled slice to ``[B, N]``, unfolds +:class:`~torchrl.envs.transforms.CatFrames` along time, keeps the last +step of every window, and -- on the inverse / write path -- excludes the +stacked ``out_keys``. Offline :class:`~torchrl.envs.transforms.CatFrames` +(``forward`` / ``unfolding``) needs a time dimension and +``("next", "done")`` to stop stacks at episode boundaries; the helper +sampler is what provides that time axis. Sampling independent transitions +and then unfolding treats the batch index as time and mixes unrelated +frames. + +Collector, ``extend`` and ``sample`` +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +A :class:`~torchrl.collectors.Collector` writes ``("collector", "traj_ids")`` +on every batch. Pass that key to the helper so the +:class:`~torchrl.data.replay_buffers.SliceSampler` agrees with the +collector on trajectory boundaries. Then ``extend`` each collected batch +and ``sample`` as usual: + +.. code-block:: python + + from torchrl.collectors import Collector, RandomPolicy + from torchrl.data import LazyTensorStorage, ReplayBuffer + from torchrl.envs import ( + CatFrames, + Compose, + GrayScale, + GymEnv, + InitTracker, + Resize, + StepCounter, + ToTensorImage, + TransformedEnv, + ) + + frame_stack = 4 + batch_size = 32 + + catframes = CatFrames( + N=frame_stack, + dim=-3, + in_keys=["pixels_trsf"], + out_keys=["pixels_trsf"], + ) + env = TransformedEnv( + GymEnv("CartPole-v1", from_pixels=True, pixels_only=True), + Compose( + InitTracker(), + ToTensorImage(in_keys=["pixels"], out_keys=["pixels_trsf"]), + GrayScale(in_keys=["pixels_trsf"]), + Resize(84, 84, in_keys=["pixels_trsf"]), + catframes, + StepCounter(), + ), + ) # doctest: +SKIP + + rb_catframes, sampler = catframes.make_rb_transform_and_sampler( + batch_size=batch_size, + traj_key=("collector", "traj_ids"), + ) + # Sample-path processing must match the env: same ToTensorImage / + # GrayScale / Resize, but applied to both the root and the "next" + # pixels so CatFrames sees the layout it saw during collection. + rb = ReplayBuffer( + storage=LazyTensorStorage(100_000), + sampler=sampler, + batch_size=batch_size, + transform=Compose( + ToTensorImage( + in_keys=["pixels", ("next", "pixels")], + out_keys=["pixels_trsf", ("next", "pixels_trsf")], + ), + GrayScale(in_keys=["pixels_trsf", ("next", "pixels_trsf")]), + Resize(84, 84, in_keys=["pixels_trsf", ("next", "pixels_trsf")]), + rb_catframes, + ), + ) # doctest: +SKIP + + collector = Collector( + env, + RandomPolicy(env.action_spec), + frames_per_batch=64, + total_frames=10_000, + ) # doctest: +SKIP + for data in collector: + # Inverse ExcludeTransform on rb_catframes already drops the + # stacked keys on write; excluding them here is equivalent and + # makes the stored tensordict obvious. + rb.extend(data.exclude("pixels_trsf", ("next", "pixels_trsf"))) + batch = rb.sample() + # batch["pixels_trsf"] has shape [32, 4, 84, 84] (grayscale, dim=-3) + +``rb.extend(data)`` is the only write API you need: the collector yields +a tensordict of shape ``[frames_per_batch]`` (or +``[batch, time]`` for a batched env -- then give the storage +``ndim=2``, see :ref:`collectors and replay buffers `). +Do not flatten away ``("collector", "traj_ids")`` or ``("next", "done")`` +before extending; the sampler uses them to keep each slice inside one +episode. + +The same ``extend`` / ``sample`` pairing without the helper looks like +this. Three extra steps are not optional: + +- ``SliceSampler(slice_len=N)``: without it, + :meth:`~torchrl.envs.transforms.CatFrames.unfolding` has no time + dimension. +- ``reshape(-1, N)``: ``SliceSampler`` returns the ``B`` windows + flattened as ``[B * N]``. Leaving that rank-1 batch as the time + axis makes later windows inherit the last frames of the previous + window. +- ``[:, -1]``: after unfolding, each window is an ``N``-step + tensordict; the last step is the completed stack. + +.. code-block:: python + + from torchrl.data import SliceSampler + from torchrl.envs import ExcludeTransform + + rb = ReplayBuffer( + storage=LazyTensorStorage(100_000), + sampler=SliceSampler( + slice_len=frame_stack, + traj_key=("collector", "traj_ids"), + ), + batch_size=batch_size * frame_stack, # B windows of length N + transform=Compose( + ToTensorImage( + in_keys=["pixels", ("next", "pixels")], + out_keys=["pixels_trsf", ("next", "pixels_trsf")], + ), + GrayScale(in_keys=["pixels_trsf", ("next", "pixels_trsf")]), + Resize(84, 84, in_keys=["pixels_trsf", ("next", "pixels_trsf")]), + # SliceSampler returns [B * N]; CatFrames needs a time axis + # per window, then only the completed stack is kept. + lambda td: td.reshape(-1, frame_stack), + CatFrames( + N=frame_stack, + dim=-3, + in_keys=["pixels_trsf", ("next", "pixels_trsf")], + out_keys=["pixels_trsf", ("next", "pixels_trsf")], + ), + lambda td: td[:, -1], + ExcludeTransform("pixels_trsf", ("next", "pixels_trsf"), inverse=True), + ), + ) # doctest: +SKIP + rb.extend(data) # stores raw "pixels" only + batch = rb.sample() # [32, 4, 84, 84] grayscale stacks + +The helper is the preferred form: it multiplies the requested +``batch_size`` by ``N`` internally and inserts the reshape / last-step +selection, so ``rb.sample()`` returns ``batch_size`` stacked +transitions rather than a flat ``batch_size * N`` window. + +Common pitfalls +^^^^^^^^^^^^^^^ + +- **Stacking twice.** Env :class:`~torchrl.envs.transforms.CatFrames` + writing ``pixels_trsf`` **and** a buffer + :class:`~torchrl.envs.transforms.CatFrames` that reads that same + already-stacked key. Store raw ``"pixels"`` and rebuild, or store the + stack and do not attach :class:`~torchrl.envs.transforms.CatFrames` to + the buffer -- not both. +- **Wrong ``dim``.** Images after + :class:`~torchrl.envs.transforms.ToTensorImage` are CHW: use + ``dim=-3``. ``dim=-1`` is for vector observations. ``dim=-4`` is only + correct after an + :class:`~torchrl.envs.transforms.UnsqueezeTransform` that inserted that + axis (the ``[N, C, H, W]`` variant in + ``examples/replay-buffers/catframes-in-buffer.py``). +- **No time axis.** Offline :class:`~torchrl.envs.transforms.CatFrames` + unfolds along time. A uniform random sample of transitions has no + time axis, so the stack is assembled from unrelated steps. Always pair + the buffer transform with a + :class:`~torchrl.data.replay_buffers.SliceSampler` (or + :meth:`~torchrl.envs.transforms.CatFrames.make_rb_transform_and_sampler`). +- **Flattened slices.** ``SliceSampler`` returns ``[B * N]``, not + ``[B, N]``. :class:`~torchrl.envs.transforms.CatFrames` then treats + the whole sample as one sequence and the first ``N - 1`` rows of each + later window mix frames from the previous window. Reshape to + ``[-1, N]`` before the transform and keep ``[:, -1]`` after it + (the helper does both). +- **Missing ``("next", ...)`` keys.** A buffer transform does not walk + into ``"next"`` on its own. List both ``"pixels"`` and + ``("next", "pixels")`` on every sample-path transform that should + apply to both. +- **Same in/out key on** + :class:`~torchrl.envs.transforms.ToTensorImage`. If the processed + tensor overwrites ``"pixels"``, there is no cheap raw frame left to + store and you cannot drop the stack on ``extend``. + +.. seealso:: + + :ref:`CatFrames for visual RL ` for the env-side + placement, reset / :class:`~torchrl.envs.transforms.InitTracker` + behaviour, and the ``dim=-3`` convention. + :ref:`Collectors and replay buffers ` for + ``ndim`` / ``traj_key`` when the collector is batched or + multi-process. A runnable script of the extra-axis (``dim=-4``) + variant is ``examples/replay-buffers/catframes-in-buffer.py``. diff --git a/docs/source/reference/envs_transforms.rst b/docs/source/reference/envs_transforms.rst index fc6b3616830..f8f17121a27 100644 --- a/docs/source/reference/envs_transforms.rst +++ b/docs/source/reference/envs_transforms.rst @@ -192,7 +192,107 @@ transform has an inverse transform that subtracts 1 from the action tensor: Using a Transform with a Replay Buffer ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -You can use a transform with a replay buffer by passing it to the ReplayBuffer constructor: +A transform passed to :class:`~torchrl.data.ReplayBuffer` (the ``transform=`` +argument, or :meth:`~torchrl.data.ReplayBuffer.append_transform`) runs on the +**sample** path: ``rb.sample()`` applies it after the storage is indexed. +If the transform implements an inverse, that inverse runs when data is +**written** (``rb.add`` / ``rb.extend``). This is how you store cheap raw +observations and reconstruct a processed view only at training time. + +The most common instance of this pattern is frame stacking. The env-side +recipe is below; the collector + replay-buffer recipe (including +``extend`` / ``sample``) lives in +:ref:`Frame stacking images with CatFrames `. + +.. _catframes-images: + +CatFrames for visual RL +~~~~~~~~~~~~~~~~~~~~~~~ + +:class:`CatFrames` concatenates the last ``N`` observations along one +existing dimension so a feed-forward policy can see motion. There are two +legitimate placements; they solve different problems and must not be +stacked on top of each other. + +1. **On the env** (this section). The policy sees a stacked observation + at every step. The transform is stateful: it keeps a rolling buffer + that is flushed on reset. +2. **On the replay buffer, at sample time** + (:ref:`catframes-collector-replay`). Raw unstacked frames are stored; + the stack is rebuilt when you call ``rb.sample()``. This is the + memory-efficient path for image off-policy training. + +For CHW images (the layout produced by :class:`ToTensorImage`) the stack +dimension is ``dim=-3`` (the channel axis). That is the DQN convention: +a grayscale frame of shape ``[1, H, W]`` becomes ``[N, H, W]``; an RGB +frame of shape ``[3, H, W]`` becomes ``[3 * N, H, W]``. Vector +observations use ``dim=-1`` instead. Using the vector default on pixels, +or stacking along ``-4`` without first inserting that axis, silently +produces the wrong layout. + +Env-side stacking (policy input) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Put :class:`CatFrames` on the env so the observation spec the policy +reads is already stacked. Write the stack to a **different** key than +the raw pixels: the collector can then persist the cheap uint8 frame +and drop the float stack (see :ref:`catframes-collector-replay`). + +:class:`CatFrames` is stateful. ``env.reset()`` (or a ``"_reset"`` flag +in the tensordict) flushes the rolling buffer and pads the new stack. +With the default ``padding="same"`` the missing history is a repeat of +the first post-reset frame; ``padding="constant"`` fills with +``padding_value`` (0 by default). :class:`InitTracker` is not required +for that flush -- :class:`CatFrames` listens to the env ``_reset`` key +-- but it should sit in the same :class:`Compose` so ``"is_init"`` marks +the steps where the stack was re-initialized. Collectors, recurrent +policies and advantage estimators all read that flag. + +.. code-block:: python + + from torchrl.envs import ( + CatFrames, + Compose, + GrayScale, + GymEnv, + InitTracker, + Resize, + StepCounter, + ToTensorImage, + TransformedEnv, + ) + + env = TransformedEnv( + GymEnv("CartPole-v1", from_pixels=True, pixels_only=True), + Compose( + InitTracker(), + ToTensorImage(in_keys=["pixels"], out_keys=["pixels_trsf"]), + GrayScale(in_keys=["pixels_trsf"]), + Resize(84, 84, in_keys=["pixels_trsf"]), + CatFrames(N=4, dim=-3, in_keys=["pixels_trsf"], out_keys=["pixels_trsf"]), + StepCounter(), + ), + ) # doctest: +SKIP + # "pixels": raw uint8 frame from the env + # "pixels_trsf": float stack of shape [4, 84, 84] (grayscale, dim=-3) + # "is_init": True on the first step after every reset + +A rollout or a :class:`~torchrl.collectors.Collector` then feeds +``pixels_trsf`` to the policy. Do **not** also attach a second +:class:`CatFrames` on the replay buffer that reads the same stacked +key -- that stacks twice (``[N, H, W]`` stored, ``[N * N, H, W]`` +sampled). Either: + +- store the already-stacked ``pixels_trsf`` and leave the buffer + transform-free (``N`` times more RAM), or +- exclude ``pixels_trsf`` on ``extend`` and rebuild the stack from + raw ``"pixels"`` at sample time (recommended; see + :ref:`catframes-collector-replay`). + +If you want a separate stack axis ``[N, C, H, W]`` rather than channel +concatenation, unsqueeze first and pass ``dim=-4`` (that is the variant +in ``examples/replay-buffers/catframes-in-buffer.py``). For a standard +visual-RL CNN the ``dim=-3`` layout above is the one to use. Cloning transforms ~~~~~~~~~~~~~~~~~~ diff --git a/torchrl/envs/transforms/_observation.py b/torchrl/envs/transforms/_observation.py index 404111470e7..e3b205a8f28 100644 --- a/torchrl/envs/transforms/_observation.py +++ b/torchrl/envs/transforms/_observation.py @@ -1089,6 +1089,11 @@ def make_rb_transform_and_sampler( .. note:: For a more complete example, refer to torchrl's github repo `examples` folder: https://github.com/pytorch/rl/tree/main/examples/replay-buffers/catframes-in-buffer.py + .. seealso:: :ref:`catframes-images` (env-side stacking, reset / + :class:`~torchrl.envs.transforms.InitTracker`) and + :ref:`catframes-collector-replay` (collector ``extend`` / buffer + ``sample``, ``dim=-3`` for CHW images). + """ from torchrl.data.replay_buffers import SliceSampler