[Minor] Remove ya gymnasium deprecation warning in vectorized envs (#…

…1573)
pytorch · Sep 27, 2023 · 7f42576 · 7f42576
1 parent 434fe58
commit 7f42576
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/torchrl/envs/libs/gym.py b/torchrl/envs/libs/gym.py
@@ -465,13 +465,22 @@ def _is_batched(self):
             self._env, tuple_of_classes + (gym_backend("vector").VectorEnv,)
         )
 
+    @implement_for("gym", None, "0.27")
     def _get_batch_size(self, env):
         if hasattr(env, "num_envs"):
             batch_size = torch.Size([env.num_envs, *self.batch_size])
         else:
             batch_size = self.batch_size
         return batch_size
 
+    @implement_for("gymnasium", "0.27", None)  # gymnasium wants the unwrapped env
+    def _get_batch_size(self, env):  # noqa: F811
+        if hasattr(env, "num_envs"):
+            batch_size = torch.Size([env.unwrapped.num_envs, *self.batch_size])
+        else:
+            batch_size = self.batch_size
+        return batch_size
+
     def _check_kwargs(self, kwargs: Dict):
         if "env" not in kwargs:
             raise TypeError("Could not find environment key 'env' in kwargs.")