syfffffjj
07027e8fde
Current code snapshot: smooth/bridge task configs, smooth random terrain, roll/pitch stability rewards, M20 URDF updates, eval profile and stability training documentation.
263 lines
11 KiB
Python
263 lines
11 KiB
Python
# Copyright (c) 2025 Deep Robotics
|
||
# SPDX-License-Identifier: BSD 3-Clause
|
||
|
||
"""指标计算：纯函数，只依赖 numpy。
|
||
|
||
输入是 :class:`~eval_kit.core.schema.Trace`、:class:`~eval_kit.core.schema.RobotProfile`
|
||
和 :class:`~eval_kit.core.schema.EvalProtocol`，输出是一个扁平的指标字典。
|
||
|
||
指标口径与 ``scripts/reinforcement_learning/rsl_rl/eval.py`` 保持一致（同样的取点、
|
||
同样的公式），因此默认协议下两套评估的数值可以直接对照；差别只在三处，且都是
|
||
有意为之：
|
||
|
||
1. **质量**不再写死在命令行默认值里，而是从 ``Trace.mass_kg`` 取（适配层从资产读），
|
||
传入的 ``profile.mass_kg`` 只作兜底。
|
||
2. **加速度阈值**以重力倍数表达，换算成 m/s² 后参与计算，使阈值可跨量级机器人复用。
|
||
3. **缺失的形态**返回 ``nan`` 而不是 0。旧实现里没有足端的机器人会得到
|
||
「支撑裕度 = 0.00」，看起来像测出来的值，实际是没测。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from typing import Any
|
||
|
||
import numpy as np
|
||
|
||
from eval_kit.core.schema import EvalProtocol, RobotProfile, Trace
|
||
|
||
#: 全部帧，用于成功判定与最大进度（含起始帧，与旧实现取点一致）。
|
||
_ALL_FRAMES = slice(None)
|
||
#: 动作执行之后的帧，用于逐帧累积类指标。
|
||
_STEP_FRAMES = slice(1, None)
|
||
|
||
_GRAVITY = 9.81
|
||
|
||
|
||
def tilt_from_projected_gravity(projected_gravity_b: np.ndarray) -> np.ndarray:
|
||
"""由机体系投影重力算倾角 [rad]。
|
||
|
||
``projected_gravity_b`` 为 ``(T, N, 3)``，返回 ``(T, N)``。
|
||
"""
|
||
z_axis = np.clip(-projected_gravity_b[..., 2], -1.0, 1.0)
|
||
return np.arccos(z_axis)
|
||
|
||
|
||
def support_margin_xy(root_xy: np.ndarray, foot_xy: np.ndarray) -> np.ndarray:
|
||
"""机身投影点到足端轴对齐包围盒的距离 [m]，在盒内为正、盒外为负。
|
||
|
||
``root_xy`` 为 ``(N, 2)``，``foot_xy`` 为 ``(N, F, 2)``，返回 ``(N,)``。
|
||
|
||
这是当前实现采用的口径（与旧 ``eval.py`` 相同），记录在
|
||
``EvalProtocol.stability_margin_method`` 里。它比凸包/ZMP 版本粗糙，
|
||
但不会在退化多边形上失败，也不需要 scipy。
|
||
"""
|
||
min_xy = np.min(foot_xy, axis=1)
|
||
max_xy = np.max(foot_xy, axis=1)
|
||
|
||
inside_x = (root_xy[:, 0] >= min_xy[:, 0]) & (root_xy[:, 0] <= max_xy[:, 0])
|
||
inside_y = (root_xy[:, 1] >= min_xy[:, 1]) & (root_xy[:, 1] <= max_xy[:, 1])
|
||
inside = inside_x & inside_y
|
||
|
||
inside_dist = np.minimum(
|
||
np.minimum(root_xy[:, 0] - min_xy[:, 0], max_xy[:, 0] - root_xy[:, 0]),
|
||
np.minimum(root_xy[:, 1] - min_xy[:, 1], max_xy[:, 1] - root_xy[:, 1]),
|
||
)
|
||
outside_dx = np.clip(np.maximum(min_xy[:, 0] - root_xy[:, 0], root_xy[:, 0] - max_xy[:, 0]), 0.0, None)
|
||
outside_dy = np.clip(np.maximum(min_xy[:, 1] - root_xy[:, 1], root_xy[:, 1] - max_xy[:, 1]), 0.0, None)
|
||
outside_dist = np.sqrt(outside_dx**2 + outside_dy**2)
|
||
|
||
return np.where(inside, inside_dist, -outside_dist)
|
||
|
||
|
||
def _nanmean_or_nan(values: np.ndarray) -> float:
|
||
"""均值；空数组返回 nan（而不是 0，也不是 numpy 的警告）。"""
|
||
if values.size == 0:
|
||
return float("nan")
|
||
return float(np.mean(values))
|
||
|
||
|
||
def compute_metrics(
|
||
trace: Trace,
|
||
profile: RobotProfile,
|
||
protocol: EvalProtocol,
|
||
*,
|
||
include_per_env: bool = False,
|
||
) -> dict[str, Any]:
|
||
"""由一次 rollout 的轨迹算出全部指标。
|
||
|
||
指标名的命名与旧 ``eval.py`` 完全一致，便于新旧结果直接对照。
|
||
"""
|
||
num_steps = trace.num_steps
|
||
num_envs = trace.num_envs
|
||
dt = float(trace.dt)
|
||
|
||
if abs(protocol.dt - dt) > 0.01 * dt:
|
||
raise ValueError(
|
||
f"协议 dt={protocol.dt} 与轨迹 dt={dt} 不一致。"
|
||
"通常意味着协议里的步长与任务的 decimation 不匹配。"
|
||
)
|
||
|
||
command = np.asarray(trace.command, dtype=np.float64)
|
||
command_xy = command[:2]
|
||
command_yaw = float(command[2])
|
||
|
||
command_dir = command_xy.copy()
|
||
norm = float(np.linalg.norm(command_dir))
|
||
command_dir = command_dir / norm if norm > 1e-6 else np.array([1.0, 0.0])
|
||
|
||
# ------------------------------------------------------------------ 成功判定
|
||
base_xy = np.asarray(trace.base_pos_w, dtype=np.float64)[..., :2]
|
||
delta_xy = base_xy - base_xy[0][None, :, :]
|
||
progress = np.einsum("tnc,c->tn", delta_xy[_ALL_FRAMES], command_dir)
|
||
max_progress = progress.max(axis=0)
|
||
reached = max_progress >= protocol.success_distance_m
|
||
|
||
tilt = tilt_from_projected_gravity(np.asarray(trace.projected_gravity_b, dtype=np.float64))
|
||
tilt_steps = tilt[_STEP_FRAMES]
|
||
max_tilt = tilt_steps.max(axis=0)
|
||
ever_fallen = max_tilt > protocol.tilt_threshold_rad
|
||
|
||
success = reached & ~ever_fallen
|
||
success_count = int(np.count_nonzero(success))
|
||
|
||
# 首次达成成功的时刻：第 i 帧（i>=1）达成则记 i*dt。
|
||
crossed = progress[_STEP_FRAMES] >= protocol.success_distance_m
|
||
complete_time = np.full((num_envs,), np.nan, dtype=np.float64)
|
||
for env_index in range(num_envs):
|
||
hit = np.flatnonzero(crossed[:, env_index])
|
||
if hit.size:
|
||
complete_time[env_index] = float(hit[0] + 1) * dt
|
||
|
||
successful_time = complete_time[success & np.isfinite(complete_time)]
|
||
avg_time = _nanmean_or_nan(successful_time)
|
||
|
||
# ------------------------------------------------------------- 速度跟踪误差
|
||
lin_vel_b = np.asarray(trace.base_lin_vel_b, dtype=np.float64)[_STEP_FRAMES]
|
||
ang_vel_b = np.asarray(trace.base_ang_vel_b, dtype=np.float64)[_STEP_FRAMES]
|
||
lin_error = np.linalg.norm(lin_vel_b[..., :2] - command_xy[None, None, :], axis=-1)
|
||
yaw_error = np.abs(ang_vel_b[..., 2] - command_yaw)
|
||
|
||
# -------------------------------------------------------------- 能耗与路程
|
||
joint_torque = np.asarray(trace.joint_torque, dtype=np.float64)[_STEP_FRAMES]
|
||
joint_vel = np.asarray(trace.joint_vel, dtype=np.float64)[_STEP_FRAMES]
|
||
mechanical_power = joint_torque * joint_vel
|
||
if protocol.energy_convention == "absolute":
|
||
integrand = np.abs(mechanical_power)
|
||
elif protocol.energy_convention == "signed":
|
||
integrand = mechanical_power
|
||
else:
|
||
raise ValueError(
|
||
f"未知能耗约定 {protocol.energy_convention!r}；可选 'absolute' 或 'signed'。"
|
||
)
|
||
energy = integrand.sum(axis=(0, 2)) * dt
|
||
|
||
lin_vel_w = np.asarray(trace.base_lin_vel_w, dtype=np.float64)
|
||
route = np.linalg.norm(lin_vel_w[_STEP_FRAMES][..., :2], axis=-1).sum(axis=0) * dt
|
||
|
||
# ------------------------------------------------------------ 运动不稳定性
|
||
accel = np.linalg.norm(np.diff(lin_vel_w, axis=0), axis=-1) / dt
|
||
instability_count = (accel > protocol.acc_threshold_mps2).sum(axis=0)
|
||
|
||
stable_time = (tilt_steps < protocol.tilt_threshold_rad).sum(axis=0) * dt
|
||
|
||
# ------------------------------------------------------------- 支撑多边形
|
||
feet_pos_w = np.asarray(trace.feet_pos_w, dtype=np.float64)[_STEP_FRAMES]
|
||
num_feet = int(feet_pos_w.shape[2])
|
||
if num_feet > 0:
|
||
margins = support_margin_xy(base_xy[_STEP_FRAMES][:, :, :].reshape(-1, 2), feet_pos_w.reshape(-1, num_feet, 2))
|
||
margins = margins.reshape(feet_pos_w.shape[0], num_envs)
|
||
avg_margin = _nanmean_or_nan(margins)
|
||
margin_method = protocol.stability_margin_method
|
||
else:
|
||
avg_margin = float("nan")
|
||
margin_method = "unavailable"
|
||
|
||
# --------------------------------------------------------------- 支撑接触
|
||
support_force = np.asarray(trace.support_force_n, dtype=np.float64)
|
||
support_contact_time = np.asarray(trace.support_contact_time_s, dtype=np.float64)
|
||
num_support_bodies = int(support_force.shape[2])
|
||
support_threshold = (
|
||
profile.support_force_threshold_n
|
||
if protocol.support_force_threshold_n is None
|
||
else protocol.support_force_threshold_n
|
||
)
|
||
if num_support_bodies > 0:
|
||
active = np.any(support_force[_STEP_FRAMES] > support_threshold, axis=2) # (T-1, N)
|
||
support_steps = active.sum(axis=0)
|
||
# 起始帧之前视为「未接触」，因此首帧接触也算一次事件（与旧实现一致）。
|
||
active_padded = np.concatenate([np.zeros((1, num_envs), dtype=bool), active], axis=0)
|
||
support_events = (active_padded[1:] & ~active_padded[:-1]).sum(axis=0)
|
||
max_support_time = support_contact_time[_STEP_FRAMES].max(axis=(0, 2))
|
||
support_contact_ratio = _nanmean_or_nan(support_steps / max(num_steps, 1))
|
||
avg_support_events = float(np.mean(support_events))
|
||
avg_max_support_time = float(np.mean(max_support_time))
|
||
if success_count > 0:
|
||
success_support_ratio = float(np.mean((support_events[success] > 0).astype(np.float64)))
|
||
else:
|
||
success_support_ratio = float("nan")
|
||
else:
|
||
support_contact_ratio = float("nan")
|
||
avg_support_events = float("nan")
|
||
avg_max_support_time = float("nan")
|
||
success_support_ratio = float("nan")
|
||
|
||
# --------------------------------------------------------- 质量与 COT
|
||
mass = trace.mass_kg if trace.mass_kg is not None else profile.mass_kg
|
||
if mass is None or mass <= 0.0:
|
||
cost_of_transport = np.full((num_envs,), np.nan)
|
||
mass_used = float("nan")
|
||
else:
|
||
cost_of_transport = energy / ((route + 1e-6) * mass * _GRAVITY)
|
||
mass_used = float(mass)
|
||
|
||
metrics: dict[str, Any] = {
|
||
"profile": profile.name,
|
||
"protocol_fingerprint": protocol.fingerprint(),
|
||
"num_envs": float(num_envs),
|
||
"num_steps": float(num_steps),
|
||
"episode_length_s": float(num_steps * dt),
|
||
"terrain_size_m": float(protocol.terrain_size_m),
|
||
"success_distance_m": float(protocol.success_distance_m),
|
||
"robot_mass_kg": mass_used,
|
||
"energy_convention": protocol.energy_convention,
|
||
"stability_margin_method": margin_method,
|
||
"tilt_threshold_rad": float(protocol.tilt_threshold_rad),
|
||
"acc_threshold_mps2": float(protocol.acc_threshold_mps2),
|
||
"support_force_threshold_n": float(support_threshold) if num_support_bodies > 0 else float("nan"),
|
||
# --- 成功率与进度
|
||
"success_count": float(success_count),
|
||
"success_rate": float(success_count / max(num_envs, 1)),
|
||
"avg_final_progress_m": float(np.mean(max_progress)),
|
||
# --- 时间
|
||
"avg_time_consumption_s": avg_time,
|
||
"time_efficiency": float(
|
||
avg_time / (protocol.success_distance_m / max(abs(float(command[0])), 1e-6))
|
||
)
|
||
if np.isfinite(avg_time)
|
||
else float("nan"),
|
||
# --- 跟踪误差
|
||
"mean_lin_vel_error": float(np.mean(lin_error)),
|
||
"mean_yaw_vel_error": float(np.mean(yaw_error)),
|
||
# --- 能耗
|
||
"avg_energy_j": float(np.mean(energy)),
|
||
"avg_cost_of_transport": _nanmean_or_nan(cost_of_transport),
|
||
# --- 稳定性
|
||
"unstability_rate": float(np.mean(instability_count / max(num_steps, 1))),
|
||
"avg_stability_margin_m": avg_margin,
|
||
"avg_max_tilt_rad": float(np.mean(max_tilt)),
|
||
"avg_stable_time_s": float(np.mean(stable_time)),
|
||
# --- 支撑接触
|
||
"avg_support_contact_ratio": support_contact_ratio,
|
||
"avg_support_events": avg_support_events,
|
||
"avg_max_support_contact_time_s": avg_max_support_time,
|
||
"successful_runs_with_support_ratio": success_support_ratio,
|
||
}
|
||
|
||
if include_per_env:
|
||
metrics["_per_env"] = {
|
||
"success": success.astype(float).tolist(),
|
||
"max_progress": max_progress.tolist(),
|
||
"max_tilt": max_tilt.tolist(),
|
||
}
|
||
return metrics
|