v1.1.4; Add ZMP margin metric

robogauge/tasks/gauge/base_gauge.py

@@ -138,13 +138,14 @@ class BaseGauge:
     
     def update_metrics(self, sim_data: SimData, goal_data: GoalData):
         if sim_data.n_step % int(self.cfg.metrics.metric_dt / sim_data.sim_dt + 1e-9) != 0:
-            return
+            return False
         metrics_results = {}
         for metric_name, metric_obj in zip(self.info['metric'], self.metrics):
             val = metric_obj(sim_data, goal_data)
             if metric_name not in ['visualization']:
                 metrics_results[metric_name] = val
         self.goals[self.goal_idx].update_metrics(metrics_results)
+        return True
     
     def reset_metrics(self):
         for metric in self.metrics:

robogauge/tasks/gauge/base_gauge_config.py

@@ -16,6 +16,7 @@ QUALITY_WEIGHTS  = {  # Weights for geometric average, to calculate quality scor
     'dof_power': 1,
     'orientation_stability': 1,
     'torque_smoothness': 1,
+    'zmp_margin': 1,
 }
 
 class BaseGaugeConfig(Config):
@@ -81,3 +82,12 @@ class BaseGaugeConfig(Config):
         class torque_smoothness:
             enabled = True
             scaling_factor = 30.0  # [Nm] scaling factor for torque smoothness metric
+
+        class zmp_margin:
+            enabled = True
+            contact_threshold = 1e-3  # [m] contact.dist <= threshold is treated as support contact
+            force_threshold = 1e-6  # [N] abs(Fz) below threshold skips ZMP and returns 1.0
+            draw_point = True  # Whether to draw the ZMP point in both viewer and offscreen render
+            draw_point_size = 0.03  # [m] ZMP point sphere radius in Mujoco visualization
+            draw_height_offset = 0.00  # [m] Lift the point slightly above the support plane for visibility
+            draw_point_rgba = [1.0, 0.85, 0.1, 1.0]  # RGBA color of the ZMP point

robogauge/tasks/gauge/goals/base_goal.py

@@ -49,10 +49,15 @@ class BaseGoal:
     def update_metrics(self, metrics: dict):
         """ Update step metrics for the current goal."""
         quality_score = 1.0
+        quality_weight_sum = 0.0
         for metric_name, value in metrics.items():
             self.goal_metrics[metric_name].append(value)
+            if metric_name not in QUALITY_WEIGHTS:
+                continue
             quality_score *= min(max(1e-9, value), 1.0) ** QUALITY_WEIGHTS[metric_name]
-        quality_score = quality_score ** (1.0 / sum(QUALITY_WEIGHTS.values()))
+            quality_weight_sum += QUALITY_WEIGHTS[metric_name]
+        if quality_weight_sum > 0.0:
+            quality_score = quality_score ** (1.0 / quality_weight_sum)
         self.goal_quality_scores.append(quality_score)
     
     @property

robogauge/tasks/gauge/metrics/__init__.py

@@ -2,4 +2,4 @@ from .base_metric import BaseMetric
 from .dof_metrics import DofLimitsMetric, DofPowerMetric
 from .visualization import VisualizationMetric
 from .vel_metrics import LinVelErrMetric, AngVelErrMetric
-from .stable_metric import OrientationStabilityMetric, TorqueSmoothnessMetric
+from .stable_metric import OrientationStabilityMetric, TorqueSmoothnessMetric, ZmpMarginMetric

robogauge/tasks/gauge/metrics/stable_metric.py

@@ -1,11 +1,11 @@
 # -*- coding: utf-8 -*-
 '''
-@File    : height_metric.py
+@File    : stable_metric.py
 @Time    : 2025/12/18 20:18:33
 @Author  : wty-yy
 @Version : 1.0
 @Blog    : https://wty-yy.github.io/
-@Desc    : Orientation Stability, Torque Smoothness Metric Implementation
+@Desc    : Stability-related metrics implementation
 '''
 import numpy as np
 
@@ -17,8 +17,8 @@ from robogauge.utils.logger import logger
 
 
 class OrientationStabilityMetric(BaseMetric):
