diff --git a/mani_skill2/utils/wrappers/record.py b/mani_skill2/utils/wrappers/record.py index ba284acbf..e19ff181b 100644 --- a/mani_skill2/utils/wrappers/record.py +++ b/mani_skill2/utils/wrappers/record.py @@ -185,7 +185,9 @@ class RecordEpisode(gym.Wrapper): save_trajectory: whether to save trajectory trajectory_name: name of trajectory file (.h5). Use timestamp if not provided. save_video: whether to save video - info_on_video: whether to write data about reward and data in the info object to the video + info_on_video: whether to write data about reward, action, and data in the info object to the video. The first video frame is generally the result + of the first env.reset() (visualizing the first observation). Text is written on frames after that, showing the action taken to get to that + environment state and reward. save_on_reset: whether to save the previous trajectory (and video of it if `save_video` is True) automatically when resetting. Not that for environments simulated on the GPU (to leverage fast parallel rendering) you must set `max_steps_per_video` to a fixed number so that every `max_steps_per_video` steps a video is saved. This is @@ -385,12 +387,14 @@ def recursive_replace(x, y): ) if self._trajectory_buffer.fail is not None: recursive_replace(self._trajectory_buffer.fail, first_step.fail) - if self.save_video: - self._render_images.append(self.capture_image()) return obs, info def step(self, action): + if self.save_video and self._video_steps == 0: + # save the first frame of the video here (s_0) instead of inside reset as user + # may call env.reset(...) multiple times but we want to ignore empty trajectories + self._render_images.append(self.capture_image()) obs, rew, terminated, truncated, info = super().step(action) if self.save_trajectory: diff --git a/manualtest/record_test.py b/manualtest/record_test.py index d5561bb9a..fc7d8c28b 100644 --- a/manualtest/record_test.py +++ b/manualtest/record_test.py @@ -10,7 +10,7 @@ if __name__ == "__main__": # sapien.set_log_level("info") # , "StackCube-v1", "PickCube-v1", "PushCube-v1", "PickSingleYCB-v1", "OpenCabinet-v1" - num_envs = 1 + num_envs = 2 for env_id in ["PickCube-v1"]: env = gym.make( env_id, @@ -39,7 +39,11 @@ ) # env = ManiSkillVectorEnv(env) - env.reset(seed=52, options=dict(reconfigure=True)) + # env.reset(seed=52, options=dict(reconfigure=True)) + env.reset() + env.reset() + env.reset() + env.reset() # for i in range(180): # env.step(env.action_space.sample()) # env.step(env.action_space.sample())