[CI] Fix doc upload (#1738)

.github/workflows/docs.yml

@@ -24,6 +24,7 @@ jobs:
  uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
  with:
  repository: pytorch/rl
+ upload-artifact: docs
  runner: "linux.g5.4xlarge.nvidia.gpu"
  docker-image: "nvidia/cudagl:11.4.0-base"
  timeout: 120
@@ -82,8 +83,11 @@ jobs:
  # bash -ic "PYOPENGL_PLATFORM=egl MUJOCO_GL=egl sphinx-build ./source _local_build" || code=$?; if [[ $code -ne 124 && $code -ne 0 ]]; then exit $code; fi
  PYOPENGL_PLATFORM=egl MUJOCO_GL=egl sphinx-build ./source _local_build
  cd ..
+
  cp -r docs/_local_build "${RUNNER_ARTIFACT_DIR}"
- cp -r docs/_local_build/* "${RUNNER_DOCS_DIR}"
+ if [[ ${{ github.event_name == 'pull_request' }} ]]; then
+ cp -r docs/_local_build/* "${RUNNER_DOCS_DIR}"
+ fi
 
  upload:
  needs: build-docs
@@ -119,8 +123,8 @@ jobs:
  # ;;
  # esac
  # fi
- # echo "Target Folder: ${TARGET_FOLDER}"
  TARGET_FOLDER="./"
+ echo "Target Folder: ${TARGET_FOLDER}"
 
  # mkdir -p "${TARGET_FOLDER}"
  # rm -rf "${TARGET_FOLDER}"/*

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.

torchrl/envs/__init__.py

@@ -20,6 +20,8 @@
  IsaacGymWrapper,
  JumanjiEnv,
  JumanjiWrapper,
+ MOGymEnv,
+ MOGymWrapper,
  MultiThreadedEnv,
  MultiThreadedEnvWrapper,
  OpenMLEnv,
@@ -78,6 +80,7 @@
  TransformedEnv,
  UnsqueezeTransform,
  VC1Transform,
+ VecGymEnvTransform,
  VecNorm,
  VIPRewardTransform,
  VIPTransform,

torchrl/envs/transforms/__init__.py

@@ -47,6 +47,7 @@
  Transform,
  TransformedEnv,
  UnsqueezeTransform,
+ VecGymEnvTransform,
  VecNorm,
 )
 from .vc1 import VC1Transform

torchrl/modules/__init__.py

@@ -17,6 +17,7 @@
  TruncatedNormal,
 )
 from .models import (
+ Conv3dNet,
  ConvNet,
  DdpgCnnActor,
  DdpgCnnQNet,