[Feature] Read storages and replay buffers through torch.utils.data

[theap06]

Description

First step of the DataLoader direction discussed for the next release: interoperate with torch.utils.data at the storage / buffer seam instead of adding a loader of our own.

• Storage is a torch.utils.data.Dataset. It gains __getitems__, so a DataLoader with any torch sampler fetches each index batch through a single get call. Multi-dimensional storages are read through storage.flatten() (the unflattened storage raises with that hint, since its __len__ and __getitem__ disagree) and StorageEnsemble is not a flat dataset and raises. Usage: DataLoader(rb.storage, batch_size=32, shuffle=True, collate_fn=tensordict_collate).
• tensordict_collate returns a fetched batch unchanged (tensor, tensordict or tuple of them) and stacks lists of samples (lazily for ragged tensordicts, element-wise for mappings and tuples), so per-item storages such as ListStorage and compositions such as ConcatDataset work through the same collate. The default torch collation iterates a TensorDict over its batch dimension and crashes.
• ReplayBuffer.as_dataset(num_batches=None) returns a ReplayBufferDataset (IterableDataset) that iterates the buffer; it requires the buffer batch_size. Under DataLoader workers each worker holds its own copy of the buffer, so the TorchRL sampler and the buffer transforms run in the worker process. num_batches is split between workers; prefetched batches are stripped when the dataset is serialized and the worker copy gets fresh locks with prefetching disabled (the DataLoader prefetches; otherwise the parent's queue is replayed by every worker, and under fork a lock held by a parent prefetch thread would deadlock the worker); a buffer built with a generator is reseeded once per worker from the worker seed (otherwise every worker, and every epoch with persistent workers, draws the same batches). RayReplayBuffer rejects as_dataset explicitly.
• Sampler.requires_shared_state (new attribute, False by default) marks samplers whose state every consumer must observe: without replacement, prioritized, consuming, staleness-aware, streaming slice, prompt group with the recency strategy, the OpenX streaming sampler, and ensembles containing one. as_dataset rejects them when num_workers > 0 instead of silently duplicating their state per process, at serialization time for spawn and in the worker for fork.
• TorchRLBufferDataset (TRL interop) now builds on ReplayBufferDataset. Its per-iterator num_batches semantics are unchanged; a shared-state sampler emits a FutureWarning at construction and workers will reject it from v0.17.
• BugFix: TensorStorage.__getstate__ called tree_map(storage, fn) with the arguments swapped, so a plain-tensor TensorStorage could not be sent to a spawned process (TypeError: 'Tensor' object is not callable). Regression test added.
• BugFix: ReplayBuffer.__getstate__ pickled the generator state as a temporary tensor. Under the spawn start method with torch's file-descriptor sharing strategy the tensor is freed right after pickling, its descriptor number is recycled for the spawn pipes, and the child fails with unable to resize file <filename not specified> to the right size: Invalid argument (22) while unpickling. This is what failed the CPU CI jobs on the first push and affects any generator-seeded buffer sent to a spawned process on Linux. The state is now pickled as bytes. Regression test added.

Docs: new sections in data_replaybuffers.rst and data_storage.rst. Benchmark: test_replay_buffer_dataset_workers in benchmarks/test_replaybuffer_benchmark.py. Example: examples/replay-buffers/dataloader_behavior_cloning.py trains a CartPole policy with BCLoss from buffer batches sampled in workers and writes a checkpoint that rlrender plays back (eval return 500/500 after 500 updates).

Throughput of DataLoader(rb.as_dataset()) with a per-frame resize on the sample path (32 x 3x96x96 uint8 to 128x128, LazyMemmapStorage, macOS 8 cores, fork):

workers	batches / s
0	178
1	139
2	281
4	483
8	611

The gain is for sample paths that are serial per item (video decoding, per-frame preprocessing). A batched torch op in an otherwise idle main process already uses every core and does not benefit.

Not in this PR, planned as follow-ups: a storage wrapping an arbitrary map-style Dataset (LeRobot's own dataset), torch Sampler adapters for TorchRL samplers, and folding torchrl/data/llm/dataset.py::get_dataloader into the same pattern.

Motivation and Context

vmoens: "let's look at how we can interact with existing PyTorch APIs perhaps? I think we should be able to hook a DL on our RB or its storage, or make the samplers+storage play with DLs." Also the reason LeRobot did not adopt torchrl's data loading.

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• Documentation (update in the documentation)
• Example (update in the folder of examples)

Checklist

• 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/4387

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit 645510b with merge base 2d258fe ():

NEW FAILURE - The following job has failed:

• Unit-tests on Linux / tests-cpu (3.10, bulk) / linux-job (gh)
  test/rb/test_rb_core.py::TestTorchDataInterop::test_as_dataset_workers_drop_prefetched_batches

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

[vmoens]

Requesting changes for worker-state correctness and PyTorch Dataset interoperability. The focused interop and storage-spawn tests pass locally, but the current head also has five failing CI jobs (two example jobs and three CPU jobs).

[vmoens]

[P1] Checking only sampler.requires_shared_state misses the existing OpenX streaming path. _StreamingSampler inherits the default False, while _StreamingStorage.get() consumes a dataset_iter created before workers are spawned. Each worker therefore receives the same cursor/state and can silently duplicate stream data. Please either mark this combination as requiring shared state or recreate and shard the stream per worker.

[vmoens]

[P2] This disables prefetching only after the worker copy has been created. With spawn, ReplayBuffer.__getstate__ has already captured and serialized the full prefetch queue into every worker; a pickle round-trip with prefetch=3 restores all three queued batches before this line clears them. For large batches this multiplies startup memory and file-descriptor traffic. Please strip prefetch state during dataset serialization instead.

[vmoens]

[P2] PyTorch's Dataset.__getitems__ contract returns a list of samples, but this returns an already-collated batch and relies on an identity collate_fn. Standard composition consequently breaks: wrapping storages in ConcatDataset falls back to scalar __getitem__ calls and tensordict_collate leaves the resulting list uncollated. Third-party Storage subclasses also fail here when _get_default_collate does not recognize them. Please preserve the Dataset batching contract or make the public collation path handle both scalar-sample lists and vectorized batches.

[vmoens]

[P1] This new example is not registered in .github/unittest/examples/scripts/test_examples.py, so both example CI shards fail with unclassified examples. Please add an ExampleSpec with smoke arguments or a documented exclusion.

[vmoens]

Thanks that looks like it's going in the right direction
Can you hold on a bit before merging this I'd like to give it some thought and run it through a couple of former colleagues!