[Environment] Fix lib CI failures

ghstack-source-id: 0c77ddc
Pull-Request-resolved: #2923

Author: Vincent Moens

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

@@ -15,7 +15,7 @@ conda activate ./env
 $DIR/install.sh
 
 # Extracted from run_test.sh to run once.
-apt-get update && apt-get install -y git wget libglew-dev libx11-dev x11proto-dev g++ cmake
+apt-get update && apt-get install -y git wget libglew-dev libx11-dev x11proto-dev g++
 
 # solves "'extras_require' must be a dictionary"
 pip install setuptools==65.3.0

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

@@ -4,7 +4,7 @@ unset PYTORCH_VERSION
 # For unittest, nightly PyTorch is used as the following section,
 # so no need to set PYTORCH_VERSION.
 # In fact, keeping PYTORCH_VERSION forces us to hardcode PyTorch version in config.
-apt-get update && apt-get install -y git wget gcc g++ cmake
+apt-get update && apt-get install -y git wget gcc g++
 
 set -e
 set -v

.github/unittest/linux_libs/scripts_gym/setup_env.sh

@@ -9,7 +9,7 @@ set -e
 
 this_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 # Avoid error: "fatal: unsafe repository"
-apt-get update && apt-get install -y git wget gcc g++ cmake
+apt-get update && apt-get install -y git wget gcc g++
 
 git config --global --add safe.directory '*'
 root_dir="$(git rev-parse --show-toplevel)"
@@ -69,6 +69,8 @@ printf "* Installing dependencies (except PyTorch)\n"
 echo "  - python=${PYTHON_VERSION}" >> "${this_dir}/environment.yml"
 cat "${this_dir}/environment.yml"
 
+conda install anaconda::cmake -y
+
 export MUJOCO_GL=egl
 conda env config vars set \
   MAX_IDLE_COUNT=1000 \

sota-implementations/a2c/utils_atari.py

@@ -74,7 +74,6 @@ def make_parallel_env(env_name, num_envs, device, gym_backend, is_test=False):
             lambda: make_base_env(env_name, gym_backend=gym_backend, is_test=is_test),
         ),
         serial_for_single=True,
-        gym_backend=gym_backend,
         device=device,
     )
     env = TransformedEnv(env)

torchrl/data/datasets/atari_dqn.py

@@ -508,6 +508,12 @@ def _download_and_preproc(self):
             if not os.listdir(tempdir):
                 os.makedirs(tempdir, exist_ok=True)
                 # get the list of runs
+                try:
+                    subprocess.run(
+                        ["gsutil", "version"], check=True, capture_output=True
+                    )
+                except subprocess.CalledProcessError:
+                    raise RuntimeError("gsutil is not installed or not found in PATH.")
                 command = f"gsutil -m ls -R gs://atari-replay-datasets/dqn/{self.dataset_id}/replay_logs"
                 output = subprocess.run(
                     command, shell=True, capture_output=True
@@ -520,9 +526,7 @@ def _download_and_preproc(self):
                 self.remote_gz_files = self._list_runs(None, files)
                 remote_gz_files = list(self.remote_gz_files)
                 if not len(remote_gz_files):
-                    raise RuntimeError(
-                        "Could not load the file list. Did you install gsutil?"
-                    )
+                    raise RuntimeError("No files in file list.")
 
                 total_runs = remote_gz_files[-1]
                 if self.num_procs == 0:

torchrl/envs/libs/smacv2.py

@@ -228,7 +228,7 @@ def _make_specs(self, env: "smacv2.env.StarCraft2Env") -> None:  # noqa: F821
             dtype=torch.bool,
             device=self.device,
         )
-        self.action_spec = self._make_action_spec()
+        self.full_action_spec = self._make_action_spec()
         self.observation_spec = self._make_observation_spec()
 
     def _init_env(self) -> None:
@@ -356,7 +356,7 @@ def _reset(
     def _step(self, tensordict: TensorDictBase) -> TensorDictBase:
         # perform actions
         action = tensordict.get(("agents", "action"))
-        action_np = self.action_spec.to_numpy(action)
+        action_np = self.full_action_spec[self.action_key].to_numpy(action)
 
         # Actions are validated by the environment.
         try:

torchrl/envs/transforms/transforms.py

@@ -8798,29 +8798,37 @@ def __init__(
     def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
         raise RuntimeError(FORWARD_NOT_IMPLEMENTED.format(type(self)))
 
+    @property
+    def action_spec(self):
+        action_spec = self.container.full_action_spec
+        keys = self.container.action_keys
+        if len(keys) == 1:
+            action_spec = action_spec[keys[0]]
+        else:
+            raise ValueError(
+                f"Too many action keys for {self.__class__.__name__}: {keys=}"
+            )
+        if not isinstance(action_spec, self.ACCEPTED_SPECS):
+            raise ValueError(
+                self.SPEC_TYPE_ERROR.format(self.ACCEPTED_SPECS, type(action_spec))
+            )
+        return action_spec
+
     def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase:
         parent = self.parent
         if parent is None:
             raise RuntimeError(
                 f"{type(self)}.parent cannot be None: make sure this transform is executed within an environment."
             )
         mask = next_tensordict.get(self.in_keys[1])
-        action_spec = self.container.action_spec
-        if not isinstance(action_spec, self.ACCEPTED_SPECS):
-            raise ValueError(
-                self.SPEC_TYPE_ERROR.format(self.ACCEPTED_SPECS, type(action_spec))
-            )
+        action_spec = self.action_spec
         action_spec.update_mask(mask.to(action_spec.device))
         return next_tensordict
 
     def _reset(
         self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
     ) -> TensorDictBase:
-        action_spec = self.container.action_spec
-        if not isinstance(action_spec, self.ACCEPTED_SPECS):
-            raise ValueError(
-                self.SPEC_TYPE_ERROR.format(self.ACCEPTED_SPECS, type(action_spec))
-            )
+        action_spec = self.action_spec
         mask = tensordict.get(self.in_keys[1], None)
         if mask is not None:
             mask = mask.to(action_spec.device)