[Environment] Fix lib CI failures

ghstack-source-id: 13a9f82b2d8fdbb1ec724743bbdd809f52018923
Pull-Request-resolved: #2923

Author: vmoens

sota-implementations/impala/utils.py

@@ -143,9 +143,7 @@ def make_ppo_modules_pixels(proof_environment):
 
 def make_ppo_models(env_name, gym_backend):
 
-    proof_environment = make_env(
-        env_name, device="cpu", gym_backend=gym_backend
-    )
+    proof_environment = make_env(env_name, device="cpu", gym_backend=gym_backend)
     common_module, policy_module, value_module = make_ppo_modules_pixels(
         proof_environment
     )

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: