[Feature] Add flat replay writes to collectors

benchmarks/test_collectors_benchmark.py

@@ -26,6 +26,7 @@
     EnvCreator,
     GymEnv,
     ParallelEnv,
+    SerialEnv,
     StepCounter,
     TransformedEnv,
 )
@@ -292,6 +293,25 @@ def single_collector_with_rb_setup_pixels():
     return ((c, rb), {})
 
 
+def flat_collector_with_rb_setup():
+    """Setup batched collection into the preferred flat replay layout."""
+    env = SerialEnv(4, functools.partial(_PayloadEnv, payload_size=4096))
+    rb = ReplayBuffer(storage=LazyTensorStorage(10000))
+    c = Collector(
+        env,
+        RandomPolicy(env.action_spec),
+        total_frames=-1,
+        frames_per_batch=128,
+        replay_buffer=rb,
+        flatten_data=True,
+    )
+    c = iter(c)
+    for i, _ in enumerate(c):
+        if i == 10:
+            break
+    return ((c, rb), {})
+
+
 def sync_collector_with_rb_setup():
     """Setup a multi-process collector whose workers write whole rollouts."""
     device = "cuda:0" if torch.cuda.device_count() else "cpu"
@@ -360,6 +380,12 @@ def test_sync_payload_with_rb(benchmark):
         collector.shutdown()
 
 
+def test_flatten_data_with_rb(benchmark):
+    """Benchmark one reshape and one extend for a batched rollout."""
+    (c, rb), _ = flat_collector_with_rb_setup()
+    benchmark(execute_collector_with_rb, c, rb)
+
+
 @pytest.mark.skipif(not torch.cuda.device_count(), reason="no rendering without cuda")
 def test_single_with_rb_pixels(benchmark):
     """Benchmark single collector with replay buffer for pixel observations."""

docs/source/reference/collectors_distributed.rst

@@ -96,11 +96,15 @@ topology.
 
 .. tip::
 
-  All distributed collectors support ``trajs_per_batch`` combined with
-  ``replay_buffer``.  When set, each remote worker assembles **complete
-  trajectories** and writes them to the shared buffer as flat 1-D sequences,
-  which is directly compatible with :class:`~torchrl.data.replay_buffers.SliceSampler`.
-  See :ref:`collectors_replay_trajs` for examples and best practices.
+  :class:`~torchrl.collectors.distributed.RayCollector` supports direct replay
+  writes through a Ray-backed ``replay_buffer``. Pass ``flatten_data=True`` for
+  regular fixed-frame collection, or set ``trajs_per_batch`` to have each
+  remote worker assemble complete trajectories. Both paths write flat 1-D
+  sequences compatible with
+  :class:`~torchrl.data.replay_buffers.SliceSampler`. The generic distributed
+  and RPC collectors return data to their caller instead of owning a shared
+  replay-buffer transport. See :ref:`collectors_replay_trajs` for examples and
+  best practices.
 
 .. autosummary::
     :toctree: generated/

docs/source/reference/collectors_replay.rst

@@ -13,79 +13,56 @@ Collectors and Replay Buffers
     recommended — see :ref:`Data layout: contiguous trajectories
     <data-layout>`.
 
