From c870d7c0edf1b7a38795bfd656308ad6ed412e54 Mon Sep 17 00:00:00 2001
From: vmoens <vmoens@meta.com>
Date: Tue, 9 Apr 2024 12:28:04 -0700
Subject: [PATCH 01/14] init

---
 test/test_libs.py             | 16 ++++++++--------
 torchrl/envs/libs/gym.py      | 19 +++++++++++++------
 torchrl/envs/libs/isaacgym.py | 30 ++++++++++++++++--------------
 3 files changed, 37 insertions(+), 28 deletions(-)

diff --git a/test/test_libs.py b/test/test_libs.py
index c72fcc562e6..c5543fe689e 100644
--- a/test/test_libs.py
+++ b/test/test_libs.py
@@ -2,7 +2,14 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
-import importlib
+import importlib.util
+_has_isaac = importlib.util.find_spec("isaacgym") is not None
+
+if _has_isaac:
+    # isaac gym asks to be imported before torch...
+    import isaacgym  # noqa
+    import isaacgymenvs  # noqa
+    from torchrl.envs.libs.isaacgym import IsaacGymEnv
 import os
 from contextlib import nullcontext
 from pathlib import Path
@@ -14,13 +21,6 @@
 from torchrl.envs.transforms import ActionMask, TransformedEnv
 from torchrl.modules import MaskedCategorical
 
-_has_isaac = importlib.util.find_spec("isaacgym") is not None
-
-if _has_isaac:
-    # isaac gym asks to be imported before torch...
-    import isaacgym  # noqa
-    import isaacgymenvs  # noqa
-    from torchrl.envs.libs.isaacgym import IsaacGymEnv
 
 import argparse
 import importlib
diff --git a/torchrl/envs/libs/gym.py b/torchrl/envs/libs/gym.py
index d8043cb9ef7..dfbf7ae6217 100644
--- a/torchrl/envs/libs/gym.py
+++ b/torchrl/envs/libs/gym.py
@@ -943,6 +943,9 @@ def _reward_space(self, env):  # noqa: F811
             return rs
 
     def _make_specs(self, env: "gym.Env", batch_size=None) -> None:  # noqa: F821
+        # If batch_size is provided, we se it to tell what batch size must be used
+        # instead of self.batch_size
+        cur_batch_size = self.batch_size if batch_size is None else torch.Size([])
         action_spec = _gym_to_torchrl_spec_transform(
             env.action_space,
             device=self.device,
@@ -953,17 +956,18 @@ def _make_specs(self, env: "gym.Env", batch_size=None) -> None:  # noqa: F821
             device=self.device,
             categorical_action_encoding=self._categorical_action_encoding,
         )
+        print("-1", observation_spec)
         if not isinstance(observation_spec, CompositeSpec):
             if self.from_pixels:
                 observation_spec = CompositeSpec(
-                    pixels=observation_spec, shape=self.batch_size
+                    pixels=observation_spec, shape=cur_batch_size
                 )
             else:
                 observation_spec = CompositeSpec(
-                    observation=observation_spec, shape=self.batch_size
+                    observation=observation_spec, shape=cur_batch_size
                 )
-        elif observation_spec.shape[: len(self.batch_size)] != self.batch_size:
-            observation_spec.shape = self.batch_size
+        elif observation_spec.shape[: len(cur_batch_size)] != cur_batch_size:
+            observation_spec.shape = cur_batch_size
 
         reward_space = self._reward_space(env)
         if reward_space is not None:
@@ -978,15 +982,18 @@ def _make_specs(self, env: "gym.Env", batch_size=None) -> None:  # noqa: F821
                 device=self.device,
             )
         if batch_size is not None:
+            print("0", observation_spec)
             action_spec = action_spec.expand(*batch_size, *action_spec.shape)
             reward_spec = reward_spec.expand(*batch_size, *reward_spec.shape)
             observation_spec = observation_spec.expand(
                 *batch_size, *observation_spec.shape
             )
+            print("1", observation_spec)
+
         self.done_spec = self._make_done_spec()
         self.action_spec = action_spec
-        if reward_spec.shape[: len(self.batch_size)] != self.batch_size:
-            self.reward_spec = reward_spec.expand(*self.batch_size, *reward_spec.shape)
+        if reward_spec.shape[: len(cur_batch_size)] != cur_batch_size:
+            self.reward_spec = reward_spec.expand(*cur_batch_size, *reward_spec.shape)
         else:
             self.reward_spec = reward_spec
         self.observation_spec = observation_spec
diff --git a/torchrl/envs/libs/isaacgym.py b/torchrl/envs/libs/isaacgym.py
index 90f90043ab7..f8f5143c9dd 100644
--- a/torchrl/envs/libs/isaacgym.py
+++ b/torchrl/envs/libs/isaacgym.py
@@ -16,7 +16,7 @@
 from tensordict import TensorDictBase
 from torchrl.envs.libs.gym import GymWrapper
 from torchrl.envs.utils import _classproperty, make_composite_from_td
-
+from torchrl.data import CompositeSpec
 _has_isaac = importlib.util.find_spec("isaacgym") is not None
 
 
@@ -51,17 +51,20 @@ def __init__(
         )
         num_envs = env.num_envs
         super().__init__(
-            env, torch.device(env.device), batch_size=torch.Size([num_envs]), **kwargs
+            env, torch.device(env.device), batch_size=torch.Size([]), **kwargs
         )
         if not hasattr(self, "task"):
             # by convention in IsaacGymEnvs
             self.task = env.__name__
 
     def _make_specs(self, env: "gym.Env") -> None:  # noqa: F821
+        print("self.batch_size", self.batch_size)
         super()._make_specs(env, batch_size=self.batch_size)
-        self.full_done_spec = {
+        self.full_done_spec = CompositeSpec({
             key: spec.squeeze(-1) for key, spec in self.full_done_spec.items(True, True)
-        }
+        }, shape=self.batch_size)
+        print(self.output_spec)
+
         self.observation_spec["obs"] = self.observation_spec["observation"]
         del self.observation_spec["observation"]
 
@@ -78,7 +81,13 @@ def _make_specs(self, env: "gym.Env") -> None:  # noqa: F821
         obs_spec.unlock_()
         obs_spec.update(specs)
         obs_spec.lock_()
-        self.__dict__["full_observation_spec"] = obs_spec
+
+    def _output_transform(self, output):
+        obs, reward, done, info = output
+        return obs, reward, done ^ done, done, done, info
+
+    def _reset_output_transform(self, reset_data):
+        return reset_data, {}
 
     @classmethod
     def _make_envs(cls, *, task, num_envs, device, seed=None, headless=True, **kwargs):
