[Feature] Substep Rendering (#724)

* Update record.py

* Update record.py

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