-Collectors and replay buffers interoperability
-----------------------------------------------
-
-In the simplest scenario where single transitions have to be sampled
-from the replay buffer, little attention has to be given to the way
-the collector is built. Flattening the data after collection will
-be a sufficient preprocessing step before populating the storage:
-
-    >>> memory = ReplayBuffer(
-    ...     storage=LazyTensorStorage(N),
-    ...     transform=lambda data: data.reshape(-1))
-    >>> for data in collector:
-    ...     memory.extend(data)
-
-If trajectory slices have to be collected, the recommended way to achieve this is to create
-a multidimensional buffer and sample using the :class:`~torchrl.data.replay_buffers.SliceSampler`
-sampler class. One must ensure that the data passed to the buffer is properly shaped, with the
-``time`` and ``batch`` dimensions clearly separated. In practice, the following configurations
-will work:
-
-    >>> # Single environment: no need for a multi-dimensional buffer
-    >>> memory = ReplayBuffer(
-    ...     storage=LazyTensorStorage(N),
-    ...     sampler=SliceSampler(num_slices=4, trajectory_key=("collector", "traj_ids"))
-    ... )
-    >>> collector = Collector(env, policy, frames_per_batch=N, total_frames=-1)
-    >>> for data in collector:
-    ...     memory.extend(data)
-    >>> # Batched environments: a multi-dim buffer is required
-    >>> memory = ReplayBuffer(
-    ...     storage=LazyTensorStorage(N, ndim=2),
-    ...     sampler=SliceSampler(num_slices=4, trajectory_key=("collector", "traj_ids"))
-    ... )
-    >>> env = ParallelEnv(4, make_env)
-    >>> collector = Collector(env, policy, frames_per_batch=N, total_frames=-1)
-    >>> for data in collector:
-    ...     memory.extend(data)
-    >>> # Synchronous process collection behaves like ParallelEnv if cat_results="stack"
-    >>> memory = ReplayBuffer(
-    ...     storage=LazyTensorStorage(N, ndim=2),
-    ...     sampler=SliceSampler(num_slices=4, trajectory_key=("collector", "traj_ids"))
-    ... )
-    >>> collector = Collector(make_env, policy,
-    ...     num_collectors=4,
-    ...     sync=True,
-    ...     frames_per_batch=N,
-    ...     total_frames=-1,
-    ...     cat_results="stack")
-    >>> for data in collector:
-    ...     memory.extend(data)
-    >>> # Process collection + parallel env: adapt ndim for both batch dimensions
-    >>> memory = ReplayBuffer(
-    ...     storage=LazyTensorStorage(N, ndim=3),
-    ...     sampler=SliceSampler(num_slices=4, trajectory_key=("collector", "traj_ids"))
-    ... )
-    >>> collector = Collector(lambda: ParallelEnv(2, make_env), policy,
-    ...     num_collectors=4,
-    ...     sync=True,
-    ...     frames_per_batch=N,
-    ...     total_frames=-1,
-    ...     cat_results="stack")
-    >>> for data in collector:
-    ...     memory.extend(data)
-
-.. important::
-
-    The ``ndim=2`` and ``ndim=3`` examples above apply to **fixed-frame
-    batches** (the default, without ``trajs_per_batch``).  When
-    ``trajs_per_batch`` is set, each trajectory is written to the buffer as a
-    **flat 1-D sequence** of variable length.  A storage with ``ndim >= 2``
-    expects a fixed second dimension that variable-length trajectories cannot
-    satisfy.  Always use the default ``ndim=1`` when combining
-    ``trajs_per_batch`` with a replay buffer.
+Preferred layout: flat 1-D storage
+----------------------------------
+
+TorchRL recommends storing replay data in a single flat, 1-D buffer
+(``ndim=1``). Collector batch dimensions describe how data was produced; they
+do not need to become replay-buffer dimensions. Trajectory boundaries remain
+available through ``("collector", "traj_ids")``, ``("next", "done")``,
+``("next", "terminated")``, and ``("next", "truncated")``, and trajectory
+slices can be reconstructed at sampling time with
+:class:`~torchrl.data.replay_buffers.SliceSampler`.
+
+When the collector writes directly to the replay buffer, set
+``flatten_data=True``. A batched rollout with shape ``[N, T]`` is then reshaped
+once to ``[N * T]`` and written with one ``extend`` call:
+
+.. code-block:: python
+
+    from torchrl.collectors import Collector
+    from torchrl.data import LazyTensorStorage, ReplayBuffer
+
+    memory = ReplayBuffer(storage=LazyTensorStorage(100_000))
+    collector = Collector(
+        env,
+        policy,
+        frames_per_batch=200,
+        total_frames=-1,
+        replay_buffer=memory,
+        flatten_data=True,
+    )
+    for _ in collector:  # yields None; transitions are written directly
+        batch = memory.sample(256)
+
+The same argument is available on :class:`~torchrl.collectors.AsyncCollector`,
+:class:`~torchrl.collectors.MultiSyncCollector`,
+:class:`~torchrl.collectors.MultiAsyncCollector`, and
+:class:`~torchrl.collectors.distributed.RayCollector`. If the caller owns the
+write instead, flatten explicitly before extending:
+
+.. code-block:: python
+
+    for data in collector:
+        memory.extend(data.reshape(-1))
+
+Higher-dimensional storage (``ndim >= 2``) remains supported for applications
+that deliberately store fixed-shape chunks. It is a specialized layout rather
+than the recommended default: the chunk dimensions become part of every replay
+item, variable-length trajectories cannot be represented directly, and shared
+multi-producer buffers require extra care around trajectory boundaries. See
+:ref:`The replay buffer ndim arg and why it doesn't multi-process well
+<data-layout-storage-ndim>` for details.
 
 .. _collectors_replay_trajs:
 

docs/source/reference/collectors_single.rst

@@ -78,6 +78,9 @@ padding) instead of being yielded.  This is the recommended pattern for
 off-policy training with :class:`~torchrl.data.replay_buffers.SliceSampler`, especially
 with multi-process collectors where fixed-frame batches can silently mix
 episodes.  See :ref:`collectors_replay_trajs` for full details and examples.
+For regular fixed-frame collection from a batched environment, pass
+``flatten_data=True`` to reshape each ``[N, T]`` rollout once and extend the
+recommended 1-D storage with ``N * T`` transitions.
 
 .. note::
     The deprecated collector aliases were removed in v0.13. Construct new

docs/source/reference/data_layout.rst

@@ -10,6 +10,9 @@ which trajectories are concatenated end-to-end and their boundaries are
 recovered from the per-step ``is_init`` / ``("next", "done")`` /
 ``("next", "truncated")`` / ``("next", "terminated")`` markers, **not**
 from a fixed ``[B, T]`` shape with a padding mask.
+For collectors that write directly to a replay buffer, pass
+``flatten_data=True`` to reshape batched rollouts to this layout before each
+``extend`` call.
 
 Two main patterns coexist in TorchRL:
 

