Add video recording to the job

mimickit/engines/engine.py

@@ -188,9 +188,22 @@ def calc_obj_mass(self, env_id, obj_id):
     @abc.abstractmethod
     def get_control_mode(self):
         return
+
+    @abc.abstractmethod
+    def create_video_recorder(self, camera_config=None):
+        return
 
     def draw_lines(self, env_id, start_verts, end_verts, cols, line_width):
         return
 
     def register_keyboard_callback(self, key_str, callback_func):
+        return
+
+    def get_video_recorder(self):
+        return None
+
+    def pre_rollout_test(self):
+        return
+
+    def post_rollout_test(self):
         return

mimickit/engines/engine_builder.py

@@ -3,15 +3,15 @@
 except:
     pass
 
-def build_engine(config, num_envs, device, visualize):
+def build_engine(config, num_envs, device, visualize, record_video=False):
     eng_name = config["engine_name"]
 
     if (eng_name == "isaac_gym"):
         import engines.isaac_gym_engine as isaac_gym_engine
         engine = isaac_gym_engine.IsaacGymEngine(config, num_envs, device, visualize)
     elif (eng_name == "isaac_lab"):
         import engines.isaac_lab_engine as isaac_lab_engine
-        engine = isaac_lab_engine.IsaacLabEngine(config, num_envs, device, visualize)
+        engine = isaac_lab_engine.IsaacLabEngine(config, num_envs, device, visualize, record_video=record_video)
     elif (eng_name == "newton"):
         import engines.newton_engine as newton_engine
         engine = newton_engine.NewtonEngine(config, num_envs, device, visualize)

mimickit/engines/isaac_gym_engine.py

@@ -919,4 +919,7 @@ def _process_gui_events(self):
                 elif (evt.action in self._keyboard_callbacks):
                     callback = self._keyboard_callbacks[evt.action]
                     callback()
-        return
+        return
+
+    def create_video_recorder(self, camera_config=None):
+        raise NotImplementedError("Video recording not supported for Isaac Gym engine")

mimickit/engines/isaac_lab_engine.py

@@ -1,4 +1,5 @@
 from isaaclab.app import AppLauncher
+from util.video_recorder import VideoRecorder
 
 import carb
 
@@ -59,9 +60,10 @@ def is_valid_clone(self, other):
 
 
 class IsaacLabEngine(engine.Engine):
-    def __init__(self, config, num_envs, device, visualize):
+    def __init__(self, config, num_envs, device, visualize, record_video=False):
         super().__init__()
 
+        self.video_recorder = None
         self._device = device
         sim_freq = config.get("sim_freq", 60)
         control_freq = config.get("control_freq", 10)
@@ -72,7 +74,7 @@ def __init__(self, config, num_envs, device, visualize):
         self._sim_steps = int(sim_freq / control_freq)
         sim_timestep = 1.0 / sim_freq
 
-        self._create_simulator(sim_timestep, visualize)
+        self._create_simulator(sim_timestep, visualize, record_video)
 
         self._env_spacing = config["env_spacing"]
         self._obj_cfgs = []
@@ -86,15 +88,35 @@ def __init__(self, config, num_envs, device, visualize):
         self._build_ground()
         self._env_offsets = self._compute_env_offsets(num_envs)
 
+        if (visualize or record_video):
+            self._build_camera()
+            self._build_lights()
+
+        # Video recorder will be created after environment is initialized
+        # so it can query environment for camera config if needed
+        self._record_video = record_video
+
         if (visualize):
             self._prev_frame_time = 0.0
-            self._build_lights()
-            self._build_camera()
             self._build_draw_interface()
             self._setup_keyboard()
 
         return
+
+    def create_video_recorder(self, camera_config={}):
+        """Create the video recorder with optional camera configuration.
+        
+        Args:
+            camera_config: Optional camera config dict. If None, uses defaults.
+        """
+        self.video_recorder = VideoRecorder(self, camera_config)
+
+        Logger.print("Video recording enabled")
+        return
 
+    def get_video_recorder(self):
+        return self.video_recorder
+
     def get_name(self):
         return "isaac_lab"
 
@@ -125,6 +147,7 @@ def initialize_sim(self):
 
     def step(self):
         self._update_reset_objs()
+        self.get_video_recorder().capture_frame()
 
         for i in range(self._sim_steps):
             self._pre_sim_step()
@@ -652,8 +675,17 @@ def _setup_keyboard(self):
         self._keyboard_callbacks = dict()
         return
 
-    def _create_simulator(self, sim_timestep, visualize):
-        self._app_launcher = AppLauncher({"headless": not visualize, "device": self._device})
+    def _create_simulator(self, sim_timestep, visualize, record_video=False):
+        # Headless rendering (record_video without a display) requires a virtual display
+        if record_video and not visualize:
+            from util.display import ensure_virtual_display
+            ensure_virtual_display()
+
+        self._app_launcher = AppLauncher({
+            "headless": not visualize,
+            "device": self._device,
+            "enable_cameras": record_video or visualize,
+        })
 
         import isaaclab.sim as sim_utils
         from isaacsim.core.utils.stage import get_current_stage