-    """ Metric to log height stability. """
-    name = 'height_std_metric'
+    """ Metric to log body orientation stability. """
+    name = 'orientation_stability_metric'
 
     def __init__(self,
         robot_cfg: RobotConfig,
@@ -59,4 +59,209 @@ class TorqueSmoothnessMetric(BaseMetric):
         rms_value = np.sqrt(np.mean(np.square(torque_diff)))
         metric_value = 1.0 - rms_value / self.scaling_factor
         logger.log(rms_value, f'stable_metric/torque_rms_diff', step=sim_data.n_step)
-        return metric_value
+        return metric_value
+
+
+def _cross_2d(o: np.ndarray, a: np.ndarray, b: np.ndarray) -> float:
+    oa = a - o
+    ob = b - o
+    return float(oa[0] * ob[1] - oa[1] * ob[0])
+
+
+def _convex_hull_2d(points: np.ndarray) -> np.ndarray:
+    if points.shape[0] <= 1:
+        return points.copy()
+    pts = np.unique(points, axis=0)
+    if pts.shape[0] <= 1:
+        return pts
+    pts = pts[np.lexsort((pts[:, 1], pts[:, 0]))]
+
+    lower = []
+    for point in pts:
+        while len(lower) >= 2 and _cross_2d(lower[-2], lower[-1], point) <= 0.0:
+            lower.pop()
+        lower.append(point)
+
+    upper = []
+    for point in pts[::-1]:
+        while len(upper) >= 2 and _cross_2d(upper[-2], upper[-1], point) <= 0.0:
+            upper.pop()
+        upper.append(point)
+
+    hull = np.array(lower[:-1] + upper[:-1], dtype=np.float32)
+    return hull if hull.size > 0 else pts[:1]
+
+
+def _polygon_area_2d(points: np.ndarray) -> float:
+    if points.shape[0] < 3:
+        return 0.0
+    x = points[:, 0]
+    y = points[:, 1]
+    return 0.5 * abs(float(np.dot(x, np.roll(y, -1)) - np.dot(y, np.roll(x, -1))))
+
+
+def _point_on_segment_2d(point: np.ndarray, start: np.ndarray, end: np.ndarray) -> bool:
+    seg = end - start
+    rel = point - start
+    seg_norm_sq = float(np.dot(seg, seg))
+    if seg_norm_sq == 0.0:
+        return float(np.linalg.norm(rel)) == 0.0
+    cross = abs(seg[0] * rel[1] - seg[1] * rel[0])
+    if cross > 0.0:
+        return False
+    dot = float(np.dot(rel, seg))
+    return 0.0 <= dot <= seg_norm_sq
+
+
+def _point_in_support_region(point: np.ndarray, support_polygon: np.ndarray) -> bool:
+    if support_polygon.shape[0] == 0:
+        return False
+    if support_polygon.shape[0] == 1:
+        return float(np.linalg.norm(point - support_polygon[0])) == 0.0
+    if support_polygon.shape[0] == 2:
+        return _point_on_segment_2d(point, support_polygon[0], support_polygon[1])
+
+    prev_sign = None
+    for idx in range(support_polygon.shape[0]):
+        start = support_polygon[idx]
+        end = support_polygon[(idx + 1) % support_polygon.shape[0]]
+        cross = _cross_2d(start, end, point)
+        if cross == 0.0:
+            continue
+        cur_sign = cross > 0.0
+        if prev_sign is None:
+            prev_sign = cur_sign
+        elif prev_sign != cur_sign:
+            return False
+    return True
+
+
+class ZmpMarginMetric(BaseMetric):
+    """ Zero Moment Point metric. """
+    name = 'zmp_margin_metric'
+    D_NORM_MIN = 1e-5
+
+    def __init__(self,
+        robot_cfg: RobotConfig,
+        contact_threshold: float = 1e-3,
+        force_threshold: float = 1e-6,
+        draw_point: bool = False,
+        draw_point_size: float = 0.03,
+        draw_height_offset: float = 0.02,
+        draw_point_rgba: list = None,
+        **kwargs
+    ):
+        super().__init__(robot_cfg)
+        self.contact_threshold = contact_threshold
+        self.force_threshold = force_threshold
+        self.draw_point = draw_point
+        self.draw_point_size = draw_point_size
+        self.draw_height_offset = draw_height_offset
+        if draw_point_rgba is None:
+            draw_point_rgba = [1.0, 0.85, 0.1, 1.0]
+        self.draw_point_rgba = np.array(draw_point_rgba, dtype=np.float32)
+
+    def _clear_zmp_visualization(self, sim_data: SimData):
+        sim_data.visual.zmp_world_pos = None
+        sim_data.visual.zmp_draw_enabled = self.draw_point
+        sim_data.visual.zmp_draw_size = self.draw_point_size
+        sim_data.visual.zmp_draw_height_offset = self.draw_height_offset
+        sim_data.visual.zmp_draw_rgba = self.draw_point_rgba.copy()
+
+    def _set_zmp_visualization(self, sim_data: SimData, zmp_world_pos: np.ndarray):
+        sim_data.visual.zmp_world_pos = np.array(zmp_world_pos, dtype=np.float32)
+        sim_data.visual.zmp_draw_enabled = self.draw_point
+        sim_data.visual.zmp_draw_size = self.draw_point_size
+        sim_data.visual.zmp_draw_height_offset = self.draw_height_offset
+        sim_data.visual.zmp_draw_rgba = self.draw_point_rgba.copy()
+
+    def _log_invalid(self, sim_data: SimData, contact_count: int = 0, metric_value: float = 1.0):
+        self._clear_zmp_visualization(sim_data)
+        logger.log(float(metric_value), 'stable_metric/zmp_margin', step=sim_data.n_step)
+        logger.log(0.0, 'stable_metric/zmp_x', step=sim_data.n_step)
+        logger.log(0.0, 'stable_metric/zmp_y', step=sim_data.n_step)
+        logger.log(0.0, 'stable_metric/zmp_norm', step=sim_data.n_step)
+        logger.log(0.0, 'stable_metric/zmp_d_norm', step=sim_data.n_step)
+        logger.log(0.0, 'stable_metric/total_force_z', step=sim_data.n_step)
+        logger.log(0.0, 'stable_metric/support_polygon_area', step=sim_data.n_step)
+        logger.log(float(contact_count), 'stable_metric/support_contact_count', step=sim_data.n_step)
+        logger.log(0.0, 'stable_metric/fzmp', step=sim_data.n_step)
+        logger.log(0.0, 'stable_metric/zmp_valid', step=sim_data.n_step)
+
+    def __call__(self, sim_data: SimData, goal_data: GoalData) -> float:
+        dynamics = sim_data.dynamics
+        if dynamics is None:
+            raise RuntimeError("ZMP metric requires sim_data.dynamics, but got None.")
+
+        rigid_bodies = dynamics.rigid_bodies
+        if rigid_bodies is None or rigid_bodies.mass.shape[0] == 0:
+            raise RuntimeError("ZMP metric requires non-empty dynamics.rigid_bodies.")
+
+        contact_mask = dynamics.contacts.distances <= self.contact_threshold
+        support_contacts = dynamics.contacts.positions[contact_mask]
+        if support_contacts.shape[0] == 0:
+            self._log_invalid(sim_data)
+            return 1.0
+
+        support_center = np.mean(support_contacts, axis=0)
+        rel_com_pos = rigid_bodies.com_pos - support_center[None, :]
+
+        gravity = np.asarray(dynamics.gravity, dtype=np.float32)
+        body_forces = rigid_bodies.mass[:, None] * (gravity[None, :] - rigid_bodies.com_lin_acc)
+        total_force = np.sum(body_forces, axis=0)
+        total_force_z = float(total_force[2])
+        if abs(total_force_z) < self.force_threshold:
+            self._log_invalid(sim_data, contact_count=support_contacts.shape[0])
+            return 1.0
+
+        inertia_alpha = np.einsum('nij,nj->ni', rigid_bodies.inertia_world, rigid_bodies.ang_acc)
+        inertia_omega = np.einsum('nij,nj->ni', rigid_bodies.inertia_world, rigid_bodies.ang_vel)
+        gyro = np.cross(rigid_bodies.ang_vel, inertia_omega)
+        body_moments = np.cross(rel_com_pos, body_forces) - (inertia_alpha + gyro)
+        total_moment = np.sum(body_moments, axis=0)
+
+        zmp_xy = np.array([
+            -total_moment[1] / total_force_z,
+            total_moment[0] / total_force_z
+        ], dtype=np.float32)
+        zmp_norm = float(np.linalg.norm(zmp_xy))
+
+        if dynamics.default_diagonal_foot_distance is None:
+            raise RuntimeError("ZMP metric requires dynamics.default_diagonal_foot_distance, but got None.")
+        d_norm = float(dynamics.default_diagonal_foot_distance)
+        if d_norm < self.D_NORM_MIN:
+            logger.warning(
+                f"ZMP metric got too small default_diagonal_foot_distance={d_norm:.8f} "
+                f"(threshold={self.D_NORM_MIN:.1e}), returning 0.0."
+            )
+            self._log_invalid(
+                sim_data,
+                contact_count=support_contacts.shape[0],
+                metric_value=0.0,
+            )
+            return 0.0
+
+        metric_value = max(0.0, 1.0 - zmp_norm / d_norm)
+        zmp_world_pos = np.array([
+            support_center[0] + zmp_xy[0],
+            support_center[1] + zmp_xy[1],
+            support_center[2],
+        ], dtype=np.float32)
+        self._set_zmp_visualization(sim_data, zmp_world_pos)
+
+        support_polygon = _convex_hull_2d(support_contacts[:, :2] - support_center[None, :2])
+        support_area = _polygon_area_2d(support_polygon)
+        is_inside_support = _point_in_support_region(zmp_xy, support_polygon)
+        fzmp = 0.0 if is_inside_support else 1.0
+
+        logger.log(metric_value, 'stable_metric/zmp_margin', step=sim_data.n_step)
+        logger.log(float(zmp_xy[0]), 'stable_metric/zmp_x', step=sim_data.n_step)
+        logger.log(float(zmp_xy[1]), 'stable_metric/zmp_y', step=sim_data.n_step)
+        logger.log(zmp_norm, 'stable_metric/zmp_norm', step=sim_data.n_step)
+        logger.log(d_norm, 'stable_metric/zmp_d_norm', step=sim_data.n_step)
+        logger.log(total_force_z, 'stable_metric/total_force_z', step=sim_data.n_step)
+        logger.log(float(support_contacts.shape[0]), 'stable_metric/support_contact_count', step=sim_data.n_step)
+        logger.log(float(support_area), 'stable_metric/support_polygon_area', step=sim_data.n_step)
+        logger.log(float(fzmp), 'stable_metric/fzmp', step=sim_data.n_step)
+        logger.log(1.0, 'stable_metric/zmp_valid', step=sim_data.n_step)
+        return metric_value

robogauge/tasks/pipeline/base_pipeline.py

@@ -107,7 +107,8 @@ class BasePipeline:
                     if self.sim_cfg.domain_rand.action_delay and i == actions_start_decimation:
                         self.sim.setup_action(action, p_gains, d_gains, control_type)
                     sim_data = self.sim.step()
-                    self.gauge.update_metrics(sim_data, goal_data)
+                    if self.gauge.update_metrics(sim_data, goal_data):
+                        self.sim.update_visuals(sim_data)
                 if not self.first_reset and self.gauge.is_reset(sim_data):
                     sim_data = self.reset_sim_and_robot(sim_data)
             except Exception as e:

robogauge/tasks/robots/go2/go2_config.py

@@ -48,7 +48,7 @@ class Go2Config(RobotConfig):
             cmd = [2.0, 2.0, 0.25]
 
     class commands(RobotConfig.commands):
-        lin_vel_x = [-2.0, 2.0]  # min max [m/s]
+        lin_vel_x = [-1.0, 1.0]  # min max [m/s]
         lin_vel_y = [-1.0, 1.0]  # min max [m/s]
         lin_vel_z = None  # min max [m/s]
         ang_vel_roll = None  # min max [rad/s]

robogauge/tasks/simulator/mujoco_simulator.py

@@ -24,7 +24,8 @@ from robogauge.utils.math_utils import get_projected_gravity, quat_rotate_invers
 from robogauge.tasks.simulator.mujoco_config import MujocoConfig
 from robogauge.tasks.simulator.sim_data import (
     SimData,
-    RobotProprioception, JointState, BaseState, IMUState
+    RobotProprioception, JointState, BaseState, IMUState,
+    RigidBodyDynamics, GroundContactState, DynamicsState, VisualState
 )
 from robogauge.tasks.gauge.goal_data import VelocityGoal
 
@@ -48,6 +49,25 @@ class MujocoSimulator:
         self.target_pos = None
         self.target_velocity: Optional[VelocityGoal] = None
         self.penetration_reset_count = 0
+        self.robot_name_prefix = None
+        self.robot_body_ids = None
+        self.robot_body_id_set = set()
+        self.dynamic_body_ids = None
+        self.dynamic_body_names = []
+        self.dynamic_body_mass = None
+        self.dynamic_body_inertia_local = None
+        self.dynamic_body_iquat = None
+        self.foot_body_ids = None
+        self.default_diagonal_foot_distance = 0.0
+        self.prev_body_com_pos = None
+        self.prev_body_rot = None
+        self.prev_body_com_lin_vel = None
+        self.prev_body_ang_vel = None
+        self.zmp_world_pos = None
+        self.zmp_draw_enabled = False
+        self.zmp_draw_size = 0.03
+        self.zmp_draw_height_offset = 0.02
+        self.zmp_draw_rgba = np.array([1.0, 0.85, 0.1, 1.0], dtype=np.float32)
 
     def load(
         self,
@@ -114,6 +134,7 @@ class MujocoSimulator:
             self.mj_data.qpos[3] = 0.0
             self.mj_data.qpos[6] = 1.0
         self.mj_data.qpos[7:] = default_dof_pos
+        self.robot_name_prefix = f'{robot_mjcf.model}/'
 
         # Domain randomization: base mass
         base_body_name = f'{robot_mjcf.model}/base_link'
@@ -184,6 +205,8 @@ class MujocoSimulator:
         self.penetration_reset_count = 0
         self.load_dof_limits()
         self.preload_sensors()
+        self.preload_dynamics_data()
+        self.reset_dynamics_cache()
 
         # Robot controller placeholders
         self.action = None
@@ -259,12 +282,16 @@ class MujocoSimulator:
             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()
+        dynamics = self.get_dynamics_data()
+        visual = VisualState()
 
         sim_data = SimData(
             n_step=self.n_step,
             sim_dt=self.sim_dt,
             sim_time=self.sim_time,
-            proprio=proprio
+            proprio=proprio,
+            dynamics=dynamics,
+            visual=visual,
         )
 
         # input("DEBUG")
@@ -288,6 +315,18 @@ class MujocoSimulator:
                 rgba=[1, 0, 0, 1]
             )
 
+        def add_zmp_sphere(geom_elem):
+            zmp_pos = np.array(self.zmp_world_pos, dtype=np.float32).copy()
+            zmp_pos[2] += self.zmp_draw_height_offset
+            mujoco.mjv_initGeom(
+                geom_elem,
+                type=mujoco.mjtGeom.mjGEOM_SPHERE,
+                size=[self.zmp_draw_size, 0, 0],
+                pos=zmp_pos,
+                mat=np.eye(3).flatten(),
+                rgba=np.array(self.zmp_draw_rgba, dtype=np.float32)
+            )
+
         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)
