Logic
957 lines
37 KiB
Python
957 lines
37 KiB
Python
"""
|
|
VLA 策略评估脚本（简化版）
|
|
|
|
该脚本使用 agent 内置的队列管理来评估训练好的 VLA 策略。
|
|
无需单独的评估器类 - agent 处理一切！
|
|
|
|
使用方法:
|
|
python roboimi/demos/eval_vla_simple.py
|
|
python roboimi/demos/eval_vla_simple.py eval.ckpt_path=checkpoints/vla_model_final.pt
|
|
python roboimi/demos/eval_vla_simple.py eval.ckpt_path=checkpoints/vla_model_best.pt
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import json
|
|
import logging
|
|
import time
|
|
import torch
|
|
import numpy as np
|
|
import hydra
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional
|
|
from tqdm import tqdm
|
|
from omegaconf import DictConfig, OmegaConf
|
|
from hydra.utils import instantiate
|
|
from einops import rearrange
|
|
|
|
from roboimi.utils.act_ex_utils import sample_transfer_pose
|
|
from roboimi.vla.eval_utils import execute_policy_action
|
|
|
|
sys.path.append(os.getcwd())
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
if not OmegaConf.has_resolver("len"):
|
|
OmegaConf.register_new_resolver("len", lambda x: len(x))
|
|
|
|
|
|
def _configure_headless_mujoco_gl(eval_cfg: DictConfig) -> None:
|
|
if not bool(eval_cfg.get('headless', False)):
|
|
return
|
|
if os.environ.get('MUJOCO_GL'):
|
|
return
|
|
os.environ['MUJOCO_GL'] = 'egl'
|
|
log.info('headless eval detected; set MUJOCO_GL=egl')
|
|
|
|
|
|
def make_sim_env(task_name: str, headless: bool = False):
|
|
from roboimi.envs.double_pos_ctrl_env import make_sim_env as _make_sim_env_impl
|
|
return _make_sim_env_impl(task_name, headless=headless)
|
|
|
|
|
|
def load_checkpoint(
|
|
ckpt_path: str,
|
|
agent_cfg: DictConfig,
|
|
device: str = 'cuda'
|
|
) -> torch.nn.Module:
|
|
"""
|
|
从检查点加载训练好的 VLA 模型，使用 Hydra agent 配置。
|
|
|
|
Args:
|
|
ckpt_path: 检查点文件路径 (.pt)
|
|
agent_cfg: Hydra agent 配置，用于实例化
|
|
device: 加载模型的设备
|
|
|
|
Returns:
|
|
加载后的 VLAAgent 模型
|
|
"""
|
|
from pathlib import Path as PathLib
|
|
|
|
ckpt_path = PathLib(ckpt_path).absolute()
|
|
if not ckpt_path.exists():
|
|
raise FileNotFoundError(f"检查点未找到: {ckpt_path}")
|
|
|
|
log.info(f"从 {ckpt_path} 加载检查点")
|
|
checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False)
|
|
log.info(f"检查点键值: {checkpoint.keys()}")
|
|
|
|
# 加载数据集统计信息用于归一化
|
|
stats = checkpoint.get('dataset_stats', None)
|
|
|
|
# 使用数据集统计信息从 Hydra 配置实例化 agent
|
|
log.info("从配置实例化 agent...")
|
|
agent = instantiate(agent_cfg, dataset_stats=stats)
|
|
|
|
# 加载模型状态
|
|
agent.load_state_dict(checkpoint['model_state_dict'])
|
|
log.info(f"✅ 模型状态已加载 (步数: {checkpoint.get('step', 'unknown')})")
|
|
|
|
if stats is not None:
|
|
log.info(f"✅ 数据集统计信息已加载 (归一化: {stats.get('normalization_type', 'gaussian')})")
|
|
else:
|
|
# 后备方案：尝试从外部 JSON 文件加载（兼容旧检查点）
|
|
stats_path = ckpt_path.parent / 'dataset_stats.json'
|
|
if stats_path.exists():
|
|
with open(stats_path, 'r') as f:
|
|
stats = json.load(f)
|
|
log.info("✅ 数据集统计信息已从外部 JSON 加载（旧版本兼容）")
|
|
else:
|
|
log.warning("⚠️ 未找到数据集统计信息。动作将无法反归一化！")
|
|
|
|
agent.eval()
|
|
agent.to(device)
|
|
|
|
log.info(f"✅ 模型已成功加载到 {device}")
|
|
return agent, stats
|
|
|
|
|
|
def prepare_observation(
|
|
obs: Dict,
|
|
camera_names: list,
|
|
image_resize_shape: Optional[tuple[int, int]] = (224, 224),
|
|
) -> Dict:
|
|
"""
|
|
将环境观测转换为 agent 格式。
|
|
|
|
Args:
|
|
obs: 环境观测字典，包含图像和 qpos
|
|
camera_names: 摄像头名称列表
|
|
|
|
Returns:
|
|
agent 格式的观测字典
|
|
"""
|
|
# 转换图像: numpy -> tensor, HWC -> CHW
|
|
images = {}
|
|
for cam_name in camera_names:
|
|
img = obs['images'][cam_name]
|
|
if image_resize_shape is not None:
|
|
import cv2
|
|
img = cv2.resize(img, tuple(image_resize_shape), interpolation=cv2.INTER_LINEAR)
|
|
img = rearrange(img, 'h w c -> c h w')
|
|
img = torch.from_numpy(img / 255.0).float()
|
|
images[cam_name] = img
|
|
|
|
# 转换 qpos: numpy -> tensor
|
|
qpos = torch.from_numpy(obs['qpos']).float()
|
|
|
|
return {'qpos': qpos, 'images': images}
|
|
|
|
|
|
def _to_numpy_action(action: Any) -> np.ndarray:
|
|
if isinstance(action, torch.Tensor):
|
|
return action.detach().cpu().numpy().astype(np.float32, copy=True)
|
|
return np.asarray(action, dtype=np.float32).copy()
|
|
|
|
|
|
def _mean_or_zero(values: list[float]) -> float:
|
|
return float(np.mean(values)) if values else 0.0
|
|
|
|
|
|
def _stats_or_zero(values: list[float]) -> dict[str, float]:
|
|
if not values:
|
|
return {
|
|
'mean': 0.0,
|
|
'std': 0.0,
|
|
'min': 0.0,
|
|
'max': 0.0,
|
|
}
|
|
array = np.asarray(values, dtype=np.float64)
|
|
return {
|
|
'mean': float(array.mean()),
|
|
'std': float(array.std()),
|
|
'min': float(array.min()),
|
|
'max': float(array.max()),
|
|
}
|
|
|
|
|
|
def _summarize_timing_breakdown(
|
|
all_timings: dict[str, list[float]],
|
|
model_forward_flags: list[bool],
|
|
) -> dict[str, Any]:
|
|
model_forward_flags = [bool(flag) for flag in model_forward_flags]
|
|
return {
|
|
'count': int(len(model_forward_flags)),
|
|
'model_forward_count': int(sum(model_forward_flags)),
|
|
'all_steps_ms': {
|
|
stage: _stats_or_zero(values)
|
|
for stage, values in all_timings.items()
|
|
},
|
|
'model_forward_steps_ms': {
|
|
stage: _stats_or_zero(
|
|
[value for value, should_keep in zip(values, model_forward_flags) if should_keep]
|
|
)
|
|
for stage, values in all_timings.items()
|
|
},
|
|
}
|
|
|
|
|
|
def _json_friendly(value: Any) -> Any:
|
|
if isinstance(value, dict):
|
|
return {str(key): _json_friendly(item) for key, item in value.items()}
|
|
if isinstance(value, (list, tuple)):
|
|
return [_json_friendly(item) for item in value]
|
|
if isinstance(value, Path):
|
|
return str(value)
|
|
if isinstance(value, np.ndarray):
|
|
return value.tolist()
|
|
if isinstance(value, (np.integer, np.floating)):
|
|
return value.item()
|
|
return value
|
|
|
|
|
|
def _resolve_artifact_paths(eval_cfg: DictConfig) -> dict[str, Optional[str]]:
|
|
save_timing = bool(eval_cfg.get('save_timing', False))
|
|
save_trajectory = bool(
|
|
eval_cfg.get('save_trajectory', False) or eval_cfg.get('save_trajectory_npz', False)
|
|
)
|
|
save_trajectory_image = bool(eval_cfg.get('save_trajectory_image', False))
|
|
wants_artifacts = any([
|
|
bool(eval_cfg.get('save_artifacts', False)),
|
|
save_timing,
|
|
save_trajectory,
|
|
save_trajectory_image,
|
|
bool(eval_cfg.get('record_video', False)),
|
|
])
|
|
output_dir: Optional[Path] = None
|
|
if wants_artifacts:
|
|
artifact_dir = eval_cfg.get('artifact_dir', None)
|
|
if artifact_dir:
|
|
output_dir = Path(str(artifact_dir)).expanduser().resolve()
|
|
else:
|
|
ckpt_stem = Path(str(eval_cfg.ckpt_path)).stem or 'rollout'
|
|
timestamp = time.strftime('%Y%m%d-%H%M%S')
|
|
output_dir = (Path.cwd() / 'rollout_artifacts' / f'{ckpt_stem}-{timestamp}').resolve()
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
video_camera_name = None
|
|
if bool(eval_cfg.get('record_video', False)):
|
|
configured_camera_name = eval_cfg.get('video_camera_name', None)
|
|
if configured_camera_name is None:
|
|
configured_camera_name = eval_cfg.get('video_camera', None)
|
|
if configured_camera_name is not None:
|
|
video_camera_name = str(configured_camera_name)
|
|
elif eval_cfg.get('camera_names'):
|
|
video_camera_name = str(eval_cfg.camera_names[0])
|
|
else:
|
|
raise ValueError('record_video=true requires eval.video_camera_name or a non-empty eval.camera_names')
|
|
|
|
trajectory_image_camera_name = None
|
|
if save_trajectory_image:
|
|
configured_camera_name = eval_cfg.get('trajectory_image_camera_name', None)
|
|
if configured_camera_name is None:
|
|
configured_camera_name = eval_cfg.get('trajectory_image_camera', None)
|
|
if configured_camera_name is not None:
|
|
trajectory_image_camera_name = str(configured_camera_name)
|
|
elif eval_cfg.get('camera_names'):
|
|
camera_names = [str(name) for name in eval_cfg.camera_names]
|
|
trajectory_image_camera_name = 'front' if 'front' in camera_names else camera_names[0]
|
|
else:
|
|
raise ValueError(
|
|
'save_trajectory_image=true requires eval.trajectory_image_camera_name '
|
|
'or a non-empty eval.camera_names'
|
|
)
|
|
|
|
return {
|
|
'output_dir': str(output_dir) if output_dir is not None else None,
|
|
'summary_json': (
|
|
str(output_dir / 'rollout_summary.json')
|
|
if output_dir is not None and bool(eval_cfg.get('save_summary_json', False))
|
|
else None
|
|
),
|
|
'timing_json': (
|
|
str(output_dir / 'timing.json')
|
|
if output_dir is not None and save_timing
|
|
else None
|
|
),
|
|
'trajectory_npz': (
|
|
str(output_dir / 'trajectory.npz')
|
|
if output_dir is not None and save_trajectory
|
|
else None
|
|
),
|
|
'video_mp4': (
|
|
str(output_dir / f'rollout_{video_camera_name}.mp4')
|
|
if output_dir is not None and bool(eval_cfg.get('record_video', False))
|
|
and video_camera_name is not None
|
|
else None
|
|
),
|
|
'video_camera_name': video_camera_name,
|
|
'trajectory_image_camera_name': trajectory_image_camera_name,
|
|
}
|
|
|
|
|
|
def _get_video_frame(obs: Dict, camera_name: Optional[str]) -> Optional[np.ndarray]:
|
|
if camera_name is None:
|
|
return None
|
|
frame = obs['images'][camera_name]
|
|
frame = np.asarray(frame)
|
|
if frame.ndim != 3 or frame.shape[2] != 3:
|
|
raise ValueError(
|
|
f'Video frame for camera {camera_name} must have shape (H, W, 3), got {frame.shape}'
|
|
)
|
|
if frame.dtype != np.uint8:
|
|
frame = np.clip(frame, 0, 255).astype(np.uint8)
|
|
return frame
|
|
|
|
|
|
def _open_video_writer(output_path: str, frame_size: tuple[int, int], fps: int):
|
|
import cv2
|
|
|
|
output_path = str(output_path)
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
|
writer = cv2.VideoWriter(output_path, fourcc, float(fps), frame_size)
|
|
if not writer.isOpened():
|
|
raise RuntimeError(f'无法打开视频输出: {output_path}')
|
|
return writer
|
|
|
|
|
|
def _episode_trajectory_image_path(
|
|
artifact_paths: dict[str, Optional[str]],
|
|
episode_idx: int,
|
|
) -> Optional[str]:
|
|
output_dir = artifact_paths.get('output_dir')
|
|
camera_name = artifact_paths.get('trajectory_image_camera_name')
|
|
if output_dir is None or camera_name is None:
|
|
return None
|
|
return str(Path(output_dir) / f'rollout_{camera_name}_ep{episode_idx + 1:02d}_trajectory.png')
|
|
|
|
|
|
def _build_action_trajectory_positions(raw_actions: list[np.ndarray]) -> dict[str, np.ndarray]:
|
|
if not raw_actions:
|
|
empty = np.zeros((0, 3), dtype=np.float32)
|
|
return {'left': empty, 'right': empty}
|
|
raw_action_array = np.asarray(raw_actions, dtype=np.float32)
|
|
return {
|
|
'left': raw_action_array[:, :3].astype(np.float32, copy=True),
|
|
'right': raw_action_array[:, 7:10].astype(np.float32, copy=True),
|
|
}
|
|
|
|
|
|
def _append_capsule_markers_to_scene(scene, markers: list[dict]) -> None:
|
|
import mujoco
|
|
|
|
for marker in markers:
|
|
if scene.ngeom >= scene.maxgeom:
|
|
break
|
|
geom = scene.geoms[scene.ngeom]
|
|
mujoco.mjv_initGeom(
|
|
geom,
|
|
mujoco.mjtGeom.mjGEOM_CAPSULE,
|
|
np.zeros(3, dtype=np.float64),
|
|
np.zeros(3, dtype=np.float64),
|
|
np.eye(3, dtype=np.float64).reshape(-1),
|
|
np.asarray(marker['rgba'], dtype=np.float32),
|
|
)
|
|
mujoco.mjv_connector(
|
|
geom,
|
|
mujoco.mjtGeom.mjGEOM_CAPSULE,
|
|
float(marker['radius']),
|
|
np.asarray(marker['from'], dtype=np.float64),
|
|
np.asarray(marker['to'], dtype=np.float64),
|
|
)
|
|
scene.ngeom += 1
|
|
|
|
|
|
def _save_rollout_trajectory_image(
|
|
env,
|
|
output_path: Optional[str],
|
|
raw_actions: list[np.ndarray],
|
|
camera_name: Optional[str],
|
|
*,
|
|
line_radius: float = 0.004,
|
|
max_markers: int = 1500,
|
|
) -> Optional[str]:
|
|
if output_path is None or camera_name is None:
|
|
return None
|
|
|
|
# IMPORTANT:
|
|
# Keep this import lazy so headless rollout can set MUJOCO_GL=egl before
|
|
# anything imports mujoco. Importing this helper at module import time would
|
|
# pull in mujoco too early on remote headless hosts and make rollout fail
|
|
# with gladLoadGL / missing DISPLAY errors.
|
|
from roboimi.utils.raw_action_trajectory_viewer import build_trajectory_capsule_markers
|
|
|
|
output_path = str(output_path)
|
|
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
frame = None
|
|
owned_renderer = None
|
|
positions = _build_action_trajectory_positions(raw_actions)
|
|
markers = build_trajectory_capsule_markers(
|
|
positions,
|
|
max_markers=max_markers,
|
|
radius=line_radius,
|
|
)
|
|
|
|
try:
|
|
renderer = None
|
|
if callable(getattr(env, '_get_or_create_offscreen_renderer', None)):
|
|
renderer = env._get_or_create_offscreen_renderer()
|
|
elif hasattr(env, 'mj_model') and hasattr(env, 'mj_data'):
|
|
import mujoco
|
|
|
|
renderer = mujoco.Renderer(env.mj_model, height=480, width=640)
|
|
owned_renderer = renderer
|
|
|
|
if renderer is not None and hasattr(env, 'mj_data'):
|
|
renderer.update_scene(env.mj_data, camera=str(camera_name))
|
|
if markers:
|
|
_append_capsule_markers_to_scene(renderer.scene, markers)
|
|
frame = renderer.render()[:, :, ::-1]
|
|
finally:
|
|
if owned_renderer is not None:
|
|
owned_renderer.close()
|
|
|
|
if frame is None and callable(getattr(env, '_get_image_obs', None)):
|
|
obs = env._get_image_obs()
|
|
frame = _get_video_frame(obs, str(camera_name))
|
|
|
|
if frame is None:
|
|
return None
|
|
|
|
import cv2
|
|
|
|
cv2.imwrite(output_path, frame)
|
|
return output_path
|
|
|
|
|
|
class _RolloutVideoRecorder:
|
|
def __init__(self, output_path: Optional[str], fps: int):
|
|
self.output_path = output_path
|
|
self.fps = int(fps)
|
|
self.writer = None
|
|
|
|
def write(self, frame: Optional[np.ndarray]):
|
|
if self.output_path is None or frame is None:
|
|
return
|
|
if self.writer is None:
|
|
frame_size = (int(frame.shape[1]), int(frame.shape[0]))
|
|
self.writer = _open_video_writer(self.output_path, frame_size, self.fps)
|
|
self.writer.write(frame)
|
|
|
|
def close(self):
|
|
if self.writer is not None:
|
|
self.writer.release()
|
|
self.writer = None
|
|
|
|
|
|
def _read_body_pose(env, body_name: str):
|
|
try:
|
|
if callable(getattr(env, 'getBodyPos', None)) and callable(getattr(env, 'getBodyQuat', None)):
|
|
pos = env.getBodyPos(body_name)
|
|
quat = env.getBodyQuat(body_name)
|
|
else:
|
|
body = env.mj_data.body(body_name)
|
|
pos = body.xpos
|
|
quat = body.xquat
|
|
except Exception:
|
|
return None
|
|
|
|
return {
|
|
'pos': np.asarray(pos, dtype=np.float32).copy(),
|
|
'quat': np.asarray(quat, dtype=np.float32).copy(),
|
|
}
|
|
|
|
|
|
def _get_executed_ee_poses(env) -> dict[str, np.ndarray]:
|
|
candidates = {
|
|
'left_link7': ('left_link7', 'eef_left'),
|
|
'right_link7': ('right_link7', 'eef_right'),
|
|
'eef_left': ('eef_left', 'left_link7'),
|
|
'eef_right': ('eef_right', 'right_link7'),
|
|
}
|
|
poses = {}
|
|
for body_key, body_names in candidates.items():
|
|
pose = None
|
|
for body_name in body_names:
|
|
pose = _read_body_pose(env, body_name)
|
|
if pose is not None:
|
|
break
|
|
if pose is None:
|
|
pose = {
|
|
'pos': np.full(3, np.nan, dtype=np.float32),
|
|
'quat': np.full(4, np.nan, dtype=np.float32),
|
|
}
|
|
poses[f'{body_key}_pos'] = pose['pos']
|
|
poses[f'{body_key}_quat'] = pose['quat']
|
|
return poses
|
|
|
|
|
|
def _empty_rollout_trajectory() -> dict[str, list]:
|
|
return {
|
|
'episode_index': [],
|
|
'step': [],
|
|
'reward': [],
|
|
'raw_action': [],
|
|
'applied_action': [],
|
|
'executed_left_link7_pos': [],
|
|
'executed_left_link7_quat': [],
|
|
'executed_right_link7_pos': [],
|
|
'executed_right_link7_quat': [],
|
|
'executed_eef_left_pos': [],
|
|
'executed_eef_left_quat': [],
|
|
'executed_eef_right_pos': [],
|
|
'executed_eef_right_quat': [],
|
|
'model_inference_triggered': [],
|
|
'obs_read_time_ms': [],
|
|
'preprocess_time_ms': [],
|
|
'inference_time_ms': [],
|
|
'env_step_time_ms': [],
|
|
'total_time_ms': [],
|
|
}
|
|
|
|
|
|
def _append_rollout_step(
|
|
storage: dict[str, list],
|
|
episode_index: int,
|
|
timestep: int,
|
|
reward: Optional[float],
|
|
raw_action: np.ndarray,
|
|
executed_action: np.ndarray,
|
|
executed_poses: dict[str, np.ndarray],
|
|
timing_ms: dict[str, float],
|
|
model_inference_triggered: bool,
|
|
):
|
|
storage['episode_index'].append(int(episode_index))
|
|
storage['step'].append(int(timestep))
|
|
storage['reward'].append(float(reward) if reward is not None else np.nan)
|
|
storage['raw_action'].append(raw_action.astype(np.float32, copy=True))
|
|
storage['applied_action'].append(executed_action.astype(np.float32, copy=True))
|
|
storage['executed_left_link7_pos'].append(executed_poses['left_link7_pos'])
|
|
storage['executed_left_link7_quat'].append(executed_poses['left_link7_quat'])
|
|
storage['executed_right_link7_pos'].append(executed_poses['right_link7_pos'])
|
|
storage['executed_right_link7_quat'].append(executed_poses['right_link7_quat'])
|
|
storage['executed_eef_left_pos'].append(executed_poses['eef_left_pos'])
|
|
storage['executed_eef_left_quat'].append(executed_poses['eef_left_quat'])
|
|
storage['executed_eef_right_pos'].append(executed_poses['eef_right_pos'])
|
|
storage['executed_eef_right_quat'].append(executed_poses['eef_right_quat'])
|
|
storage['model_inference_triggered'].append(bool(model_inference_triggered))
|
|
for key, value in timing_ms.items():
|
|
storage[key].append(float(value))
|
|
|
|
|
|
def _save_rollout_trajectory_npz(output_path: str, storage: dict[str, list]):
|
|
step = np.asarray(storage['step'], dtype=np.int32)
|
|
raw_action = np.asarray(storage['raw_action'], dtype=np.float32)
|
|
applied_action = np.asarray(storage['applied_action'], dtype=np.float32)
|
|
executed_left_link7_pos = np.asarray(storage['executed_left_link7_pos'], dtype=np.float32)
|
|
executed_left_link7_quat = np.asarray(storage['executed_left_link7_quat'], dtype=np.float32)
|
|
executed_right_link7_pos = np.asarray(storage['executed_right_link7_pos'], dtype=np.float32)
|
|
executed_right_link7_quat = np.asarray(storage['executed_right_link7_quat'], dtype=np.float32)
|
|
executed_eef_left_pos = np.asarray(storage['executed_eef_left_pos'], dtype=np.float32)
|
|
executed_eef_left_quat = np.asarray(storage['executed_eef_left_quat'], dtype=np.float32)
|
|
executed_eef_right_pos = np.asarray(storage['executed_eef_right_pos'], dtype=np.float32)
|
|
executed_eef_right_quat = np.asarray(storage['executed_eef_right_quat'], dtype=np.float32)
|
|
np.savez_compressed(
|
|
output_path,
|
|
episode_index=np.asarray(storage['episode_index'], dtype=np.int32),
|
|
step=step,
|
|
timestep=step,
|
|
reward=np.asarray(storage['reward'], dtype=np.float32),
|
|
raw_action=raw_action,
|
|
raw_predicted_ee_action=raw_action,
|
|
applied_action=applied_action,
|
|
executed_ee_action=applied_action,
|
|
executed_left_link7_pos=executed_left_link7_pos,
|
|
executed_left_link7_quat=executed_left_link7_quat,
|
|
executed_right_link7_pos=executed_right_link7_pos,
|
|
executed_right_link7_quat=executed_right_link7_quat,
|
|
executed_eef_left_pos=executed_eef_left_pos,
|
|
executed_eef_left_quat=executed_eef_left_quat,
|
|
executed_eef_right_pos=executed_eef_right_pos,
|
|
executed_eef_right_quat=executed_eef_right_quat,
|
|
left_ee_pos=executed_eef_left_pos,
|
|
left_ee_quat=executed_eef_left_quat,
|
|
right_ee_pos=executed_eef_right_pos,
|
|
right_ee_quat=executed_eef_right_quat,
|
|
model_inference_triggered=np.asarray(storage['model_inference_triggered'], dtype=bool),
|
|
obs_read_time_ms=np.asarray(storage['obs_read_time_ms'], dtype=np.float32),
|
|
preprocess_time_ms=np.asarray(storage['preprocess_time_ms'], dtype=np.float32),
|
|
inference_time_ms=np.asarray(storage['inference_time_ms'], dtype=np.float32),
|
|
env_step_time_ms=np.asarray(storage['env_step_time_ms'], dtype=np.float32),
|
|
total_time_ms=np.asarray(storage['total_time_ms'], dtype=np.float32),
|
|
)
|
|
|
|
|
|
def _save_summary_json(output_path: str, summary: dict[str, Any]):
|
|
with open(output_path, 'w', encoding='utf-8') as f:
|
|
json.dump(_json_friendly(summary), f, ensure_ascii=False, indent=2)
|
|
|
|
|
|
class ActionSmoother:
|
|
"""
|
|
动作平滑器（指数移动平均）
|
|
用于平滑执行动作以获得更稳定的控制
|
|
"""
|
|
|
|
def __init__(self, alpha: float = 0.3):
|
|
"""
|
|
Args:
|
|
alpha: 平滑系数 (0-1)，值越大越重视当前动作
|
|
"""
|
|
self.alpha = alpha
|
|
self.prev_action = None
|
|
|
|
def smooth(self, action: np.ndarray) -> np.ndarray:
|
|
"""
|
|
平滑动作
|
|
|
|
Args:
|
|
action: 当前动作
|
|
|
|
Returns:
|
|
平滑后的动作
|
|
"""
|
|
if self.prev_action is None:
|
|
smoothed = action
|
|
else:
|
|
smoothed = self.alpha * action + (1 - self.alpha) * self.prev_action
|
|
self.prev_action = smoothed
|
|
return smoothed
|
|
|
|
def reset(self):
|
|
"""重置平滑器状态"""
|
|
self.prev_action = None
|
|
|
|
|
|
def _close_env(env):
|
|
if env is None:
|
|
return
|
|
|
|
if hasattr(env, 'exit_flag'):
|
|
env.exit_flag = True
|
|
|
|
cam_thread = getattr(env, 'cam_thread', None)
|
|
if cam_thread is not None and hasattr(cam_thread, 'join'):
|
|
cam_thread.join(timeout=1.0)
|
|
|
|
viewer = getattr(env, 'viewer', None)
|
|
if viewer is not None and hasattr(viewer, 'close'):
|
|
viewer.close()
|
|
|
|
|
|
def _run_eval(cfg: DictConfig):
|
|
"""
|
|
使用 agent 内置队列管理的简化版 VLA 评估
|
|
|
|
所有评估参数来自 vla/conf/eval.yaml，合并到 cfg 中。
|
|
命令行覆盖: python eval_vla_simple.py eval.ckpt_path=... eval.num_episodes=5
|
|
"""
|
|
|
|
# 打印配置
|
|
print("=" * 80)
|
|
print("VLA 评估配置:")
|
|
print("=" * 80)
|
|
print(OmegaConf.to_yaml(cfg))
|
|
print("=" * 80)
|
|
|
|
eval_cfg = cfg.eval
|
|
_configure_headless_mujoco_gl(eval_cfg)
|
|
device = eval_cfg.device
|
|
camera_names = list(eval_cfg.camera_names)
|
|
artifact_paths = _resolve_artifact_paths(eval_cfg)
|
|
video_recorder = _RolloutVideoRecorder(
|
|
output_path=artifact_paths['video_mp4'],
|
|
fps=int(eval_cfg.get('video_fps', 30)),
|
|
)
|
|
rollout_trajectory = _empty_rollout_trajectory()
|
|
global_obs_read_times_ms = []
|
|
global_preprocess_times_ms = []
|
|
global_inference_times_ms = []
|
|
global_env_step_times_ms = []
|
|
global_total_times_ms = []
|
|
global_model_forward_flags = []
|
|
|
|
# =========================================================================
|
|
# 加载模型
|
|
# =========================================================================
|
|
log.info(f"🚀 从 {eval_cfg.ckpt_path} 加载模型...")
|
|
agent, dataset_stats = load_checkpoint(
|
|
ckpt_path=eval_cfg.ckpt_path,
|
|
agent_cfg=cfg.agent,
|
|
device=device
|
|
)
|
|
vision_encoder = getattr(agent, 'vision_encoder', None)
|
|
image_resize_shape = getattr(vision_encoder, 'eval_image_resize_shape', (224, 224))
|
|
|
|
# 重置 agent 的队列
|
|
agent.reset()
|
|
|
|
# 可选：动作平滑器
|
|
smoother = ActionSmoother(alpha=eval_cfg.smooth_alpha) if eval_cfg.use_smoothing else None
|
|
|
|
# =========================================================================
|
|
# 创建环境
|
|
# =========================================================================
|
|
env = make_sim_env(eval_cfg.task_name, headless=eval_cfg.headless)
|
|
|
|
# =========================================================================
|
|
# 运行评估回合
|
|
# =========================================================================
|
|
all_stats = []
|
|
episode_rewards = []
|
|
episode_max_rewards = []
|
|
try:
|
|
for episode_idx in range(eval_cfg.num_episodes):
|
|
print(f"\n{'='*60}")
|
|
print(f"回合 {episode_idx + 1}/{eval_cfg.num_episodes}")
|
|
print(f"{'='*60}\n")
|
|
|
|
box_pos = sample_transfer_pose()
|
|
env.reset(box_pos)
|
|
|
|
# 为新回合重置 agent 队列
|
|
agent.reset()
|
|
if smoother:
|
|
smoother.reset()
|
|
|
|
# 计时统计
|
|
obs_read_times_ms = []
|
|
preprocess_times_ms = []
|
|
inference_times_ms = []
|
|
env_step_times_ms = []
|
|
total_times_ms = []
|
|
model_forward_flags = []
|
|
episode_reward = 0.0
|
|
episode_max_reward = float('-inf')
|
|
episode_raw_actions: list[np.ndarray] = []
|
|
|
|
with torch.inference_mode():
|
|
for t in tqdm(range(eval_cfg.max_timesteps), desc=f"回合 {episode_idx + 1}"):
|
|
start_total = time.perf_counter()
|
|
|
|
# 从环境获取观测
|
|
obs = env._get_image_obs()
|
|
qpos_obs = env._get_qpos_obs()
|
|
obs['qpos'] = qpos_obs['qpos']
|
|
end_obs_read = time.perf_counter()
|
|
|
|
video_frame = _get_video_frame(obs, artifact_paths['video_camera_name'])
|
|
video_recorder.write(video_frame)
|
|
|
|
# 准备给 agent 的观测
|
|
observation = prepare_observation(
|
|
obs,
|
|
camera_names,
|
|
image_resize_shape=image_resize_shape,
|
|
)
|
|
end_preprocess = time.perf_counter()
|
|
|
|
# 选择动作（agent 内部处理队列管理）
|
|
action_queue = getattr(agent, '_queues', {}).get('action', None)
|
|
model_inference_triggered = len(action_queue) == 0 if action_queue is not None else True
|
|
start_inference = time.perf_counter()
|
|
action = agent.select_action(observation)
|
|
|
|
if str(device).startswith('cuda') and torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
end_inference = time.perf_counter()
|
|
|
|
# 转换为 numpy
|
|
raw_action = _to_numpy_action(action)
|
|
episode_raw_actions.append(raw_action.astype(np.float32, copy=True))
|
|
|
|
# 调试：打印当前时间步的动作（由配置控制）
|
|
if eval_cfg.get('verbose_action', False):
|
|
print(f"\n[Step {t:3d}] 预测动作: {raw_action}")
|
|
print(f" - 动作形状: {raw_action.shape}")
|
|
print(f" - 动作范围: [{raw_action.min():.4f}, {raw_action.max():.4f}]")
|
|
print(f" - 动作均值: {raw_action.mean():.4f}, 标准差: {raw_action.std():.4f}")
|
|
|
|
# 可选：平滑动作
|
|
executed_action = raw_action.copy()
|
|
if smoother:
|
|
executed_action = smoother.smooth(executed_action)
|
|
|
|
# 执行动作
|
|
start_env_step = time.perf_counter()
|
|
execute_policy_action(env, executed_action)
|
|
end_env_step = time.perf_counter()
|
|
executed_poses = _get_executed_ee_poses(env)
|
|
reward = getattr(env, 'rew', None)
|
|
if reward is not None:
|
|
reward = float(reward)
|
|
episode_reward += reward
|
|
episode_max_reward = max(episode_max_reward, reward)
|
|
if not eval_cfg.headless:
|
|
env.render()
|
|
|
|
end_total = time.perf_counter()
|
|
|
|
step_timing_ms = {
|
|
'obs_read_time_ms': (end_obs_read - start_total) * 1000.0,
|
|
'preprocess_time_ms': (end_preprocess - end_obs_read) * 1000.0,
|
|
'inference_time_ms': (end_inference - start_inference) * 1000.0,
|
|
'env_step_time_ms': (end_env_step - start_env_step) * 1000.0,
|
|
'total_time_ms': (end_total - start_total) * 1000.0,
|
|
}
|
|
|
|
# 记录计时
|
|
obs_read_times_ms.append(step_timing_ms['obs_read_time_ms'])
|
|
preprocess_times_ms.append(step_timing_ms['preprocess_time_ms'])
|
|
inference_times_ms.append(step_timing_ms['inference_time_ms'])
|
|
env_step_times_ms.append(step_timing_ms['env_step_time_ms'])
|
|
total_times_ms.append(step_timing_ms['total_time_ms'])
|
|
model_forward_flags.append(bool(model_inference_triggered))
|
|
global_obs_read_times_ms.append(step_timing_ms['obs_read_time_ms'])
|
|
global_preprocess_times_ms.append(step_timing_ms['preprocess_time_ms'])
|
|
global_inference_times_ms.append(step_timing_ms['inference_time_ms'])
|
|
global_env_step_times_ms.append(step_timing_ms['env_step_time_ms'])
|
|
global_total_times_ms.append(step_timing_ms['total_time_ms'])
|
|
global_model_forward_flags.append(bool(model_inference_triggered))
|
|
|
|
if artifact_paths['trajectory_npz'] is not None:
|
|
_append_rollout_step(
|
|
rollout_trajectory,
|
|
episode_index=episode_idx,
|
|
timestep=t,
|
|
reward=reward,
|
|
raw_action=raw_action,
|
|
executed_action=executed_action,
|
|
executed_poses=executed_poses,
|
|
timing_ms=step_timing_ms,
|
|
model_inference_triggered=model_inference_triggered,
|
|
)
|
|
|
|
# =========================================================================
|
|
# 打印回合统计
|
|
# =========================================================================
|
|
avg_obs_read_time_ms = _mean_or_zero(obs_read_times_ms)
|
|
avg_preprocess_time_ms = _mean_or_zero(preprocess_times_ms)
|
|
avg_inference_time_ms = _mean_or_zero(inference_times_ms)
|
|
avg_env_step_time_ms = _mean_or_zero(env_step_times_ms)
|
|
avg_total_time_ms = _mean_or_zero(total_times_ms)
|
|
timing_breakdown = _summarize_timing_breakdown(
|
|
{
|
|
'obs_read': obs_read_times_ms,
|
|
'preprocess': preprocess_times_ms,
|
|
'inference': inference_times_ms,
|
|
'env_step': env_step_times_ms,
|
|
'loop_total': total_times_ms,
|
|
},
|
|
model_forward_flags,
|
|
)
|
|
episode_artifact_paths = {
|
|
'video': artifact_paths['video_mp4'],
|
|
'trajectory': artifact_paths['trajectory_npz'],
|
|
'trajectory_image': _save_rollout_trajectory_image(
|
|
env,
|
|
_episode_trajectory_image_path(artifact_paths, episode_idx),
|
|
episode_raw_actions,
|
|
artifact_paths['trajectory_image_camera_name'],
|
|
),
|
|
'timing': artifact_paths['timing_json'] or artifact_paths['summary_json'],
|
|
}
|
|
|
|
stats = {
|
|
'inference_fps': 1000.0 / avg_inference_time_ms if avg_inference_time_ms > 0 else 0.0,
|
|
'control_fps': 1000.0 / avg_total_time_ms if avg_total_time_ms > 0 else 0.0,
|
|
'avg_obs_read_time_ms': avg_obs_read_time_ms,
|
|
'avg_preprocess_time_ms': avg_preprocess_time_ms,
|
|
'avg_inference_time_ms': avg_inference_time_ms,
|
|
'avg_env_step_time_ms': avg_env_step_time_ms,
|
|
'avg_total_time_ms': avg_total_time_ms,
|
|
'num_inferences': int(sum(model_forward_flags)),
|
|
'num_model_forwards': int(sum(model_forward_flags)),
|
|
'num_steps': len(total_times_ms),
|
|
'episode_reward': float(episode_reward),
|
|
'episode_max_reward': (
|
|
float(episode_max_reward) if episode_max_reward != float('-inf') else None
|
|
),
|
|
'artifact_paths': episode_artifact_paths,
|
|
'timing_breakdown_ms': timing_breakdown['all_steps_ms'],
|
|
'timing_summary': timing_breakdown,
|
|
}
|
|
all_stats.append(stats)
|
|
episode_rewards.append(float(episode_reward))
|
|
if episode_max_reward != float('-inf'):
|
|
episode_max_rewards.append(float(episode_max_reward))
|
|
|
|
print(f"\n回合 {episode_idx + 1} 完成 ({eval_cfg.max_timesteps} 时间步)")
|
|
print(f" 模型推理 FPS: {stats['inference_fps']:.2f} Hz")
|
|
print(f" 控制循环 FPS: {stats['control_fps']:.2f} Hz")
|
|
print(f" 平均读观测时间: {stats['avg_obs_read_time_ms']:.2f} ms")
|
|
print(f" 平均预处理时间: {stats['avg_preprocess_time_ms']:.2f} ms")
|
|
print(f" 平均推理时间: {stats['avg_inference_time_ms']:.2f} ms")
|
|
print(f" 平均环境步进时间: {stats['avg_env_step_time_ms']:.2f} ms")
|
|
print(f" 平均总时间: {stats['avg_total_time_ms']:.2f} ms")
|
|
print(f" 总推理次数: {stats['num_inferences']}")
|
|
print(f" 回合累计奖励: {stats['episode_reward']:.2f}")
|
|
|
|
# =========================================================================
|
|
# 总体统计
|
|
# =========================================================================
|
|
print(f"\n{'='*60}")
|
|
print("评估完成!")
|
|
print(f"{'='*60}")
|
|
|
|
summary = {
|
|
'num_episodes': int(eval_cfg.num_episodes),
|
|
'episode_rewards': episode_rewards,
|
|
'episode_max_rewards': episode_max_rewards,
|
|
'avg_reward': float(np.mean(episode_rewards)) if episode_rewards else 0.0,
|
|
'avg_max_reward': float(np.mean(episode_max_rewards)) if episode_max_rewards else 0.0,
|
|
'episodes': all_stats,
|
|
'artifact_dir': artifact_paths['output_dir'],
|
|
'artifacts': artifact_paths,
|
|
}
|
|
|
|
if all_stats:
|
|
avg_inference_fps = np.mean([s['inference_fps'] for s in all_stats])
|
|
avg_control_fps = np.mean([s['control_fps'] for s in all_stats])
|
|
avg_obs_read_time = _mean_or_zero(global_obs_read_times_ms)
|
|
avg_preprocess_time = _mean_or_zero(global_preprocess_times_ms)
|
|
avg_inference_time = _mean_or_zero(global_inference_times_ms)
|
|
avg_env_step_time = _mean_or_zero(global_env_step_times_ms)
|
|
avg_total_time = _mean_or_zero(global_total_times_ms)
|
|
summary.update({
|
|
'avg_inference_fps': float(avg_inference_fps),
|
|
'avg_control_fps': float(avg_control_fps),
|
|
'avg_obs_read_time_ms': float(avg_obs_read_time),
|
|
'avg_preprocess_time_ms': float(avg_preprocess_time),
|
|
'avg_inference_time_ms': float(avg_inference_time),
|
|
'avg_env_step_time_ms': float(avg_env_step_time),
|
|
'avg_total_time_ms': float(avg_total_time),
|
|
'timing_summary': _summarize_timing_breakdown(
|
|
{
|
|
'obs_read': global_obs_read_times_ms,
|
|
'preprocess': global_preprocess_times_ms,
|
|
'inference': global_inference_times_ms,
|
|
'env_step': global_env_step_times_ms,
|
|
'loop_total': global_total_times_ms,
|
|
},
|
|
global_model_forward_flags,
|
|
),
|
|
})
|
|
|
|
print(f"\n总体统计 ({eval_cfg.num_episodes} 个回合):")
|
|
print(f" 平均模型推理 FPS: {avg_inference_fps:.2f} Hz")
|
|
print(f" 平均控制循环 FPS: {avg_control_fps:.2f} Hz")
|
|
print(f" 平均读观测时间: {avg_obs_read_time:.2f} ms")
|
|
print(f" 平均预处理时间: {avg_preprocess_time:.2f} ms")
|
|
print(f" 平均推理时间: {avg_inference_time:.2f} ms")
|
|
print(f" 平均环境步进时间: {avg_env_step_time:.2f} ms")
|
|
print(f" 平均总时间: {avg_total_time:.2f} ms")
|
|
print(f" 平均累计奖励: {summary['avg_reward']:.2f}")
|
|
|
|
if artifact_paths['trajectory_npz'] is not None:
|
|
_save_rollout_trajectory_npz(artifact_paths['trajectory_npz'], rollout_trajectory)
|
|
if artifact_paths['summary_json'] is not None:
|
|
_save_summary_json(artifact_paths['summary_json'], summary)
|
|
if artifact_paths['timing_json'] is not None:
|
|
_save_summary_json(artifact_paths['timing_json'], summary.get('timing_summary', {}))
|
|
print()
|
|
return _json_friendly(summary)
|
|
finally:
|
|
video_recorder.close()
|
|
_close_env(env)
|
|
|
|
|
|
@hydra.main(version_base=None, config_path="../../vla/conf", config_name="config")
|
|
def main(cfg: DictConfig):
|
|
return _run_eval(cfg)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|