v0.1.2

robogauge/tasks/simulator/mujoco_simulator.py

@@ -0,0 +1,143 @@
+# -*- coding: utf-8 -*-
+'''
+@File    : mujoco_simulator.py
+@Time    : 2025/11/27 15:54:20
+@Author  : wty-yy
+@Version : 1.0
+@Blog    : https://wty-yy.github.io/
+@Desc    : None
+'''
+import mujoco
+import mujoco.viewer
+from dm_control import mjcf
+
+import time
+import imageio
+import numpy as np
+
+from robogauge.utils.logger import logger
+from robogauge.utils.helpers import pares_path
+from robogauge.tasks.simulator.mujoco_config import MujocoConfig
+
+class MujocoSimulator:
+    def __init__(self, sim_cfg: MujocoConfig):
+        self.cfg = sim_cfg
+        self.terrain_xml = None
+        self.robot_xml = None
+        self.terrain_spawn_xy = None
+        self.robot_spawn_height = None
+        self.viewer = None
+        self.renderer = None
+        self.vid_writer = None
+        self.vid_count = 0
+        self._pause = True
+    
+    def load(
+        self,
+        terrain_xml: str = None,
+        robot_xml: str = None,
+        terrain_spawn_xy: list = None,
+        robot_spawn_height: float = None,
+    ):
+        """ Load terrain and robot into the simulator, support re-loading. """
+        if terrain_xml is not None:
+            self.terrain_xml = pares_path(terrain_xml)
+        if robot_xml is not None:
+            self.robot_xml = pares_path(robot_xml)
+        if terrain_spawn_xy is not None:
+            self.terrain_spawn_xy = terrain_spawn_xy
+        if robot_spawn_height is not None:
+            self.robot_spawn_height = robot_spawn_height
+
+        terrain_xml = self.terrain_xml
+        robot_xml = self.robot_xml
+        terrain_spawn_xy = self.terrain_spawn_xy
+        robot_spawn_height = self.robot_spawn_height
+        if terrain_xml is None or robot_xml is None:
+            raise ValueError("Terrain and robot XML paths must be provided.")
+        
+        robot_mjcf = mjcf.from_path(robot_xml)
+        terrain_mjcf = mjcf.from_path(terrain_xml)
+        for j in robot_mjcf.find_all('joint'):
+            if j.tag == 'freejoint':
+                j.remove()
+        attachment_frame = terrain_mjcf.attach(robot_mjcf)
+        attachment_frame.add('freejoint')
+        attachment_frame.pos = [*terrain_spawn_xy, robot_spawn_height]
+
+        if self.viewer is not None:
+            self.close_viewer()
+        self.mj_physics = mjcf.Physics.from_mjcf_model(terrain_mjcf)
+        self.mj_model = self.mj_physics.model.ptr
+        self.mj_data = self.mj_physics.data.ptr
+        self.mj_model.opt.timestep = self.cfg.physics.simulation_dt
+        self.headless = self.cfg.viewer.headless
+        if self.cfg.render.save_video and self.headless:
+            logger.warning("Cannot save video in headless mode, disabling video saving.")
+            self.cfg.render.save_video = False
+        if not self.headless:
+            self.viewer = mujoco.viewer.launch_passive(self.mj_model, self.mj_data, key_callback=self.key_callback)
+            self.last_render_time = time.time()
+            if self.cfg.render.save_video:
+                self.renderer = mujoco.Renderer(self.mj_model, height=self.cfg.render.height, width=self.cfg.render.width)
+
+                vid_dir = logger.log_dir / "videos"
+                vid_dir.mkdir(parents=True, exist_ok=True)
+                vid_path = str(vid_dir / f"sim_video_{self.vid_count:03d}.mp4")
+                self.vid_writer = imageio.get_writer(
+                    vid_path,
+                    fps=int(1 / self.cfg.physics.simulation_dt),
+                )
+                logger.info(f"Saving simulation video to: {vid_path}")
+                self.vid_count += 1
+        self._pause = False
+    
+    def key_callback(self, keycode):
+        if keycode == 32:
+            self._pause = not self._pause
+            logger.info(f"Pause toggled: {self._pause}")
+    
+    def step(self) -> dict:
+        """ Simulation step, pause will block thread. """
+        while self._pause:
+            time.sleep(0.1)
+        self.mj_physics.step()
+        if self.viewer is not None:
+            if self.viewer.is_running():
+                time_untile_next_render = self.cfg.physics.simulation_dt - (time.time() - self.last_render_time)
+                if time_untile_next_render > 0:
+                    time.sleep(time_untile_next_render)
+                self.viewer.sync()
+                if self.vid_writer is not None:
+                    self.renderer.update_scene(self.mj_data, camera=self.viewer.cam)
+                    frame = self.renderer.render()
+                    self.vid_writer.append_data(frame)
+                self.last_render_time = time.time()
+            else:
+                logger.warning("Viewer closed by user, stop video recording.")
+                self.close_viewer()
+
+        info = {}
+        return info
+    
+    def reset(self):
+        """ Reset the simulator to initial state. """
+        self.mj_physics.reset()
+        if self.viewer is not None:
+            self.viewer.sync()
+
+    def apply_action(self, action: np.ndarray):
+        """ Apply action to the simulator. """
+        self.mj_data.ctrl[:] = action
+    
+    def close_viewer(self):
+        """ Close the viewer and video writer. """
+        if self.viewer is not None:
+            self.viewer.close()
+            self.viewer = None
+            logger.info("Closing viewer.")
+        if self.vid_writer is not None:
+            self.vid_writer.close()
+            self.vid_writer = None
+            self.renderer = None
+            logger.info("Closing video writer.")