From 4cc0687114cfe714ec8ea5d699dd879a01ce399d Mon Sep 17 00:00:00 2001 From: wty-yy <993660140@qq.com> Date: Thu, 27 Nov 2025 17:58:37 +0800 Subject: [PATCH] v0.1.2 --- UPDATE.md | 1 + robogauge/scripts/run.py | 1 + robogauge/tasks/simulator/mujoco_simulator.py | 163 +++++++++++++++++- robogauge/tasks/simulator/sim_data.py | 29 ++++ robogauge/utils/logger.py | 13 +- 5 files changed, 194 insertions(+), 13 deletions(-) create mode 100644 robogauge/tasks/simulator/sim_data.py diff --git a/UPDATE.md b/UPDATE.md index aefd6df..c9c584e 100644 --- a/UPDATE.md +++ b/UPDATE.md @@ -3,6 +3,7 @@ ## 20251127 ### v0.1.2 1. 修改logger +2. 完成mujoco的sensor数据, sim_data数据类处理 ## 20251124-25 ### v0.1.1 diff --git a/robogauge/scripts/run.py b/robogauge/scripts/run.py index 249380d..979c1ec 100644 --- a/robogauge/scripts/run.py +++ b/robogauge/scripts/run.py @@ -15,5 +15,6 @@ from robogauge.utils.logger import logger if __name__ == '__main__': args = parse_args() logger.create(args.experiment_name) + logger.info(f"Starting experiment: {args.experiment_name}") pipeline: BasePipeline = task_register.make_pipeline(args.task_name, args=args) pipeline.run() diff --git a/robogauge/tasks/simulator/mujoco_simulator.py b/robogauge/tasks/simulator/mujoco_simulator.py index 5b1b0f9..df710cb 100644 --- a/robogauge/tasks/simulator/mujoco_simulator.py +++ b/robogauge/tasks/simulator/mujoco_simulator.py @@ -11,6 +11,7 @@ import mujoco import mujoco.viewer from dm_control import mjcf +import re import time import imageio import numpy as np @@ -18,6 +19,7 @@ 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 +from robogauge.tasks.simulator.sim_data import RobotProprioception, JointState, BaseState, IMUState class MujocoSimulator: def __init__(self, sim_cfg: MujocoConfig): @@ -31,7 +33,8 @@ class MujocoSimulator: self.vid_writer = None self.vid_count = 0 self._pause = True - + self.n_step = 0 + def load( self, terrain_xml: str = None, @@ -76,10 +79,15 @@ class MujocoSimulator: 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.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) + 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) @@ -91,7 +99,9 @@ class MujocoSimulator: logger.info(f"Saving simulation video to: {vid_path}") self.vid_count += 1 self._pause = False - + self.n_step = 0 + self.preload_sensors() + def key_callback(self, keycode): if keycode == 32: self._pause = not self._pause @@ -104,7 +114,9 @@ class MujocoSimulator: 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) + 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() @@ -117,8 +129,33 @@ class MujocoSimulator: logger.warning("Viewer closed by user, stop video recording.") self.close_viewer() - info = {} - return info + n_sensor = self.mj_model.nsensor + + proprio = RobotProprioception( + joint=JointState( + pos=self.get_sensor_data('joint_pos'), + vel=self.get_sensor_data('joint_vel'), + force=self.get_sensor_data('joint_eff'), + ), + imu=IMUState( + quat=self.get_sensor_data('imu_quat'), + ang_vel=self.get_sensor_data('imu_ang_vel'), + acc=self.get_sensor_data('imu_acc'), + pos=self.get_sensor_data('imu_pos'), + lin_vel=self.get_sensor_data('imu_lin_vel'), + ), + base=BaseState( + pos=self.mj_data.qpos[:3], # world frame + quat=self.mj_data.qpos[3:7], # world frame + vel=self.mj_data.qvel[:3], # body frame + ang_vel=self.mj_data.qvel[3:6], # body frame + ) + ) + if self.n_step == 0: + self.debug_print_proprio_shapes(proprio) + self.n_step += 1 + + return proprio def reset(self): """ Reset the simulator to initial state. """ @@ -141,3 +178,115 @@ class MujocoSimulator: 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") + logger.info( + f"\n{'='*20} XML SENSOR NAMES {'='*20}\n" + f"""Joint Position Sensors [{len(self.joint_pos_sensor_names)}]: {self.joint_pos_sensor_names}\n""" + f"""Joint Velocity Sensors [{len(self.joint_vel_sensor_names)}]: {self.joint_vel_sensor_names}\n""" + f"""Joint Effort Sensors [{len(self.joint_eff_sensor_names)}]: {self.joint_eff_sensor_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}""" + ) + + # 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) + 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, proprio: RobotProprioception): + """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 = proprio.joint + bs = proprio.base + imu = 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.force: { _shape(jp.force) }") + + logger.info(f" base.pos: { _shape(bs.pos) }") + logger.info(f" base.quat: { _shape(bs.quat) }") + logger.info(f" base.vel: { _shape(bs.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) }") diff --git a/robogauge/tasks/simulator/sim_data.py b/robogauge/tasks/simulator/sim_data.py new file mode 100644 index 0000000..2ee311d --- /dev/null +++ b/robogauge/tasks/simulator/sim_data.py @@ -0,0 +1,29 @@ +import numpy as np +from dataclasses import dataclass + +@dataclass +class JointState: + pos: np.ndarray + vel: np.ndarray + force: np.ndarray + +@dataclass +class BaseState: + pos: np.ndarray + quat: np.ndarray + vel: np.ndarray + ang_vel: np.ndarray + +@dataclass +class IMUState: + quat: np.ndarray + ang_vel: np.ndarray + acc: np.ndarray + pos: np.ndarray + lin_vel: np.ndarray + +@dataclass +class RobotProprioception: + joint: JointState + base: BaseState + imu: IMUState diff --git a/robogauge/utils/logger.py b/robogauge/utils/logger.py index 9076819..5e5b984 100644 --- a/robogauge/utils/logger.py +++ b/robogauge/utils/logger.py @@ -69,6 +69,7 @@ class Logger: """ self.logger = logging.getLogger(experiment_name + "_logger") self.logger.setLevel(log_level) + self.logger.propagate = False console_formatter = ColorFormatter( # console output format fmt="%(asctime)s - %(color_level)s - %(filename)s:%(lineno)d - %(message)s", @@ -92,21 +93,21 @@ class Logger: fh = logging.FileHandler(path_log_file, mode=save_file_mode, encoding='utf-8') fh.setFormatter(file_formatter) self.logger.addHandler(fh) - + def debug(self, msg, *args, **kwargs): - self.logger.debug(msg, *args, **kwargs) + self.logger.debug(msg, *args, **kwargs, stacklevel=2) def info(self, msg, *args, **kwargs): - self.logger.info(msg, *args, **kwargs) + self.logger.info(msg, *args, **kwargs, stacklevel=2) def warning(self, msg, *args, **kwargs): - self.logger.warning(msg, *args, **kwargs) + self.logger.warning(msg, *args, **kwargs, stacklevel=2) def error(self, msg, *args, **kwargs): - self.logger.error(msg, *args, **kwargs) + self.logger.error(msg, *args, **kwargs, stacklevel=2) def critical(self, msg, *args, **kwargs): - self.logger.critical(msg, *args, **kwargs) + self.logger.critical(msg, *args, **kwargs, stacklevel=2) logger = Logger()