@@ -329,6 +368,14 @@ class MujocoSimulator:
             else:
                 handle.scene.ngeom += 1
                 add_target_sphere(handle.scene.geoms[self.renderer.scene.ngeom - 1])
+
+        if self.zmp_draw_enabled and self.zmp_world_pos is not None:
+            if ctype == 'viewer':
+                add_zmp_sphere(handle.user_scn.geoms[viewer_geom_idx])
+                viewer_geom_idx += 1
+            else:
+                handle.scene.ngeom += 1
+                add_zmp_sphere(handle.scene.geoms[self.renderer.scene.ngeom - 1])
         
         if self.target_velocity is not None:
             base_pos_world = self.mj_data.qpos[:3]
@@ -413,8 +460,10 @@ class MujocoSimulator:
             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.reset_dynamics_cache()
 
         self.action = None
+        self.zmp_world_pos = None
         if self.viewer is not None:
             self.viewer.sync()
 
@@ -544,6 +593,273 @@ class MujocoSimulator:
             data_list.append(self.mj_data.sensordata[adr:adr+dim])
         return np.concatenate(data_list)
 
+    def update_visuals(self, sim_data: SimData):
+        visual = sim_data.visual
+        if visual is None:
+            self.zmp_world_pos = None
+            self.zmp_draw_enabled = False
+            return
+
+        self.zmp_world_pos = visual.zmp_world_pos
+        self.zmp_draw_enabled = bool(visual.zmp_draw_enabled)
+        self.zmp_draw_size = float(visual.zmp_draw_size)
+        self.zmp_draw_height_offset = float(visual.zmp_draw_height_offset)
+        if visual.zmp_draw_rgba is not None:
+            self.zmp_draw_rgba = np.array(visual.zmp_draw_rgba, dtype=np.float32)
+
+    @staticmethod
+    def _quat_to_rotmat(quat: np.ndarray) -> np.ndarray:
+        quat = np.asarray(quat, dtype=np.float64)
+        if quat.ndim == 1:
+            quat = quat[None, :]
+        quat_norm = np.linalg.norm(quat, axis=1, keepdims=True)
+        quat_norm = np.maximum(quat_norm, 1e-12)
+        quat = quat / quat_norm
+        w, x, y, z = quat[:, 0], quat[:, 1], quat[:, 2], quat[:, 3]
+        rot = np.empty((quat.shape[0], 3, 3), dtype=np.float32)
+        rot[:, 0, 0] = 1.0 - 2.0 * (y * y + z * z)
+        rot[:, 0, 1] = 2.0 * (x * y - z * w)
+        rot[:, 0, 2] = 2.0 * (x * z + y * w)
+        rot[:, 1, 0] = 2.0 * (x * y + z * w)
+        rot[:, 1, 1] = 1.0 - 2.0 * (x * x + z * z)
+        rot[:, 1, 2] = 2.0 * (y * z - x * w)
+        rot[:, 2, 0] = 2.0 * (x * z - y * w)
+        rot[:, 2, 1] = 2.0 * (y * z + x * w)
+        rot[:, 2, 2] = 1.0 - 2.0 * (x * x + y * y)
+        return rot
+
+    @staticmethod
+    def _rotation_delta_to_angular_velocity(
+            prev_rot: np.ndarray,
+            cur_rot: np.ndarray,
+            dt: float
+        ) -> np.ndarray:
+        ang_vel = np.zeros((cur_rot.shape[0], 3), dtype=np.float32)
+        if prev_rot is None or dt <= 0.0:
+            return ang_vel
+
+        rot_delta = cur_rot @ np.transpose(prev_rot, (0, 2, 1))
+        traces = np.trace(rot_delta, axis1=1, axis2=2)
+        cos_theta = np.clip((traces - 1.0) * 0.5, -1.0, 1.0)
+        theta = np.arccos(cos_theta)
+        skew = np.stack([
+            rot_delta[:, 2, 1] - rot_delta[:, 1, 2],
+            rot_delta[:, 0, 2] - rot_delta[:, 2, 0],
+            rot_delta[:, 1, 0] - rot_delta[:, 0, 1],
+        ], axis=1)
+
+        small_mask = theta < 1e-6
+        if np.any(small_mask):
+            ang_vel[small_mask] = 0.5 * skew[small_mask] / dt
+
+        normal_mask = ~small_mask
+        if np.any(normal_mask):
+            idx = np.where(normal_mask)[0]
+            sin_theta = np.sin(theta[idx])
+            singular_mask = np.abs(sin_theta) < 1e-6
+            if np.any(singular_mask):
+                singular_idx = idx[singular_mask]
+                ang_vel[singular_idx] = 0.5 * skew[singular_idx] / dt
+            regular_idx = idx[~singular_mask]
+            if regular_idx.size > 0:
+                axis = skew[regular_idx] / (2.0 * sin_theta[~singular_mask, None])
+                ang_vel[regular_idx] = axis * theta[regular_idx, None] / dt
+        return ang_vel
+
+    @staticmethod
+    def _max_pairwise_xy_distance(positions: np.ndarray) -> float:
+        if positions.shape[0] < 2:
+            return 0.0
+        xy = positions[:, :2]
+        diffs = xy[:, None, :] - xy[None, :, :]
+        dist = np.linalg.norm(diffs, axis=-1)
+        return float(np.max(dist))
+
+    def _get_body_name(self, body_id: int) -> str:
+        if body_id == 0:
+            return "worldbody"
+        name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, body_id)
+        return name if name is not None else f"body_{body_id}"
+
+    def _get_geom_name(self, geom_id: int) -> str:
+        name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, geom_id)
+        return name if name is not None else f"geom_{geom_id}"
+
+    def reset_dynamics_cache(self):
+        self.prev_body_com_pos = None
+        self.prev_body_rot = None
+        self.prev_body_com_lin_vel = None
+        self.prev_body_ang_vel = None
+
+    def preload_dynamics_data(self):
+        self.robot_body_ids = []
+        robot_body_names = []
+        for body_id in range(self.mj_model.nbody):
+            body_name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, body_id)
+            if body_name is not None and body_name.startswith(self.robot_name_prefix):
+                self.robot_body_ids.append(body_id)
+                robot_body_names.append(body_name)
+        self.robot_body_ids = np.array(self.robot_body_ids, dtype=np.int32)
+        self.robot_body_id_set = set(int(body_id) for body_id in self.robot_body_ids.tolist())
+
+        if self.robot_body_ids.size == 0:
+            logger.warning("No robot bodies found for dynamics preprocessing.")
+            self.dynamic_body_ids = np.array([], dtype=np.int32)
+            self.dynamic_body_names = []
+            self.dynamic_body_mass = np.zeros((0,), dtype=np.float32)
+            self.dynamic_body_inertia_local = np.zeros((0, 3), dtype=np.float32)
+            self.dynamic_body_iquat = np.zeros((0, 4), dtype=np.float32)
+            self.foot_body_ids = np.array([], dtype=np.int32)
+            self.default_diagonal_foot_distance = 0.0
+            return
+
+        robot_body_mass = np.array(self.mj_model.body_mass[self.robot_body_ids], dtype=np.float32)
+        dynamic_mask = robot_body_mass > 0.0
+        self.dynamic_body_ids = self.robot_body_ids[dynamic_mask]
+        self.dynamic_body_names = [
+            robot_body_names[idx]
+            for idx, is_dynamic in enumerate(dynamic_mask.tolist())
+            if is_dynamic
+        ]
+        self.dynamic_body_mass = np.array(self.mj_model.body_mass[self.dynamic_body_ids], dtype=np.float32)
+        self.dynamic_body_inertia_local = np.array(self.mj_model.body_inertia[self.dynamic_body_ids], dtype=np.float32)
+        if hasattr(self.mj_model, 'body_iquat'):
+            self.dynamic_body_iquat = np.array(self.mj_model.body_iquat[self.dynamic_body_ids], dtype=np.float32)
+        else:
+            self.dynamic_body_iquat = np.tile(np.array([[1.0, 0.0, 0.0, 0.0]], dtype=np.float32), (len(self.dynamic_body_ids), 1))
+
+        self.foot_body_ids = np.array([
+            body_id for body_id, body_name in zip(self.robot_body_ids.tolist(), robot_body_names)
+            if 'foot' in body_name.rsplit('/', 1)[-1].lower()
+        ], dtype=np.int32)
+        self.default_diagonal_foot_distance = self.compute_default_diagonal_foot_distance()
+        logger.info(
+            f"Dynamics preprocessing: robot_bodies={len(self.robot_body_ids)}, "
+            f"dynamic_bodies={len(self.dynamic_body_ids)}, foot_bodies={len(self.foot_body_ids)}, "
+            f"default_diagonal_foot_distance={self.default_diagonal_foot_distance:.4f}"
+        )
+
+    def compute_default_diagonal_foot_distance(self) -> float:
+        if self.foot_body_ids is not None and self.foot_body_ids.size >= 2:
+            if hasattr(self.mj_data, 'xpos'):
+                foot_pos = np.array(self.mj_data.xpos[self.foot_body_ids], dtype=np.float32)
+            else:
+                foot_pos = np.array(self.mj_data.xipos[self.foot_body_ids], dtype=np.float32)
+            return self._max_pairwise_xy_distance(foot_pos)
+
+        contacts = self.get_ground_contact_state()
+        if contacts.positions.shape[0] >= 2:
+            return self._max_pairwise_xy_distance(contacts.positions)
+        return 0.0
+
+    def get_ground_contact_state(self) -> GroundContactState:
+        positions = []
+        distances = []
+        robot_geom_names = []
+        other_geom_names = []
+        robot_body_names = []
+        other_body_names = []
+
+        for i in range(self.mj_data.ncon):
+            contact = self.mj_data.contact[i]
+            geom1 = int(contact.geom1)
+            geom2 = int(contact.geom2)
+            body1 = int(self.mj_model.geom_bodyid[geom1])
+            body2 = int(self.mj_model.geom_bodyid[geom2])
+            is_robot1 = body1 in self.robot_body_id_set
+            is_robot2 = body2 in self.robot_body_id_set
+            if is_robot1 == is_robot2:
+                continue
+
+            robot_geom = geom1 if is_robot1 else geom2
+            other_geom = geom2 if is_robot1 else geom1
+            robot_body = body1 if is_robot1 else body2
+            other_body = body2 if is_robot1 else body1
+
+            positions.append(np.array(contact.pos, dtype=np.float32))
+            distances.append(float(contact.dist))
+            robot_geom_names.append(self._get_geom_name(robot_geom))
+            other_geom_names.append(self._get_geom_name(other_geom))
+            robot_body_names.append(self._get_body_name(robot_body))
+            other_body_names.append(self._get_body_name(other_body))
+
+        positions = np.array(positions, dtype=np.float32).reshape(-1, 3)
+        distances = np.array(distances, dtype=np.float32)
+        return GroundContactState(
+            positions=positions,
+            distances=distances,
+            robot_geom_names=robot_geom_names,
+            other_geom_names=other_geom_names,
+            robot_body_names=robot_body_names,
+            other_body_names=other_body_names,
+        )
+
+    def get_dynamics_data(self) -> DynamicsState:
+        if self.dynamic_body_ids is None or self.dynamic_body_ids.size == 0:
+            empty_body = RigidBodyDynamics(
+                names=[],
+                mass=np.zeros((0,), dtype=np.float32),
+                com_pos=np.zeros((0, 3), dtype=np.float32),
+                com_lin_vel=np.zeros((0, 3), dtype=np.float32),
+                com_lin_acc=np.zeros((0, 3), dtype=np.float32),
+                ang_vel=np.zeros((0, 3), dtype=np.float32),
+                ang_acc=np.zeros((0, 3), dtype=np.float32),
+                inertia_world=np.zeros((0, 3, 3), dtype=np.float32),
+            )
+            return DynamicsState(
+                gravity=np.array(self.mj_model.opt.gravity, dtype=np.float32),
+                rigid_bodies=empty_body,
+                contacts=self.get_ground_contact_state(),
+                default_diagonal_foot_distance=float(self.default_diagonal_foot_distance),
+            )
+
+        com_pos = np.array(self.mj_data.xipos[self.dynamic_body_ids], dtype=np.float32)
+        body_rot = np.array(self.mj_data.xmat[self.dynamic_body_ids], dtype=np.float32).reshape(-1, 3, 3)
+        com_lin_vel = np.zeros_like(com_pos)
+        com_lin_acc = np.zeros_like(com_pos)
+        ang_vel = np.zeros_like(com_pos)
+        ang_acc = np.zeros_like(com_pos)
+
+        if self.prev_body_com_pos is not None:
+            com_lin_vel = (com_pos - self.prev_body_com_pos) / self.sim_dt
+            if self.prev_body_com_lin_vel is not None:
+                com_lin_acc = (com_lin_vel - self.prev_body_com_lin_vel) / self.sim_dt
+        if self.prev_body_rot is not None:
+            ang_vel = self._rotation_delta_to_angular_velocity(self.prev_body_rot, body_rot, self.sim_dt)
+            if self.prev_body_ang_vel is not None:
+                ang_acc = (ang_vel - self.prev_body_ang_vel) / self.sim_dt
+
+        if hasattr(self.mj_data, 'ximat'):
+            inertial_rot = np.array(self.mj_data.ximat[self.dynamic_body_ids], dtype=np.float32).reshape(-1, 3, 3)
+        else:
+            inertial_rot = body_rot @ self._quat_to_rotmat(self.dynamic_body_iquat)
+        inertia_local = np.zeros((self.dynamic_body_ids.size, 3, 3), dtype=np.float32)
+        inertia_local[:, 0, 0] = self.dynamic_body_inertia_local[:, 0]
+        inertia_local[:, 1, 1] = self.dynamic_body_inertia_local[:, 1]
+        inertia_local[:, 2, 2] = self.dynamic_body_inertia_local[:, 2]
+        inertia_world = inertial_rot @ inertia_local @ np.transpose(inertial_rot, (0, 2, 1))
+
+        self.prev_body_com_pos = com_pos.copy()
+        self.prev_body_rot = body_rot.copy()
+        self.prev_body_com_lin_vel = com_lin_vel.copy()
+        self.prev_body_ang_vel = ang_vel.copy()
+
+        return DynamicsState(
+            gravity=np.array(self.mj_model.opt.gravity, dtype=np.float32),
+            rigid_bodies=RigidBodyDynamics(
+                names=self.dynamic_body_names,
+                mass=self.dynamic_body_mass.copy(),
+                com_pos=com_pos,
+                com_lin_vel=com_lin_vel,
+                com_lin_acc=com_lin_acc,
+                ang_vel=ang_vel,
+                ang_acc=ang_acc,
+                inertia_world=inertia_world,
+            ),
+            contacts=self.get_ground_contact_state(),
+            default_diagonal_foot_distance=float(self.default_diagonal_foot_distance),
+        )
+
     def debug_print_proprio_shapes(self):
         """Log shapes (or lengths) of each numpy vector inside a RobotProprioception.
 

robogauge/tasks/simulator/sim_data.py

@@ -5,7 +5,7 @@ from dataclasses import dataclass
 class JointState:
     pos: np.ndarray     # [rad] shape (n_dof,)
     vel: np.ndarray     # [rad/s] shape (n_dof,)
-    torque: np.ndarray   # [N*m] shape (n_dof,)
+    torque: np.ndarray  # [N*m] shape (n_dof,)
     limits: np.ndarray  # [rad] shape (n_dof, 2), lower and upper limits
     names: list         # list of joint names
 
@@ -30,9 +30,46 @@ class RobotProprioception:
     base: BaseState
     imu: IMUState
 
+@dataclass
+class RigidBodyDynamics:
+    names: list                   # list of rigid body names
+    mass: np.ndarray              # [kg] shape (n_body,)
+    com_pos: np.ndarray           # [m] world frame, shape (n_body, 3)
+    com_lin_vel: np.ndarray       # [m/s] world frame, shape (n_body, 3)
+    com_lin_acc: np.ndarray       # [m/s^2] world frame, shape (n_body, 3)
+    ang_vel: np.ndarray           # [rad/s] world frame, shape (n_body, 3)
+    ang_acc: np.ndarray           # [rad/s^2] world frame, shape (n_body, 3)
+    inertia_world: np.ndarray     # [kg*m^2] world frame, shape (n_body, 3, 3)
+
+@dataclass
+class GroundContactState:
+    positions: np.ndarray         # [m] world frame, shape (n_contact, 3)
+    distances: np.ndarray         # [m] shape (n_contact,)
+    robot_geom_names: list        # list of robot geom names
+    other_geom_names: list        # list of non-robot geom names
+    robot_body_names: list        # list of robot body names
+    other_body_names: list        # list of non-robot body names
+
+@dataclass
+class DynamicsState:
+    gravity: np.ndarray
+    rigid_bodies: RigidBodyDynamics
+    contacts: GroundContactState
+    default_diagonal_foot_distance: float
+
+@dataclass
+class VisualState:
+    zmp_world_pos: np.ndarray = None
+    zmp_draw_enabled: bool = False
+    zmp_draw_size: float = 0.03
+    zmp_draw_height_offset: float = 0.02
+    zmp_draw_rgba: np.ndarray = None
+
 @dataclass
 class SimData:
     n_step: int
     sim_dt: float
     sim_time: float
     proprio: RobotProprioception
+    dynamics: DynamicsState
+    visual: VisualState