syfffffjj
7d290ca6ff
- eval_kit/core: 零框架依赖的指标内核（Trace/metrics/aggregation/terrain_spec/report/naming） - eval_kit/isaaclab: 适配层（env_cfg 改造、rollout 采集、checkpoint 加载） - eval_kit/profiles.py / protocols.py: 机器人形态描述与评估口径预设 - eval_kit/tests: 24 项内核自检 + 21 项命令行自检（都不需要仿真） - eval_suite.py: 编排入口，每个（阶段, 种子）一个子进程；与 eval.py 指标口径一致 - WebRTC 可视化：--livestream 透传、视口相机跟随/自由（--view-env / --free-camera）、 连接前停顿（--pause-s）、自动放开客户端改画幅、打印连接地址 - 诊断与替代路径：rollout 进度打印、--record-frames 抓帧
258 lines
10 KiB
Python
258 lines
10 KiB
Python
# Copyright (c) 2025 Deep Robotics
|
|
# SPDX-License-Identifier: BSD 3-Clause
|
|
|
|
"""跑 rollout 并把仿真信号采集成框架无关的 :class:`~eval_kit.core.schema.Trace`。
|
|
|
|
这是唯一读取 ``env`` 内部状态的地方。采集完之后环境就可以关掉了，后续所有
|
|
指标计算都在内核里对纯数组进行。
|
|
|
|
形态相关的下标（谁是轮、谁是支撑体、哪些关节耗电）全部由
|
|
:class:`~eval_kit.core.schema.RobotProfile` 的正则解析得到，代码里没有硬编码的
|
|
link/joint 名字。解析不到时记录警告并跳过对应指标，而不是抛异常。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from eval_kit.core.naming import resolve_matching_indices
|
|
from eval_kit.core.schema import EvalProtocol, RobotProfile, Trace
|
|
|
|
#: 默认的接触力传感器名字，与 ``velocity_env_cfg.py`` 的场景定义一致。
|
|
DEFAULT_CONTACT_SENSOR = "contact_forces"
|
|
|
|
|
|
@dataclass
|
|
class RolloutRecorder:
|
|
"""一次评估里固定不变的东西：形态下标与质量。建一次，跑多次。"""
|
|
|
|
feet_ids: list[int] = field(default_factory=list)
|
|
feet_names: list[str] = field(default_factory=list)
|
|
support_ids: list[int] = field(default_factory=list)
|
|
support_names: list[str] = field(default_factory=list)
|
|
joint_ids: list[int] = field(default_factory=list)
|
|
joint_names: list[str] = field(default_factory=list)
|
|
mass_kg: float | None = None
|
|
contact_sensor_name: str = DEFAULT_CONTACT_SENSOR
|
|
warnings: list[str] = field(default_factory=list)
|
|
|
|
def describe(self) -> str:
|
|
lines = [
|
|
f"足端 {len(self.feet_ids)} 个: {self.feet_names or '（无）'}",
|
|
f"支撑体 {len(self.support_ids)} 个: {self.support_names or '（无）'}",
|
|
f"驱动关节 {len(self.joint_ids)} 个: {self.joint_names or '（无）'}",
|
|
f"整机质量: {'未知' if self.mass_kg is None else f'{self.mass_kg:.3f} kg'}",
|
|
]
|
|
for warning in self.warnings:
|
|
lines.append(f"[警告] {warning}")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def read_total_mass(robot: Any) -> float | None:
|
|
"""从仿真资产读整机质量。
|
|
|
|
优先读 PhysX 运行时的真实质量（会反映任何质量随机化），读不到再退回资产默认质量。
|
|
"""
|
|
try:
|
|
masses = robot.root_physx_view.get_masses()
|
|
return float(masses[0].sum().item())
|
|
except Exception: # noqa: BLE001 - 不同版本 API 差异，读不到就退回
|
|
pass
|
|
try:
|
|
return float(robot.data.default_mass[0].sum().item())
|
|
except Exception: # noqa: BLE001
|
|
return None
|
|
|
|
|
|
def build_recorder(
|
|
env: Any,
|
|
profile: RobotProfile,
|
|
*,
|
|
contact_sensor_name: str = DEFAULT_CONTACT_SENSOR,
|
|
) -> RolloutRecorder:
|
|
"""按机器人描述解析出足端/支撑体/关节下标，并读一次质量。"""
|
|
base_env = env.unwrapped
|
|
robot = base_env.scene["robot"]
|
|
recorder = RolloutRecorder(contact_sensor_name=contact_sensor_name)
|
|
|
|
feet = resolve_matching_indices(profile.feet_body_patterns, robot.body_names)
|
|
recorder.feet_ids = feet.indices
|
|
recorder.feet_names = feet.names
|
|
for pattern in feet.unmatched_patterns:
|
|
recorder.warnings.append(
|
|
f"机器人 {profile.name!r} 的足端正则 {pattern!r} 在 body_names 里没有命中，"
|
|
"支撑裕度指标将被跳过。"
|
|
)
|
|
|
|
robot_bodies = set(robot.body_names)
|
|
sensor = base_env.scene.sensors.get(contact_sensor_name) if base_env.scene.sensors else None
|
|
if sensor is None:
|
|
recorder.warnings.append(
|
|
f"场景里没有接触力传感器 {contact_sensor_name!r}，支撑接触类指标将被跳过。"
|
|
)
|
|
else:
|
|
support = resolve_matching_indices(profile.support_body_patterns, sensor.body_names)
|
|
recorder.support_ids = support.indices
|
|
recorder.support_names = support.names
|
|
missing = [name for name in support.names if name not in robot_bodies]
|
|
if missing:
|
|
recorder.warnings.append(
|
|
f"接触力传感器上的 {missing} 不在机器人的 body 列表里，指标可能对不上。"
|
|
)
|
|
for pattern in support.unmatched_patterns:
|
|
if pattern in profile.support_body_patterns:
|
|
recorder.warnings.append(
|
|
f"机器人 {profile.name!r} 的支撑体正则 {pattern!r} 没有命中，"
|
|
"支撑接触类指标将被跳过。"
|
|
)
|
|
|
|
joints = resolve_matching_indices(profile.actuated_joint_patterns, robot.joint_names)
|
|
recorder.joint_ids = joints.indices
|
|
recorder.joint_names = joints.names
|
|
if not joints.indices:
|
|
recorder.warnings.append(
|
|
"驱动关节正则一个都没命中，能耗与 COT 将无法计算（会返回 nan）。"
|
|
)
|
|
|
|
recorder.mass_kg = read_total_mass(robot)
|
|
if recorder.mass_kg is None and profile.mass_kg is None:
|
|
recorder.warnings.append("既读不到资产质量、描述里也没写兜底值，COT 将返回 nan。")
|
|
|
|
return recorder
|
|
|
|
|
|
class _FrameCollector:
|
|
"""逐帧收集张量，最后一次性转成 numpy——每帧都转会明显变慢。"""
|
|
|
|
def __init__(self) -> None:
|
|
self.frames: dict[str, list[np.ndarray]] = {}
|
|
|
|
def add(self, **signals: np.ndarray) -> None:
|
|
for key, value in signals.items():
|
|
self.frames.setdefault(key, []).append(value)
|
|
|
|
def stack(self) -> dict[str, np.ndarray]:
|
|
return {key: np.stack(values, axis=0) for key, values in self.frames.items()}
|
|
|
|
|
|
def _capture(robot: Any, sensor: Any, recorder: RolloutRecorder, num_envs: int) -> dict[str, np.ndarray]:
|
|
"""读一帧信号。所有张量都要 clone，否则会被后续仿真步就地覆写。"""
|
|
|
|
def to_numpy(tensor: torch.Tensor) -> np.ndarray:
|
|
return tensor.detach().to("cpu", torch.float64).numpy()
|
|
|
|
signals = {
|
|
"base_pos_w": to_numpy(robot.data.root_pos_w.clone()),
|
|
"base_lin_vel_w": to_numpy(robot.data.root_lin_vel_w.clone()),
|
|
"base_lin_vel_b": to_numpy(robot.data.root_lin_vel_b.clone()),
|
|
"base_ang_vel_b": to_numpy(robot.data.root_ang_vel_b.clone()),
|
|
"projected_gravity_b": to_numpy(robot.data.projected_gravity_b.clone()),
|
|
"joint_torque": to_numpy(robot.data.applied_torque[:, recorder.joint_ids].clone())
|
|
if recorder.joint_ids
|
|
else np.zeros((num_envs, 0)),
|
|
"joint_vel": to_numpy(robot.data.joint_vel[:, recorder.joint_ids].clone())
|
|
if recorder.joint_ids
|
|
else np.zeros((num_envs, 0)),
|
|
"feet_pos_w": to_numpy(robot.data.body_pos_w[:, recorder.feet_ids, :2].clone())
|
|
if recorder.feet_ids
|
|
else np.zeros((num_envs, 0, 2)),
|
|
}
|
|
|
|
if sensor is not None and recorder.support_ids:
|
|
forces = sensor.data.net_forces_w
|
|
contact_time = sensor.data.current_contact_time
|
|
if forces is None:
|
|
# 传感器还没更新过（第一帧），按无接触处理
|
|
signals["support_force_n"] = np.zeros((num_envs, len(recorder.support_ids)))
|
|
signals["support_contact_time_s"] = np.zeros((num_envs, len(recorder.support_ids)))
|
|
else:
|
|
selected = forces[:, recorder.support_ids, :]
|
|
signals["support_force_n"] = to_numpy(torch.linalg.norm(selected, dim=-1).clone())
|
|
signals["support_contact_time_s"] = to_numpy(
|
|
(contact_time[:, recorder.support_ids].clone() if contact_time is not None else torch.zeros_like(selected[..., 0]))
|
|
)
|
|
else:
|
|
signals["support_force_n"] = np.zeros((num_envs, 0))
|
|
signals["support_contact_time_s"] = np.zeros((num_envs, 0))
|
|
|
|
return signals
|
|
|
|
|
|
def record_rollout(
|
|
env: Any,
|
|
policy: Any,
|
|
protocol: EvalProtocol,
|
|
profile: RobotProfile,
|
|
recorder: RolloutRecorder | None = None,
|
|
) -> Trace:
|
|
"""跑 ``protocol.steps`` 步，采集一条 :class:`Trace`。
|
|
|
|
第 0 帧记录动作执行之前的起始状态，之后每步记录一次动作之后的状态，
|
|
因此帧数比步数多一。
|
|
"""
|
|
base_env = env.unwrapped
|
|
robot = base_env.scene["robot"]
|
|
if recorder is None:
|
|
recorder = build_recorder(env, profile)
|
|
sensors = base_env.scene.sensors
|
|
sensor = sensors.get(recorder.contact_sensor_name) if sensors else None
|
|
|
|
num_envs = int(base_env.num_envs)
|
|
collector = _FrameCollector()
|
|
initial = _capture(robot, sensor, recorder, num_envs)
|
|
# 起始帧还没有任何 episode 结束，置零以保证与后续帧形状一致
|
|
initial["done"] = np.zeros(num_envs, dtype=np.float64)
|
|
collector.add(**initial)
|
|
|
|
obs = env.get_observations()
|
|
# 终端侧进度：可视化时推流画面可能因为渲染/编码出问题而冻住，这时只有终端能
|
|
# 说明仿真到底有没有在推进（推进则打印继续，冻结则这里也停住）。每 10% 打一行。
|
|
progress_every = max(1, protocol.steps // 10)
|
|
started = time.monotonic()
|
|
for step in range(protocol.steps):
|
|
with torch.inference_mode():
|
|
actions = policy(obs)
|
|
if isinstance(actions, tuple):
|
|
actions = actions[0]
|
|
obs, _, dones, _ = env.step(actions)
|
|
frame = _capture(robot, sensor, recorder, num_envs)
|
|
frame["done"] = dones.detach().to("cpu", torch.float64).numpy().reshape(num_envs)
|
|
collector.add(**frame)
|
|
if (step + 1) % progress_every == 0 or step + 1 == protocol.steps:
|
|
print(
|
|
f"[eval_kit] rollout {step + 1}/{protocol.steps} 步，用时 {time.monotonic() - started:.1f} s",
|
|
flush=True,
|
|
)
|
|
|
|
arrays = collector.stack()
|
|
return Trace(
|
|
dt=float(base_env.step_dt),
|
|
command=np.asarray(protocol.command, dtype=np.float64),
|
|
mass_kg=recorder.mass_kg,
|
|
**arrays,
|
|
)
|
|
|
|
|
|
def summarize_early_resets(trace: Trace) -> dict[str, Any]:
|
|
"""统计 episode 中途被重置的环境数量。
|
|
|
|
中途重置会让后续帧来自新 episode，路程/进度类指标会被污染。旧
|
|
``eval.py`` 没有检查这一点，所以出现过「指标看着正常其实混了两次 rollout」
|
|
的情况不容易发现。
|
|
"""
|
|
# 最后一帧的 done 属于 episode 正常结束，不算早退
|
|
early = trace.done[:-1] > 0.5
|
|
per_env = early.sum(axis=0)
|
|
affected = int(np.count_nonzero(per_env > 0))
|
|
return {
|
|
"envs_with_early_reset": affected,
|
|
"total_early_resets": int(per_env.sum()),
|
|
"max_resets_single_env": int(per_env.max()) if per_env.size else 0,
|
|
"polluted": affected > 0,
|
|
}
|