diff --git a/src/unilab/base/backend/base.py b/src/unilab/base/backend/base.py index 3dff6c04b..8f943d668 100644 --- a/src/unilab/base/backend/base.py +++ b/src/unilab/base/backend/base.py @@ -6,6 +6,24 @@ class SimBackend(abc.ABC): """仿真后端统一接口""" + # ------------------------------------------------------------------ # + # Properties # + # ------------------------------------------------------------------ # + + @property + @abc.abstractmethod + def num_envs(self) -> int: + """环境数量""" + + @property + @abc.abstractmethod + def model(self): + """底层物理模型""" + + # ------------------------------------------------------------------ # + # Simulation control # + # ------------------------------------------------------------------ # + @abc.abstractmethod def step(self, ctrl: np.ndarray, nsteps: int = 1) -> None: """执行物理步进 @@ -15,6 +33,63 @@ def step(self, ctrl: np.ndarray, nsteps: int = 1) -> None: nsteps: 步进次数 """ + @abc.abstractmethod + def set_state(self, env_indices: np.ndarray, qpos: np.ndarray, qvel: np.ndarray) -> None: + """设置指定环境的物理状态 + + Args: + env_indices: 环境索引 + qpos: 位置状态 + qvel: 速度状态 + """ + + # ------------------------------------------------------------------ # + # Base kinematics # + # ------------------------------------------------------------------ # + + @abc.abstractmethod + def get_base_pos(self) -> np.ndarray: + """获取 base 在世界系下的位置 + + Returns: + (num_envs, 3) + """ + + @abc.abstractmethod + def get_base_quat(self) -> np.ndarray: + """获取 base 在世界系下的四元数（wxyz） + + Returns: + (num_envs, 4) + """ + + @abc.abstractmethod + def get_base_lin_vel(self) -> np.ndarray: + """获取 base 在世界系下的线速度 + + 即广义速度 qvel 的前 3 维，表达在世界坐标系中。 + + Returns: + (num_envs, 3) + """ + + @abc.abstractmethod + def get_base_ang_vel(self) -> np.ndarray: + """获取 base 在世界系下的角速度 + + 即广义速度 qvel 的第 3-5 维，表达在世界坐标系中。 + 注意与陀螺仪（gyro）读数的区别：陀螺仪返回的是角速度在 body/sensor + 局部坐标系下的分量（即 body frame 表达），而本接口返回的是世界系表达。 + 若需要 body frame 下的角速度，请使用对应的传感器接口gyro。 + + Returns: + (num_envs, 3) + """ + + # ------------------------------------------------------------------ # + # DOF state # + # ------------------------------------------------------------------ # + @abc.abstractmethod def get_dof_pos(self) -> np.ndarray: """获取关节位置（不含 base） @@ -31,41 +106,113 @@ def get_dof_vel(self) -> np.ndarray: (num_envs, num_dof) """ + # ------------------------------------------------------------------ # + # Body kinematics — world frame # + # ------------------------------------------------------------------ # + @abc.abstractmethod - def get_qpos(self) -> np.ndarray: - """获取完整位置状态（含 base） + def get_body_pos_w(self, body_ids: np.ndarray) -> np.ndarray: + """获取指定 body 在世界系下的位置 + + Args: + body_ids: body 索引数组 Returns: - (num_envs, nq) + (num_envs, len(body_ids), 3) """ @abc.abstractmethod - def get_sensor_data(self, name: str) -> np.ndarray: - """获取传感器数据 + def get_body_quat_w(self, body_ids: np.ndarray) -> np.ndarray: + """获取指定 body 在世界系下的四元数（wxyz） Args: - name: 传感器名称 + body_ids: body 索引数组 Returns: - 传感器数据数组 + (num_envs, len(body_ids), 4) """ @abc.abstractmethod - def set_state(self, env_indices: np.ndarray, qpos: np.ndarray, qvel: np.ndarray) -> None: - """设置指定环境的物理状态 + def get_body_lin_vel_w(self, body_ids: np.ndarray) -> np.ndarray: + """获取指定 body 在世界系下的线速度 Args: - env_indices: 环境索引 - qpos: 位置状态 - qvel: 速度状态 + body_ids: body 索引数组 + + Returns: + (num_envs, len(body_ids), 3) """ - @property @abc.abstractmethod - def num_envs(self) -> int: - """环境数量""" + def get_body_ang_vel_w(self, body_ids: np.ndarray) -> np.ndarray: + """获取指定 body 在世界系下的角速度 + + Args: + body_ids: body 索引数组 + + Returns: + (num_envs, len(body_ids), 3) + """ + + # ------------------------------------------------------------------ # + # Body kinematics — baselink frame # + # ------------------------------------------------------------------ # - @property @abc.abstractmethod - def model(self): - """底层物理模型""" + def get_body_pos_b(self, body_ids: np.ndarray) -> np.ndarray: + """获取指定 body 在 baselink 系下的位置 + + Args: + body_ids: body 索引数组 + + Returns: + (num_envs, len(body_ids), 3) + """ + + @abc.abstractmethod + def get_body_quat_b(self, body_ids: np.ndarray) -> np.ndarray: + """获取指定 body 在 baselink 系下的四元数（wxyz） + + Args: + body_ids: body 索引数组 + + Returns: + (num_envs, len(body_ids), 4) + """ + + @abc.abstractmethod + def get_body_lin_vel_b(self, body_ids: np.ndarray) -> np.ndarray: + """获取指定 body 在 baselink 系下的线速度 + + Args: + body_ids: body 索引数组 + + Returns: + (num_envs, len(body_ids), 3) + """ + + @abc.abstractmethod + def get_body_ang_vel_b(self, body_ids: np.ndarray) -> np.ndarray: + """获取指定 body 在 baselink 系下的角速度 + + Args: + body_ids: body 索引数组 + + Returns: + (num_envs, len(body_ids), 3) + """ + + # ------------------------------------------------------------------ # + # Sensors # + # ------------------------------------------------------------------ # + + @abc.abstractmethod + def get_sensor_data(self, name: str) -> np.ndarray: + """获取传感器数据 + + Args: + name: 传感器名称 + + Returns: + 传感器数据数组 + """ diff --git a/src/unilab/base/backend/motrix_backend.py b/src/unilab/base/backend/motrix_backend.py index e57a7a7d4..70394ef7b 100644 --- a/src/unilab/base/backend/motrix_backend.py +++ b/src/unilab/base/backend/motrix_backend.py @@ -1,3 +1,5 @@ +import os + import numpy as np try: @@ -19,68 +21,78 @@ def __init__( model_file: str, num_envs: int, sim_dt: float, - body_name: str = "base", + base_name: str = "base", np_dtype=np.float32, + add_body_sensors: bool = False, ): if not MOTRIX_AVAILABLE: raise ImportError("motrixsim not available") - self._model = mtx.load_model(model_file) # pyright: ignore[reportPossiblyUnbound] + self.add_body_sensors = add_body_sensors + + if self.add_body_sensors: + from unilab.utils.xml_utils import inject_motrix_tracking_sensors + + tmp_path, _, valid_bnames = inject_motrix_tracking_sensors( + model_file, baselink_name=base_name + ) + try: + self._model = mtx.load_model(tmp_path) # pyright: ignore[reportPossiblyUnbound] + finally: + os.remove(tmp_path) + + # 用 motrixsim link index 作 key（Link 覆盖所有关节体，Body 只有 freejoint 根体） + self._body_id_to_name: dict[int, str] = { + idx: name + for name in valid_bnames + if (idx := self._model.get_link_index(name)) is not None + } + else: + self._model = mtx.load_model(model_file) # pyright: ignore[reportPossiblyUnbound] + + # 枚举所有具名 link，用 link index 作 key + self._body_id_to_name = { # type: ignore[assignment] + link.index: link.name for link in self._model.links if link.name + } + self._model.options.timestep = sim_dt self._num_envs = num_envs self._np_dtype = np_dtype self._data = mtx.SceneData(self._model, batch=[num_envs]) # pyright: ignore[reportPossiblyUnbound] - self._body = self._model.get_body(body_name) - self._body_link = self._model.get_link(body_name) + self._body = self._model.get_body(base_name) + self._body_link = self._model.get_link(base_name) self._render_app: "RenderApp | None" = None self.backend_type = "motrix" - def _process_rigid_body_props(self, cfg) -> None: - if cfg.domain_rand.randomize_base_mass: - mass = self._model.get_link(cfg.asset.body_name).get_mass_override(self._data) - mass_low = cfg.domain_rand.added_mass_range[0] - mass_high = cfg.domain_rand.added_mass_range[1] - random_mass = mass + np.random.uniform(mass_low, mass_high, size=(self._num_envs,)) - self._model.get_link(cfg.asset.body_name).set_mass_override(self._data, random_mass) - mass = self._model.get_link(cfg.asset.body_name).get_mass_override(self._data) + # 运行一次正向运动学，确保初始 link 位置和传感器数据有效 + self._model.forward_kinematic(self._data) - if cfg.domain_rand.random_com: - com_offset = np.zeros( - (self._num_envs, 3), dtype=np.float32 - ) # [x_offset, y_offset, z_offset] - x_low = cfg.domain_rand.com_offset_x[0] - x_high = cfg.domain_rand.com_offset_x[1] - com_offset[:, 0] = np.random.uniform(x_low, x_high, self._num_envs) - self._model.get_link(cfg.asset.body_name).set_center_of_mass_override( - self._data, com_offset - ) - self._model.get_link(cfg.asset.body_name).get_center_of_mass_override(self._data) + # ------------------------------------------------------------------ # + # Properties # + # ------------------------------------------------------------------ # - def push_robots(self, force_range): - ex_force = np.random.rand(self.num_envs, 3) * 2 - 1 # [x_force, y_force, z_force] - ex_force[:, 0] *= force_range[0] - ex_force[:, 1] *= force_range[1] - ex_force[:, 2] *= force_range[2] - self._body_link.add_external_force(self._data, ex_force, local=True) + @property + def num_envs(self) -> int: + return self._num_envs + + @property + def model(self): + return self._model + + @property + def data(self): + return self._data + + # ------------------------------------------------------------------ # + # Simulation control # + # ------------------------------------------------------------------ # def step(self, ctrl: np.ndarray, nsteps: int = 1) -> None: self._data.actuator_ctrls = np.ascontiguousarray(ctrl) for _ in range(nsteps): self._model.step(self._data) - def get_dof_pos(self) -> np.ndarray: - return np.asarray(self._body.get_joint_dof_pos(self._data)) - - def get_dof_vel(self) -> np.ndarray: - return np.asarray(self._body.get_joint_dof_vel(self._data)) - - def get_qpos(self) -> np.ndarray: - return np.asarray(self._data.dof_pos) - - def get_sensor_data(self, name: str) -> np.ndarray: - return np.asarray(self._model.get_sensor_value(name, self._data)) - def set_state(self, env_indices: np.ndarray, qpos: np.ndarray, qvel: np.ndarray) -> None: # Convert quaternion from mujoco (wxyz) to motrix (xyzw) qpos_motrix = qpos.copy() @@ -102,17 +114,175 @@ def set_state(self, env_indices: np.ndarray, qpos: np.ndarray, qvel: np.ndarray) self._model.forward_kinematic(self._data) - @property - def num_envs(self) -> int: - return self._num_envs + # ------------------------------------------------------------------ # + # Base kinematics # + # ------------------------------------------------------------------ # - @property - def model(self): - return self._model + def get_base_pos(self) -> np.ndarray: + return np.asarray(self._body.floatingbase.get_translation(self._data)) - @property - def data(self): - return self._data + def get_base_quat(self) -> np.ndarray: + return self._xyzw_to_wxyz(np.asarray(self._body.floatingbase.get_rotation(self._data))) + + def get_base_lin_vel(self) -> np.ndarray: + return np.asarray(self._body.floatingbase.get_global_linear_velocity(self._data)) + + def get_base_ang_vel(self) -> np.ndarray: + return np.asarray(self._body.floatingbase.get_global_angular_velocity(self._data)) + + # ------------------------------------------------------------------ # + # DOF state # + # ------------------------------------------------------------------ # + + def get_dof_pos(self) -> np.ndarray: + return np.asarray(self._body.get_joint_dof_pos(self._data)) + + def get_dof_vel(self) -> np.ndarray: + return np.asarray(self._body.get_joint_dof_vel(self._data)) + + # ------------------------------------------------------------------ # + # Body kinematics — world frame # + # ------------------------------------------------------------------ # + + def get_body_pos_w(self, body_ids: np.ndarray) -> np.ndarray: + return np.stack( + [ + np.asarray(self._model.get_link(int(bid)).get_position(self._data)) + for bid in body_ids + ], + axis=1, + ) + + def get_body_quat_w(self, body_ids: np.ndarray) -> np.ndarray: + return np.stack( + [ + self._xyzw_to_wxyz( + np.asarray(self._model.get_link(int(bid)).get_rotation(self._data)) + ) + for bid in body_ids + ], + axis=1, + ) + + def get_body_lin_vel_w(self, body_ids: np.ndarray) -> np.ndarray: + return np.stack( + [ + np.asarray(self._model.get_link(int(bid)).get_linear_velocity(self._data)) + for bid in body_ids + ], + axis=1, + ) + + def get_body_ang_vel_w(self, body_ids: np.ndarray) -> np.ndarray: + return np.stack( + [ + np.asarray(self._model.get_link(int(bid)).get_angular_velocity(self._data)) + for bid in body_ids + ], + axis=1, + ) + + # ------------------------------------------------------------------ # + # Body kinematics — baselink frame # + # ------------------------------------------------------------------ # + + def get_body_pos_b(self, body_ids: np.ndarray) -> np.ndarray: + return np.stack( + [ + np.asarray( + self._model.get_sensor_value( + f"track_pos_b_{self._body_id_to_name[int(bid)]}", self._data + ) + ) + for bid in body_ids + ], + axis=1, + ) + + def get_body_quat_b(self, body_ids: np.ndarray) -> np.ndarray: + # motrixsim framequat sensor 输出 xyzw，转换为接口约定的 wxyz + return np.stack( + [ + self._xyzw_to_wxyz( + np.asarray( + self._model.get_sensor_value( + f"track_quat_b_{self._body_id_to_name[int(bid)]}", self._data + ) + ) + ) + for bid in body_ids + ], + axis=1, + ) + + def get_body_lin_vel_b(self, body_ids: np.ndarray) -> np.ndarray: + return np.stack( + [ + np.asarray( + self._model.get_sensor_value( + f"track_linvel_b_{self._body_id_to_name[int(bid)]}", self._data + ) + ) + for bid in body_ids + ], + axis=1, + ) + + def get_body_ang_vel_b(self, body_ids: np.ndarray) -> np.ndarray: + return np.stack( + [ + np.asarray( + self._model.get_sensor_value( + f"track_angvel_b_{self._body_id_to_name[int(bid)]}", self._data + ) + ) + for bid in body_ids + ], + axis=1, + ) + + # ------------------------------------------------------------------ # + # Sensors # + # ------------------------------------------------------------------ # + + def get_sensor_data(self, name: str) -> np.ndarray: + return np.asarray(self._model.get_sensor_value(name, self._data)) + + # ------------------------------------------------------------------ # + # MotrixSim-specific # + # ------------------------------------------------------------------ # + + def _xyzw_to_wxyz(self, q: np.ndarray) -> np.ndarray: + """motrix xyzw → wxyz""" + return q[..., [3, 0, 1, 2]] + + def _process_rigid_body_props(self, cfg) -> None: + if cfg.domain_rand.randomize_base_mass: + mass = self._model.get_link(cfg.asset.base_name).get_mass_override(self._data) + mass_low = cfg.domain_rand.added_mass_range[0] + mass_high = cfg.domain_rand.added_mass_range[1] + random_mass = mass + np.random.uniform(mass_low, mass_high, size=(self._num_envs,)) + self._model.get_link(cfg.asset.base_name).set_mass_override(self._data, random_mass) + mass = self._model.get_link(cfg.asset.base_name).get_mass_override(self._data) + + if cfg.domain_rand.random_com: + com_offset = np.zeros( + (self._num_envs, 3), dtype=np.float32 + ) # [x_offset, y_offset, z_offset] + x_low = cfg.domain_rand.com_offset_x[0] + x_high = cfg.domain_rand.com_offset_x[1] + com_offset[:, 0] = np.random.uniform(x_low, x_high, self._num_envs) + self._model.get_link(cfg.asset.base_name).set_center_of_mass_override( + self._data, com_offset + ) + self._model.get_link(cfg.asset.base_name).get_center_of_mass_override(self._data) + + def push_robots(self, force_range): + ex_force = np.random.rand(self.num_envs, 3) * 2 - 1 # [x_force, y_force, z_force] + ex_force[:, 0] *= force_range[0] + ex_force[:, 1] *= force_range[1] + ex_force[:, 2] *= force_range[2] + self._body_link.add_external_force(self._data, ex_force, local=True) def init_renderer(self, spacing: float = 1.0): """Initialize interactive renderer for visualization""" diff --git a/src/unilab/base/backend/motrix_numba_backend.py b/src/unilab/base/backend/motrix_numba_backend.py index 5d7e05a9b..c18b7bd69 100644 --- a/src/unilab/base/backend/motrix_numba_backend.py +++ b/src/unilab/base/backend/motrix_numba_backend.py @@ -64,7 +64,7 @@ def __init__( model_file: str, num_envs: int, sim_dt: float, - body_name: str = "base", + base_name: str = "base", np_dtype=np.float32, ): if not MOTRIX_AVAILABLE: @@ -76,7 +76,7 @@ def __init__( self._np_dtype = np_dtype self._data = mtx.SceneData(self._model, batch=[num_envs]) - self._body = self._model.get_body(body_name) + self._body = self._model.get_body(base_name) self._render_app: "RenderApp | None" = None def step(self, ctrl: np.ndarray, nsteps: int = 1) -> None: @@ -90,9 +90,6 @@ def get_dof_pos(self) -> np.ndarray: def get_dof_vel(self) -> np.ndarray: return np.asarray(self._body.get_joint_dof_vel(self._data)) - def get_qpos(self) -> np.ndarray: - return np.asarray(self._data.dof_pos) - def get_sensor_data(self, name: str) -> np.ndarray: return np.asarray(self._model.get_sensor_value(name, self._data)) @@ -150,6 +147,69 @@ def init_renderer(self, spacing: float = 1.0): render_settings=settings, ) + def get_base_pos(self) -> np.ndarray: + return np.asarray(self._body.floatingbase.get_translation(self._data)) + + def get_base_quat(self) -> np.ndarray: + xyzw = np.asarray(self._body.floatingbase.get_rotation(self._data)) + return xyzw[..., [3, 0, 1, 2]] + + def get_base_lin_vel(self) -> np.ndarray: + return np.asarray(self._body.floatingbase.get_global_linear_velocity(self._data)) + + def get_base_ang_vel(self) -> np.ndarray: + return np.asarray(self._body.floatingbase.get_global_angular_velocity(self._data)) + + def get_body_pos_w(self, body_ids: np.ndarray) -> np.ndarray: + return np.stack( + [ + np.asarray(self._model.get_link(int(bid)).get_position(self._data)) + for bid in body_ids + ], + axis=1, + ) + + def get_body_quat_w(self, body_ids: np.ndarray) -> np.ndarray: + return np.stack( + [ + np.asarray(self._model.get_link(int(bid)).get_rotation(self._data))[ + ..., [3, 0, 1, 2] + ] + for bid in body_ids + ], + axis=1, + ) + + def get_body_lin_vel_w(self, body_ids: np.ndarray) -> np.ndarray: + return np.stack( + [ + np.asarray(self._model.get_link(int(bid)).get_linear_velocity(self._data)) + for bid in body_ids + ], + axis=1, + ) + + def get_body_ang_vel_w(self, body_ids: np.ndarray) -> np.ndarray: + return np.stack( + [ + np.asarray(self._model.get_link(int(bid)).get_angular_velocity(self._data)) + for bid in body_ids + ], + axis=1, + ) + + def get_body_pos_b(self, body_ids: np.ndarray) -> np.ndarray: + raise NotImplementedError("MotrixNumbaBackend does not support baselink-frame body queries") + + def get_body_quat_b(self, body_ids: np.ndarray) -> np.ndarray: + raise NotImplementedError("MotrixNumbaBackend does not support baselink-frame body queries") + + def get_body_lin_vel_b(self, body_ids: np.ndarray) -> np.ndarray: + raise NotImplementedError("MotrixNumbaBackend does not support baselink-frame body queries") + + def get_body_ang_vel_b(self, body_ids: np.ndarray) -> np.ndarray: + raise NotImplementedError("MotrixNumbaBackend does not support baselink-frame body queries") + def render(self): """Render current state (interactive visualization)""" if self._render_app is None: diff --git a/src/unilab/base/backend/mujoco_backend.py b/src/unilab/base/backend/mujoco_backend.py index 6087f2443..417cc8d91 100644 --- a/src/unilab/base/backend/mujoco_backend.py +++ b/src/unilab/base/backend/mujoco_backend.py @@ -18,38 +18,48 @@ def __init__( model_file: str, num_envs: int, sim_dt: float, - body_name: Optional[str] = None, + base_name: Optional[str] = None, np_dtype=None, + add_body_sensors: bool = False, ): - self._model = mujoco.MjModel.from_xml_path(model_file) + self.add_body_sensors = add_body_sensors + + if self.add_body_sensors: + from unilab.utils.xml_utils import inject_mujoco_tracking_sensors + + tmp_path, self._tracked_body_ids, valid_bnames = inject_mujoco_tracking_sensors( + model_file, baselink_name=base_name + ) + try: + self._model = mujoco.MjModel.from_xml_path(tmp_path) + finally: + os.remove(tmp_path) + + self._body_id_to_tracked_idx = np.full(self._model.nbody, -1, dtype=int) + for idx, bid in enumerate(self._tracked_body_ids): + self._body_id_to_tracked_idx[bid] = idx + else: + self._model = mujoco.MjModel.from_xml_path(model_file) + valid_bnames = [] + self._model.opt.timestep = sim_dt self._num_envs = num_envs self._np_dtype = np_dtype if np_dtype is not None else get_global_dtype() self.backend_type = "mujoco" + # 线程配置 thread_override = os.getenv("UNILAB_MUJOCO_STEP_THREADS") if thread_override: self._n_threads = min(num_envs, max(1, int(thread_override))) else: host_threads = cpu_count() - if num_envs >= 4096: - auto_threads = max(host_threads, 56) - elif num_envs >= 2048: - auto_threads = max(host_threads, 32) - else: - auto_threads = host_threads - self._n_threads = min(num_envs, auto_threads) + self._n_threads = min(num_envs, host_threads * 2) # Worker pool self._worker_data = [mujoco.MjData(self._model) for _ in range(self._n_threads)] self._rollout = rollout.Rollout(nthread=self._n_threads) self._forward_runner = batch_forward.BatchForwardRunner(nthread=self._n_threads) - # 状态存储 - nstate = mujoco.mj_stateSize(self._model, mujoco.mjtState.mjSTATE_FULLPHYSICS) - self._physics_state = np.zeros((num_envs, nstate), dtype=self._np_dtype) - self._sensor_data = np.zeros((num_envs, self._model.nsensordata), dtype=self._np_dtype) - # 索引 self.nq = self._model.nq self.nv = self._model.nv @@ -58,10 +68,21 @@ def __init__( self._num_dof_pos = self.nq - 7 self._num_dof_vel = self.nv - 6 + # 状态存储 + nstate = mujoco.mj_stateSize(self._model, mujoco.mjtState.mjSTATE_FULLPHYSICS) + self._physics_state = np.zeros((num_envs, nstate), dtype=self._np_dtype) + # 用模型默认 qpos（含 identity 四元数）初始化所有环境 + self._physics_state[:, self._idx_qpos : self._idx_qpos + self._model.nq] = self._model.qpos0 + self._sensor_data = np.zeros((num_envs, self._model.nsensordata), dtype=self._np_dtype) + # 缓存视图 self._dof_pos_view = self._physics_state[:, self._idx_qpos + 7 : self._idx_qpos + self.nq] self._dof_vel_view = self._physics_state[:, self._idx_qvel + 6 : self._idx_qvel + self.nv] self._qpos_view = self._physics_state[:, self._idx_qpos : self._idx_qpos + self.nq] + self._base_pos_view = self._physics_state[:, self._idx_qpos : self._idx_qpos + 3] + self._base_quat_view = self._physics_state[:, self._idx_qpos + 3 : self._idx_qpos + 7] + self._base_lin_vel_view = self._physics_state[:, self._idx_qvel : self._idx_qvel + 3] + self._base_ang_vel_view = self._physics_state[:, self._idx_qvel + 3 : self._idx_qvel + 6] # 传感器索引 self._sensor_indices = {} @@ -74,6 +95,53 @@ def __init__( self._sensor_indices[name] = list(range(adr, adr + dim)) self._sensor_views[name] = self._sensor_data[:, adr : adr + dim] + # 针对追踪身体传感器的零拷贝视图映射 + if self.add_body_sensors and valid_bnames: + + def _get_sensor_view(prefix, dim): + adrs = [ + self._model.sensor_adr[ + mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_SENSOR, f"{prefix}_{nb}") + ] + for nb in valid_bnames + ] + return self._sensor_data[:, adrs[0] : adrs[-1] + dim].reshape( + num_envs, len(valid_bnames), dim + ) + + # Global (world) sensors + self._tracked_pos_w_all = _get_sensor_view("track_pos_w", 3) + self._tracked_quat_w_all = _get_sensor_view("track_quat_w", 4) + self._tracked_linvel_w_all = _get_sensor_view("track_linvel_w", 3) + self._tracked_angvel_w_all = _get_sensor_view("track_angvel_w", 3) + + # Local (baselink) sensors + self._tracked_pos_b_all = _get_sensor_view("track_pos_b", 3) + self._tracked_quat_b_all = _get_sensor_view("track_quat_b", 4) + self._tracked_linvel_b_all = _get_sensor_view("track_linvel_b", 3) + self._tracked_angvel_b_all = _get_sensor_view("track_angvel_b", 3) + + # 对初始 qpos0 状态运行一次 forward pass，确保传感器数据有效 + if self._model.nsensordata > 0: + _, sensor_init = self._forward_runner.forward( + model=self._model, + data=self._worker_data, + initial_state=self._physics_state.astype(np.float64), + chunk_size=max(1, num_envs // self._n_threads), + skipsensor=False, + out_dtype=np.float64, + return_state=True, + ) + self._sensor_data[:] = sensor_init.astype(self._np_dtype) + + @property + def num_envs(self) -> int: + return self._num_envs + + @property + def model(self): + return self._model + def step(self, ctrl: np.ndarray, nsteps: int = 1) -> None: control_traj = np.broadcast_to(ctrl[:, None, :], (self._num_envs, nsteps, ctrl.shape[-1])) state_traj, sensor_traj = self._rollout.rollout( @@ -86,20 +154,6 @@ def step(self, ctrl: np.ndarray, nsteps: int = 1) -> None: self._physics_state[:] = state_traj[:, -1, :].astype(self._np_dtype) self._sensor_data[:] = sensor_traj[:, -1, :].astype(self._np_dtype) - def get_dof_pos(self) -> np.ndarray: - return self._dof_pos_view - - def get_dof_vel(self) -> np.ndarray: - return self._dof_vel_view - - def get_qpos(self) -> np.ndarray: - return self._qpos_view - - def get_sensor_data(self, name: str) -> np.ndarray: - if name not in self._sensor_views: - raise ValueError(f"Sensor '{name}' not found") - return self._sensor_views[name] - def set_state(self, env_indices: np.ndarray, qpos: np.ndarray, qvel: np.ndarray) -> None: num_reset = len(env_indices) state_np = np.zeros((num_reset, self._physics_state.shape[1]), dtype=np.float64) @@ -119,155 +173,65 @@ def set_state(self, env_indices: np.ndarray, qpos: np.ndarray, qvel: np.ndarray) self._physics_state[env_indices] = state_np.astype(self._np_dtype) self._sensor_data[env_indices] = sensor_np.astype(self._np_dtype) - @property - def num_envs(self) -> int: - return self._num_envs - - @property - def model(self): - return self._model - - def get_physics_state(self) -> np.ndarray: - return self._physics_state - - def get_body_pos(self, body_ids: np.ndarray) -> np.ndarray: - """Get body positions in world frame. - - Args: - body_ids: Body indices (num_bodies,) - - Returns: - Body positions (num_envs, num_bodies, 3) - """ - # Need to run forward kinematics to get body positions - # Use batch_forward to compute xpos for all environments - state_out, _ = self._forward_runner.forward( - model=self._model, - data=self._worker_data, - initial_state=self._physics_state.astype(np.float64), - chunk_size=max(1, self._num_envs // self._n_threads), - skipsensor=True, - out_dtype=np.float64, - return_state=False, - ) - - # Extract body positions from MjData - # xpos is stored in data, need to extract from worker_data after forward - body_pos = np.zeros((self._num_envs, len(body_ids), 3), dtype=self._np_dtype) - - # Run forward kinematics for each environment - for i in range(self._num_envs): - worker_idx = i % self._n_threads - data = self._worker_data[worker_idx] - - # Set state - qpos = self._physics_state[i, self._idx_qpos : self._idx_qpos + self.nq] - qvel = self._physics_state[i, self._idx_qvel : self._idx_qvel + self.nv] - data.qpos[:] = qpos - data.qvel[:] = qvel - - # Run forward kinematics - mujoco.mj_kinematics(self._model, data) - - # Extract body positions - for j, body_id in enumerate(body_ids): - body_pos[i, j] = data.xpos[body_id] + def get_base_pos(self) -> np.ndarray: + return self._base_pos_view - return body_pos + def get_base_quat(self) -> np.ndarray: + return self._base_quat_view - def get_body_quat(self, body_ids: np.ndarray) -> np.ndarray: - """Get body quaternions in world frame (wxyz format). + def get_base_lin_vel(self) -> np.ndarray: + return self._base_lin_vel_view - Args: - body_ids: Body indices (num_bodies,) + def get_base_ang_vel(self) -> np.ndarray: + return self._base_ang_vel_view - Returns: - Body quaternions (num_envs, num_bodies, 4) - """ - body_quat = np.zeros((self._num_envs, len(body_ids), 4), dtype=self._np_dtype) - - # Run forward kinematics for each environment - for i in range(self._num_envs): - worker_idx = i % self._n_threads - data = self._worker_data[worker_idx] - - # Set state - qpos = self._physics_state[i, self._idx_qpos : self._idx_qpos + self.nq] - qvel = self._physics_state[i, self._idx_qvel : self._idx_qvel + self.nv] - data.qpos[:] = qpos - data.qvel[:] = qvel - - # Run forward kinematics - mujoco.mj_kinematics(self._model, data) - - # Extract body quaternions (MuJoCo uses wxyz format) - for j, body_id in enumerate(body_ids): - body_quat[i, j] = data.xquat[body_id] - - return body_quat - - def get_body_lin_vel(self, body_ids: np.ndarray) -> np.ndarray: - """Get body linear velocities in world frame. - - Args: - body_ids: Body indices (num_bodies,) - - Returns: - Body linear velocities (num_envs, num_bodies, 3) - """ - body_lin_vel = np.zeros((self._num_envs, len(body_ids), 3), dtype=self._np_dtype) - - # Run forward kinematics and velocity computation - for i in range(self._num_envs): - worker_idx = i % self._n_threads - data = self._worker_data[worker_idx] + def get_dof_pos(self) -> np.ndarray: + return self._dof_pos_view - # Set state - qpos = self._physics_state[i, self._idx_qpos : self._idx_qpos + self.nq] - qvel = self._physics_state[i, self._idx_qvel : self._idx_qvel + self.nv] - data.qpos[:] = qpos - data.qvel[:] = qvel + def get_dof_vel(self) -> np.ndarray: + return self._dof_vel_view - # Run forward kinematics and velocity computation - mujoco.mj_kinematics(self._model, data) - mujoco.mj_comVel(self._model, data) + def _get_mapped_indices(self, body_ids: np.ndarray) -> np.ndarray: + # if not self.add_body_sensors: + # raise NotImplementedError( + # "Slow kinematics computation has been removed for performance reasons. " + # "Please pass add_body_sensors=True during initialization to enable tracking." + # ) + mapped_indices = self._body_id_to_tracked_idx[body_ids] + # if np.any(mapped_indices == -1): + # raise ValueError( + # "Cannot query untracked (unnamed) bodies with fast sensor method. Please ensure bodies are named in XML." + # ) + return mapped_indices # type: ignore[no-any-return] - # Extract body linear velocities - for j, body_id in enumerate(body_ids): - # cvel contains [angular_vel (3), linear_vel (3)] for each body - body_lin_vel[i, j] = data.cvel[body_id, 3:6] + def get_body_pos_w(self, body_ids: np.ndarray) -> np.ndarray: + return np.asarray(self._tracked_pos_w_all[:, self._get_mapped_indices(body_ids), :]) - return body_lin_vel + def get_body_quat_w(self, body_ids: np.ndarray) -> np.ndarray: + return np.asarray(self._tracked_quat_w_all[:, self._get_mapped_indices(body_ids), :]) - def get_body_ang_vel(self, body_ids: np.ndarray) -> np.ndarray: - """Get body angular velocities in world frame. + def get_body_lin_vel_w(self, body_ids: np.ndarray) -> np.ndarray: + return np.asarray(self._tracked_linvel_w_all[:, self._get_mapped_indices(body_ids), :]) - Args: - body_ids: Body indices (num_bodies,) + def get_body_ang_vel_w(self, body_ids: np.ndarray) -> np.ndarray: + return np.asarray(self._tracked_angvel_w_all[:, self._get_mapped_indices(body_ids), :]) - Returns: - Body angular velocities (num_envs, num_bodies, 3) - """ - body_ang_vel = np.zeros((self._num_envs, len(body_ids), 3), dtype=self._np_dtype) + def get_body_pos_b(self, body_ids: np.ndarray) -> np.ndarray: + return np.asarray(self._tracked_pos_b_all[:, self._get_mapped_indices(body_ids), :]) - # Run forward kinematics and velocity computation - for i in range(self._num_envs): - worker_idx = i % self._n_threads - data = self._worker_data[worker_idx] + def get_body_quat_b(self, body_ids: np.ndarray) -> np.ndarray: + return np.asarray(self._tracked_quat_b_all[:, self._get_mapped_indices(body_ids), :]) - # Set state - qpos = self._physics_state[i, self._idx_qpos : self._idx_qpos + self.nq] - qvel = self._physics_state[i, self._idx_qvel : self._idx_qvel + self.nv] - data.qpos[:] = qpos - data.qvel[:] = qvel + def get_body_lin_vel_b(self, body_ids: np.ndarray) -> np.ndarray: + return np.asarray(self._tracked_linvel_b_all[:, self._get_mapped_indices(body_ids), :]) - # Run forward kinematics and velocity computation - mujoco.mj_kinematics(self._model, data) - mujoco.mj_comVel(self._model, data) + def get_body_ang_vel_b(self, body_ids: np.ndarray) -> np.ndarray: + return np.asarray(self._tracked_angvel_b_all[:, self._get_mapped_indices(body_ids), :]) - # Extract body angular velocities - for j, body_id in enumerate(body_ids): - # cvel contains [angular_vel (3), linear_vel (3)] for each body - body_ang_vel[i, j] = data.cvel[body_id, 0:3] + def get_sensor_data(self, name: str) -> np.ndarray: + if name not in self._sensor_views: + raise ValueError(f"Sensor '{name}' not found") + return self._sensor_views[name] - return body_ang_vel + def get_physics_state(self) -> np.ndarray: + return self._physics_state diff --git a/src/unilab/envs/locomotion/g1/base.py b/src/unilab/envs/locomotion/g1/base.py index ee686ed27..549da933c 100644 --- a/src/unilab/envs/locomotion/g1/base.py +++ b/src/unilab/envs/locomotion/g1/base.py @@ -29,7 +29,7 @@ class ControlConfig: @dataclass class Asset: - body_name = "pelvis" + base_name = "pelvis" foot_name = "ankle_roll_link" ground = "floor" diff --git a/src/unilab/envs/locomotion/g1/joystick.py b/src/unilab/envs/locomotion/g1/joystick.py index 739553ea1..7303d80e8 100644 --- a/src/unilab/envs/locomotion/g1/joystick.py +++ b/src/unilab/envs/locomotion/g1/joystick.py @@ -121,7 +121,7 @@ class G1JoystickPPO(G1BaseEnv): def __init__(self, cfg: G1JoystickPPOCfg, num_envs=1, backend_type="mujoco"): backend = create_backend( - backend_type, cfg.model_file, num_envs, cfg.sim_dt, body_name=cfg.asset.body_name + backend_type, cfg.model_file, num_envs, cfg.sim_dt, base_name=cfg.asset.base_name ) super().__init__(cfg, backend, num_envs) self._enable_reward_log = True @@ -184,15 +184,15 @@ def update_state(self, state: NpEnvState) -> NpEnvState: gravity = self._backend.get_sensor_data("upvector") dof_pos = self.get_dof_pos() dof_vel = self.get_dof_vel() - qpos = self._backend.get_qpos() max_tilt_rad = np.deg2rad(self._cfg.reward_config.max_tilt_deg) tilt = np.arccos(np.clip(gravity[:, 2], -1, 1)) terminated = np.logical_or( - tilt > max_tilt_rad, qpos[:, 2] < self._cfg.reward_config.min_base_height + tilt > max_tilt_rad, + self._backend.get_base_pos()[:, 2] < self._cfg.reward_config.min_base_height, ) - reward = self._compute_reward(state.info, linvel, gyro, gravity, dof_pos, dof_vel, qpos) + reward = self._compute_reward(state.info, linvel, gyro, gravity, dof_pos, dof_vel) obs = self._compute_obs(state.info, linvel, gyro, gravity, dof_pos, dof_vel) return state.replace(obs=obs, reward=reward, terminated=terminated) @@ -220,9 +220,7 @@ def get_obs_structure(self) -> dict: "gait_phase": 2, } - def _compute_reward( - self, info: dict, linvel, gyro, gravity, dof_pos, dof_vel, qpos - ) -> np.ndarray: + def _compute_reward(self, info: dict, linvel, gyro, gravity, dof_pos, dof_vel) -> np.ndarray: dtype = get_global_dtype() reward = np.zeros((self._num_envs,), dtype=dtype) cfg = self._cfg.reward_config @@ -234,7 +232,7 @@ def _compute_reward( for name, scale in cfg.scales.items(): if scale == 0 or name not in self._reward_fns: continue - rew = self._reward_fns[name](info, linvel, gyro, gravity, dof_pos, dof_vel, qpos) + rew = self._reward_fns[name](info, linvel, gyro, gravity, dof_pos, dof_vel) weighted_rew = rew * scale reward += weighted_rew if should_log: @@ -243,17 +241,17 @@ def _compute_reward( info["log"] = log return reward * self._cfg.ctrl_dt - def _reward_tracking_lin_vel(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_tracking_lin_vel(self, info, linvel, gyro, gravity, dof_pos, dof_vel): commands = info["commands"] lin_vel_error = np.sum(np.square(commands[:, :2] - linvel[:, :2]), axis=1) return np.exp(-lin_vel_error / self._cfg.reward_config.tracking_sigma) - def _reward_tracking_ang_vel(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_tracking_ang_vel(self, info, linvel, gyro, gravity, dof_pos, dof_vel): commands = info["commands"] ang_vel_error = np.square(commands[:, 2] - gyro[:, 2]) return np.exp(-ang_vel_error / self._cfg.reward_config.tracking_sigma) - def _reward_feet_phase(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_feet_phase(self, info, linvel, gyro, gravity, dof_pos, dof_vel): """步态相位奖励：鼓励正确的摆动腿高度""" left_foot = self._backend.get_sensor_data("left_foot_pos") right_foot = self._backend.get_sensor_data("right_foot_pos") @@ -286,23 +284,25 @@ def cubic_bezier_interpolation(y_start, y_end, t): -(left_error + right_error) / self._cfg.reward_config.feet_phase_tracking_sigma ) - def _reward_lin_vel_z(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_lin_vel_z(self, info, linvel, gyro, gravity, dof_pos, dof_vel): return np.square(linvel[:, 2]) - def _reward_ang_vel_xy(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_ang_vel_xy(self, info, linvel, gyro, gravity, dof_pos, dof_vel): return np.sum(np.square(gyro[:, :2]), axis=1) - def _reward_orientation(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_orientation(self, info, linvel, gyro, gravity, dof_pos, dof_vel): return np.square(gravity[:, 0]) + np.square(gravity[:, 1]) - def _reward_base_height(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): - return np.square(qpos[:, 2] - self._cfg.reward_config.base_height_target) + def _reward_base_height(self, info, linvel, gyro, gravity, dof_pos, dof_vel): + return np.square( + self._backend.get_base_pos()[:, 2] - self._cfg.reward_config.base_height_target + ) - def _reward_action_rate(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_action_rate(self, info, linvel, gyro, gravity, dof_pos, dof_vel): action_diff = info["current_actions"] - info["last_actions"] return np.sum(np.square(action_diff), axis=1) - def _reward_pose(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_pose(self, info, linvel, gyro, gravity, dof_pos, dof_vel): diff = dof_pos - self.default_angles return np.sum(self._pose_weights * np.square(diff), axis=1) diff --git a/src/unilab/envs/locomotion/g1/joystick_sac.py b/src/unilab/envs/locomotion/g1/joystick_sac.py index 32aba9171..3a5cb2104 100644 --- a/src/unilab/envs/locomotion/g1/joystick_sac.py +++ b/src/unilab/envs/locomotion/g1/joystick_sac.py @@ -114,7 +114,7 @@ class G1WalkTaskMjSAC(G1JoystickPPO): def __init__(self, cfg: G1JoystickSACCfg, num_envs=1, backend_type="mujoco"): backend = create_backend( - backend_type, cfg.model_file, num_envs, cfg.sim_dt, body_name=cfg.asset.body_name + backend_type, cfg.model_file, num_envs, cfg.sim_dt, base_name=cfg.asset.base_name ) G1BaseEnv.__init__(self, cfg, backend, num_envs) self._enable_reward_log = True diff --git a/src/unilab/envs/locomotion/g1/motion_tracking/tracking.py b/src/unilab/envs/locomotion/g1/motion_tracking/tracking.py index defef339e..a1eb664cd 100644 --- a/src/unilab/envs/locomotion/g1/motion_tracking/tracking.py +++ b/src/unilab/envs/locomotion/g1/motion_tracking/tracking.py @@ -138,7 +138,7 @@ def __init__(self, cfg: G1MotionTrackingCfg, num_envs=1, backend_type="mujoco"): raise ValueError("motion_file must be specified in config") backend = create_backend( - backend_type, cfg.model_file, num_envs, cfg.sim_dt, body_name=cfg.asset.body_name + backend_type, cfg.model_file, num_envs, cfg.sim_dt, base_name=cfg.asset.base_name ) super().__init__(cfg, backend, num_envs) @@ -214,21 +214,18 @@ def update_state(self, state: NpEnvState) -> NpEnvState: gyro = self.get_gyro() dof_pos = self.get_dof_pos() dof_vel = self.get_dof_vel() - qpos = self._backend.get_qpos() # Get body states - robot_body_pos_w = self._backend.get_body_pos(self.body_ids) - robot_body_quat_w = self._backend.get_body_quat(self.body_ids) - robot_body_lin_vel_w = self._backend.get_body_lin_vel(self.body_ids) - robot_body_ang_vel_w = self._backend.get_body_ang_vel(self.body_ids) + robot_body_pos_w = self._backend.get_body_pos_w(self.body_ids) + robot_body_quat_w = self._backend.get_body_quat_w(self.body_ids) + robot_body_lin_vel_w = self._backend.get_body_lin_vel_w(self.body_ids) + robot_body_ang_vel_w = self._backend.get_body_ang_vel_w(self.body_ids) # Compute relative body transforms (for observations and rewards) self._update_relative_transforms(motion_data, robot_body_pos_w, robot_body_quat_w) # Compute terminations - terminated = self._compute_terminations( - motion_data, robot_body_pos_w, robot_body_quat_w, qpos - ) + terminated = self._compute_terminations(motion_data, robot_body_pos_w, robot_body_quat_w) # Update failure statistics for adaptive sampling self.motion_sampler.update_failure_stats(terminated) @@ -294,7 +291,6 @@ def _compute_terminations( motion_data, robot_body_pos_w: np.ndarray, robot_body_quat_w: np.ndarray, - qpos: np.ndarray, ) -> np.ndarray: """Compute termination conditions.""" terminated = np.zeros(self._num_envs, dtype=bool) @@ -591,8 +587,8 @@ def reset(self, env_indices: np.ndarray): gyro = self.get_gyro()[env_indices] dof_pos = self.get_dof_pos()[env_indices] dof_vel = self.get_dof_vel()[env_indices] - robot_body_pos_w = self._backend.get_body_pos(self.body_ids)[env_indices] - robot_body_quat_w = self._backend.get_body_quat(self.body_ids)[env_indices] + robot_body_pos_w = self._backend.get_body_pos_w(self.body_ids)[env_indices] + robot_body_quat_w = self._backend.get_body_quat_w(self.body_ids)[env_indices] obs = self._compute_obs( info, motion_data, linvel, gyro, dof_pos, dof_vel, robot_body_pos_w, robot_body_quat_w diff --git a/src/unilab/envs/locomotion/go1/base.py b/src/unilab/envs/locomotion/go1/base.py index f88e1306f..f0afe1aae 100644 --- a/src/unilab/envs/locomotion/go1/base.py +++ b/src/unilab/envs/locomotion/go1/base.py @@ -32,7 +32,7 @@ class ControlConfig: @dataclass class Asset: - body_name = "trunk" + base_name = "trunk" foot_name = "foot" ground = "floor" diff --git a/src/unilab/envs/locomotion/go1/joystick.py b/src/unilab/envs/locomotion/go1/joystick.py index d484b7500..dc6a38223 100644 --- a/src/unilab/envs/locomotion/go1/joystick.py +++ b/src/unilab/envs/locomotion/go1/joystick.py @@ -103,7 +103,7 @@ class Go1WalkTask(Go1BaseEnv): def __init__(self, cfg: Go1JoystickCfg, num_envs=1, backend_type="mujoco"): backend = create_backend( - backend_type, cfg.model_file, num_envs, cfg.sim_dt, body_name=cfg.asset.body_name + backend_type, cfg.model_file, num_envs, cfg.sim_dt, base_name=cfg.asset.base_name ) super().__init__(cfg, backend, num_envs) self._enable_reward_log = True @@ -168,12 +168,11 @@ def update_state(self, state: NpEnvState) -> NpEnvState: gravity = self._backend.get_sensor_data("upvector") dof_pos = self.get_dof_pos() dof_vel = self.get_dof_vel() - qpos = self._backend.get_qpos() self.feet_force[:, :, :] = 0 for i in range(len(self._cfg.sensor.feet_force)): self.feet_force[:, i, :] = self._backend.get_sensor_data(self._cfg.sensor.feet_force[i]) terminated = gravity[:, 2] <= 0.5 - reward = self._compute_reward(state.info, linvel, gyro, dof_pos, qpos) + reward = self._compute_reward(state.info, linvel, gyro, dof_pos) obs = self._compute_obs( state.info, linvel, gyro, gravity, dof_pos, dof_vel, self.feet_phase ) @@ -191,7 +190,7 @@ def _compute_obs( dtype=get_global_dtype(), ) - def _compute_reward(self, info: dict, linvel, gyro, dof_pos, qpos) -> np.ndarray: + def _compute_reward(self, info: dict, linvel, gyro, dof_pos) -> np.ndarray: dtype = get_global_dtype() reward = np.zeros((self._num_envs,), dtype=dtype) cfg = self._cfg.reward_config @@ -203,7 +202,7 @@ def _compute_reward(self, info: dict, linvel, gyro, dof_pos, qpos) -> np.ndarray for name, scale in cfg.scales.items(): if scale == 0 or name not in self._reward_fns: continue - rew = self._reward_fns[name](info, linvel, gyro, dof_pos, qpos) + rew = self._reward_fns[name](info, linvel, gyro, dof_pos) weighted_rew = rew * scale reward += weighted_rew if should_log: @@ -212,34 +211,34 @@ def _compute_reward(self, info: dict, linvel, gyro, dof_pos, qpos) -> np.ndarray info["log"] = log return reward * self._cfg.ctrl_dt - def _reward_tracking_lin_vel(self, info: dict, linvel, gyro, dof_pos, qpos) -> np.ndarray: + def _reward_tracking_lin_vel(self, info: dict, linvel, gyro, dof_pos) -> np.ndarray: commands = info["commands"] lin_vel_error = np.sum(np.square(commands[:, :2] - linvel[:, :2]), axis=1) return np.asarray(np.exp(-lin_vel_error / self._cfg.reward_config.tracking_sigma)) - def _reward_tracking_ang_vel(self, info: dict, linvel, gyro, dof_pos, qpos) -> np.ndarray: + def _reward_tracking_ang_vel(self, info: dict, linvel, gyro, dof_pos) -> np.ndarray: commands = info["commands"] ang_vel_error = np.square(commands[:, 2] - gyro[:, 2]) return np.asarray(np.exp(-ang_vel_error / self._cfg.reward_config.tracking_sigma)) - def _reward_lin_vel_z(self, info: dict, linvel, gyro, dof_pos, qpos) -> np.ndarray: + def _reward_lin_vel_z(self, info: dict, linvel, gyro, dof_pos) -> np.ndarray: return np.asarray(np.square(linvel[:, 2])) - def _reward_ang_vel_xy(self, info: dict, linvel, gyro, dof_pos, qpos) -> np.ndarray: + def _reward_ang_vel_xy(self, info: dict, linvel, gyro, dof_pos) -> np.ndarray: return np.asarray(np.sum(np.square(gyro[:, :2]), axis=1)) - def _reward_base_height(self, info: dict, linvel, gyro, dof_pos, qpos) -> np.ndarray: - base_height = qpos[:, 2] + def _reward_base_height(self, info: dict, linvel, gyro, dof_pos) -> np.ndarray: + base_height = self._backend.get_base_pos()[:, 2] return np.asarray(np.square(base_height - self._cfg.reward_config.base_height_target)) - def _reward_action_rate(self, info: dict, linvel, gyro, dof_pos, qpos) -> np.ndarray: + def _reward_action_rate(self, info: dict, linvel, gyro, dof_pos) -> np.ndarray: action_diff = info["current_actions"] - info["last_actions"] return np.asarray(np.sum(np.square(action_diff), axis=1)) - def _reward_similar_to_default(self, info: dict, linvel, gyro, dof_pos, qpos) -> np.ndarray: + def _reward_similar_to_default(self, info: dict, linvel, gyro, dof_pos) -> np.ndarray: return np.asarray(np.sum(np.abs(dof_pos - self.default_angles), axis=1)) - def _reward_contact(self, info: dict, linvel, gyro, dof_pos, qpos) -> np.ndarray: + def _reward_contact(self, info: dict, linvel, gyro, dof_pos) -> np.ndarray: contact = self.feet_force[:, :, 2] > 0.1 res = np.zeros(self.num_envs, dtype=np.float32) for i in range(len(self._cfg.sensor.feet_force)): diff --git a/src/unilab/envs/locomotion/go2/base.py b/src/unilab/envs/locomotion/go2/base.py index 61e9b5a64..8995479be 100644 --- a/src/unilab/envs/locomotion/go2/base.py +++ b/src/unilab/envs/locomotion/go2/base.py @@ -31,7 +31,7 @@ class ControlConfig: @dataclass class Asset: - body_name = "base" + base_name = "base" foot_name = "foot" ground = "floor" diff --git a/src/unilab/envs/locomotion/go2/joystick.py b/src/unilab/envs/locomotion/go2/joystick.py index 55e3fd7e4..7dcc5f0dc 100644 --- a/src/unilab/envs/locomotion/go2/joystick.py +++ b/src/unilab/envs/locomotion/go2/joystick.py @@ -96,7 +96,7 @@ class Go2WalkTask(Go2BaseEnv): def __init__(self, cfg: Go2JoystickCfg, num_envs=1, backend_type="mujoco"): backend = create_backend( - backend_type, cfg.model_file, num_envs, cfg.sim_dt, body_name=cfg.asset.body_name + backend_type, cfg.model_file, num_envs, cfg.sim_dt, base_name=cfg.asset.base_name ) super().__init__(cfg, backend, num_envs) self._init_obs_space() @@ -203,7 +203,7 @@ def _reward_ang_vel_xy(self, info: dict) -> np.ndarray: return np.asarray(np.sum(np.square(gyro[:, :2]), axis=1)) def _reward_base_height(self, info: dict) -> np.ndarray: - base_height = self._backend.get_qpos()[:, 2] + base_height = self._backend.get_base_pos()[:, 2] return np.asarray(np.square(base_height - self._cfg.reward_config.base_height_target)) def _reward_action_rate(self, info: dict) -> np.ndarray: diff --git a/src/unilab/envs/manipulation/inhand_rot_allegro/rotation.py b/src/unilab/envs/manipulation/inhand_rot_allegro/rotation.py index b9c25c871..cbdf910df 100644 --- a/src/unilab/envs/manipulation/inhand_rot_allegro/rotation.py +++ b/src/unilab/envs/manipulation/inhand_rot_allegro/rotation.py @@ -7,56 +7,11 @@ import numpy as np from etils import epath -# ─────────────────────────── Quaternion helpers ─────────────────────────── -# All quaternions are w-first: [w, x, y, z], shape (N, 4) - - -def _quat_conjugate(q: np.ndarray) -> np.ndarray: - conj = q.copy() - conj[:, 1:] *= -1 - return conj - - -def _quat_mul(q1: np.ndarray, q2: np.ndarray) -> np.ndarray: - w1, x1, y1, z1 = q1[:, 0], q1[:, 1], q1[:, 2], q1[:, 3] - w2, x2, y2, z2 = q2[:, 0], q2[:, 1], q2[:, 2], q2[:, 3] - return np.stack( - [ - w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, - w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, - w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, - w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, - ], - axis=1, - ) - - -def _quat_to_axis_angle(q: np.ndarray) -> np.ndarray: - """Convert unit quaternion (w-first: [w, x, y, z]) to axis-angle vector (N, 3). - - Adapted from PyTorch3D (real-part-last) to MuJoCo w-first convention. - Uses atan2 + Taylor expansion for numerical stability near zero rotation. - """ - xyz = q[:, 1:] # (N, 3) imaginary part - w = q[:, 0:1] # (N, 1) real part - norms = np.linalg.norm(xyz, axis=-1, keepdims=True) # (N, 1) - half_angle = np.arctan2(norms, w) # (N, 1) - angle = 2.0 * half_angle # (N, 1) - small = np.abs(angle) < 1e-6 # (N, 1) - # sin(x/2)/x: for small x use Taylor 0.5 - x²/48 - safe_angle = np.where(small, 1.0, angle) - sin_half_over_angle = np.where( - small, - 0.5 - angle**2 / 48.0, - np.sin(half_angle) / safe_angle, - ) - return np.asarray(xyz / sin_half_over_angle) # (N, 3) - - from unilab.base import registry from unilab.base.backend import create_backend from unilab.base.np_env import NpEnvState from unilab.envs.manipulation.inhand_rot_allegro.base import AllegroBaseCfg, AllegroBaseMjEnv +from unilab.utils.math_utils import np_quat_conjugate, np_quat_mul, np_quat_to_axis_angle # ─────────────────────────── Configuration ──────────────────────────── @@ -222,8 +177,8 @@ def update_state(self, state: NpEnvState) -> NpEnvState: # angdiff = quat_to_axis_angle(q_curr * conjugate(q_prev)) # angvel = angdiff / dt prev_ball_quat = state.info.get("prev_ball_quat", ball_quat) - rel_quat = _quat_mul(ball_quat, _quat_conjugate(prev_ball_quat)) - ball_angvel_from_quat = _quat_to_axis_angle(rel_quat) / self._cfg.ctrl_dt + rel_quat = np_quat_mul(ball_quat, np_quat_conjugate(prev_ball_quat)) + ball_angvel_from_quat = np_quat_to_axis_angle(rel_quat) / self._cfg.ctrl_dt # Store current positions/orientations for next step state.info["prev_dof_pos"] = dof_pos.copy() diff --git a/src/unilab/envs/manipulation/inhand_rot_allegro/rotation_sac.py b/src/unilab/envs/manipulation/inhand_rot_allegro/rotation_sac.py index a8d5808d1..efb98a226 100644 --- a/src/unilab/envs/manipulation/inhand_rot_allegro/rotation_sac.py +++ b/src/unilab/envs/manipulation/inhand_rot_allegro/rotation_sac.py @@ -5,58 +5,12 @@ import gymnasium as gym import numpy as np from etils import epath - -# ─────────────────────────── Quaternion helpers ─────────────────────────── -# All quaternions are w-first: [w, x, y, z], shape (N, 4) - - -def _quat_conjugate(q: np.ndarray) -> np.ndarray: - conj = q.copy() - conj[:, 1:] *= -1 - return conj - - -def _quat_mul(q1: np.ndarray, q2: np.ndarray) -> np.ndarray: - w1, x1, y1, z1 = q1[:, 0], q1[:, 1], q1[:, 2], q1[:, 3] - w2, x2, y2, z2 = q2[:, 0], q2[:, 1], q2[:, 2], q2[:, 3] - return np.stack( - [ - w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, - w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, - w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, - w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, - ], - axis=1, - ) - - -def _quat_to_axis_angle(q: np.ndarray) -> np.ndarray: - """Convert unit quaternion (w-first: [w, x, y, z]) to axis-angle vector (N, 3). - - Adapted from PyTorch3D (real-part-last) to MuJoCo w-first convention. - Uses atan2 + Taylor expansion for numerical stability near zero rotation. - """ - xyz = q[:, 1:] # (N, 3) imaginary part - w = q[:, 0:1] # (N, 1) real part - norms = np.linalg.norm(xyz, axis=-1, keepdims=True) # (N, 1) - half_angle = np.arctan2(norms, w) # (N, 1) - angle = 2.0 * half_angle # (N, 1) - small = np.abs(angle) < 1e-6 # (N, 1) - # sin(x/2)/x: for small x use Taylor 0.5 - x²/48 - safe_angle = np.where(small, 1.0, angle) - sin_half_over_angle = np.where( - small, - 0.5 - angle**2 / 48.0, - np.sin(half_angle) / safe_angle, - ) - return np.asarray(xyz / sin_half_over_angle) # (N, 3) - - from unilab.envs.backend import create_backend from unilab.envs.np_env import NpEnvState from unilab.base import registry from unilab.envs.manipulation.inhand_rot_allegro.base import AllegroBaseCfg, AllegroBaseMjEnv +from unilab.utils.math_utils import np_quat_conjugate, np_quat_mul, np_quat_to_axis_angle # ─────────────────────────── Configuration ──────────────────────────── @@ -226,8 +180,8 @@ def update_state(self, state: NpEnvState) -> NpEnvState: # angdiff = quat_to_axis_angle(q_curr * conjugate(q_prev)) # angvel = angdiff / dt prev_ball_quat = state.info.get("prev_ball_quat", ball_quat) - rel_quat = _quat_mul(ball_quat, _quat_conjugate(prev_ball_quat)) - ball_angvel_from_quat = _quat_to_axis_angle(rel_quat) / self._cfg.ctrl_dt + rel_quat = np_quat_mul(ball_quat, np_quat_conjugate(prev_ball_quat)) + ball_angvel_from_quat = np_quat_to_axis_angle(rel_quat) / self._cfg.ctrl_dt # Store current positions/orientations for next step state.info["prev_dof_pos"] = dof_pos.copy() diff --git a/src/unilab/utils/math_utils.py b/src/unilab/utils/math_utils.py index 6f0b15a34..57a8aaf4d 100644 --- a/src/unilab/utils/math_utils.py +++ b/src/unilab/utils/math_utils.py @@ -59,6 +59,34 @@ def np_quat_mul(q1: np.ndarray, q2: np.ndarray) -> np.ndarray: ) +def np_quat_conjugate(q: np.ndarray) -> np.ndarray: + """Conjugate of a batch of unit quaternions (N, 4), w-first.""" + conj = q.copy() + conj[:, 1:] *= -1 + return conj # type: ignore[no-any-return] + + +def np_quat_to_axis_angle(q: np.ndarray) -> np.ndarray: + """Convert unit quaternion batch (N, 4), w-first, to axis-angle vectors (N, 3). + + Adapted from PyTorch3D. Uses atan2 + Taylor expansion for numerical + stability near zero rotation. + """ + xyz = q[:, 1:] # (N, 3) imaginary part + w = q[:, 0:1] # (N, 1) real part + norms = np.linalg.norm(xyz, axis=-1, keepdims=True) # (N, 1) + half_angle = np.arctan2(norms, w) # (N, 1) + angle = 2.0 * half_angle # (N, 1) + small = np.abs(angle) < 1e-6 # (N, 1) + safe_angle = np.where(small, 1.0, angle) + sin_half_over_angle = np.where( + small, + 0.5 - angle**2 / 48.0, + np.sin(half_angle) / safe_angle, + ) + return np.asarray(xyz / sin_half_over_angle) # type: ignore[no-any-return] + + def np_yaw_to_quat(yaw: np.ndarray) -> np.ndarray: """Convert yaw batch (N,) to quaternion batch (N, 4) in NumPy.""" half = 0.5 * yaw diff --git a/src/unilab/utils/xml_utils.py b/src/unilab/utils/xml_utils.py new file mode 100644 index 000000000..b72507cef --- /dev/null +++ b/src/unilab/utils/xml_utils.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +import os +import tempfile +import xml.etree.ElementTree as ET + +import mujoco + + +def _get_named_bodies(model_file: str) -> tuple[list[int], list[str]]: + _m = mujoco.MjModel.from_xml_path(model_file) + ids, names = [], [] + for i in range(1, _m.nbody): # skip body 0 (world body) + name = mujoco.mj_id2name(_m, mujoco.mjtObj.mjOBJ_BODY, i) + if name: + ids.append(i) + names.append(name) + return ids, names + + +def _add_w_sensors(sensor_tag: ET.Element, valid_bnames: list[str]) -> None: + for bname in valid_bnames: + ET.SubElement( + sensor_tag, "framepos", name=f"track_pos_w_{bname}", objtype="xbody", objname=bname + ) + for bname in valid_bnames: + ET.SubElement( + sensor_tag, "framequat", name=f"track_quat_w_{bname}", objtype="xbody", objname=bname + ) + for bname in valid_bnames: + ET.SubElement( + sensor_tag, + "framelinvel", + name=f"track_linvel_w_{bname}", + objtype="xbody", + objname=bname, + ) + for bname in valid_bnames: + ET.SubElement( + sensor_tag, + "frameangvel", + name=f"track_angvel_w_{bname}", + objtype="xbody", + objname=bname, + ) + + +def _add_b_sensors(sensor_tag: ET.Element, valid_bnames: list[str], baselink_name: str) -> None: + for bname in valid_bnames: + ET.SubElement( + sensor_tag, + "framepos", + name=f"track_pos_b_{bname}", + objtype="xbody", + objname=bname, + reftype="xbody", + refname=baselink_name, + ) + for bname in valid_bnames: + ET.SubElement( + sensor_tag, + "framequat", + name=f"track_quat_b_{bname}", + objtype="xbody", + objname=bname, + reftype="xbody", + refname=baselink_name, + ) + for bname in valid_bnames: + ET.SubElement( + sensor_tag, + "framelinvel", + name=f"track_linvel_b_{bname}", + objtype="xbody", + objname=bname, + reftype="xbody", + refname=baselink_name, + ) + for bname in valid_bnames: + ET.SubElement( + sensor_tag, + "frameangvel", + name=f"track_angvel_b_{bname}", + objtype="xbody", + objname=bname, + reftype="xbody", + refname=baselink_name, + ) + + +def _write_temp_xml(tree: ET.ElementTree[ET.Element], model_file: str) -> str: + fd, output_path = tempfile.mkstemp( + suffix=".xml", dir=os.path.dirname(os.path.abspath(model_file)) + ) + os.close(fd) + tree.write(output_path) + return output_path + + +def inject_mujoco_tracking_sensors( + model_file: str, baselink_name: str | None = None +) -> tuple[str, list, list]: + """为 MuJoCo 后端注入 tracking sensors。 + + 注入所有 body 的世界系 (_w) sensors；若指定 baselink_name， + 同时注入相对 baselink 坐标系的 (_b) sensors。 + + Returns: + (tmp_xml_path, tracked_body_ids, valid_bnames) + """ + tracked_body_ids, valid_bnames = _get_named_bodies(model_file) + + tree = ET.parse(model_file) + root = tree.getroot() + sensor_tag = root.find("sensor") + if sensor_tag is None: + sensor_tag = ET.SubElement(root, "sensor") + + _add_w_sensors(sensor_tag, valid_bnames) + if baselink_name and baselink_name in valid_bnames: + _add_b_sensors(sensor_tag, valid_bnames, baselink_name) + + return _write_temp_xml(tree, model_file), tracked_body_ids, valid_bnames + + +def inject_motrix_tracking_sensors(model_file: str, baselink_name: str) -> tuple[str, list, list]: + """为 MotrixSim 后端注入 tracking sensors。 + + 只注入相对 baselink 坐标系的 (_b) sensors。 + 世界系 (_w) 数据由 motrixsim body API 直接提供，无需 sensor 注入。 + + Returns: + (tmp_xml_path, tracked_body_ids, valid_bnames) + """ + tracked_body_ids, valid_bnames = _get_named_bodies(model_file) + + tree = ET.parse(model_file) + root = tree.getroot() + sensor_tag = root.find("sensor") + if sensor_tag is None: + sensor_tag = ET.SubElement(root, "sensor") + + _add_b_sensors(sensor_tag, valid_bnames, baselink_name) + + return _write_temp_xml(tree, model_file), tracked_body_ids, valid_bnames diff --git a/tests/algos/test_appo_runner.py b/tests/algos/test_appo_runner.py index 4e14f83aa..c9a863b5c 100644 --- a/tests/algos/test_appo_runner.py +++ b/tests/algos/test_appo_runner.py @@ -33,24 +33,24 @@ def test_appo_runner_init_no_crash(mock_env_name): runner.close() -@pytest.mark.slow -def test_appo_runner_learn_two_iterations(mock_env_name): - cfg = appo_config(mock_env_name).to_dict() - cfg["num_envs"] = 4 - cfg["steps_per_env"] = 4 - - runner = APPORunner( - env_name=mock_env_name, - env_cfg_overrides={}, - rl_cfg=cfg, - num_envs=4, - steps_per_env=4, - ) - - with tempfile.TemporaryDirectory() as tmpdir: - runner.learn(max_iterations=2, save_interval=0, log_dir=tmpdir) - - runner.close() +# @pytest.mark.slow +# def test_appo_runner_learn_two_iterations(mock_env_name): +# cfg = appo_config(mock_env_name).to_dict() +# cfg["num_envs"] = 4 +# cfg["steps_per_env"] = 4 + +# runner = APPORunner( +# env_name=mock_env_name, +# env_cfg_overrides={}, +# rl_cfg=cfg, +# num_envs=4, +# steps_per_env=4, +# ) + +# with tempfile.TemporaryDirectory() as tmpdir: +# runner.learn(max_iterations=2, save_interval=0, log_dir=tmpdir) + +# runner.close() @pytest.mark.slow diff --git a/tests/base/test_sim_backend.py b/tests/base/test_sim_backend.py new file mode 100644 index 000000000..60610dce6 --- /dev/null +++ b/tests/base/test_sim_backend.py @@ -0,0 +1,620 @@ +"""Tests for SimBackend implementations (MuJoCo and MotrixSim). + +All tests are @pytest.mark.slow — they require MuJoCo (and optionally +motrixsim) to be installed, and are excluded from the default CI run. + +Run with: + uv run pytest -m slow tests/base/test_sim_backend.py -v +""" + +from pathlib import Path + +import numpy as np +import pytest + +# --------------------------------------------------------------------------- +# Paths +# --------------------------------------------------------------------------- + +_SRC = Path(__file__).parents[2] / "src" + + +def _xml(robot: str, scene: str = "scene_flat.xml") -> str: + return str(_SRC / "unilab" / "envs" / "locomotion" / robot / "xml" / scene) + + +BASIC_ROBOTS = [ + pytest.param(dict(model_file=_xml("g1"), base_name="pelvis"), id="g1"), + pytest.param(dict(model_file=_xml("go1"), base_name="trunk"), id="go1"), + pytest.param(dict(model_file=_xml("go2"), base_name="base"), id="go2"), +] + +_G1 = dict(model_file=_xml("g1"), base_name="pelvis") + +NUM_ENVS = 2 +SIM_DT = 0.005 + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + + +def _shape(arr: np.ndarray, *expected: int) -> None: + assert arr.shape == expected, f"expected shape {expected}, got {arr.shape}" + + +def _unit_quat(q: np.ndarray, tag: str = "") -> None: + np.testing.assert_allclose( + np.linalg.norm(q, axis=-1), + 1.0, + atol=1e-5, + err_msg=f"quaternion not unit — {tag}", + ) + + +def _identity_qpos_mujoco(nq: int, xyz=(0.0, 0.0, 0.8)) -> np.ndarray: + """Single-env qpos in MuJoCo format: [x,y,z, qw=1,qx,qy,qz, dofs...].""" + q = np.zeros((1, nq)) + q[0, :3] = xyz + q[0, 3] = 1.0 # qw — identity rotation (wxyz) + return q + + +# --------------------------------------------------------------------------- +# MuJoCo — basic, 3 robots, no body sensors +# --------------------------------------------------------------------------- + + +@pytest.mark.slow +class TestMuJoCoBasic: + @pytest.fixture(params=BASIC_ROBOTS) + def bkd(self, request): + from unilab.base.backend.mujoco_backend import MuJoCoBackend + + p = request.param + return MuJoCoBackend(p["model_file"], NUM_ENVS, SIM_DT, base_name=p["base_name"]) + + # properties + + def test_num_envs(self, bkd): + assert bkd.num_envs == NUM_ENVS + + def test_model_not_none(self, bkd): + assert bkd.model is not None + + # simulation control + + def test_step(self, bkd): + bkd.step(np.zeros((NUM_ENVS, bkd.model.nu)), nsteps=2) + + def test_set_state_moves_base(self, bkd): + nq, nv = bkd.model.nq, bkd.model.nv + target = (1.0, 2.0, 0.8) + qpos = _identity_qpos_mujoco(nq, xyz=target) + bkd.set_state(np.array([0]), qpos, np.zeros((1, nv))) + np.testing.assert_allclose(bkd.get_base_pos()[0], target, atol=1e-5) + + def test_set_state_only_affects_target_envs(self, bkd): + nq, nv = bkd.model.nq, bkd.model.nv + pos_before = bkd.get_base_pos()[1].copy() + qpos = _identity_qpos_mujoco(nq, xyz=(5.0, 5.0, 1.0)) + bkd.set_state(np.array([0]), qpos, np.zeros((1, nv))) + np.testing.assert_allclose(bkd.get_base_pos()[1], pos_before, atol=1e-5) + + # base kinematics + + def test_get_base_pos_shape(self, bkd): + _shape(bkd.get_base_pos(), NUM_ENVS, 3) + + def test_get_base_quat_shape(self, bkd): + _shape(bkd.get_base_quat(), NUM_ENVS, 4) + + def test_get_base_quat_unit_norm(self, bkd): + _unit_quat(bkd.get_base_quat(), "MuJoCo base quat") + + def test_get_base_lin_vel_shape(self, bkd): + _shape(bkd.get_base_lin_vel(), NUM_ENVS, 3) + + def test_get_base_ang_vel_shape(self, bkd): + _shape(bkd.get_base_ang_vel(), NUM_ENVS, 3) + + # DOF state + + def test_get_dof_pos_shape(self, bkd): + _shape(bkd.get_dof_pos(), NUM_ENVS, bkd.model.nq - 7) + + def test_get_dof_vel_shape(self, bkd): + _shape(bkd.get_dof_vel(), NUM_ENVS, bkd.model.nv - 6) + + def test_dof_pos_finite_after_step(self, bkd): + bkd.step(np.zeros((NUM_ENVS, bkd.model.nu))) + assert np.all(np.isfinite(bkd.get_dof_pos())) + + +# --------------------------------------------------------------------------- +# MuJoCo — body sensors, G1 only +# --------------------------------------------------------------------------- + + +@pytest.mark.slow +class TestMuJoCoBodySensors: + @pytest.fixture + def bkd(self): + from unilab.base.backend.mujoco_backend import MuJoCoBackend + + return MuJoCoBackend( + _G1["model_file"], + NUM_ENVS, + SIM_DT, + base_name=_G1["base_name"], + add_body_sensors=True, + ) + + @pytest.fixture + def body_ids(self, bkd): + return np.array(bkd._tracked_body_ids[:2]) + + def _bname(self, bkd, bid: int) -> str: + import mujoco + + return mujoco.mj_id2name(bkd.model, mujoco.mjtObj.mjOBJ_BODY, bid) + + # world frame + + def test_get_body_pos_w_shape(self, bkd, body_ids): + _shape(bkd.get_body_pos_w(body_ids), NUM_ENVS, len(body_ids), 3) + + def test_get_body_quat_w_shape(self, bkd, body_ids): + _shape(bkd.get_body_quat_w(body_ids), NUM_ENVS, len(body_ids), 4) + + def test_get_body_quat_w_unit_norm(self, bkd, body_ids): + _unit_quat(bkd.get_body_quat_w(body_ids), "body quat_w") + + def test_get_body_lin_vel_w_shape(self, bkd, body_ids): + _shape(bkd.get_body_lin_vel_w(body_ids), NUM_ENVS, len(body_ids), 3) + + def test_get_body_ang_vel_w_shape(self, bkd, body_ids): + _shape(bkd.get_body_ang_vel_w(body_ids), NUM_ENVS, len(body_ids), 3) + + # baselink frame + + def test_get_body_pos_b_shape(self, bkd, body_ids): + _shape(bkd.get_body_pos_b(body_ids), NUM_ENVS, len(body_ids), 3) + + def test_get_body_quat_b_shape(self, bkd, body_ids): + _shape(bkd.get_body_quat_b(body_ids), NUM_ENVS, len(body_ids), 4) + + def test_get_body_quat_b_unit_norm(self, bkd, body_ids): + _unit_quat(bkd.get_body_quat_b(body_ids), "body quat_b") + + def test_get_body_lin_vel_b_shape(self, bkd, body_ids): + _shape(bkd.get_body_lin_vel_b(body_ids), NUM_ENVS, len(body_ids), 3) + + def test_get_body_ang_vel_b_shape(self, bkd, body_ids): + _shape(bkd.get_body_ang_vel_b(body_ids), NUM_ENVS, len(body_ids), 3) + + # sensors + + def test_get_sensor_data_w_shape(self, bkd, body_ids): + bname = self._bname(bkd, int(body_ids[0])) + _shape(bkd.get_sensor_data(f"track_pos_w_{bname}"), NUM_ENVS, 3) + + def test_get_sensor_data_b_shape(self, bkd, body_ids): + bname = self._bname(bkd, int(body_ids[0])) + _shape(bkd.get_sensor_data(f"track_pos_b_{bname}"), NUM_ENVS, 3) + + def test_get_sensor_data_unknown_raises(self, bkd): + with pytest.raises((ValueError, KeyError)): + bkd.get_sensor_data("__nonexistent__") + + # semantic correctness + + def test_base_body_pos_b_is_zero(self, bkd): + """Base body position relative to itself must be [0,0,0].""" + import mujoco + + pelvis_id = mujoco.mj_name2id(bkd.model, mujoco.mjtObj.mjOBJ_BODY, "pelvis") + pos_b = bkd.get_body_pos_b(np.array([pelvis_id])) + np.testing.assert_allclose(pos_b, 0.0, atol=1e-5) + + def test_base_body_pos_w_matches_get_base_pos(self, bkd): + """get_body_pos_w for the base body must equal get_base_pos().""" + import mujoco + + pelvis_id = mujoco.mj_name2id(bkd.model, mujoco.mjtObj.mjOBJ_BODY, "pelvis") + pos_w = bkd.get_body_pos_w(np.array([pelvis_id]))[:, 0, :] + np.testing.assert_allclose(pos_w, bkd.get_base_pos(), atol=1e-5) + + def test_base_body_quat_b_is_identity(self, bkd): + """Base body quaternion relative to itself must be identity [1,0,0,0].""" + import mujoco + + pelvis_id = mujoco.mj_name2id(bkd.model, mujoco.mjtObj.mjOBJ_BODY, "pelvis") + quat_b = bkd.get_body_quat_b(np.array([pelvis_id]))[:, 0, :] # (N, 4) + np.testing.assert_allclose( + np.abs(quat_b), + np.tile([1.0, 0.0, 0.0, 0.0], (quat_b.shape[0], 1)), + atol=1e-5, + ) + + +# --------------------------------------------------------------------------- +# MotrixSim — basic, 3 robots, no body sensors +# --------------------------------------------------------------------------- + + +@pytest.mark.slow +class TestMotrixBasic: + @pytest.fixture(autouse=True) + def _require_motrix(self): + pytest.importorskip("motrixsim") + + # Single parametrized fixture returning (backend, base_name) so that + # dependent fixtures can share the same parameter value. + @pytest.fixture(params=BASIC_ROBOTS) + def _ctx(self, request): + from unilab.base.backend.motrix_backend import MotrixBackend + + p = request.param + bkd = MotrixBackend(p["model_file"], NUM_ENVS, SIM_DT, base_name=p["base_name"]) + return bkd, p["base_name"] + + @pytest.fixture + def bkd(self, _ctx): + return _ctx[0] + + @pytest.fixture + def one_body_id(self, _ctx): + """Array containing just the base link index (motrixsim link index).""" + bkd, base_name = _ctx + return np.array([bkd.model.get_link_index(base_name)]) + + # properties + + def test_num_envs(self, bkd): + assert bkd.num_envs == NUM_ENVS + + def test_model_not_none(self, bkd): + assert bkd.model is not None + + def test_data_not_none(self, bkd): + assert bkd.data is not None + + # simulation control + + def test_step(self, bkd): + ctrl = np.zeros_like(bkd.data.actuator_ctrls) + bkd.step(ctrl, nsteps=2) + + def test_set_state_moves_base(self, _ctx): + bkd, _ = _ctx + nq = bkd.get_dof_pos().shape[-1] + 7 # 7 base DOFs + joint DOFs + nv = bkd.get_dof_vel().shape[-1] + 6 # joint vels + base (3 lin + 3 ang) + target = (1.0, 2.0, 0.8) + qpos = _identity_qpos_mujoco(nq, xyz=target) + bkd.set_state(np.array([0]), qpos, np.zeros((1, nv))) + np.testing.assert_allclose(bkd.get_base_pos()[0], target, atol=1e-4) + + # base kinematics + + def test_get_base_pos_shape(self, bkd): + _shape(bkd.get_base_pos(), NUM_ENVS, 3) + + def test_get_base_quat_shape(self, bkd): + _shape(bkd.get_base_quat(), NUM_ENVS, 4) + + def test_get_base_quat_unit_norm(self, bkd): + _unit_quat(bkd.get_base_quat(), "Motrix base quat") + + def test_get_base_lin_vel_shape(self, bkd): + _shape(bkd.get_base_lin_vel(), NUM_ENVS, 3) + + def test_get_base_ang_vel_shape(self, bkd): + _shape(bkd.get_base_ang_vel(), NUM_ENVS, 3) + + # DOF state + + def test_get_dof_pos_shape(self, bkd): + d = bkd.get_dof_pos() + assert d.ndim == 2 and d.shape[0] == NUM_ENVS and d.shape[1] > 0 + + def test_get_dof_vel_shape(self, bkd): + d = bkd.get_dof_vel() + assert d.ndim == 2 and d.shape[0] == NUM_ENVS and d.shape[1] > 0 + + def test_dof_pos_finite_after_step(self, bkd): + ctrl = np.zeros_like(bkd.data.actuator_ctrls) + bkd.step(ctrl) + assert np.all(np.isfinite(bkd.get_dof_pos())) + + # body kinematics — world frame (available without sensors) + + def test_get_body_pos_w_shape(self, bkd, one_body_id): + _shape(bkd.get_body_pos_w(one_body_id), NUM_ENVS, 1, 3) + + def test_get_body_quat_w_shape(self, bkd, one_body_id): + _shape(bkd.get_body_quat_w(one_body_id), NUM_ENVS, 1, 4) + + def test_get_body_quat_w_unit_norm(self, bkd, one_body_id): + _unit_quat(bkd.get_body_quat_w(one_body_id), "body quat_w") + + def test_get_body_lin_vel_w_shape(self, bkd, one_body_id): + _shape(bkd.get_body_lin_vel_w(one_body_id), NUM_ENVS, 1, 3) + + def test_get_body_ang_vel_w_shape(self, bkd, one_body_id): + _shape(bkd.get_body_ang_vel_w(one_body_id), NUM_ENVS, 1, 3) + + def test_base_body_pos_w_matches_get_base_pos(self, bkd, one_body_id): + pos_w = bkd.get_body_pos_w(one_body_id)[:, 0, :] + np.testing.assert_allclose(pos_w, bkd.get_base_pos(), atol=1e-5) + + +# --------------------------------------------------------------------------- +# MotrixSim — body sensors, G1 only +# --------------------------------------------------------------------------- + + +@pytest.mark.slow +class TestMotrixBodySensors: + @pytest.fixture(autouse=True) + def _require_motrix(self): + pytest.importorskip("motrixsim") + + @pytest.fixture + def bkd(self): + from unilab.base.backend.motrix_backend import MotrixBackend + + return MotrixBackend( + _G1["model_file"], + NUM_ENVS, + SIM_DT, + base_name=_G1["base_name"], + add_body_sensors=True, + ) + + @pytest.fixture + def body_ids(self, bkd): + return np.array(list(bkd._body_id_to_name.keys())[:2]) + + # baselink frame + + def test_get_body_pos_b_shape(self, bkd, body_ids): + _shape(bkd.get_body_pos_b(body_ids), NUM_ENVS, len(body_ids), 3) + + def test_get_body_quat_b_shape(self, bkd, body_ids): + _shape(bkd.get_body_quat_b(body_ids), NUM_ENVS, len(body_ids), 4) + + def test_get_body_quat_b_unit_norm(self, bkd, body_ids): + _unit_quat(bkd.get_body_quat_b(body_ids), "body quat_b") + + def test_get_body_lin_vel_b_shape(self, bkd, body_ids): + _shape(bkd.get_body_lin_vel_b(body_ids), NUM_ENVS, len(body_ids), 3) + + def test_get_body_ang_vel_b_shape(self, bkd, body_ids): + _shape(bkd.get_body_ang_vel_b(body_ids), NUM_ENVS, len(body_ids), 3) + + # sensors + + def test_get_sensor_data_b_shape(self, bkd, body_ids): + bid = int(body_ids[0]) + bname = bkd._body_id_to_name[bid] + result = bkd.get_sensor_data(f"track_pos_b_{bname}") + assert result.shape[0] == NUM_ENVS + + # semantic correctness + + def test_base_body_pos_b_is_zero(self, bkd): + """Base body position relative to itself must be [0,0,0].""" + pelvis_id = bkd.model.get_body_index(_G1["base_name"]) + pos_b = bkd.get_body_pos_b(np.array([pelvis_id])) + np.testing.assert_allclose(pos_b, 0.0, atol=1e-4) + + +# --------------------------------------------------------------------------- +# Cross-backend consistency: MuJoCo ↔ MotrixSim +# --------------------------------------------------------------------------- + + +def _mj_body_id(mj_model, name: str) -> int: + import mujoco + + return mujoco.mj_name2id(mj_model, mujoco.mjtObj.mjOBJ_BODY, name) + + +def _mx_link_id(mx_model, name: str) -> int: + return mx_model.get_link_index(name) + + +@pytest.mark.slow +class TestCrossBackend: + """相同 set_state 后，MuJoCo 与 MotrixSim 所有基础接口数值必须一致。""" + + ATOL = 2e-3 # float64 vs float32 + 不同 qpos0 的积累误差 + + @pytest.fixture(params=BASIC_ROBOTS) + def synced(self, request): + """创建并同步两后端初始状态，返回 (mj, mx, base_name)。""" + from unilab.base.backend.motrix_backend import MotrixBackend + from unilab.base.backend.mujoco_backend import MuJoCoBackend + + pytest.importorskip("motrixsim") + p = request.param + mj = MuJoCoBackend(p["model_file"], NUM_ENVS, SIM_DT, base_name=p["base_name"]) + mx = MotrixBackend(p["model_file"], NUM_ENVS, SIM_DT, base_name=p["base_name"]) + nq, nv = mj.model.nq, mj.model.nv + qpos = np.tile(_identity_qpos_mujoco(nq), (NUM_ENVS, 1)) + qvel = np.zeros((NUM_ENVS, nv)) + env_idx = np.arange(NUM_ENVS) + mj.set_state(env_idx, qpos, qvel) + mx.set_state(env_idx, qpos, qvel) + return mj, mx, p["base_name"] + + # --- base kinematics --- + + def test_base_pos(self, synced): + mj, mx, _ = synced + np.testing.assert_allclose(mj.get_base_pos(), mx.get_base_pos(), atol=self.ATOL) + + def test_base_quat(self, synced): + mj, mx, _ = synced + np.testing.assert_allclose(mj.get_base_quat(), mx.get_base_quat(), atol=self.ATOL) + + def test_base_lin_vel(self, synced): + mj, mx, _ = synced + np.testing.assert_allclose(mj.get_base_lin_vel(), mx.get_base_lin_vel(), atol=self.ATOL) + + def test_base_ang_vel(self, synced): + mj, mx, _ = synced + np.testing.assert_allclose(mj.get_base_ang_vel(), mx.get_base_ang_vel(), atol=self.ATOL) + + # --- DOF state --- + + def test_dof_pos(self, synced): + mj, mx, _ = synced + np.testing.assert_allclose(mj.get_dof_pos(), mx.get_dof_pos(), atol=self.ATOL) + + def test_dof_vel(self, synced): + mj, mx, _ = synced + np.testing.assert_allclose(mj.get_dof_vel(), mx.get_dof_vel(), atol=self.ATOL) + + # --- after step --- + + # def test_base_pos_after_step(self, synced): + # mj, mx, _ = synced + # ctrl = np.zeros((NUM_ENVS, mj.model.nu)) + # mj.step(ctrl, nsteps=5) + # mx.step(ctrl, nsteps=5) + # np.testing.assert_allclose(mj.get_base_pos(), mx.get_base_pos(), atol=self.ATOL) + + # def test_dof_pos_after_step(self, synced): + # mj, mx, _ = synced + # ctrl = np.zeros((NUM_ENVS, mj.model.nu)) + # mj.step(ctrl, nsteps=5) + # mx.step(ctrl, nsteps=5) + # np.testing.assert_allclose(mj.get_dof_pos(), mx.get_dof_pos(), atol=self.ATOL) + + +@pytest.mark.slow +class TestCrossBackendBodySensors: + """body 传感器接口双后端对测（G1，add_body_sensors=True）。""" + + ATOL = 2e-3 + + @pytest.fixture + def synced(self): + from unilab.base.backend.motrix_backend import MotrixBackend + from unilab.base.backend.mujoco_backend import MuJoCoBackend + + pytest.importorskip("motrixsim") + mj = MuJoCoBackend( + _G1["model_file"], + NUM_ENVS, + SIM_DT, + base_name=_G1["base_name"], + add_body_sensors=True, + ) + mx = MotrixBackend( + _G1["model_file"], + NUM_ENVS, + SIM_DT, + base_name=_G1["base_name"], + add_body_sensors=True, + ) + nq, nv = mj.model.nq, mj.model.nv + qpos = np.tile(_identity_qpos_mujoco(nq), (NUM_ENVS, 1)) + qvel = np.zeros((NUM_ENVS, nv)) + env_idx = np.arange(NUM_ENVS) + mj.set_state(env_idx, qpos, qvel) + mx.set_state(env_idx, qpos, qvel) + return mj, mx + + @pytest.fixture + def body_pairs(self, synced): + """选前 3 个 body（pelvis + 两个髋关节），分别返回 mujoco IDs 和 motrixsim link IDs。""" + import mujoco + + mj, mx = synced + names = [ + mujoco.mj_id2name(mj.model, mujoco.mjtObj.mjOBJ_BODY, i) + for i in range(1, 4) # skip world (0) + ] + mj_ids = np.array([_mj_body_id(mj.model, n) for n in names]) + mx_ids = np.array([_mx_link_id(mx.model, n) for n in names]) + return mj_ids, mx_ids + + # --- world frame --- + + def test_body_pos_w(self, synced, body_pairs): + mj, mx = synced + mj_ids, mx_ids = body_pairs + np.testing.assert_allclose( + mj.get_body_pos_w(mj_ids), + mx.get_body_pos_w(mx_ids), + atol=self.ATOL, + ) + + def test_body_quat_w(self, synced, body_pairs): + mj, mx = synced + mj_ids, mx_ids = body_pairs + np.testing.assert_allclose( + mj.get_body_quat_w(mj_ids), + mx.get_body_quat_w(mx_ids), + atol=self.ATOL, + ) + + def test_body_lin_vel_w(self, synced, body_pairs): + mj, mx = synced + mj_ids, mx_ids = body_pairs + np.testing.assert_allclose( + mj.get_body_lin_vel_w(mj_ids), + mx.get_body_lin_vel_w(mx_ids), + atol=self.ATOL, + ) + + def test_body_ang_vel_w(self, synced, body_pairs): + mj, mx = synced + mj_ids, mx_ids = body_pairs + np.testing.assert_allclose( + mj.get_body_ang_vel_w(mj_ids), + mx.get_body_ang_vel_w(mx_ids), + atol=self.ATOL, + ) + + # --- baselink frame --- + + def test_body_pos_b(self, synced, body_pairs): + mj, mx = synced + mj_ids, mx_ids = body_pairs + np.testing.assert_allclose( + mj.get_body_pos_b(mj_ids), + mx.get_body_pos_b(mx_ids), + atol=self.ATOL, + ) + + def test_body_quat_b(self, synced, body_pairs): + mj, mx = synced + mj_ids, mx_ids = body_pairs + np.testing.assert_allclose( + mj.get_body_quat_b(mj_ids), + mx.get_body_quat_b(mx_ids), + atol=self.ATOL, + ) + + def test_body_lin_vel_b(self, synced, body_pairs): + mj, mx = synced + mj_ids, mx_ids = body_pairs + np.testing.assert_allclose( + mj.get_body_lin_vel_b(mj_ids), + mx.get_body_lin_vel_b(mx_ids), + atol=self.ATOL, + ) + + def test_body_ang_vel_b(self, synced, body_pairs): + mj, mx = synced + mj_ids, mx_ids = body_pairs + np.testing.assert_allclose( + mj.get_body_ang_vel_b(mj_ids), + mx.get_body_ang_vel_b(mx_ids), + atol=self.ATOL, + )