Update record.py

haosulab · Mar 6, 2024 · 16e5a07 · 16e5a07
1 parent faf62dc
commit 16e5a07
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/mani_skill2/utils/wrappers/record.py b/mani_skill2/utils/wrappers/record.py
@@ -299,9 +299,6 @@ def reset(
             else:
                 self.flush_trajectory(env_idxs_to_flush=to_numpy(options["env_idx"]))
 
-        self.last_reset_kwargs = copy.deepcopy(
-            dict(seed=seed, options=options, **kwargs)
-        )
         obs, info = super().reset(*args, seed=seed, options=options, **kwargs)
 
         if self.save_trajectory:
@@ -365,7 +362,11 @@ def recursive_replace(x, y):
                     )
                 if self._trajectory_buffer.fail is not None:
                     recursive_replace(self._trajectory_buffer.fail, first_step.fail)
-
+        if "env_idx" in options:
+            options["env_idx"] = to_numpy(options["env_idx"])
+        self.last_reset_kwargs = copy.deepcopy(
+            dict(seed=seed, options=options, **kwargs)
+        )
         return obs, info
 
     def step(self, action):