# -*- 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 : Mujoco Simulator for Robogauge ''' import mujoco import mujoco.viewer from dm_control import mjcf import re import time import imageio import numpy as np from pathlib import Path from typing import Literal, List, Optional, Union from robogauge.utils.logger import logger from robogauge.utils.helpers import parse_path from robogauge.utils.math_utils import get_projected_gravity, quat_rotate_inverse from robogauge.tasks.simulator.mujoco_config import MujocoConfig from robogauge.tasks.simulator.sim_data import ( SimData, RobotProprioception, JointState, BaseState, IMUState ) from robogauge.tasks.gauge.goal_data import VelocityGoal class MujocoSimulator: def __init__(self, sim_cfg: MujocoConfig): self.cfg = sim_cfg self.terrain_xmls = None self.robot_xml = None self.terrain_spawn_pos = None self.robot_spawn_height = None self.default_dof_pos = None self.invert_yaw = None self.viewer = None self.offscreen_cam = mujoco.MjvCamera() self.renderer = None self.vid_writer = None self.vid_count = 0 self._pause = True self.n_step = 0 self.sim_time = 0.0 self.target_pos = None self.target_velocity: Optional[VelocityGoal] = None self.penetration_reset_count = 0 def load( self, terrain_xmls: List[str] = None, robot_xml: str = None, terrain_spawn_pos: list = None, default_dof_pos: list = None, invert_yaw: bool = None, ): """ Load terrain and robot into the simulator, support re-loading. """ if terrain_xmls is not None: self.terrain_xmls = [parse_path(xml) for xml in terrain_xmls] if robot_xml is not None: self.robot_xml = parse_path(robot_xml) if terrain_spawn_pos is not None: self.terrain_spawn_pos = terrain_spawn_pos if default_dof_pos is not None: self.default_dof_pos = default_dof_pos if invert_yaw is not None: self.invert_yaw = invert_yaw terrain_xmls = self.terrain_xmls robot_xml = self.robot_xml terrain_spawn_pos = self.terrain_spawn_pos if terrain_xmls is None or robot_xml is None: raise ValueError("Terrain and robot XML paths must be provided.") if default_dof_pos is None: raise ValueError("Default DOF positions must be provided.") # Create MJCF models robot_mjcf = mjcf.from_path(robot_xml) terrain_mjcf = mjcf.from_path(terrain_xmls[0]) visual_elem = terrain_mjcf.visual global_elem = visual_elem.get_children('global') global_elem.offwidth = 1920 global_elem.offheight = 1080 for path in terrain_xmls[1:]: next_terrain = mjcf.from_path(path) terrain_mjcf.attach(next_terrain) for j in robot_mjcf.find_all('joint'): if j.tag == 'freejoint': j.remove() robot_base = robot_mjcf.find('body', 'base_link') if robot_base is not None: origin_robot_height = robot_base.pos.copy() if robot_base.pos is not None else None robot_base.pos = [0, 0, 0] # move base_link translation to terrain_spawn_pos else: raise ValueError("Robot base_link body not found in the robot MJCF model.") attachment_frame = terrain_mjcf.attach(robot_mjcf) attachment_frame.add('freejoint', name='root') if origin_robot_height is not None: terrain_spawn_pos = np.array(terrain_spawn_pos) + origin_robot_height attachment_frame.pos = terrain_spawn_pos self.close_viewer() self.close_video_writer() 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.sim_dt = self.cfg.physics.simulation_dt if self.invert_yaw: self.mj_data.qpos[3] = 0.0 self.mj_data.qpos[6] = 1.0 self.mj_data.qpos[7:] = default_dof_pos # Domain randomization: base mass base_body_name = f'{robot_mjcf.model}/base_link' body_id = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, base_body_name) assert body_id != -1, f"Body '{base_body_name}' not found in the model." if self.cfg.domain_rand.base_mass != 0.0: original_mass = self.mj_model.body_mass[body_id] new_mass = max(0.01, original_mass + self.cfg.domain_rand.base_mass) self.mj_model.body_mass[body_id] = new_mass logger.info(f"Randomized base mass: {original_mass:.3f} -> {new_mass:.3f} kg") # Domain randomization: friction if self.cfg.domain_rand.friction != 1.0: for i in range(self.mj_model.ngeom): # Both change robot friction and terrain friction, usually robot friction < 1.0 # Mujoco friction calculation takes the *max* between two contacting geoms geom_friction = self.mj_model.geom_friction[i] geom_friction[0] *= self.cfg.domain_rand.friction self.mj_model.geom_friction[i] = geom_friction logger.info(f"Scaled geom friction by factor: {self.cfg.domain_rand.friction:.3f}") mujoco.mj_forward(self.mj_model, self.mj_data) # Setup offscreen camera self.offscreen_cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING self.offscreen_cam.trackbodyid = body_id self.offscreen_cam.distance = self.cfg.viewer.camera_distance self.offscreen_cam.elevation = self.cfg.viewer.camera_elevation self.offscreen_cam.azimuth = self.cfg.viewer.camera_azimuth self.offscreen_cam.lookat = np.array([0.0, 0.0, 0.0]) # Setup viewer self.headless = self.cfg.viewer.headless if not self.headless: self.viewer = mujoco.viewer.launch_passive( self.mj_model, self.mj_data, key_callback=self.key_callback ) # set viewer.camera to follow robot self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING self.viewer.cam.trackbodyid = body_id self.viewer.cam.distance = self.cfg.viewer.camera_distance self.viewer.cam.elevation = self.cfg.viewer.camera_elevation self.viewer.cam.azimuth = self.cfg.viewer.camera_azimuth self.last_render_time = time.time() # Setup video writer 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=self.cfg.render.video_fps, ) self.vid_frame_skip = int(1 / (self.cfg.render.video_fps * self.sim_dt)) logger.info(f"Simulation video saved at: {vid_path}") self.vid_count += 1 # Initialize simulation state self._pause = False self.n_step = 0 self.sim_time = 0.0 self.penetration_reset_count = 0 self.load_dof_limits() self.preload_sensors() # Robot controller placeholders self.action = None self.p_gains = None self.d_gains = None self.control_type = None def key_callback(self, keycode): if keycode == 32: self._pause = not self._pause logger.info(f"Pause toggled: {self._pause}") def step(self) -> SimData: """ Simulation step, pause will block thread. """ while self._pause: time.sleep(0.1) self.update_torque() self.mj_physics.step() # Viewer sync if self.viewer is not None: if self.viewer.is_running(): self.update_external_rendering(self.viewer, ctype='viewer') self.viewer.sync() 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.last_render_time = time.time() else: logger.warning("Viewer closed by user.") self.close_viewer() # Video recording if self.vid_writer is not None and self.n_step % self.vid_frame_skip == 0: render_cam = self.viewer.cam if self.viewer is not None else self.offscreen_cam # mujoco.mjv_updateCamera(render_cam) self.renderer.update_scene(self.mj_data, camera=render_cam) self.update_external_rendering(self.renderer, ctype='renderer') frame = self.renderer.render() self.vid_writer.append_data(frame) self.proprio = proprio = RobotProprioception( joint=JointState( pos=self.get_sensor_data('joint_pos'), vel=self.get_sensor_data('joint_vel'), torque=self.get_sensor_data('joint_eff'), limits=self.dof_limits, names=self.dof_names, ), imu=IMUState( pos=self.get_sensor_data('imu_pos'), quat=self.get_sensor_data('imu_quat'), acc=self.get_sensor_data('imu_acc'), lin_vel=self.get_sensor_data('imu_lin_vel'), # body frame, check direction, go2 is inverted ang_vel=self.get_sensor_data('imu_ang_vel'), # body frame, check direction, go2 is inverted ), base=BaseState( pos=self.mj_data.qpos[:3], # world frame quat=self.mj_data.qpos[3:7], # world frame lin_vel=quat_rotate_inverse(self.mj_data.qpos[3:7], self.mj_data.qvel[:3]), # body frame ang_vel=quat_rotate_inverse(self.mj_data.qpos[3:7], self.mj_data.qvel[3:6]), # body frame ) ) if self.n_step % int(0.1 / self.sim_dt) == 0: logger.log(value=np.mean(proprio.imu.quat - proprio.base.quat), tag="sim/delta_quat", step=self.n_step) logger.log(value=np.mean(proprio.imu.ang_vel - proprio.base.ang_vel), tag="sim/delta_ang_vel", step=self.n_step) logger.log(value=np.mean(proprio.imu.lin_vel - proprio.base.lin_vel), tag="sim/delta_lin_vel", step=self.n_step) logger.log(value=proprio.imu.lin_vel[0], tag="sim/imu_lin_vel_x", step=self.n_step) logger.log(value=proprio.imu.lin_vel[1], tag="sim/imu_lin_vel_y", step=self.n_step) logger.log(value=proprio.base.lin_vel[0], tag="sim/base_lin_vel_x", step=self.n_step) logger.log(value=proprio.base.lin_vel[1], tag="sim/base_lin_vel_y", step=self.n_step) if self.n_step == 0: self.debug_print_proprio_shapes() sim_data = SimData( n_step=self.n_step, sim_dt=self.sim_dt, sim_time=self.sim_time, proprio=proprio ) # input("DEBUG") self.n_step += 1 self.sim_time = self.n_step * self.sim_dt self.check_truncation(sim_data) return sim_data def update_external_rendering(self, handle: Union[mujoco.viewer.Handle, mujoco.Renderer], ctype: Literal['viewer', 'renderer'], ): """ Update external rendering handle (viewer or renderer). """ def add_target_sphere(geom_elem): mujoco.mjv_initGeom( geom_elem, type=mujoco.mjtGeom.mjGEOM_SPHERE, size=[0.1, 0, 0], pos=self.target_pos, mat=np.eye(3).flatten(), rgba=[1, 0, 0, 1] ) def add_thick_arrow(geom_elem, pos, vec, rgba, scale=0.7): vel_norm = np.linalg.norm(vec) display_norm = min(vel_norm * scale, 1.0) if display_norm < 0.10: mujoco.mjv_initGeom( geom_elem, type=mujoco.mjtGeom.mjGEOM_NONE, size=[0,0,0], pos=pos, mat=np.eye(3).flatten(), rgba=[0,0,0,0] ) return mat = np.zeros(9) target_quat = np.zeros(4) vec_normalized = vec / vel_norm mujoco.mju_quatZ2Vec(target_quat, vec_normalized) mujoco.mju_quat2Mat(mat, target_quat) mat = mat.reshape(3, 3) mat[:, 2] *= display_norm mujoco.mjv_initGeom( geom_elem, type=mujoco.mjtGeom.mjGEOM_ARROW, size=[0.02, 0.02, display_norm], # [height, width, length] pos=pos, mat=mat.flatten(), rgba=rgba ) viewer_geom_idx = 0 if ctype == 'viewer': handle.user_scn.ngeom = 0 # reset user scene geometry if self.target_pos is not None: if ctype == 'viewer': add_target_sphere(handle.user_scn.geoms[viewer_geom_idx]) viewer_geom_idx += 1 else: handle.scene.ngeom += 1 add_target_sphere(handle.scene.geoms[self.renderer.scene.ngeom - 1]) if self.target_velocity is not None: base_pos_world = self.mj_data.qpos[:3] base_quat = self.mj_data.qpos[3:7] # rendering arrows start position offset_body = np.array([0.0, 0.0, 0.2]) offset_world = np.zeros(3) mujoco.mju_rotVecQuat(offset_world, offset_body, base_quat) start_pos = base_pos_world + offset_world tgt_vel_body = np.array([self.target_velocity.lin_vel_x, self.target_velocity.lin_vel_y, 0.0]) raw_cur_vel = self.proprio.base.lin_vel cur_vel_body = np.array([raw_cur_vel[0], raw_cur_vel[1], 0.0]) tgt_vel_world = np.zeros(3) cur_vel_world = np.zeros(3) mujoco.mju_rotVecQuat(tgt_vel_world, tgt_vel_body, base_quat) if ctype == 'viewer': mujoco.mju_rotVecQuat(cur_vel_world, cur_vel_body, base_quat) else: mujoco.mju_rotVecQuat(cur_vel_world, cur_vel_body, base_quat) COLOR_CMD = [0, 1, 0, 1] # Green 0x00ff00 COLOR_REAL = [0, 0, 1, 1] # Blue 0x0000ff if ctype == 'viewer': # Cmd Arrow add_thick_arrow(handle.user_scn.geoms[viewer_geom_idx], start_pos, tgt_vel_world, COLOR_CMD) viewer_geom_idx += 1 # Real Arrow add_thick_arrow(handle.user_scn.geoms[viewer_geom_idx], start_pos, cur_vel_world, COLOR_REAL) viewer_geom_idx += 1 else: # Renderer Append handle.scene.ngeom += 1 add_thick_arrow(handle.scene.geoms[handle.scene.ngeom - 1], start_pos, tgt_vel_world, COLOR_CMD) handle.scene.ngeom += 1 add_thick_arrow(handle.scene.geoms[handle.scene.ngeom - 1], start_pos, cur_vel_world, COLOR_REAL) if ctype == 'viewer': handle.user_scn.ngeom = viewer_geom_idx def check_penetration(self, threshold: float = -0.02): if self.penetration_reset_count >= self.cfg.truncation.penetration_max_reset_num: return False, None, None, None for i in range(self.mj_data.ncon): contact = self.mj_data.contact[i] if contact.dist < threshold: geom1_name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, contact.geom1) geom2_name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, contact.geom2) return True, geom1_name, geom2_name, contact.dist return False, None, None, None def check_truncation(self, sim_data: SimData): if self.cfg.truncation.enabled: projected_gravity = get_projected_gravity(sim_data.proprio.base.quat) if -projected_gravity[2] < np.cos(self.cfg.truncation.projected_gravity_rad): raise RuntimeError(f"[Roll Error] Episode truncated due to excessive projected gravity, angle: {np.arccos(-projected_gravity[2]):.3f} rad, projected: {projected_gravity}") is_penetrated, geom1, geom2, dist = self.check_penetration(self.cfg.truncation.penetration_threshold) if is_penetrated: is_err = True if self.cfg.truncation.skip_penetration_geoms is not None and ( any(skip_geom in geom1.lower() for skip_geom in self.cfg.truncation.skip_penetration_geoms) or any(skip_geom in geom2.lower() for skip_geom in self.cfg.truncation.skip_penetration_geoms) ): is_err = False if self.cfg.truncation.skip_self_penetration: if geom1.split('/')[0] == geom2.split('/')[0]: is_err = False if is_err: self.penetration_reset_count += 1 raise RuntimeError(f"[Penetration Error] Episode truncated: Penetration ({geom1} <-> {geom2}), distance: {dist}") def reset(self): """ Reset the simulator to initial state. """ self.mj_physics.reset() if self.invert_yaw: self.mj_data.qpos[3] = 0.0 self.mj_data.qpos[6] = 1.0 self.mj_data.qpos[7:] = self.default_dof_pos mujoco.mj_forward(self.mj_model, self.mj_data) self.action = None if self.viewer is not None: self.viewer.sync() def setup_action(self, action: np.ndarray, p_gains: np.ndarray = None, d_gains: np.ndarray = None, control_type: Literal['P'] = 'P' ): """ Setup action to the simulator. """ self.action = action self.p_gains = p_gains self.d_gains = d_gains self.control_type = control_type def update_torque(self): if self.action is None: return dof_pos = self.proprio.joint.pos dof_vel = self.proprio.joint.vel if self.control_type == 'P': torques = self.p_gains * (self.action - dof_pos) - self.d_gains * dof_vel else: raise NotImplementedError(f"Control type '{self.control_type}' not implemented.") self.mj_data.ctrl[:] = torques 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.") def close_video_writer(self): """ Close the video writer if exists. """ if self.vid_writer is not None: self.vid_writer.close() self.vid_writer = None self.renderer = None logger.info("Closing video writer.") def preload_sensors(self): # Preload sensor names self.joint_pos_sensor_names = self.find_sensors(tag_name="jointpos") self.joint_vel_sensor_names = self.find_sensors(tag_name="jointvel") self.joint_eff_sensor_names = self.find_sensors(tag_name="jointactuatorfrc") self.imu_quat = self.find_sensors(tag_name="framequat") self.imu_ang_vel = self.find_sensors(tag_name="gyro") self.imu_acc = self.find_sensors(tag_name="accelerometer") self.imu_pos = self.find_sensors(tag_name="framepos") self.imu_lin_vel = self.find_sensors(tag_name="framelinvel") actuator_names = [mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_ACTUATOR, i) for i in range(self.mj_model.nu)] logger.info( f"""\nRobot XML: {self.robot_xml}\n""" f"""Robot joint names: {[x.rsplit('/')[-1] for x in self.dof_names]}\n""" f"""{'='*20} XML SENSOR NAMES {'='*20}\n""" f"""Joint Position Sensors [{len(self.joint_pos_sensor_names)}]: {[x.rsplit('/')[-1] for x in self.joint_pos_sensor_names]}\n""" f"""Joint Velocity Sensors [{len(self.joint_vel_sensor_names)}]: {[x.rsplit('/')[-1] for x in self.joint_vel_sensor_names]}\n""" f"""Joint Effort Sensors [{len(self.joint_eff_sensor_names)}]: {[x.rsplit('/')[-1] for x in self.joint_eff_sensor_names]}\n""" f"""Actuators [{len(actuator_names)}]: {[x.rsplit('/')[-1] for x in actuator_names]}\n""" f"""IMU Sensors: Quat{self.imu_quat}, AngVel{self.imu_ang_vel}, Acc{self.imu_acc}, Pos{self.imu_pos}, LinVel{self.imu_lin_vel}\n""" f"""!!!Checkout actuators order is consistent with joint sensors!!!\n""" f"{'='*58}" ) # Cache sensor indices self.sensor_cache = {} all_lists = { 'joint_pos': self.joint_pos_sensor_names, 'joint_vel': self.joint_vel_sensor_names, 'joint_eff': self.joint_eff_sensor_names, 'imu_quat': self.imu_quat, 'imu_ang_vel': self.imu_ang_vel, 'imu_acc': self.imu_acc, 'imu_pos': self.imu_pos, 'imu_lin_vel': self.imu_lin_vel } for key, name_list in all_lists.items(): indices = [] for name in name_list: sid = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_SENSOR, name) if sid == -1: continue adr = int(self.mj_model.sensor_adr[sid]) dim = int(self.mj_model.sensor_dim[sid]) indices.append((adr, dim)) self.sensor_cache[key] = indices def find_sensors(self, *, pattern: re.Pattern = None, tag_name: str = None) -> list: model = self.mj_model found = [] tag_map = { "jointpos": mujoco.mjtSensor.mjSENS_JOINTPOS, "jointvel": mujoco.mjtSensor.mjSENS_JOINTVEL, "jointactuatorfrc": mujoco.mjtSensor.mjSENS_JOINTACTFRC, "accelerometer": mujoco.mjtSensor.mjSENS_ACCELEROMETER, "gyro": mujoco.mjtSensor.mjSENS_GYRO, "framepos": mujoco.mjtSensor.mjSENS_FRAMEPOS, "framequat": mujoco.mjtSensor.mjSENS_FRAMEQUAT, "framelinvel": mujoco.mjtSensor.mjSENS_FRAMELINVEL, "frameangvel": mujoco.mjtSensor.mjSENS_FRAMEANGVEL, } tag_type_id = None if tag_name: if tag_name not in tag_map: logger.warning(f"Unknown tag_name '{tag_name}', ignoring tag filter.") return [] tag_type_id = tag_map[tag_name] for i in range(model.nsensor): if tag_type_id is not None and model.sensor_type[i] != tag_type_id: continue name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_SENSOR, i) if pattern is None or name and pattern.search(name): found.append(name) if len(found) == 0: logger.warning(f"No sensors found for pattern='{pattern}' tag_name='{tag_name}'") return found def get_sensor_data(self, cache_key: str) -> np.ndarray: ids = self.sensor_cache.get(cache_key, []) if not ids: return np.array([]) data_list = [] for adr, dim in ids: data_list.append(self.mj_data.sensordata[adr:adr+dim]) return np.concatenate(data_list) def debug_print_proprio_shapes(self): """Log shapes (or lengths) of each numpy vector inside a RobotProprioception. This helps debug mismatched sensor sizes between robots. """ def _shape(x): try: arr = np.asarray(x) return arr.shape except Exception: return None jp = self.proprio.joint bs = self.proprio.base imu = self.proprio.imu logger.info("Proprioception shapes:") logger.info(f" joint.pos: { _shape(jp.pos) }") logger.info(f" joint.vel: { _shape(jp.vel) }") logger.info(f" joint.torque: { _shape(jp.torque) }") logger.info(f" base.pos: { _shape(bs.pos) }") logger.info(f" base.quat: { _shape(bs.quat) }") logger.info(f" base.vel: { _shape(bs.lin_vel) }") logger.info(f" base.ang_vel: { _shape(bs.ang_vel) }") logger.info(f" imu.quat: { _shape(imu.quat) }") logger.info(f" imu.ang_vel: { _shape(imu.ang_vel) }") logger.info(f" imu.acc: { _shape(imu.acc) }") logger.info(f" imu.pos: { _shape(imu.pos) }") logger.info(f" imu.lin_vel: { _shape(imu.lin_vel) }") def load_dof_limits(self): self.dof_limits = [] self.dof_names = [] for i in range(self.mj_model.njnt): name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_JOINT, i) jnt_type = self.mj_model.jnt_type[i] if jnt_type == mujoco.mjtJoint.mjJNT_FREE: continue limits = self.mj_model.jnt_range[i] self.dof_limits.append(limits) self.dof_names.append(name) self.dof_limits = np.array(self.dof_limits, np.float32)