Update

[ghstack-poisoned]

.github/unittest/linux_libs/scripts_gym/batch_scripts.sh

@@ -6,6 +6,7 @@
 DIR="$(cd "$(dirname "$0")" && pwd)"
 
 set -e
+set -v
 
 eval "$(./conda/bin/conda shell.bash hook)"
 conda activate ./env

.github/unittest/linux_libs/scripts_gym/install.sh

@@ -7,6 +7,7 @@ unset PYTORCH_VERSION
 apt-get update && apt-get install -y git wget gcc g++
 
 set -e
+set -v
 
 eval "$(./conda/bin/conda shell.bash hook)"
 conda activate ./env
@@ -39,7 +40,7 @@ printf "Installing PyTorch with %s\n" "${CU_VERSION}"
 if [ "${CU_VERSION:-}" == cpu ] ; then
     conda install pytorch==2.0 torchvision==0.15 cpuonly -c pytorch -y
 else
-    conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.8 -c pytorch -c nvidia -y
+    conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.8 numpy==1.26 numpy-base==1.26 -c pytorch -c nvidia -y
 fi
 
 # Solving circular import: https://stackoverflow.com/questions/75501048/how-to-fix-attributeerror-partially-initialized-module-charset-normalizer-has

.github/unittest/linux_libs/scripts_habitat/setup_env.sh

@@ -67,8 +67,9 @@ pip install pip --upgrade
 
 conda env update --file "${this_dir}/environment.yml" --prune
 
-#conda install habitat-sim withbullet headless -c conda-forge -c aihabitat -y
 conda install habitat-sim withbullet headless -c conda-forge -c aihabitat -y
-conda run python -m pip install git+https://github.com/facebookresearch/habitat-lab.git@stable#subdirectory=habitat-lab
-#conda run python -m pip install git+https://github.com/facebookresearch/habitat-lab.git#subdirectory=habitat-baselines
+git clone https://github.com/facebookresearch/habitat-lab.git
+cd habitat-lab
+pip3 install -e habitat-lab
+pip3 install -e habitat-baselines  # install habitat_baselines
 conda run python -m pip install "gym[atari,accept-rom-license]" pygame

.github/unittest/linux_libs/scripts_minari/environment.yml

@@ -17,4 +17,4 @@ dependencies:
     - pyyaml
     - scipy
     - hydra-core
-    - minari
+    - minari[gcs]

.github/unittest/linux_olddeps/scripts_gym_0_13/install.sh

@@ -39,7 +39,7 @@ printf "Installing PyTorch with %s\n" "${CU_VERSION}"
 if [ "${CU_VERSION:-}" == cpu ] ; then
     conda install pytorch==2.0 torchvision==0.15 cpuonly -c pytorch -y
 else
-    conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.8 -c pytorch -c nvidia -y
+    conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.8 numpy==1.26 numpy-base==1.26 -c pytorch -c nvidia -y
 fi
 
 # Solving circular import: https://stackoverflow.com/questions/75501048/how-to-fix-attributeerror-partially-initialized-module-charset-normalizer-has

.github/workflows/build-wheels-aarch64-linux.yml

@@ -0,0 +1,51 @@
+name: Build Aarch64 Linux Wheels
+
+on:
+  pull_request:
+  push:
+    branches:
+      - nightly
+      - main
+      - release/*
+    tags:
+        # NOTE: Binary build pipelines should only get triggered on release candidate builds
+        # Release candidate tags look like: v1.11.0-rc1
+        - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
+  workflow_dispatch:
+
+permissions:
+  id-token: write
+  contents: read
+
+jobs:
+  generate-matrix:
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
+    with:
+      package-type: wheel
+      os: linux-aarch64
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      with-cuda: disable
+  build:
+    needs: generate-matrix
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - repository: pytorch/rl
+            smoke-test-script: test/smoke_test.py
+            package-name: torchrl
+    name: pytorch/rl
+    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
+    with:
+      repository: ${{ matrix.repository }}
+      ref: ""
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      package-name: ${{ matrix.package-name }}
+      smoke-test-script: ${{ matrix.smoke-test-script }}
+      trigger-event: ${{ github.event_name }}
+      env-var-script: .github/scripts/td_script.sh
+      architecture: aarch64
+      setup-miniconda: false

.github/workflows/build-wheels-windows.yml

@@ -32,6 +32,8 @@ jobs:
       matrix:
         include:
           - repository: pytorch/rl
+            pre-script: .github/scripts/td_script.sh
+            env-script: .github/scripts/version_script.bat
             post-script: "python packaging/wheel/relocate.py"
             smoke-test-script: test/smoke_test.py
             package-name: torchrl
@@ -43,8 +45,9 @@ jobs:
       test-infra-repository: pytorch/test-infra
       test-infra-ref: main
       build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      pre-script: ${{ matrix.pre-script }}
+      env-script: ${{ matrix.env-script }}
+      post-script: ${{ matrix.post-script }}
       package-name: ${{ matrix.package-name }}
       smoke-test-script: ${{ matrix.smoke-test-script }}
       trigger-event: ${{ github.event_name }}
-      pre-script: .github/scripts/td_script.sh
-      env-script: .github/scripts/version_script.bat

.github/workflows/docs.yml

@@ -50,7 +50,7 @@ jobs:
         conda activate "${env_dir}"
         
         # 2. upgrade pip, ninja and packaging
-        apt-get install python3.9 python3-pip -y
+        # apt-get install python3.9 python3-pip -y
         python3 -m pip install --upgrade pip
         python3 -m pip install setuptools ninja packaging -U
         

docs/source/reference/modules.rst

@@ -57,23 +57,29 @@ projected (in a L1-manner) into the desired domain.
     SafeSequential
     TanhModule
 
-Exploration wrappers
-~~~~~~~~~~~~~~~~~~~~
+Exploration wrappers and modules
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-To efficiently explore the environment, TorchRL proposes a series of wrappers
+To efficiently explore the environment, TorchRL proposes a series of modules
 that will override the action sampled by the policy by a noisier version.
 Their behavior is controlled by :func:`~torchrl.envs.utils.exploration_mode`:
 if the exploration is set to ``"random"``, the exploration is active. In all
 other cases, the action written in the tensordict is simply the network output.
 
-.. currentmodule:: torchrl.modules.tensordict_module
+.. note:: Unlike other exploration modules, :class:`~torchrl.modules.ConsistentDropoutModule`
+  uses the ``train``/``eval`` mode to comply with the regular `Dropout` API in PyTorch.
+  The :func:`~torchrl.envs.utils.set_exploration_mode` context manager will have no effect on
+  this module.
+
+.. currentmodule:: torchrl.modules
 
 .. autosummary::
     :toctree: generated/
     :template: rl_template_noinherit.rst
 
     AdditiveGaussianModule
     AdditiveGaussianWrapper
+    ConsistentDropoutModule
     EGreedyModule
     EGreedyWrapper
     OrnsteinUhlenbeckProcessModule
@@ -438,12 +444,13 @@ Regular modules
     :toctree: generated/
     :template: rl_template_noinherit.rst
 
-    MLP
-    ConvNet
+    BatchRenorm1d
+    ConsistentDropout
     Conv3dNet
-    SqueezeLayer
+    ConvNet
+    MLP
     Squeeze2dLayer
-    BatchRenorm1d
+    SqueezeLayer
 
 Algorithm-specific modules
 ~~~~~~~~~~~~~~~~~~~~~~~~~~