From 399d4fa9cea043f4da5e39c5d1f94ac4fe40c888 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 11 Oct 2024 08:56:36 +0100
Subject: [PATCH 01/24] init

---
 .github/workflows/benchmarks.yml    | 30 ++++++++++++++++++++---------
 .github/workflows/benchmarks_pr.yml | 14 +++++++-------
 torchrl/objectives/common.py        |  2 +-
 3 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index 7d8b714ad4d..4c9f5540786 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -24,10 +24,17 @@ jobs:
     name: CPU Pytest benchmark
     runs-on: ubuntu-20.04
     steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - name: Who triggered this?
+        run: |
+          echo "Action triggered by ${{ github.event.pull_request.html_url }}"
+      - name: Checkout
+        uses: actions/checkout@v3
         with:
-          python-version: 3.9
+          fetch-depth: 50 # this is to make sure we obtain the target base commit
+      - name: Python Setup
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
       - name: Setup Environment
         run: |
           python3 -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U 
@@ -36,15 +43,15 @@ jobs:
           python3 -m pip install pytest pytest-benchmark
           python3 -m pip install "gym[accept-rom-license,atari]"
           python3 -m pip install "dm_control" "mujoco"
-          export TD_GET_DEFAULTS_TO_NONE=1
       - name: Run benchmarks
         run: |
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
-          python -m pytest --benchmark-json output.json
+          export TD_GET_DEFAULTS_TO_NONE=1
+          python -m pytest -vvv --rank 0 --benchmark-json output.json
       - name: Store benchmark results
