Fix gymnasium vectorized environments reset returns when called twice

Toni-SM · Aug 3, 2024 · ebb2095 · ebb2095
1 parent 7e4d1cd
commit ebb2095
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/skrl/envs/wrappers/torch/gymnasium_envs.py b/skrl/envs/wrappers/torch/gymnasium_envs.py
@@ -134,13 +134,14 @@ def reset(self) -> Tuple[torch.Tensor, Any]:
         :return: Observation, info
         :rtype: torch.Tensor and any other info
         """
-        # handle vectorized envs
+        # handle vectorized environments (vector environments are autoreset)
         if self._vectorized:
-            if not self._reset_once:
-                return self._observation, self._info
-            self._reset_once = False
+            if self._reset_once:
+                observation, self._info = self._env.reset()
+                self._observation = self._observation_to_tensor(observation)
+                self._reset_once = False
+            return self._observation, self._info
 
-        # reset the env/envs
         observation, info = self._env.reset()
         return self._observation_to_tensor(observation), info