@@ -125,15 +134,8 @@ def read_done(
             done = done.bool()
         return terminated, truncated, done, done.any()
 
-    def read_reward(self, total_reward, step_reward):
-        """Reads a reward and the total reward so far (in the frame skip loop) and returns a sum of the two.
-
-        Args:
-            total_reward (torch.Tensor or TensorDict): total reward so far in the step
-            step_reward (reward in the format provided by the inner env): reward of this particular step
-
-        """
-        return total_reward + step_reward
+    def read_reward(self, total_reward):
+        return total_reward
 
     def read_obs(
         self, observations: Union[Dict[str, Any], torch.Tensor, np.ndarray]

From b4ca1b6a0ddce40ae2bfad1446c318bb785464a7 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Tue, 9 Apr 2024 21:42:06 +0200
Subject: [PATCH 02/14] amend

---
 test/test_libs.py             | 31 +++++++++++++++++--------------
 torchrl/envs/libs/gym.py      |  3 ---
 torchrl/envs/libs/isaacgym.py | 16 +++++++++-------
 3 files changed, 26 insertions(+), 24 deletions(-)

diff --git a/test/test_libs.py b/test/test_libs.py
index c5543fe689e..f8617187d3c 100644
--- a/test/test_libs.py
+++ b/test/test_libs.py
@@ -3,6 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 import importlib.util
+
 _has_isaac = importlib.util.find_spec("isaacgym") is not None
 
 if _has_isaac:
@@ -10,22 +11,13 @@
     import isaacgym  # noqa
     import isaacgymenvs  # noqa
     from torchrl.envs.libs.isaacgym import IsaacGymEnv
-import os
-from contextlib import nullcontext
-from pathlib import Path
-
-from torchrl._utils import logger as torchrl_logger
-
-from torchrl.data.datasets.gen_dgrl import GenDGRLExperienceReplay
-
-from torchrl.envs.transforms import ActionMask, TransformedEnv
-from torchrl.modules import MaskedCategorical
-
-
 import argparse
 import importlib
+import os
 
 import time
+from contextlib import nullcontext
+from pathlib import Path
 from sys import platform
 from typing import Optional, Union
 
@@ -52,7 +44,8 @@
     TensorDictSequential,
 )
 from torch import nn
-from torchrl._utils import implement_for
+
+from torchrl._utils import implement_for, logger as torchrl_logger
 from torchrl.collectors.collectors import SyncDataCollector
 from torchrl.data import (
     BinaryDiscreteTensorSpec,
@@ -69,6 +62,8 @@
 )
 from torchrl.data.datasets.atari_dqn import AtariDQNExperienceReplay
 from torchrl.data.datasets.d4rl import D4RLExperienceReplay
+
+from torchrl.data.datasets.gen_dgrl import GenDGRLExperienceReplay
 from torchrl.data.datasets.minari_data import MinariExperienceReplay
 from torchrl.data.datasets.openml import OpenMLExperienceReplay
 from torchrl.data.datasets.openx import OpenXExperienceReplay
@@ -108,13 +103,21 @@
 from torchrl.envs.libs.robohive import _has_robohive, RoboHiveEnv
 from torchrl.envs.libs.smacv2 import _has_smacv2, SMACv2Env
 from torchrl.envs.libs.vmas import _has_vmas, VmasEnv, VmasWrapper
+
+from torchrl.envs.transforms import ActionMask, TransformedEnv
 from torchrl.envs.utils import (
     check_env_specs,
     ExplorationType,
     MarlGroupMapType,
     RandomPolicy,
 )
-from torchrl.modules import ActorCriticOperator, MLP, SafeModule, ValueOperator
+from torchrl.modules import (
+    ActorCriticOperator,
+    MaskedCategorical,
+    MLP,
+    SafeModule,
+    ValueOperator,
+)
 
 _has_d4rl = importlib.util.find_spec("d4rl") is not None
 
diff --git a/torchrl/envs/libs/gym.py b/torchrl/envs/libs/gym.py
index dfbf7ae6217..50655004f52 100644
--- a/torchrl/envs/libs/gym.py
+++ b/torchrl/envs/libs/gym.py
@@ -956,7 +956,6 @@ def _make_specs(self, env: "gym.Env", batch_size=None) -> None:  # noqa: F821
             device=self.device,
             categorical_action_encoding=self._categorical_action_encoding,
         )
-        print("-1", observation_spec)
         if not isinstance(observation_spec, CompositeSpec):
             if self.from_pixels:
                 observation_spec = CompositeSpec(
@@ -982,13 +981,11 @@ def _make_specs(self, env: "gym.Env", batch_size=None) -> None:  # noqa: F821
                 device=self.device,
             )
         if batch_size is not None:
-            print("0", observation_spec)
             action_spec = action_spec.expand(*batch_size, *action_spec.shape)
             reward_spec = reward_spec.expand(*batch_size, *reward_spec.shape)
             observation_spec = observation_spec.expand(
                 *batch_size, *observation_spec.shape
             )
-            print("1", observation_spec)
 
         self.done_spec = self._make_done_spec()
         self.action_spec = action_spec
diff --git a/torchrl/envs/libs/isaacgym.py b/torchrl/envs/libs/isaacgym.py
index f8f5143c9dd..e130775bf5b 100644
--- a/torchrl/envs/libs/isaacgym.py
+++ b/torchrl/envs/libs/isaacgym.py
@@ -14,9 +14,10 @@
 import torch
 
 from tensordict import TensorDictBase
+from torchrl.data import CompositeSpec
 from torchrl.envs.libs.gym import GymWrapper
 from torchrl.envs.utils import _classproperty, make_composite_from_td
-from torchrl.data import CompositeSpec
+
 _has_isaac = importlib.util.find_spec("isaacgym") is not None
 
 