-        if: ${{ github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' }}
         uses: benchmark-action/github-action-benchmark@v1
+        if: ${{ github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' }}
         with:
           name: CPU Benchmark Results
           tool: 'pytest'
@@ -68,6 +75,9 @@ jobs:
       image: nvidia/cuda:12.3.0-base-ubuntu22.04
       options: --gpus all
     steps:
+      - name: Who triggered this?
+        run: |
+          echo "Action triggered by ${{ github.event.pull_request.html_url }}"
       - name: Install deps
         run: |
           export TZ=Europe/London
@@ -82,10 +92,12 @@ jobs:
         run: ldd --version
       - name: Checkout
         uses: actions/checkout@v3
+        with:
+          fetch-depth: 50 # this is to make sure we obtain the target base commit
       - name: Python Setup
         uses: actions/setup-python@v4
         with:
-          python-version: 3.9
+          python-version: '3.10'
       - name: Setup git
         run: git config --global --add safe.directory /__w/rl/rl
       - name: setup Path
@@ -99,7 +111,6 @@ jobs:
           python3 -m pip install pytest pytest-benchmark
           python3 -m pip install "gym[accept-rom-license,atari]"
           python3 -m pip install "dm_control" "mujoco"
-          export TD_GET_DEFAULTS_TO_NONE=1
       - name: check GPU presence
         run: |
           python -c """import torch
@@ -109,7 +120,8 @@ jobs:
         run: |
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
-          python3 -m pytest --benchmark-json output.json
+          export TD_GET_DEFAULTS_TO_NONE=1
+          python -m pytest -vvv --rank 0 --benchmark-json output.json
       - name: Store benchmark results
         uses: benchmark-action/github-action-benchmark@v1
         if: ${{ github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' }}
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index fa1b8037ecb..6fd9e5a38df 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -1,5 +1,4 @@
 name: Continuous Benchmark (PR)
-
 on:
   pull_request:
 
@@ -12,6 +11,7 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+
   benchmark_cpu:
     name: CPU Pytest benchmark
     runs-on: ubuntu-20.04
@@ -26,7 +26,7 @@ jobs:
       - name: Python Setup
         uses: actions/setup-python@v4
         with:
-          python-version: 3.9
+          python-version: '3.10'
       - name: Setup Environment
         run: |
           python3 -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U 
@@ -35,7 +35,6 @@ jobs:
           python3 -m pip install pytest pytest-benchmark
           python3 -m pip install "gym[accept-rom-license,atari]"
           python3 -m pip install "dm_control" "mujoco"
-          export TD_GET_DEFAULTS_TO_NONE=1
       - name: Setup benchmarks
         run: |
           echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV
@@ -47,7 +46,8 @@ jobs:
         run: |
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
-          RUN_BENCHMARK="pytest --rank 0 --benchmark-json "
+          export TD_GET_DEFAULTS_TO_NONE=1
+          RUN_BENCHMARK="pytest -vvv --rank 0 --benchmark-json"
           git checkout ${{ github.event.pull_request.base.sha }}
           $RUN_BENCHMARK ${{ env.BASELINE_JSON }}
           git checkout ${{ github.event.pull_request.head.sha }}
@@ -96,7 +96,7 @@ jobs:
       - name: Python Setup
         uses: actions/setup-python@v4
         with:
-          python-version: 3.9
+          python-version: '3.10'
       - name: Setup git
         run: git config --global --add safe.directory /__w/rl/rl
       - name: setup Path
@@ -110,7 +110,6 @@ jobs:
           python3 -m pip install pytest pytest-benchmark
           python3 -m pip install "gym[accept-rom-license,atari]"
           python3 -m pip install "dm_control" "mujoco"
-          export TD_GET_DEFAULTS_TO_NONE=1
       - name: check GPU presence
         run: |
           python -c """import torch
@@ -127,7 +126,8 @@ jobs:
         run: |
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
-          RUN_BENCHMARK="pytest --rank 0 --benchmark-json "
+          export TD_GET_DEFAULTS_TO_NONE=1
+          RUN_BENCHMARK="pytest -vvv --rank 0 --benchmark-json"
           git checkout ${{ github.event.pull_request.base.sha }}
           $RUN_BENCHMARK ${{ env.BASELINE_JSON }}
           git checkout ${{ github.event.pull_request.head.sha }}
diff --git a/torchrl/objectives/common.py b/torchrl/objectives/common.py
index a1c70612484..f6935ceae82 100644
--- a/torchrl/objectives/common.py
+++ b/torchrl/objectives/common.py
@@ -26,7 +26,7 @@
 
 try:
     from torch.compiler import is_dynamo_compiling
-except ModuleNotFoundError:
+except ImportError:
     from torch._dynamo import is_compiling as is_dynamo_compiling
 
 

From 263507804aa9e9ef19cf8202c172d95bb5dc368a Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 11 Oct 2024 09:14:34 +0100
Subject: [PATCH 02/24] amend

---
 .github/workflows/benchmarks.yml    | 6 +++---
 .github/workflows/benchmarks_pr.yml | 6 +++---
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index 4c9f5540786..ff2fe163b13 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -48,7 +48,7 @@ jobs:
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
-          python -m pytest -vvv --rank 0 --benchmark-json output.json
+          python3 -m pytest -vvv --rank 0 --benchmark-json output.json
       - name: Store benchmark results
         uses: benchmark-action/github-action-benchmark@v1
         if: ${{ github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' }}
@@ -113,7 +113,7 @@ jobs:
           python3 -m pip install "dm_control" "mujoco"
       - name: check GPU presence
         run: |
-          python -c """import torch
+          python3 -c """import torch
           assert torch.cuda.device_count()
           """
       - name: Run benchmarks
@@ -121,7 +121,7 @@ jobs:
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
-          python -m pytest -vvv --rank 0 --benchmark-json output.json
+          python3 -m pytest -vvv --rank 0 --benchmark-json output.json
       - name: Store benchmark results
         uses: benchmark-action/github-action-benchmark@v1
         if: ${{ github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' }}
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index 6fd9e5a38df..9891142631c 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -47,7 +47,7 @@ jobs:
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
-          RUN_BENCHMARK="pytest -vvv --rank 0 --benchmark-json"
+          RUN_BENCHMARK="python3 -m pytest -vvv --rank 0 --benchmark-json"
           git checkout ${{ github.event.pull_request.base.sha }}
           $RUN_BENCHMARK ${{ env.BASELINE_JSON }}
           git checkout ${{ github.event.pull_request.head.sha }}
@@ -112,7 +112,7 @@ jobs:
           python3 -m pip install "dm_control" "mujoco"
       - name: check GPU presence
         run: |
-          python -c """import torch
+          python3 -c """import torch
           assert torch.cuda.device_count()
           """
       - name: Setup benchmarks
@@ -127,7 +127,7 @@ jobs:
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
-          RUN_BENCHMARK="pytest -vvv --rank 0 --benchmark-json"
+          RUN_BENCHMARK="python3 -m pytest -vvv --rank 0 --benchmark-json"
           git checkout ${{ github.event.pull_request.base.sha }}
           $RUN_BENCHMARK ${{ env.BASELINE_JSON }}
           git checkout ${{ github.event.pull_request.head.sha }}

From 405d99cbab32cf8d6759e3b86c4e9b1e41244df1 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 11 Oct 2024 09:29:35 +0100
Subject: [PATCH 03/24] amend

---
 .github/workflows/benchmarks_pr.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index 9891142631c..287b68d4161 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -71,7 +71,7 @@ jobs:
       run:
         shell: bash -l {0}
     container:
-      image: nvidia/cuda:12.3.0-base-ubuntu22.04
+      image: nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
       options: --gpus all
     steps:
       - name: Who triggered this?

From d45b6694a578ec1c48df3ce73b73714f92f67430 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 11 Oct 2024 10:16:56 +0100
Subject: [PATCH 04/24] amend

---
 .github/workflows/benchmarks_pr.yml | 22 +++++++++-------------
 test/test_utils.py                  |  2 +-
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index 287b68d4161..ce7b8abb7bc 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -27,7 +27,7 @@ jobs:
         uses: actions/setup-python@v4
         with:
           python-version: '3.10'
-      - name: Setup Environment
+      - name: Setup Environment + benchmarks and tests
         run: |
           python3 -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U 
           python3 -m pip install git+https://github.com/pytorch/tensordict 
@@ -35,15 +35,13 @@ jobs:
           python3 -m pip install pytest pytest-benchmark
           python3 -m pip install "gym[accept-rom-license,atari]"
           python3 -m pip install "dm_control" "mujoco"
-      - name: Setup benchmarks
-        run: |
+          
           echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV
           echo "HEAD_SHA=$(echo ${{ github.event.pull_request.head.sha }} | cut -c1-8)" >> $GITHUB_ENV
           echo "BASELINE_JSON=$(mktemp)" >> $GITHUB_ENV
           echo "CONTENDER_JSON=$(mktemp)" >> $GITHUB_ENV
           echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
-      - name: Run benchmarks
-        run: |
+          
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
@@ -102,28 +100,26 @@ jobs:
       - name: setup Path
         run: |
           echo /usr/local/bin >> $GITHUB_PATH
-      - name: Setup Environment
+      - name: Setup Environment + benchmarks and tests
         run: |
           python3 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
           python3 -m pip install git+https://github.com/pytorch/tensordict 
-          python3 setup.py develop
           python3 -m pip install pytest pytest-benchmark
           python3 -m pip install "gym[accept-rom-license,atari]"
           python3 -m pip install "dm_control" "mujoco"
-      - name: check GPU presence
-        run: |
+          python3 setup.py develop
+
+          # test import
           python3 -c """import torch
           assert torch.cuda.device_count()
           """
-      - name: Setup benchmarks
-        run: |
+
           echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV
           echo "HEAD_SHA=$(echo ${{ github.event.pull_request.head.sha }} | cut -c1-8)" >> $GITHUB_ENV
           echo "BASELINE_JSON=$(mktemp)" >> $GITHUB_ENV
           echo "CONTENDER_JSON=$(mktemp)" >> $GITHUB_ENV
           echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
-      - name: Run benchmarks
-        run: |
+
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
diff --git a/test/test_utils.py b/test/test_utils.py
index f94b776a31b..4224a36b54f 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -191,7 +191,7 @@ def test_implement_for_check_versions(
     [
         ("0.27.0", None, "1.0.0"),
         ("0.27.2", None, "1.0.0"),
-        ("1.0.1", "1.0.0", None),
+        # ("1.0.1", "1.0.0", None),
     ],
 )
 @pytest.mark.parametrize(

From 19f09fed4707f09477dfeb5ed224504231a5527b Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 11 Oct 2024 10:39:40 +0100
Subject: [PATCH 05/24] amend

---
 .github/workflows/benchmarks_pr.yml | 30 +++++++++++++++--------------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index ce7b8abb7bc..a01f1b5bb43 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -27,7 +27,14 @@ jobs:
         uses: actions/setup-python@v4
         with:
           python-version: '3.10'
-      - name: Setup Environment + benchmarks and tests
+      - name: Setup benchmarks
+        run: |
+          echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV
+          echo "HEAD_SHA=$(echo ${{ github.event.pull_request.head.sha }} | cut -c1-8)" >> $GITHUB_ENV
+          echo "BASELINE_JSON=$(mktemp)" >> $GITHUB_ENV
+          echo "CONTENDER_JSON=$(mktemp)" >> $GITHUB_ENV
+          echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
+      - name: Setup Environment and tests
         run: |
           python3 -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U 
           python3 -m pip install git+https://github.com/pytorch/tensordict 
@@ -36,12 +43,6 @@ jobs:
           python3 -m pip install "gym[accept-rom-license,atari]"
           python3 -m pip install "dm_control" "mujoco"
           
-          echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV
-          echo "HEAD_SHA=$(echo ${{ github.event.pull_request.head.sha }} | cut -c1-8)" >> $GITHUB_ENV
-          echo "BASELINE_JSON=$(mktemp)" >> $GITHUB_ENV
-          echo "CONTENDER_JSON=$(mktemp)" >> $GITHUB_ENV
-          echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
-          
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
@@ -100,7 +101,14 @@ jobs:
       - name: setup Path
         run: |
           echo /usr/local/bin >> $GITHUB_PATH
-      - name: Setup Environment + benchmarks and tests
+      - name: Setup benchmarks
+        run: |
+          echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV
+          echo "HEAD_SHA=$(echo ${{ github.event.pull_request.head.sha }} | cut -c1-8)" >> $GITHUB_ENV
+          echo "BASELINE_JSON=$(mktemp)" >> $GITHUB_ENV
+          echo "CONTENDER_JSON=$(mktemp)" >> $GITHUB_ENV
+          echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
+      - name: Setup Environment and tests
         run: |
           python3 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
           python3 -m pip install git+https://github.com/pytorch/tensordict 
@@ -114,12 +122,6 @@ jobs:
           assert torch.cuda.device_count()
           """
 
-          echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV
-          echo "HEAD_SHA=$(echo ${{ github.event.pull_request.head.sha }} | cut -c1-8)" >> $GITHUB_ENV
-          echo "BASELINE_JSON=$(mktemp)" >> $GITHUB_ENV
-          echo "CONTENDER_JSON=$(mktemp)" >> $GITHUB_ENV
-          echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
-
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1

From f07f86cf7af9ed2b383c4bbb51f3d992e9238a82 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 11 Oct 2024 10:41:05 +0100
Subject: [PATCH 06/24] amend

---
 .github/workflows/benchmarks.yml | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index ff2fe163b13..ee40744de01 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -43,8 +43,7 @@ jobs:
           python3 -m pip install pytest pytest-benchmark
           python3 -m pip install "gym[accept-rom-license,atari]"
           python3 -m pip install "dm_control" "mujoco"
-      - name: Run benchmarks
-        run: |
+
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
@@ -111,13 +110,11 @@ jobs:
           python3 -m pip install pytest pytest-benchmark
           python3 -m pip install "gym[accept-rom-license,atari]"
           python3 -m pip install "dm_control" "mujoco"
-      - name: check GPU presence
-        run: |
+
           python3 -c """import torch
           assert torch.cuda.device_count()
           """
-      - name: Run benchmarks
-        run: |
+
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1

From 5a1f1c6ac43b520b0bb5952137f2e2917453d266 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 11 Oct 2024 11:22:03 +0100
Subject: [PATCH 07/24] amend

---
 .github/workflows/benchmarks.yml    | 6 ++++++
 .github/workflows/benchmarks_pr.yml | 6 ++++++
 2 files changed, 12 insertions(+)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index ee40744de01..ae16d14870d 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -37,6 +37,9 @@ jobs:
           python-version: '3.10'
       - name: Setup Environment
         run: |
+          python3.10 -m venv ./py310
+          activate ./py310/bin/activate
+
           python3 -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U 
           python3 -m pip install git+https://github.com/pytorch/tensordict 
           python3 setup.py develop
@@ -104,6 +107,9 @@ jobs:
           echo /usr/local/bin >> $GITHUB_PATH
       - name: Setup Environment
         run: |
+          python3.10 -m venv ./py310
+          activate ./py310/bin/activate
+
           python3 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
           python3 -m pip install git+https://github.com/pytorch/tensordict 
           python3 setup.py develop
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index a01f1b5bb43..fee2048fd8d 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -36,6 +36,9 @@ jobs:
           echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
       - name: Setup Environment and tests
         run: |
+          python3.10 -m venv ./py310
+          activate ./py310/bin/activate
+
           python3 -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U 
           python3 -m pip install git+https://github.com/pytorch/tensordict 
           python3 setup.py develop
@@ -110,6 +113,9 @@ jobs:
           echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
       - name: Setup Environment and tests
         run: |
+          python3.10 -m venv ./py310
+          activate ./py310/bin/activate
+
           python3 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
           python3 -m pip install git+https://github.com/pytorch/tensordict 
           python3 -m pip install pytest pytest-benchmark

From cd9d49ff45abee7e05c88c343b433af017e100a9 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 11 Oct 2024 11:24:58 +0100
Subject: [PATCH 08/24] amend

---
 .github/workflows/benchmarks.yml    | 4 ++--
 .github/workflows/benchmarks_pr.yml | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index ae16d14870d..d2d3119f945 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -38,7 +38,7 @@ jobs:
       - name: Setup Environment
         run: |
           python3.10 -m venv ./py310
-          activate ./py310/bin/activate
+          source ./py310/bin/activate
 
           python3 -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U 
           python3 -m pip install git+https://github.com/pytorch/tensordict 
@@ -108,7 +108,7 @@ jobs:
       - name: Setup Environment
         run: |
           python3.10 -m venv ./py310
-          activate ./py310/bin/activate
+          source ./py310/bin/activate
 
           python3 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
           python3 -m pip install git+https://github.com/pytorch/tensordict 
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index fee2048fd8d..6cbb00e31f7 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -37,7 +37,7 @@ jobs:
       - name: Setup Environment and tests
         run: |
           python3.10 -m venv ./py310
-          activate ./py310/bin/activate
+          source ./py310/bin/activate
 
           python3 -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U 
           python3 -m pip install git+https://github.com/pytorch/tensordict 
@@ -114,7 +114,7 @@ jobs:
       - name: Setup Environment and tests
         run: |
           python3.10 -m venv ./py310
-          activate ./py310/bin/activate
+          source ./py310/bin/activate
 
           python3 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
           python3 -m pip install git+https://github.com/pytorch/tensordict 

From b213370ba42c8167e682266fee6cdfd725e751c7 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 11 Oct 2024 11:43:37 +0100
Subject: [PATCH 09/24] amend

---
 .github/workflows/benchmarks.yml    | 23 +++++++++++++++++------
 .github/workflows/benchmarks_pr.yml | 23 +++++++++++++++++------
 2 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index d2d3119f945..31e88546562 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -77,6 +77,17 @@ jobs:
       image: nvidia/cuda:12.3.0-base-ubuntu22.04
       options: --gpus all
     steps:
+      - name: Set GITHUB_BRANCH environment variable
+        run: |
+          if [ "${{ github.event_name }}" == "push" ]; then
+            export GITHUB_BRANCH=${{ github.event.branch }}
+          elif [ "${{ github.event_name }}" == "pull_request" ]; then
+            export GITHUB_BRANCH=${{ github.event.pull_request.head.ref }}
+          else
+            echo "Unsupported event type"
+            exit 1
+          fi
+          echo "GITHUB_BRANCH=$GITHUB_BRANCH" >> $GITHUB_ENV
       - name: Who triggered this?
         run: |
           echo "Action triggered by ${{ github.event.pull_request.html_url }}"
@@ -110,12 +121,12 @@ jobs:
           python3.10 -m venv ./py310
           source ./py310/bin/activate
 
-          python3 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
-          python3 -m pip install git+https://github.com/pytorch/tensordict 
-          python3 setup.py develop
-          python3 -m pip install pytest pytest-benchmark
-          python3 -m pip install "gym[accept-rom-license,atari]"
-          python3 -m pip install "dm_control" "mujoco"
+          python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
+          python3.10 -m pip install git+https://github.com/pytorch/tensordict 
+          python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
+          python3.10 -m pip install pytest pytest-benchmark
+          python3.10 -m pip install "gym[accept-rom-license,atari]"
+          python3.10 -m pip install "dm_control" "mujoco"
 
           python3 -c """import torch
           assert torch.cuda.device_count()
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index 6cbb00e31f7..883afd5628c 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -76,6 +76,17 @@ jobs:
       image: nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
       options: --gpus all
     steps:
+      - name: Set GITHUB_BRANCH environment variable
+        run: |
+          if [ "${{ github.event_name }}" == "push" ]; then
+            export GITHUB_BRANCH=${{ github.event.branch }}
+          elif [ "${{ github.event_name }}" == "pull_request" ]; then
+            export GITHUB_BRANCH=${{ github.event.pull_request.head.ref }}
+          else
+            echo "Unsupported event type"
+            exit 1
+          fi
+          echo "GITHUB_BRANCH=$GITHUB_BRANCH" >> $GITHUB_ENV
       - name: Who triggered this?
         run: |
           echo "Action triggered by ${{ github.event.pull_request.html_url }}"
@@ -116,12 +127,12 @@ jobs:
           python3.10 -m venv ./py310
           source ./py310/bin/activate
 
-          python3 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
-          python3 -m pip install git+https://github.com/pytorch/tensordict 
-          python3 -m pip install pytest pytest-benchmark
-          python3 -m pip install "gym[accept-rom-license,atari]"
-          python3 -m pip install "dm_control" "mujoco"
-          python3 setup.py develop
+          python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
+          python3.10 -m pip install git+https://github.com/pytorch/tensordict 
+          python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
+          python3.10 -m pip install pytest pytest-benchmark
+          python3.10 -m pip install "gym[accept-rom-license,atari]"
+          python3.10 -m pip install "dm_control" "mujoco"
 
           # test import
           python3 -c """import torch

From 389959eafe709a4e1eda3d8fc1f0965c2e19030e Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 11 Oct 2024 13:23:42 +0100
Subject: [PATCH 10/24] amend

---
 .github/workflows/benchmarks.yml    | 4 ++--
 .github/workflows/benchmarks_pr.yml | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index 31e88546562..d18ea29fa8d 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -50,7 +50,7 @@ jobs:
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
-          python3 -m pytest -vvv --rank 0 --benchmark-json output.json
+          python3 -m pytest -vvv --rank 0 --benchmark-json output.json --ignore test_objectives_benchmarks.py
       - name: Store benchmark results
         uses: benchmark-action/github-action-benchmark@v1
         if: ${{ github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' }}
@@ -135,7 +135,7 @@ jobs:
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
-          python3 -m pytest -vvv --rank 0 --benchmark-json output.json
+          python3 -m pytest -vvv --rank 0 --benchmark-json output.json --ignore test_objectives_benchmarks.py
       - name: Store benchmark results
         uses: benchmark-action/github-action-benchmark@v1
         if: ${{ github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' }}
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index 883afd5628c..fc4f0d2c941 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -49,7 +49,7 @@ jobs:
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
-          RUN_BENCHMARK="python3 -m pytest -vvv --rank 0 --benchmark-json"
+          RUN_BENCHMARK="python3 -m pytest -vvv --rank 0 --ignore test_objectives_benchmarks.py --benchmark-json "
           git checkout ${{ github.event.pull_request.base.sha }}
           $RUN_BENCHMARK ${{ env.BASELINE_JSON }}
           git checkout ${{ github.event.pull_request.head.sha }}
@@ -142,7 +142,7 @@ jobs:
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
-          RUN_BENCHMARK="python3 -m pytest -vvv --rank 0 --benchmark-json"
+          RUN_BENCHMARK="python3 -m pytest -vvv --rank 0 --ignore test_objectives_benchmarks.py --benchmark-json "
           git checkout ${{ github.event.pull_request.base.sha }}
           $RUN_BENCHMARK ${{ env.BASELINE_JSON }}
           git checkout ${{ github.event.pull_request.head.sha }}

From d7c69819b36258ecc55f9c83ec506655aab91c49 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 11 Oct 2024 13:35:45 +0100
Subject: [PATCH 11/24] amend

---
 .github/workflows/benchmarks.yml    | 4 ++--
 .github/workflows/benchmarks_pr.yml | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index d18ea29fa8d..2b7402657f6 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -50,7 +50,7 @@ jobs:
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
-          python3 -m pytest -vvv --rank 0 --benchmark-json output.json --ignore test_objectives_benchmarks.py
+          python3 -m pytest -vvv --rank 0 --benchmark-json output.json --ignore test_collectors_benchmark.py
       - name: Store benchmark results
         uses: benchmark-action/github-action-benchmark@v1
         if: ${{ github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' }}
@@ -135,7 +135,7 @@ jobs:
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
-          python3 -m pytest -vvv --rank 0 --benchmark-json output.json --ignore test_objectives_benchmarks.py
+          python3 -m pytest -vvv --rank 0 --benchmark-json output.json --ignore test_collectors_benchmark.py
       - name: Store benchmark results
         uses: benchmark-action/github-action-benchmark@v1
         if: ${{ github.ref == 'refs/heads/main' || github.event_name == 'workflow_dispatch' }}
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index fc4f0d2c941..e4ad721620c 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -49,7 +49,7 @@ jobs:
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
-          RUN_BENCHMARK="python3 -m pytest -vvv --rank 0 --ignore test_objectives_benchmarks.py --benchmark-json "
+          RUN_BENCHMARK="python3 -m pytest -vvv --rank 0 --ignore test_collectors_benchmark.py --benchmark-json "
           git checkout ${{ github.event.pull_request.base.sha }}
           $RUN_BENCHMARK ${{ env.BASELINE_JSON }}
           git checkout ${{ github.event.pull_request.head.sha }}
@@ -142,7 +142,7 @@ jobs:
           cd benchmarks/
           export TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1
           export TD_GET_DEFAULTS_TO_NONE=1
-          RUN_BENCHMARK="python3 -m pytest -vvv --rank 0 --ignore test_objectives_benchmarks.py --benchmark-json "
+          RUN_BENCHMARK="python3 -m pytest -vvv --rank 0 --ignore test_collectors_benchmark.py --benchmark-json "
           git checkout ${{ github.event.pull_request.base.sha }}
           $RUN_BENCHMARK ${{ env.BASELINE_JSON }}
           git checkout ${{ github.event.pull_request.head.sha }}

From 4a707625665a8103ebf1d25522fa1d0f82bb7013 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 11 Oct 2024 14:23:27 +0100
Subject: [PATCH 12/24] amend

---
 .github/workflows/benchmarks.yml    | 1 +
 .github/workflows/benchmarks_pr.yml | 1 +
 2 files changed, 2 insertions(+)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index 2b7402657f6..3eacc615108 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -122,6 +122,7 @@ jobs:
           source ./py310/bin/activate
 
           python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
+          python3.10 -m pip install cmake ninja
           python3.10 -m pip install git+https://github.com/pytorch/tensordict 
           python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
           python3.10 -m pip install pytest pytest-benchmark
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index e4ad721620c..f4085e72791 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -128,6 +128,7 @@ jobs:
           source ./py310/bin/activate
 
           python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
+          python3.10 -m pip install cmake ninja
           python3.10 -m pip install git+https://github.com/pytorch/tensordict 
           python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
           python3.10 -m pip install pytest pytest-benchmark

From 8b8d1bc3ccb967e2c1ddf90f161a5e0fc53bcf52 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 11 Oct 2024 15:05:36 +0100
Subject: [PATCH 13/24] amend

---
 .github/workflows/benchmarks.yml    | 3 ++-
 .github/workflows/benchmarks_pr.yml | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index 3eacc615108..f0ea3ee9d69 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -124,7 +124,8 @@ jobs:
           python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
           python3.10 -m pip install cmake ninja
           python3.10 -m pip install git+https://github.com/pytorch/tensordict 
-          python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
+#          python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
+          python3.10 setup.py develop
           python3.10 -m pip install pytest pytest-benchmark
           python3.10 -m pip install "gym[accept-rom-license,atari]"
           python3.10 -m pip install "dm_control" "mujoco"
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index f4085e72791..d889a69f5f4 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -130,7 +130,8 @@ jobs:
           python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
           python3.10 -m pip install cmake ninja
           python3.10 -m pip install git+https://github.com/pytorch/tensordict 
-          python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
+#          python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
+          python3.10 setup.py develop
           python3.10 -m pip install pytest pytest-benchmark
           python3.10 -m pip install "gym[accept-rom-license,atari]"
           python3.10 -m pip install "dm_control" "mujoco"

From b6ca4fb7eaaf0d82a9493d20884cae981d8495f2 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 11 Oct 2024 15:11:57 +0100
Subject: [PATCH 14/24] amend

---
 .github/workflows/benchmarks.yml    | 2 +-
 .github/workflows/benchmarks_pr.yml | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index f0ea3ee9d69..b064e6db857 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -124,7 +124,7 @@ jobs:
           python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
           python3.10 -m pip install cmake ninja
           python3.10 -m pip install git+https://github.com/pytorch/tensordict 
-#          python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
+          # python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
           python3.10 setup.py develop
           python3.10 -m pip install pytest pytest-benchmark
           python3.10 -m pip install "gym[accept-rom-license,atari]"
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index d889a69f5f4..5bb6c42b2be 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -130,7 +130,7 @@ jobs:
           python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
           python3.10 -m pip install cmake ninja
           python3.10 -m pip install git+https://github.com/pytorch/tensordict 
-#          python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
+          # python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
           python3.10 setup.py develop
           python3.10 -m pip install pytest pytest-benchmark
           python3.10 -m pip install "gym[accept-rom-license,atari]"

From cfe8ca60144d3d8f9abc2969481079c7050869ad Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Fri, 11 Oct 2024 16:01:04 +0100
Subject: [PATCH 15/24] amend

---
 .github/workflows/benchmarks.yml    | 29 ++++++++++++-----------------
 .github/workflows/benchmarks_pr.yml | 29 ++++++++++++-----------------
 2 files changed, 24 insertions(+), 34 deletions(-)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index b064e6db857..c834a301d1b 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -91,16 +91,6 @@ jobs:
       - name: Who triggered this?
         run: |
           echo "Action triggered by ${{ github.event.pull_request.html_url }}"
-      - name: Install deps
-        run: |
-          export TZ=Europe/London
-          export DEBIAN_FRONTEND=noninteractive  # tzdata bug
-          apt-get update -y
-          apt-get install software-properties-common -y
-          add-apt-repository ppa:git-core/candidate -y
-          apt-get update -y
-          apt-get upgrade -y
-          apt-get -y install libglu1-mesa libgl1-mesa-glx libosmesa6 gcc curl g++ unzip wget libglfw3-dev libgles2-mesa-dev libglew-dev sudo git cmake libz-dev
       - name: Check ldd --version
         run: ldd --version
       - name: Checkout
@@ -118,17 +108,22 @@ jobs:
           echo /usr/local/bin >> $GITHUB_PATH
       - name: Setup Environment
         run: |
+          export TZ=Europe/London
+          export DEBIAN_FRONTEND=noninteractive  # tzdata bug
+          apt-get update -y
+          apt-get install software-properties-common -y
+          add-apt-repository ppa:git-core/candidate -y
+          apt-get update -y
+          apt-get upgrade -y
+          apt-get -y install libglu1-mesa libgl1-mesa-glx libosmesa6 gcc curl g++ unzip wget libglfw3-dev libgles2-mesa-dev libglew-dev sudo git cmake libz-dev
+
           python3.10 -m venv ./py310
           source ./py310/bin/activate
 
-          python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
-          python3.10 -m pip install cmake ninja
+          python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu124 -U
+          python3.10 -m pip install cmake ninja pytest pytest-benchmark mujoco dm_control "gym[accept-rom-license,atari]"
           python3.10 -m pip install git+https://github.com/pytorch/tensordict 
-          # python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
-          python3.10 setup.py develop
-          python3.10 -m pip install pytest pytest-benchmark
-          python3.10 -m pip install "gym[accept-rom-license,atari]"
-          python3.10 -m pip install "dm_control" "mujoco"
+          python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
 
           python3 -c """import torch
           assert torch.cuda.device_count()
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index 5bb6c42b2be..cf32864fd80 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -90,16 +90,6 @@ jobs:
       - name: Who triggered this?
         run: |
           echo "Action triggered by ${{ github.event.pull_request.html_url }}"
-      - name: Install deps
-        run: |
-          export TZ=Europe/London
-          export DEBIAN_FRONTEND=noninteractive  # tzdata bug
-          apt-get update -y
-          apt-get install software-properties-common -y
-          add-apt-repository ppa:git-core/candidate -y
-          apt-get update -y
-          apt-get upgrade -y
-          apt-get -y install libglu1-mesa libgl1-mesa-glx libosmesa6 gcc curl g++ unzip wget libglfw3-dev libgles2-mesa-dev libglew-dev sudo git cmake libz-dev
       - name: Check ldd --version
         run: ldd --version
       - name: Checkout
@@ -124,17 +114,22 @@ jobs:
           echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
       - name: Setup Environment and tests
         run: |
+          export TZ=Europe/London
+          export DEBIAN_FRONTEND=noninteractive  # tzdata bug
+          apt-get update -y
+          apt-get install software-properties-common -y
+          add-apt-repository ppa:git-core/candidate -y
+          apt-get update -y
+          apt-get upgrade -y
+          apt-get -y install libglu1-mesa libgl1-mesa-glx libosmesa6 gcc curl g++ unzip wget libglfw3-dev libgles2-mesa-dev libglew-dev sudo git cmake libz-dev
+
           python3.10 -m venv ./py310
           source ./py310/bin/activate
 
-          python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121 -U
-          python3.10 -m pip install cmake ninja
+          python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu124 -U
+          python3.10 -m pip install cmake ninja pytest pytest-benchmark mujoco dm_control "gym[accept-rom-license,atari]"
           python3.10 -m pip install git+https://github.com/pytorch/tensordict 
-          # python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
-          python3.10 setup.py develop
-          python3.10 -m pip install pytest pytest-benchmark
-          python3.10 -m pip install "gym[accept-rom-license,atari]"
-          python3.10 -m pip install "dm_control" "mujoco"
+          python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
 
           # test import
           python3 -c """import torch

From b112c75ed50e96b3d0dcdbb3420dde6dcd7a7c19 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Mon, 14 Oct 2024 08:46:24 +0100
Subject: [PATCH 16/24] amend

---
 .github/workflows/benchmarks.yml    | 13 +++++++------
 .github/workflows/benchmarks_pr.yml | 28 ++++++++++++++--------------
 2 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index c834a301d1b..dbda7e7c1a9 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -101,11 +101,6 @@ jobs:
         uses: actions/setup-python@v4
         with:
           python-version: '3.10'
-      - name: Setup git
-        run: git config --global --add safe.directory /__w/rl/rl
-      - name: setup Path
-        run: |
-          echo /usr/local/bin >> $GITHUB_PATH
       - name: Setup Environment
         run: |
           export TZ=Europe/London
@@ -116,7 +111,13 @@ jobs:
           apt-get update -y
           apt-get upgrade -y
           apt-get -y install libglu1-mesa libgl1-mesa-glx libosmesa6 gcc curl g++ unzip wget libglfw3-dev libgles2-mesa-dev libglew-dev sudo git cmake libz-dev
-
+      - name: Setup git
+        run: git config --global --add safe.directory /__w/rl/rl
+      - name: setup Path
+        run: |
+          echo /usr/local/bin >> $GITHUB_PATH
+      - name: Run
+        run: |
           python3.10 -m venv ./py310
           source ./py310/bin/activate
 
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index cf32864fd80..e1d772e1f14 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -100,19 +100,7 @@ jobs:
         uses: actions/setup-python@v4
         with:
           python-version: '3.10'
-      - name: Setup git
-        run: git config --global --add safe.directory /__w/rl/rl
-      - name: setup Path
-        run: |
-          echo /usr/local/bin >> $GITHUB_PATH
-      - name: Setup benchmarks
-        run: |
-          echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV
-          echo "HEAD_SHA=$(echo ${{ github.event.pull_request.head.sha }} | cut -c1-8)" >> $GITHUB_ENV
-          echo "BASELINE_JSON=$(mktemp)" >> $GITHUB_ENV
-          echo "CONTENDER_JSON=$(mktemp)" >> $GITHUB_ENV
-          echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
-      - name: Setup Environment and tests
+      - name: Setup Environment
         run: |
           export TZ=Europe/London
           export DEBIAN_FRONTEND=noninteractive  # tzdata bug
@@ -122,7 +110,13 @@ jobs:
           apt-get update -y
           apt-get upgrade -y
           apt-get -y install libglu1-mesa libgl1-mesa-glx libosmesa6 gcc curl g++ unzip wget libglfw3-dev libgles2-mesa-dev libglew-dev sudo git cmake libz-dev
-
+      - name: Setup git
+        run: git config --global --add safe.directory /__w/rl/rl
+      - name: setup Path
+        run: |
+          echo /usr/local/bin >> $GITHUB_PATH
+      - name: Run
+        run: |
           python3.10 -m venv ./py310
           source ./py310/bin/activate
 
@@ -131,6 +125,12 @@ jobs:
           python3.10 -m pip install git+https://github.com/pytorch/tensordict 
           python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
 
+          echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV
+          echo "HEAD_SHA=$(echo ${{ github.event.pull_request.head.sha }} | cut -c1-8)" >> $GITHUB_ENV
+          echo "BASELINE_JSON=$(mktemp)" >> $GITHUB_ENV
+          echo "CONTENDER_JSON=$(mktemp)" >> $GITHUB_ENV
+          echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
+
           # test import
           python3 -c """import torch
           assert torch.cuda.device_count()

From b23d0a68058802508ad7241f5830c6af133c62c2 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Mon, 14 Oct 2024 09:05:35 +0100
Subject: [PATCH 17/24] amend

---
 .github/workflows/benchmarks.yml    | 11 ++++++++++-
 .github/workflows/benchmarks_pr.yml | 16 +++++++++-------
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index dbda7e7c1a9..52ea5b899af 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -116,6 +116,13 @@ jobs:
       - name: setup Path
         run: |
           echo /usr/local/bin >> $GITHUB_PATH
+      - name: Setup benchmarks
+        run: |
+          echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV
+          echo "HEAD_SHA=$(echo ${{ github.event.pull_request.head.sha }} | cut -c1-8)" >> $GITHUB_ENV
+          echo "BASELINE_JSON=$(mktemp)" >> $GITHUB_ENV
+          echo "CONTENDER_JSON=$(mktemp)" >> $GITHUB_ENV
+          echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
       - name: Run
         run: |
           python3.10 -m venv ./py310
@@ -124,8 +131,10 @@ jobs:
           python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu124 -U
           python3.10 -m pip install cmake ninja pytest pytest-benchmark mujoco dm_control "gym[accept-rom-license,atari]"
           python3.10 -m pip install git+https://github.com/pytorch/tensordict 
-          python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
+          python3.10 setup.py develop
+          # python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
 
+          # test import
           python3 -c """import torch
           assert torch.cuda.device_count()
           """
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index e1d772e1f14..e5c53d54d9a 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -115,6 +115,13 @@ jobs:
       - name: setup Path
         run: |
           echo /usr/local/bin >> $GITHUB_PATH
+      - name: Setup benchmarks
+        run: |
+          echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV
+          echo "HEAD_SHA=$(echo ${{ github.event.pull_request.head.sha }} | cut -c1-8)" >> $GITHUB_ENV
+          echo "BASELINE_JSON=$(mktemp)" >> $GITHUB_ENV
+          echo "CONTENDER_JSON=$(mktemp)" >> $GITHUB_ENV
+          echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
       - name: Run
         run: |
           python3.10 -m venv ./py310
@@ -123,13 +130,8 @@ jobs:
           python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu124 -U
           python3.10 -m pip install cmake ninja pytest pytest-benchmark mujoco dm_control "gym[accept-rom-license,atari]"
           python3.10 -m pip install git+https://github.com/pytorch/tensordict 
-          python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
-
-          echo "BASE_SHA=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-8)" >> $GITHUB_ENV
-          echo "HEAD_SHA=$(echo ${{ github.event.pull_request.head.sha }} | cut -c1-8)" >> $GITHUB_ENV
-          echo "BASELINE_JSON=$(mktemp)" >> $GITHUB_ENV
-          echo "CONTENDER_JSON=$(mktemp)" >> $GITHUB_ENV
-          echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
+          python3.10 setup.py develop
+          # python3.10 -m pip install git+https://github.com/pytorch/rl@$GITHUB_BRANCH
 
           # test import
           python3 -c """import torch

From f81d48a7ee9edd597dd3d3286ec28ca4daad6b38 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Mon, 14 Oct 2024 10:00:55 +0100
Subject: [PATCH 18/24] amend

---
 .github/workflows/benchmarks.yml    | 2 +-
 .github/workflows/benchmarks_pr.yml | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index 52ea5b899af..16d013c4b6e 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -110,7 +110,7 @@ jobs:
           add-apt-repository ppa:git-core/candidate -y
           apt-get update -y
           apt-get upgrade -y
-          apt-get -y install libglu1-mesa libgl1-mesa-glx libosmesa6 gcc curl g++ unzip wget libglfw3-dev libgles2-mesa-dev libglew-dev sudo git cmake libz-dev
+          apt-get -y install libglu1-mesa libgl1-mesa-glx libosmesa6 gcc curl g++ unzip wget libglfw3-dev libgles2-mesa-dev libglew-dev sudo git cmake libz-dev libpython3.10-dev
       - name: Setup git
         run: git config --global --add safe.directory /__w/rl/rl
       - name: setup Path
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index e5c53d54d9a..9f30a65a057 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -109,7 +109,7 @@ jobs:
           add-apt-repository ppa:git-core/candidate -y
           apt-get update -y
           apt-get upgrade -y
-          apt-get -y install libglu1-mesa libgl1-mesa-glx libosmesa6 gcc curl g++ unzip wget libglfw3-dev libgles2-mesa-dev libglew-dev sudo git cmake libz-dev
+          apt-get -y install libglu1-mesa libgl1-mesa-glx libosmesa6 gcc curl g++ unzip wget libglfw3-dev libgles2-mesa-dev libglew-dev sudo git cmake libz-dev libpython3.10-dev
       - name: Setup git
         run: git config --global --add safe.directory /__w/rl/rl
       - name: setup Path

From 43ff46d3c7373de6c54bc40d44d1739bd2239922 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Mon, 14 Oct 2024 10:08:52 +0100
Subject: [PATCH 19/24] amend

---
 .github/workflows/benchmarks.yml    | 2 +-
 .github/workflows/benchmarks_pr.yml | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index 16d013c4b6e..dc76865f890 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -125,7 +125,7 @@ jobs:
           echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
       - name: Run
         run: |
-          python3.10 -m venv ./py310
+          python3.10 -m venv --system-site-packages ./py310
           source ./py310/bin/activate
 
           python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu124 -U
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index 9f30a65a057..cfd73456cb3 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -124,7 +124,7 @@ jobs:
           echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
       - name: Run
         run: |
-          python3.10 -m venv ./py310
+          python3.10 -m venv --system-site-packages ./py310
           source ./py310/bin/activate
 
           python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu124 -U

From 734582bc9069a73a866bf616c1dda8b8134a3252 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Mon, 14 Oct 2024 10:17:01 +0100
Subject: [PATCH 20/24] amend

---
 .github/workflows/benchmarks.yml    | 1 +
 .github/workflows/benchmarks_pr.yml | 1 +
 2 files changed, 2 insertions(+)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index dc76865f890..28832d5229b 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -127,6 +127,7 @@ jobs:
         run: |
           python3.10 -m venv --system-site-packages ./py310
           source ./py310/bin/activate
+          export PYTHON_INCLUDE_DIR=/usr/include/python3.10
 
           python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu124 -U
           python3.10 -m pip install cmake ninja pytest pytest-benchmark mujoco dm_control "gym[accept-rom-license,atari]"
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index cfd73456cb3..dfd8850a6f7 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -126,6 +126,7 @@ jobs:
         run: |
           python3.10 -m venv --system-site-packages ./py310
           source ./py310/bin/activate
+          export PYTHON_INCLUDE_DIR=/usr/include/python3.10
 
           python3.10 -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu124 -U
           python3.10 -m pip install cmake ninja pytest pytest-benchmark mujoco dm_control "gym[accept-rom-license,atari]"

From 2b9041a21f0ceb71ccd676204b7b532803c43d05 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Mon, 14 Oct 2024 10:39:20 +0100
Subject: [PATCH 21/24] amend

---
 .github/workflows/benchmarks.yml    | 2 ++
 .github/workflows/benchmarks_pr.yml | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index 28832d5229b..dce44a35719 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -125,6 +125,8 @@ jobs:
           echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
       - name: Run
         run: |
+          echo $(find /usr -name Python.h)
+
           python3.10 -m venv --system-site-packages ./py310
           source ./py310/bin/activate
           export PYTHON_INCLUDE_DIR=/usr/include/python3.10
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index dfd8850a6f7..6e29d7a6c3b 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -124,6 +124,8 @@ jobs:
           echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
       - name: Run
         run: |
+          echo $(find /usr -name Python.h)
+
           python3.10 -m venv --system-site-packages ./py310
           source ./py310/bin/activate
           export PYTHON_INCLUDE_DIR=/usr/include/python3.10

From 4fb81ed8ada1ef51190c2dff48e545d766740349 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Mon, 14 Oct 2024 10:58:22 +0100
Subject: [PATCH 22/24] amend

---
 setup.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index d37c179600f..2cb3bd78871 100644
--- a/setup.py
+++ b/setup.py
@@ -152,11 +152,15 @@ def get_extensions():
     }
     sources = list(extension_sources)
 
+    include_dirs = [this_dir]
+    python_include_dir = os.getenv('PYTHON_INCLUDE_DIR')
+    if python_include_dir is not None:
+        include_dirs.append(python_include_dir)
     ext_modules = [
         extension(
             "torchrl._torchrl",
             sources,
-            include_dirs=[this_dir],
+            include_dirs=include_dirs,
             extra_compile_args=extra_compile_args,
             extra_link_args=extra_link_args,
         )

From 5e2da2273abc31f6c97f2ac56eb61cf461a35bb3 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Mon, 14 Oct 2024 10:58:37 +0100
Subject: [PATCH 23/24] amend

---
 .github/workflows/benchmarks.yml    | 2 --
 .github/workflows/benchmarks_pr.yml | 2 --
 2 files changed, 4 deletions(-)

diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
index dce44a35719..28832d5229b 100644
--- a/.github/workflows/benchmarks.yml
+++ b/.github/workflows/benchmarks.yml
@@ -125,8 +125,6 @@ jobs:
           echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
       - name: Run
         run: |
-          echo $(find /usr -name Python.h)
-
           python3.10 -m venv --system-site-packages ./py310
           source ./py310/bin/activate
           export PYTHON_INCLUDE_DIR=/usr/include/python3.10
diff --git a/.github/workflows/benchmarks_pr.yml b/.github/workflows/benchmarks_pr.yml
index 6e29d7a6c3b..dfd8850a6f7 100644
--- a/.github/workflows/benchmarks_pr.yml
+++ b/.github/workflows/benchmarks_pr.yml
@@ -124,8 +124,6 @@ jobs:
           echo "PR_COMMENT=$(mktemp)" >>  $GITHUB_ENV
       - name: Run
         run: |
-          echo $(find /usr -name Python.h)
-
           python3.10 -m venv --system-site-packages ./py310
           source ./py310/bin/activate
           export PYTHON_INCLUDE_DIR=/usr/include/python3.10

From b957a401bb8fbbeebe97c801d4fd1bc0b0e3d3a4 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Mon, 14 Oct 2024 11:17:03 +0100
Subject: [PATCH 24/24] amend

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 2cb3bd78871..a52711b0ed5 100644
--- a/setup.py
+++ b/setup.py
@@ -153,7 +153,7 @@ def get_extensions():
     sources = list(extension_sources)
 
     include_dirs = [this_dir]
-    python_include_dir = os.getenv('PYTHON_INCLUDE_DIR')
+    python_include_dir = os.getenv("PYTHON_INCLUDE_DIR")
     if python_include_dir is not None:
         include_dirs.append(python_include_dir)
     ext_modules = [