docs/source/reference/data_replaybuffers.rst

@@ -20,6 +20,16 @@ restrictions, and expected performance trade-offs, and
 :ref:`ref_distributed_transport_layouts` for the per-operation layout
 discovery and buffer lifecycle.
 
+.. important::
+
+    TorchRL recommends a flat, 1-D replay-buffer layout (the default
+    ``ndim=1`` storage). Flatten batched collector output before insertion and
+    recover trajectory structure from per-step boundary markers with
+    :class:`~torchrl.data.replay_buffers.SliceSampler`. Higher-dimensional
+    storage is supported when fixed-shape chunks are intentionally part of each
+    replay item, but it is not the preferred general-purpose layout. See
+    :ref:`Data layout: contiguous trajectories <data-layout>`.
+
 .. code-block:: python
 
     from functools import partial

test/test_collectors.py

@@ -5683,6 +5683,114 @@ def test_ray_map_fn_and_get_distant_attr(self):
 
 
 class TestCollectorRB:
+    @pytest.mark.parametrize(
+        "collector_class",
+        [Collector, AsyncCollector, MultiSyncCollector, MultiAsyncCollector],
+    )
+    def test_flatten_data_replay_buffer(self, collector_class):
+        """Batched collector rollouts can be written to flat 1-D storage."""
+
+        def make_env():
+            return SerialEnv(2, CountingEnv)
+
+        probe = make_env()
+        policy = RandomPolicy(probe.action_spec)
+        probe.close(raise_if_closed=False)
+        rb = ReplayBuffer(storage=LazyTensorStorage(64), batch_size=4)
+        if collector_class is Collector:
+            create_env_fn = make_env()
+        elif collector_class is AsyncCollector:
+            create_env_fn = make_env
+        else:
+            create_env_fn = [make_env]
+        collector = collector_class(
+            create_env_fn,
+            policy,
+            replay_buffer=rb,
+            flatten_data=True,
+            total_frames=8,
+            frames_per_batch=8,
+        )
+        try:
+            assert all(data is None for data in collector)
+        finally:
+            collector.shutdown()
+
+        assert rb.write_count == 8
+        assert rb[:].shape == (8,)
+        traj_ids = rb["collector", "traj_ids"]
+        torch.testing.assert_close(traj_ids[:4], torch.zeros(4, dtype=torch.long))
+        torch.testing.assert_close(traj_ids[4:], torch.ones(4, dtype=torch.long))
+
+    def test_flatten_data_replay_buffer_requires_extend(self):
+        env = CountingEnv()
+        rb = ReplayBuffer(storage=LazyTensorStorage(64), batch_size=4)
+        try:
+            with pytest.raises(TypeError, match="requires extend_buffer=True"):
+                Collector(
+                    env,
+                    RandomPolicy(env.action_spec),
+                    replay_buffer=rb,
+                    flatten_data=True,
+                    extend_buffer=False,
+                    total_frames=8,
+                    frames_per_batch=8,
+                )
+        finally:
+            env.close(raise_if_closed=False)
+
+    @pytest.mark.parametrize(
+        "collector_class", [Collector, MultiSyncCollector, MultiAsyncCollector]
+    )
+    def test_flatten_data_requires_replay_buffer(self, collector_class):
+        """flatten_data without a replay buffer is an error, not a silent no-op."""
+        env = CountingEnv()
+        try:
+            policy = RandomPolicy(env.action_spec)
+            if collector_class is Collector:
+                create_env_fn = env
+            else:
+                create_env_fn = [CountingEnv]
+            with pytest.raises(TypeError, match="requires a replay buffer"):
+                collector_class(
+                    create_env_fn,
+                    policy,
+                    flatten_data=True,
+                    total_frames=8,
+                    frames_per_batch=8,
+                )
+        finally:
+            env.close(raise_if_closed=False)
+
+    def test_flatten_data_replay_buffer_postproc(self):
+        """postproc runs on the batched [N, T] rollout, before flattening."""
+        env = SerialEnv(2, CountingEnv)
+        rb = ReplayBuffer(storage=LazyTensorStorage(64), batch_size=4)
+        postproc_shapes = []
+
+        def postproc(data):
+            postproc_shapes.append(tuple(data.shape))
+            data["postproc_flag"] = torch.ones(data.shape, dtype=torch.bool)
+            return data
+
+        collector = Collector(
+            env,
+            RandomPolicy(env.action_spec),
+            replay_buffer=rb,
+            flatten_data=True,
+            postproc=postproc,
+            total_frames=8,
+            frames_per_batch=8,
+        )
+        try:
+            assert all(data is None for data in collector)
+        finally:
+            collector.shutdown()
+
+        assert postproc_shapes == [(2, 4)]
+        assert rb[:].shape == (8,)
+        assert rb["postproc_flag"].all()
+
     @pytest.mark.skipif(not _has_gym, reason="requires gym.")
     def test_collector_rb_sync(self):
         env = SerialEnv(8, lambda cp=CARTPOLE_VERSIONED(): GymEnv(cp))