268 lines
11 KiB
Python
268 lines
11 KiB
Python
# -*- 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
|