[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

[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!

[vmoens]

Thanks for addressing the original inline comments. After thinking more about the API direction, I think the storage side of this interop should use an explicit adapter rather than changing the public base class.

For this PR, please keep Storage's inheritance unchanged and expose Storage.as_dataset() returning a small map-style adapter (for example, StorageDataset) that owns the PyTorch Dataset and __getitems__ behavior. DataLoader does not require Dataset inheritance: Storage already has __len__ and __getitem__. The current Storage(Dataset) change has several compatibility costs:

• it changes the MRO and isinstance/issubclass results for a public extension point, and can break downstream multiple-inheritance classes;
• it inherits adjacent behavior such as Dataset.__add__, so storage addition gains new semantics unrelated to replay storage;
• more importantly, adding Storage.__getitems__ silently changes existing DataLoader(storage, ...) calls from N scalar reads followed by a list passed to the user's collator into one batched get(indices) result. Existing custom collators may therefore break even if they worked before this PR.

An adapter gives us a clean opt-in and BC boundary while retaining the optimized batched fetch. ReplayBuffer.as_dataset() / ReplayBufferDataset is likewise the right shape for the buffer side, since that is explicitly an iterable view of replay-buffer sampling rather than a new identity for ReplayBuffer itself.

Two useful follow-ups, explicitly out of scope for this PR:

1. The inverse adapter, DatasetStorage(dataset), could let a fixed-length map-style PyTorch dataset back a read-only replay buffer together with the existing ImmutableDatasetWriter. The initial contract should probably be deliberately narrow: scalar and batched index forwarding, uniform and without-replacement sampling, explicit collation and checkpoint semantics, and no IterableDataset or trajectory-aware samplers until those contracts are designed. Stochastic __getitem__ also needs documentation because a priority would apply to an index, not necessarily to a stable realized sample.
2. Utilities to export and upload offline replay data to the Hugging Face Hub would be valuable. I would make an interoperable dataset representation that can be consumed outside TorchRL the primary exchange format (including LeRobot-oriented conversion where the schema fits), while optionally supporting a native TorchRL storage artifact for lossless, fast reload. The portable schema, trajectory metadata, transforms, and native-storage versioning deserve a separate design rather than being coupled to this DataLoader adapter.

Separately, the current head still has a red tests-cpu (3.10, bulk) job in test_as_dataset_workers_drop_prefetched_batches (ValueError: bad value(s) in fds_to_keep), which also needs to be resolved or shown to be unrelated before approval.