@@ -727,7 +759,7 @@ def _post_sim_step(self):
 
     def _clear_forces(self):
         for obj in self._objs:
-            if (obj.has_external_wrench):
+            if (getattr(obj, 'has_external_wrentch', False)):
                 forces = torch.zeros([1, 3], dtype=torch.float, device=self._device)
                 torques = torch.zeros([1, 3], dtype=torch.float, device=self._device)
                 obj.set_external_force_and_torque(forces=forces, torques=torques,
@@ -1158,3 +1190,13 @@ def _on_keyboard_event(self, event):
                 callback = self._keyboard_callbacks[event.input]
                 callback()
         return
+
+    def pre_rollout_test(self):
+        if self._record_video:
+            self.video_recorder.start_recording()
+        return
+
+    def post_rollout_test(self):
+        if self._record_video:
+            self.video_recorder.stop_recording()
+        return

mimickit/engines/newton_engine.py

@@ -1075,6 +1075,9 @@ def _on_keyboard_event(self, symbol, modifiers):
             callback()
         return
 
+    def create_video_recorder(self, camera_config=None):
+        raise NotImplementedError("Video recording not supported for Newton engine")
+
 
 @wp.kernel
 def clamp_arrays(x: wp.array(dtype=float),

mimickit/envs/add_env.py

@@ -5,9 +5,9 @@
 import util.torch_util as torch_util
 
 class ADDEnv(amp_env.AMPEnv):
-    def __init__(self, env_config, engine_config, num_envs, device, visualize):
+    def __init__(self, env_config, engine_config, num_envs, device, visualize, record_video=False):
         super().__init__(env_config=env_config, engine_config=engine_config, 
-                         num_envs=num_envs, device=device, visualize=visualize)
+                         num_envs=num_envs, device=device, visualize=visualize, record_video=record_video)
         return
 
     def _build_disc_obs_buffers(self):

mimickit/envs/amp_env.py

@@ -7,11 +7,11 @@
 import util.torch_util as torch_util
 
 class AMPEnv(deepmimic_env.DeepMimicEnv):
-    def __init__(self, env_config, engine_config, num_envs, device, visualize):
+    def __init__(self, env_config, engine_config, num_envs, device, visualize, record_video=False):
         self._num_disc_obs_steps = env_config["num_disc_obs_steps"]
 
         super().__init__(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device,
-                         visualize=visualize)
+                         visualize=visualize, record_video=record_video)
         return
 
     def get_disc_obs_space(self):

mimickit/envs/ase_env.py

@@ -4,11 +4,11 @@
 import envs.char_env as char_env
 
 class ASEEnv(amp_env.AMPEnv):
-    def __init__(self, env_config, engine_config, num_envs, device, visualize):
+    def __init__(self, env_config, engine_config, num_envs, device, visualize, record_video=False):
         self._default_reset_prob = env_config["default_reset_prob"]
 
         super().__init__(env_config=env_config, engine_config=engine_config, 
-                         num_envs=num_envs, device=device, visualize=visualize)
+                         num_envs=num_envs, device=device, visualize=visualize, record_video=record_video)
         return
 
     def _reset_char(self, env_ids):

mimickit/envs/base_env.py

@@ -66,4 +66,14 @@ def get_env_time(self, env_ids=None):
         return 0.0
 
     def record_diagnostics(self):
-        return self._diagnostics
+        return self._diagnostics
+
+    def get_video_camera_config(self) -> dict:
+        """Optional method for environments to provide camera config for video recording.
+        
+        Returns:
+            dict with optional keys:
+                - cam_pos: np.ndarray, camera position [x, y, z]
+                - cam_target: np.ndarray, camera target [x, y, z]
+        """
+        return {}

mimickit/envs/char_dof_test_env.py

@@ -6,11 +6,11 @@
 import engines.engine as engine
 
 class CharDofTestEnv(char_env.CharEnv):
-    def __init__(self, env_config, engine_config, num_envs, device, visualize):
+    def __init__(self, env_config, engine_config, num_envs, device, visualize, record_video=False):
         self._time_per_dof = 4.0
 
         super().__init__(env_config=env_config, engine_config=engine_config,
-                         num_envs=num_envs, device=device, visualize=visualize)
+                         num_envs=num_envs, device=device, visualize=visualize, record_video=record_video)
 
         self._episode_length = self._time_per_dof * self._pd_low.shape[0]
         return

mimickit/envs/char_env.py

@@ -15,13 +15,13 @@
 import engines.engine as engine
 
 class CharEnv(sim_env.SimEnv):
-    def __init__(self, env_config, engine_config, num_envs, device, visualize):
+    def __init__(self, env_config, engine_config, num_envs, device, visualize, record_video=False):
         self._global_obs = env_config["global_obs"]
         self._root_height_obs = env_config.get("root_height_obs", True)
         self._zero_center_action = env_config.get("zero_center_action", False)
 
         super().__init__(env_config=env_config, engine_config=engine_config,
-                         num_envs=num_envs, device=device, visualize=visualize)
+                         num_envs=num_envs, device=device, visualize=visualize, record_video=record_video)
 
         char_id = self._get_char_id()
         self._print_char_prop(0, char_id)
