# 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, }