From bec2b94e23dbd6a44684419b9a810ace2e73ac7c Mon Sep 17 00:00:00 2001 From: Stone Tao Date: Thu, 28 Nov 2024 11:32:21 -0800 Subject: [PATCH] [Feature] Substep Rendering (#724) * Update record.py * Update record.py --- mani_skill/utils/wrappers/record.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/mani_skill/utils/wrappers/record.py b/mani_skill/utils/wrappers/record.py index baa5ade30..709240a0a 100644 --- a/mani_skill/utils/wrappers/record.py +++ b/mani_skill/utils/wrappers/record.py @@ -206,6 +206,8 @@ class RecordEpisode(gym.Wrapper): record_reward: whether to record the reward in the trajectory data record_env_state: whether to record the environment state in the trajectory data video_fps (int): The FPS of the video to generate if save_video is True + render_substeps (bool): Whether to render substeps for video. This is captures an image of the environment after each physics step. This runs slower but generates more image frames + per environment step which when coupled with a higher video FPS can yield a smoother video. avoid_overwriting_video (bool): If true, the wrapper will iterate over possible video names to avoid overwriting existing videos in the output directory. Useful for resuming training runs. source_type (Optional[str]): a word to describe the source of the actions used to record episodes (e.g. RL, motionplanning, teleoperation) source_desc (Optional[str]): A longer description describing how the demonstrations are collected @@ -226,6 +228,7 @@ def __init__( record_reward: bool = True, record_env_state: bool = True, video_fps: int = 30, + render_substeps: bool = False, avoid_overwriting_video: bool = False, source_type: Optional[str] = None, source_desc: Optional[str] = None, @@ -302,6 +305,19 @@ def __init__( else: break + self.render_substeps = render_substeps + if self.render_substeps: + _original_after_simulation_step = self.base_env._after_simulation_step + + def wrapped_after_simulation_step(): + _original_after_simulation_step() + if self.save_video: + if self.base_env.gpu_sim_enabled: + self.base_env.scene._gpu_fetch_all() + self.render_images.append(self.capture_image()) + + self.base_env._after_simulation_step = wrapped_after_simulation_step + @property def num_envs(self): return self.base_env.num_envs