@@ -49,7 +50,6 @@ def __init__(
         warnings.warn(
             "IsaacGym environment support is an experimental feature that may change in the future."
         )
-        num_envs = env.num_envs
         super().__init__(
             env, torch.device(env.device), batch_size=torch.Size([]), **kwargs
         )
@@ -58,12 +58,14 @@ def __init__(
             self.task = env.__name__
 
     def _make_specs(self, env: "gym.Env") -> None:  # noqa: F821
-        print("self.batch_size", self.batch_size)
         super()._make_specs(env, batch_size=self.batch_size)
-        self.full_done_spec = CompositeSpec({
-            key: spec.squeeze(-1) for key, spec in self.full_done_spec.items(True, True)
-        }, shape=self.batch_size)
-        print(self.output_spec)
+        self.full_done_spec = CompositeSpec(
+            {
+                key: spec.squeeze(-1)
+                for key, spec in self.full_done_spec.items(True, True)
+            },
+            shape=self.batch_size,
+        )
 
         self.observation_spec["obs"] = self.observation_spec["observation"]
         del self.observation_spec["observation"]

From 96501797101dfdb536259933014da9e9e62394f2 Mon Sep 17 00:00:00 2001
From: vmoens <vmoens@meta.com>
Date: Wed, 10 Apr 2024 02:49:37 -0700
Subject: [PATCH 03/14] amend

---
 test/test_libs.py             | 12 +++++++-----
 torchrl/envs/libs/isaacgym.py | 10 ++++++++--
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/test/test_libs.py b/test/test_libs.py
index f8617187d3c..234cca7d25a 100644
--- a/test/test_libs.py
+++ b/test/test_libs.py
@@ -3077,22 +3077,24 @@ def test_data(self, dataset):
 )
 @pytest.mark.parametrize("num_envs", [10, 20])
 @pytest.mark.parametrize("device", get_default_devices())
+@pytest.mark.parametrize("from_pixels", [True, False])
 class TestIsaacGym:
     @classmethod
-    def _run_on_proc(cls, q, task, num_envs, device):
+    def _run_on_proc(cls, q, task, num_envs, device, from_pixels):
         try:
-            env = IsaacGymEnv(task=task, num_envs=num_envs, device=device)
+            env = IsaacGymEnv(task=task, num_envs=num_envs, device=device, from_pixels=from_pixels)
+            print(env.rollout(3))
             check_env_specs(env)
             q.put(("succeeded!", None))
         except Exception as err:
             q.put(("failed!", err))
             raise err
 
-    def test_env(self, task, num_envs, device):
+    def test_env(self, task, num_envs, device, from_pixels):
         from torch import multiprocessing as mp
-
         q = mp.Queue(1)
-        proc = mp.Process(target=self._run_on_proc, args=(q, task, num_envs, device))
+        self._run_on_proc(q, task, num_envs, device, from_pixels)
+        proc = mp.Process(target=self._run_on_proc, args=(q, task, num_envs, device, from_pixels))
         try:
             proc.start()
             msg, error = q.get()
diff --git a/torchrl/envs/libs/isaacgym.py b/torchrl/envs/libs/isaacgym.py
index e130775bf5b..ad0d79d8d12 100644
--- a/torchrl/envs/libs/isaacgym.py
+++ b/torchrl/envs/libs/isaacgym.py
@@ -86,9 +86,14 @@ def _make_specs(self, env: "gym.Env") -> None:  # noqa: F821
 
     def _output_transform(self, output):
         obs, reward, done, info = output
+        if self.from_pixels:
+            obs["pixels"] = self._env.render(mode='rgb_array')
         return obs, reward, done ^ done, done, done, info
 
     def _reset_output_transform(self, reset_data):
+        reset_data.pop("reward", None)
+        if self.from_pixels:
+            reset_data["pixels"] = self._env.render(mode='rgb_array')
         return reset_data, {}
 
     @classmethod
@@ -187,6 +192,7 @@ def __init__(self, task=None, *, env=None, num_envs, device, **kwargs):
             raise RuntimeError("Cannot provide both `task` and `env` arguments.")
         elif env is not None:
             task = env
-        envs = self._make_envs(task=task, num_envs=num_envs, device=device, **kwargs)
+        from_pixels = kwargs.pop("from_pixels", False)
+        envs = self._make_envs(task=task, num_envs=num_envs, device=device, virtual_screen_capture=True, **kwargs)
         self.task = task
-        super().__init__(envs, **kwargs)
+        super().__init__(envs, from_pixels=from_pixels, **kwargs)

From f2f8f3787d62ba9294914248f23557799a8405f4 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 24 Apr 2024 09:23:01 +0100
Subject: [PATCH 04/14] amend

---
 .github/workflows/test-linux-libs.yml | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/.github/workflows/test-linux-libs.yml b/.github/workflows/test-linux-libs.yml
index cc0e4a4f54e..17a37eb65b8 100644
--- a/.github/workflows/test-linux-libs.yml
+++ b/.github/workflows/test-linux-libs.yml
@@ -219,6 +219,22 @@ jobs:
         ./.github/unittest/linux_libs/scripts_gym/batch_scripts.sh
         ./.github/unittest/linux_libs/scripts_gym/post_process.sh
 
+  unittests-isaac:
+    strategy:
+      matrix:
+        python_version: ["3.9"]
+        cuda_arch_version: ["12.1"]
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    with:
+      docker-image: "nvcr.io/nvidia/isaac-sim:2023.1.1"
+      repository: pytorch/rl
+      runner: "linux.g5.4xlarge.nvidia.gpu"
+      gpu-arch-type: cuda
+      gpu-arch-version: "11.7"
+      timeout: 120
+      script: |
+        echo "Hello world"
+
   unittests-jumanji:
     strategy:
       matrix:

From 97481e9afa5ca46214bbf7564cbbfca12ccfb302 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 24 Apr 2024 09:24:56 +0100
Subject: [PATCH 05/14] amend

---
 .github/workflows/test-linux-libs.yml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/test-linux-libs.yml b/.github/workflows/test-linux-libs.yml
index 07d94231c89..ddb888d836b 100644
--- a/.github/workflows/test-linux-libs.yml
+++ b/.github/workflows/test-linux-libs.yml
@@ -226,9 +226,10 @@ jobs:
       matrix:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
+    if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
     with:
-      docker-image: "nvcr.io/nvidia/isaac-sim:2023.1.1"
+      docker-image: "nvidia/isaac-sim:2023.1.1"
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
       gpu-arch-type: cuda

From e0c0a57814f22964c15229ee0991c854f1e6c902 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 24 Apr 2024 10:52:58 +0100
Subject: [PATCH 06/14] amend

---
 .github/workflows/test-linux-libs.yml |  2 +-
 test/test_libs.py                     | 10 +++++++---
 torchrl/envs/libs/isaacgym.py         | 12 +++++++++---
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/test-linux-libs.yml b/.github/workflows/test-linux-libs.yml
index ddb888d836b..9850dff4876 100644
--- a/.github/workflows/test-linux-libs.yml
+++ b/.github/workflows/test-linux-libs.yml
@@ -227,7 +227,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
-    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    uses: vmoens/test-infra/.github/workflows/linux_job.yml@secret-api-docker
     with:
       docker-image: "nvidia/isaac-sim:2023.1.1"
       repository: pytorch/rl
diff --git a/test/test_libs.py b/test/test_libs.py
index b97aef47fe3..57f7c77a6dc 100644
--- a/test/test_libs.py
+++ b/test/test_libs.py
@@ -3088,8 +3088,9 @@ class TestIsaacGym:
     @classmethod
     def _run_on_proc(cls, q, task, num_envs, device, from_pixels):
         try:
-            env = IsaacGymEnv(task=task, num_envs=num_envs, device=device, from_pixels=from_pixels)
-            print(env.rollout(3))
+            env = IsaacGymEnv(
+                task=task, num_envs=num_envs, device=device, from_pixels=from_pixels
+            )
             check_env_specs(env)
             q.put(("succeeded!", None))
         except Exception as err:
@@ -3098,9 +3099,12 @@ def _run_on_proc(cls, q, task, num_envs, device, from_pixels):
 
     def test_env(self, task, num_envs, device, from_pixels):
         from torch import multiprocessing as mp
+
         q = mp.Queue(1)
         self._run_on_proc(q, task, num_envs, device, from_pixels)
-        proc = mp.Process(target=self._run_on_proc, args=(q, task, num_envs, device, from_pixels))
+        proc = mp.Process(
+            target=self._run_on_proc, args=(q, task, num_envs, device, from_pixels)
+        )
         try:
             proc.start()
             msg, error = q.get()
diff --git a/torchrl/envs/libs/isaacgym.py b/torchrl/envs/libs/isaacgym.py
index ad0d79d8d12..b4bf594e022 100644
--- a/torchrl/envs/libs/isaacgym.py
+++ b/torchrl/envs/libs/isaacgym.py
@@ -87,13 +87,13 @@ def _make_specs(self, env: "gym.Env") -> None:  # noqa: F821
     def _output_transform(self, output):
         obs, reward, done, info = output
         if self.from_pixels:
-            obs["pixels"] = self._env.render(mode='rgb_array')
+            obs["pixels"] = self._env.render(mode="rgb_array")
         return obs, reward, done ^ done, done, done, info
 
     def _reset_output_transform(self, reset_data):
         reset_data.pop("reward", None)
         if self.from_pixels:
-            reset_data["pixels"] = self._env.render(mode='rgb_array')
+            reset_data["pixels"] = self._env.render(mode="rgb_array")
         return reset_data, {}
 
     @classmethod
@@ -193,6 +193,12 @@ def __init__(self, task=None, *, env=None, num_envs, device, **kwargs):
         elif env is not None:
             task = env
         from_pixels = kwargs.pop("from_pixels", False)
-        envs = self._make_envs(task=task, num_envs=num_envs, device=device, virtual_screen_capture=True, **kwargs)
+        envs = self._make_envs(
+            task=task,
+            num_envs=num_envs,
+            device=device,
+            virtual_screen_capture=True,
+            **kwargs,
+        )
         self.task = task
         super().__init__(envs, from_pixels=from_pixels, **kwargs)

From 8e3563f3cd47917a484ed38e2bfeacbf48c91a6a Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 24 Apr 2024 10:54:33 +0100
Subject: [PATCH 07/14] amend

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

diff --git a/.github/workflows/test-linux-libs.yml b/.github/workflows/test-linux-libs.yml
index 9850dff4876..0b73b1bd74e 100644
--- a/.github/workflows/test-linux-libs.yml
+++ b/.github/workflows/test-linux-libs.yml
@@ -227,7 +227,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job.yml@secret-api-docker
+    uses: vmoens/test-infra/.github/workflows/linux_job.yml@main
     with:
       docker-image: "nvidia/isaac-sim:2023.1.1"
       repository: pytorch/rl

From bb93bd783e708b58ea394170aa0d51e90462e290 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Thu, 25 Apr 2024 14:28:38 +0100
Subject: [PATCH 08/14] amend

---
 .github/workflows/test-linux-libs.yml | 17 -----------------
 1 file changed, 17 deletions(-)

diff --git a/.github/workflows/test-linux-libs.yml b/.github/workflows/test-linux-libs.yml
index 0b73b1bd74e..9e1875cac18 100644
--- a/.github/workflows/test-linux-libs.yml
+++ b/.github/workflows/test-linux-libs.yml
@@ -221,23 +221,6 @@ jobs:
         ./.github/unittest/linux_libs/scripts_gym/batch_scripts.sh
         ./.github/unittest/linux_libs/scripts_gym/post_process.sh
 
-  unittests-isaac:
-    strategy:
-      matrix:
-        python_version: ["3.9"]
-        cuda_arch_version: ["12.1"]
-    if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job.yml@main
-    with:
-      docker-image: "nvidia/isaac-sim:2023.1.1"
-      repository: pytorch/rl
-      runner: "linux.g5.4xlarge.nvidia.gpu"
-      gpu-arch-type: cuda
-      gpu-arch-version: "11.7"
-      timeout: 120
-      script: |
-        echo "Hello world"
-
   unittests-jumanji:
     strategy:
       matrix:

From 70c3cc708735273812f95f48338415ef1b21bbf0 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Thu, 25 Apr 2024 15:30:38 +0100
Subject: [PATCH 09/14] amend

---
 torchrl/envs/libs/jumanji.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/torchrl/envs/libs/jumanji.py b/torchrl/envs/libs/jumanji.py
index 3c51acead4d..5e7866864cd 100644
--- a/torchrl/envs/libs/jumanji.py
+++ b/torchrl/envs/libs/jumanji.py
@@ -67,8 +67,8 @@ def _jumanji_to_torchrl_spec_transform(
             dtype = numpy_to_torch_dtype_dict[spec.dtype]
         return BoundedTensorSpec(
             shape=shape,
-            low=np.asarray(spec.low),
-            high=np.asarray(spec.high),
+            low=np.asarray(spec.minimum),
+            high=np.asarray(spec.maximum),
             dtype=dtype,
             device=device,
         )

From 2449bf3eccf78f2625d34eb5d0511d94bcd26f2f Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 8 May 2024 14:34:51 +0200
Subject: [PATCH 10/14] amend

---
 test/test_libs.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/test_libs.py b/test/test_libs.py
index 8dab5f6139f..67fe74f9f6e 100644
--- a/test/test_libs.py
+++ b/test/test_libs.py
@@ -3221,7 +3221,7 @@ def test_data(self, dataset):
 )
 @pytest.mark.parametrize("num_envs", [10, 20])
 @pytest.mark.parametrize("device", get_default_devices())
-@pytest.mark.parametrize("from_pixels", [True, False])
+@pytest.mark.parametrize("from_pixels", [False])
 class TestIsaacGym:
     @classmethod
     def _run_on_proc(cls, q, task, num_envs, device, from_pixels):

From 5354114e0f435e3d9818d5c48364ead5b39426f7 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 8 May 2024 14:39:22 +0200
Subject: [PATCH 11/14] amend

---
 torchrl/envs/libs/isaacgym.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchrl/envs/libs/isaacgym.py b/torchrl/envs/libs/isaacgym.py
index b4bf594e022..25725dbc0e4 100644
--- a/torchrl/envs/libs/isaacgym.py
+++ b/torchrl/envs/libs/isaacgym.py
@@ -100,7 +100,7 @@ def _reset_output_transform(self, reset_data):
     def _make_envs(cls, *, task, num_envs, device, seed=None, headless=True, **kwargs):
         import isaacgym  # noqa
         import isaacgymenvs  # noqa
-
+        _ = kwargs.pop("from_pixels", None)
         envs = isaacgymenvs.make(
             seed=seed,
             task=task,

From 56ba3467f7a8672d67399f1d6e62d1d544f93f23 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 8 May 2024 14:45:07 +0200
Subject: [PATCH 12/14] amend

---
 torchrl/envs/libs/isaacgym.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchrl/envs/libs/isaacgym.py b/torchrl/envs/libs/isaacgym.py
index 25725dbc0e4..2c28aad3f59 100644
--- a/torchrl/envs/libs/isaacgym.py
+++ b/torchrl/envs/libs/isaacgym.py
@@ -97,7 +97,7 @@ def _reset_output_transform(self, reset_data):
         return reset_data, {}
 
     @classmethod
-    def _make_envs(cls, *, task, num_envs, device, seed=None, headless=True, **kwargs):
+    def _make_envs(cls, *, task, num_envs, device, seed=None, headless=False, **kwargs):
         import isaacgym  # noqa
         import isaacgymenvs  # noqa
         _ = kwargs.pop("from_pixels", None)

From 5fde20a7b5388d07126547dab10d19e741368948 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 8 May 2024 14:49:03 +0200
Subject: [PATCH 13/14] amend

---
 torchrl/envs/libs/isaacgym.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchrl/envs/libs/isaacgym.py b/torchrl/envs/libs/isaacgym.py
index 2c28aad3f59..1a9a14299c3 100644
--- a/torchrl/envs/libs/isaacgym.py
+++ b/torchrl/envs/libs/isaacgym.py
@@ -197,7 +197,7 @@ def __init__(self, task=None, *, env=None, num_envs, device, **kwargs):
             task=task,
             num_envs=num_envs,
             device=device,
-            virtual_screen_capture=True,
+            virtual_screen_capture=False,
             **kwargs,
         )
         self.task = task

From fb9b27367f4fdf9118c9664c398e4a8eac213f23 Mon Sep 17 00:00:00 2001
From: Vincent Moens <vmoens@meta.com>
Date: Wed, 8 May 2024 14:51:38 +0200
Subject: [PATCH 14/14] amend

---
 torchrl/envs/libs/isaacgym.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/torchrl/envs/libs/isaacgym.py b/torchrl/envs/libs/isaacgym.py
index 1a9a14299c3..4c56bea304a 100644
--- a/torchrl/envs/libs/isaacgym.py
+++ b/torchrl/envs/libs/isaacgym.py
@@ -100,6 +100,7 @@ def _reset_output_transform(self, reset_data):
     def _make_envs(cls, *, task, num_envs, device, seed=None, headless=False, **kwargs):
         import isaacgym  # noqa
         import isaacgymenvs  # noqa
+
         _ = kwargs.pop("from_pixels", None)
         envs = isaacgymenvs.make(
             seed=seed,