init

.github/workflows/docs.yml

@@ -88,7 +88,7 @@ jobs:
  upload:
  needs: build-docs
  if: github.repository == 'pytorch/rl' && github.event_name == 'push' && 
- ((github.ref_type == 'branch' && github.ref_name == 'main') || github.ref_type == 'tag')
+ ((github.ref_type == 'branch' && github.ref_name == 'main') || github.ref_type == 'tag' || contains(github.event.pull_request.labels.*.name, 'ciflow/docs'))
  permissions:
  contents: write
  uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
@@ -124,7 +124,7 @@ jobs:
 
  # mkdir -p "${TARGET_FOLDER}"
  # rm -rf "${TARGET_FOLDER}"/*
- mv "${RUNNER_ARTIFACT_DIR}"/html/* "${TARGET_FOLDER}"
+ cp -r "${RUNNER_ARTIFACT_DIR}"/html/* "${TARGET_FOLDER}"
  git add "${TARGET_FOLDER}" || true
 
  # if [[ "${TARGET_FOLDER}" == main ]]; then

docs/source/reference/data.rst

@@ -60,33 +60,6 @@ The following mean sampling latency improvements over using ListStorage were fou
 | :class:`LazyMemmapStorage` | 3.44x |
 +-------------------------------+-----------+
 
-Replay buffers with a shared storage and regular (RoundRobin) writers can also
-be shared between processes on a single node. This allows each worker to read and
-write onto the storage. The following code snippet examplifies this feature:
-
- >>> from torchrl.data import TensorDictReplayBuffer, LazyMemmapStorage
- >>> import torch
- >>> from torch import multiprocessing as mp
- >>> from tensordict import TensorDict
- >>>
- >>> def worker(rb):
- ... # Updates the replay buffer with new data
- ... td = TensorDict({"a": torch.ones(10)}, [10])
- ... rb.extend(td)
- ...
- >>> if __name__ == "__main__":
- ... rb = TensorDictReplayBuffer(storage=LazyMemmapStorage(21))
- ... td = TensorDict({"a": torch.zeros(10)}, [10])
- ... rb.extend(td)
- ...
- ... proc = mp.Process(target=worker, args=(rb,))
- ... proc.start()
- ... proc.join()
- ... # the replay buffer now has a length of 20, since the worker updated it
- ... assert len(rb) == 20
- ... assert (rb["_data", "a"][:10] == 0).all() # data from main process
- ... assert (rb["_data", "a"][10:20] == 1).all() # data from remote process
-
 Sharing replay buffers across processes
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -105,30 +78,28 @@ are currently not sharable, and the same goes for stateful samplers such as
 A shared replay buffer can be read and extended on any process that has access
 to it, as the following example shows:
 
- >>> import pickle
- >>>
  >>> from torchrl.data import TensorDictReplayBuffer, LazyMemmapStorage
  >>> import torch
  >>> from torch import multiprocessing as mp
  >>> from tensordict import TensorDict
  >>>
  >>> def worker(rb):
+ ... # Updates the replay buffer with new data
  ... td = TensorDict({"a": torch.ones(10)}, [10])
- ... # Extends the shared replay buffer on a subprocess
  ... rb.extend(td)
- >>>
+ ...
  >>> if __name__ == "__main__":
  ... rb = TensorDictReplayBuffer(storage=LazyMemmapStorage(21))
  ... td = TensorDict({"a": torch.zeros(10)}, [10])
- .. # extends the replay buffer on the main process
  ... rb.extend(td)
  ...
  ... proc = mp.Process(target=worker, args=(rb,))
  ... proc.start()
  ... proc.join()
- ... # Checks that the length of the buffer equates the length of both
- ... # extensions (local and remote)
+ ... # the replay buffer now has a length of 20, since the worker updated it
  ... assert len(rb) == 20
+ ... assert (rb["_data", "a"][:10] == 0).all() # data from main process
+ ... assert (rb["_data", "a"][10:20] == 1).all() # data from remote process
 
 
 Storing trajectories
@@ -175,7 +146,7 @@ can be used:
  is_shared=False)
 
 Checkpointing Replay Buffers
-----------------------------
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 Each component of the replay buffer can potentially be stateful and, as such,
 require a dedicated way of being serialized.