@@ -344,6 +344,21 @@ def _build_body_ids_tensor(self, body_names):
     def _has_key_bodies(self):
         return len(self._key_body_ids) > 0
 
+    def get_video_camera_config(self) -> dict:
+        """Provide camera config for video recording."""
+        env_id = 0
+        char_id = self._get_char_id()
+        char_root_pos = self._engine.get_root_pos(char_id)
+        char_pos = char_root_pos[env_id].cpu().numpy()
+
+        cam_pos = np.array([char_pos[0], char_pos[1] - 5.0, 3.0])
+        cam_target = np.array([char_pos[0], char_pos[1], 1.0])
+
+        return {
+            "cam_pos": cam_pos,
+            "cam_target": cam_target,
+        }
+
     def _build_camera(self, env_config):
         env_id = 0
         char_id = self._get_char_id()

mimickit/envs/deepmimic_env.py

@@ -10,7 +10,7 @@
 import util.torch_util as torch_util
 
 class DeepMimicEnv(char_env.CharEnv):
-    def __init__(self, env_config, engine_config, num_envs, device, visualize):
+    def __init__(self, env_config, engine_config, num_envs, device, visualize, record_video=False):
         self._enable_early_termination = env_config["enable_early_termination"]
         self._num_phase_encoding = env_config.get("num_phase_encoding", 0)
 
@@ -40,7 +40,7 @@ def __init__(self, env_config, engine_config, num_envs, device, visualize):
         self._visualize_ref_char = env_config.get("visualize_ref_char", True)
 
         super().__init__(env_config=env_config, engine_config=engine_config,
-                         num_envs=num_envs, device=device, visualize=visualize)
+                         num_envs=num_envs, device=device, visualize=visualize, record_video=record_video)
         return
 
     def get_reward_succ(self):

mimickit/envs/env_builder.py

@@ -5,47 +5,54 @@
 
 from util.logger import Logger
 
-def build_env(env_file, engine_file, num_envs, device, visualize):
+def build_env(env_file, engine_file, num_envs, device, visualize, record_video=False):
     env_config, engine_config = load_configs(env_file, engine_file)
 
     env_name = env_config["env_name"]
     Logger.print("Building {} env".format(env_name))
 
+    # Build environment based on env_name
     if (env_name == "char"):
         import envs.char_env as char_env
-        env = char_env.CharEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize)
+        env = char_env.CharEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize, record_video=record_video)
     elif (env_name == "deepmimic"):
         import envs.deepmimic_env as deepmimic_env
-        env = deepmimic_env.DeepMimicEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize)
+        env = deepmimic_env.DeepMimicEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize, record_video=record_video)
     elif (env_name == "amp"):
         import envs.amp_env as amp_env
-        env = amp_env.AMPEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize)
+        env = amp_env.AMPEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize, record_video=record_video)
     elif (env_name == "ase"):
         import envs.ase_env as ase_env
-        env = ase_env.ASEEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize)
+        env = ase_env.ASEEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize, record_video=record_video)
     elif (env_name == "add"):
         import envs.add_env as add_env
-        env = add_env.ADDEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize)
+        env = add_env.ADDEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize, record_video=record_video)
     elif (env_name == "char_dof_test"):
         import envs.char_dof_test_env as char_dof_test_env
-        env = char_dof_test_env.CharDofTestEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize)
+        env = char_dof_test_env.CharDofTestEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize, record_video=record_video)
     elif (env_name == "view_motion"):
         import envs.view_motion_env as view_motion_env
-        env = view_motion_env.ViewMotionEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize)
+        env = view_motion_env.ViewMotionEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize, record_video=record_video)
     elif (env_name == "task_location"):
         import envs.task_location_env as task_location_env
-        env = task_location_env.TaskLocationEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize)
+        env = task_location_env.TaskLocationEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize, record_video=record_video)
     elif (env_name == "task_steering"):
         import envs.task_steering_env as task_steering_env
-        env = task_steering_env.TaskSteeringEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize)
+        env = task_steering_env.TaskSteeringEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize, record_video=record_video)
     elif (env_name == "static_objects"):
         import envs.static_objects_env as static_objects_env
-        env = static_objects_env.StaticObjectsEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize)
+        env = static_objects_env.StaticObjectsEnv(env_config=env_config, engine_config=engine_config, num_envs=num_envs, device=device, visualize=visualize, record_video=record_video)
     else:
         assert(False), "Unsupported env: {}".format(env_name)
 
     return env
 
+def get_engine_name(engine_file):
+    engine_config = load_config(engine_file)
+    if engine_config is not None:
+        return engine_config.get("engine_name", "")
+    return ""
+
 def load_config(file):
     if (file is not None and file != ""):
         with open(file, "r") as stream: