# -*- coding: utf-8 -*- ''' @File : stable_metric.py @Time : 2025/12/18 20:18:33 @Author : wty-yy @Version : 1.0 @Blog : https://wty-yy.github.io/ @Desc : Stability-related metrics implementation ''' import numpy as np from robogauge.tasks.robots import RobotConfig from robogauge.tasks.gauge.metrics.base_metric import BaseMetric, GoalData, SimData from robogauge.utils.math_utils import get_projected_gravity from robogauge.utils.logger import logger class OrientationStabilityMetric(BaseMetric): """ Metric to log body orientation stability. """ name = 'orientation_stability_metric' def __init__(self, robot_cfg: RobotConfig, **kwargs ): super().__init__(robot_cfg) def __call__(self, sim_data: SimData, goal_data: GoalData) -> float: projected_gravity = get_projected_gravity(sim_data.proprio.base.quat) projected_y = projected_gravity[1] metric_value = 1 - abs(projected_y) # consider roll only logger.log(abs(projected_y), f'stable_metric/projected_y_abs', step=sim_data.n_step) return metric_value class TorqueSmoothnessMetric(BaseMetric): """ Metric to log torque smoothness. """ name = 'torque_smoothness_metric' def __init__(self, robot_cfg: RobotConfig, scaling_factor: float = 30.0, **kwargs ): super().__init__(robot_cfg) self.last_torque = None self.scaling_factor = scaling_factor def reset(self): self.last_torque = None def __call__(self, sim_data: SimData, goal_data: GoalData) -> float: current_torque = np.array(sim_data.proprio.joint.torque, np.float32) if self.last_torque is None: self.last_torque = current_torque return 1.0 # No change at first step torque_diff = current_torque - self.last_torque self.last_torque = current_torque 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 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