Logic
1973 lines
74 KiB
Python
1973 lines
74 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 queue
|
||
import concurrent.futures
|
||
import multiprocessing
|
||
import torch
|
||
import numpy as np
|
||
import hydra
|
||
from pathlib import Path
|
||
from collections import deque
|
||
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.envs.double_pos_ctrl_env import make_sim_env
|
||
from roboimi.utils.act_ex_utils import (
|
||
sample_air_insert_socket_peg_state,
|
||
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 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),
|
||
task_description: Optional[str] = None,
|
||
) -> Dict:
|
||
"""
|
||
将环境观测转换为 agent 格式。
|
||
|
||
Args:
|
||
obs: 环境观测字典，包含图像和 qpos
|
||
camera_names: 摄像头名称列表
|
||
|
||
Returns:
|
||
agent 格式的观测字典
|
||
"""
|
||
import cv2
|
||
|
||
# 转换图像: numpy -> tensor, HWC -> CHW
|
||
images = {}
|
||
for cam_name in camera_names:
|
||
img = obs['images'][cam_name]
|
||
if image_resize_shape is not None:
|
||
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()
|
||
|
||
observation = {'qpos': qpos, 'images': images}
|
||
if 'task' in obs:
|
||
observation['task'] = obs['task']
|
||
elif task_description is not None:
|
||
observation['task'] = task_description
|
||
return observation
|
||
|
||
|
||
def _normalize_resize_shape(shape) -> Optional[tuple[int, int]]:
|
||
if shape is None:
|
||
return None
|
||
normalized = tuple(int(v) for v in shape)
|
||
if len(normalized) != 2:
|
||
raise ValueError(f'image resize shape must contain exactly two values, got {normalized}')
|
||
return normalized
|
||
|
||
|
||
def _resolve_eval_image_resize_shape(cfg: DictConfig) -> Optional[tuple[int, int]]:
|
||
image_resize_shape = cfg.get('data', {}).get('image_resize_shape', (224, 224))
|
||
agent_cfg = cfg.agent
|
||
if 'eval_image_resize_shape' in agent_cfg:
|
||
return _normalize_resize_shape(agent_cfg.get('eval_image_resize_shape'))
|
||
for backbone_key in ('vision_backbone', 'condition_encoder'):
|
||
backbone_cfg = agent_cfg.get(backbone_key, None)
|
||
if backbone_cfg is not None and 'eval_image_resize_shape' in backbone_cfg:
|
||
image_resize_shape = backbone_cfg.get('eval_image_resize_shape')
|
||
break
|
||
return _normalize_resize_shape(image_resize_shape)
|
||
|
||
|
||
def _resolve_policy_camera_names(cfg: DictConfig) -> list[str]:
|
||
agent_cfg = cfg.agent
|
||
eval_cfg = cfg.eval
|
||
camera_names = agent_cfg.get('camera_names', None)
|
||
if camera_names is not None:
|
||
return list(camera_names)
|
||
|
||
agent_target = str(agent_cfg.get('_target_', ''))
|
||
if agent_target.endswith('VLAAgentGr00tDiT'):
|
||
return list(eval_cfg.camera_names)
|
||
return sorted(eval_cfg.camera_names)
|
||
|
||
|
||
def _new_local_policy_queues(obs_horizon: int) -> dict[str, deque]:
|
||
return {
|
||
'qpos': deque(maxlen=int(obs_horizon)),
|
||
'images': deque(maxlen=int(obs_horizon)),
|
||
'task': deque(maxlen=int(obs_horizon)),
|
||
'action': deque(),
|
||
}
|
||
|
||
|
||
def _populate_local_policy_queues(
|
||
queues: dict[str, deque],
|
||
observation: Dict[str, torch.Tensor],
|
||
) -> None:
|
||
if 'qpos' in observation:
|
||
queues['qpos'].append(observation['qpos'].detach().clone())
|
||
if 'images' in observation:
|
||
queues['images'].append({
|
||
camera_name: image.detach().clone()
|
||
for camera_name, image in observation['images'].items()
|
||
})
|
||
if 'task' in observation:
|
||
queues['task'].append(observation['task'])
|
||
|
||
|
||
def _prepare_local_policy_batch(
|
||
queues: dict[str, deque],
|
||
obs_horizon: int,
|
||
camera_names: list[str],
|
||
) -> Dict[str, torch.Tensor]:
|
||
qpos_list = list(queues['qpos'])
|
||
if not qpos_list:
|
||
raise ValueError('observation queue is empty.')
|
||
while len(qpos_list) < int(obs_horizon):
|
||
qpos_list.append(qpos_list[-1])
|
||
batch_qpos = torch.stack(qpos_list, dim=0).unsqueeze(0)
|
||
|
||
images_list = list(queues['images'])
|
||
if not images_list:
|
||
raise ValueError('image queue is empty.')
|
||
while len(images_list) < int(obs_horizon):
|
||
images_list.append(images_list[-1])
|
||
|
||
ordered_camera_names = list(camera_names) if camera_names else sorted(images_list[0].keys())
|
||
batch_images = {
|
||
camera_name: torch.stack(
|
||
[image_history[camera_name] for image_history in images_list],
|
||
dim=0,
|
||
).unsqueeze(0)
|
||
for camera_name in ordered_camera_names
|
||
}
|
||
batch = {'qpos': batch_qpos, 'images': batch_images}
|
||
if queues.get('task'):
|
||
batch['task'] = [list(queues['task'])[-1]]
|
||
return batch
|
||
|
||
|
||
def _enqueue_predicted_actions(
|
||
queues: dict[str, deque],
|
||
predicted_actions: Any,
|
||
obs_horizon: int,
|
||
num_action_steps: int,
|
||
action_chunk_start: Optional[int] = None,
|
||
) -> None:
|
||
if isinstance(predicted_actions, np.ndarray):
|
||
predicted_actions = torch.from_numpy(predicted_actions)
|
||
if predicted_actions.ndim == 2:
|
||
predicted_actions = predicted_actions.unsqueeze(0)
|
||
|
||
start = int(obs_horizon) - 1 if action_chunk_start is None else int(action_chunk_start)
|
||
end = start + int(num_action_steps)
|
||
executable_actions = predicted_actions[:, start:end]
|
||
for action_index in range(executable_actions.shape[1]):
|
||
queues['action'].append(
|
||
executable_actions[:, action_index].squeeze(0).detach().cpu().clone()
|
||
)
|
||
|
||
|
||
def _serialize_policy_batch(batch: Dict[str, torch.Tensor]) -> dict[str, Any]:
|
||
serialized = {
|
||
'qpos': batch['qpos'].detach().cpu().numpy().astype(np.float32, copy=True),
|
||
'images': {
|
||
camera_name: image.detach().cpu().numpy().astype(np.float32, copy=True)
|
||
for camera_name, image in batch['images'].items()
|
||
},
|
||
}
|
||
if 'task' in batch:
|
||
serialized['task'] = batch['task']
|
||
return serialized
|
||
|
||
|
||
def _deserialize_policy_batch(batch: dict[str, Any], device: str) -> Dict[str, torch.Tensor]:
|
||
deserialized = {
|
||
'qpos': torch.as_tensor(batch['qpos'], dtype=torch.float32, device=device),
|
||
'images': {
|
||
camera_name: torch.as_tensor(image, dtype=torch.float32, device=device)
|
||
for camera_name, image in batch['images'].items()
|
||
},
|
||
}
|
||
if 'task' in batch:
|
||
deserialized['task'] = batch['task']
|
||
return deserialized
|
||
|
||
|
||
class _LocalPolicyRunner:
|
||
def __init__(self, agent: torch.nn.Module):
|
||
self.agent = agent
|
||
self.uses_local_model = True
|
||
|
||
def reset(self):
|
||
self.agent.reset()
|
||
|
||
def select_action(
|
||
self,
|
||
observation: Dict[str, torch.Tensor],
|
||
*,
|
||
episode_index: int,
|
||
timestep: int,
|
||
) -> tuple[Any, bool]:
|
||
del episode_index, timestep
|
||
action_queue = getattr(self.agent, '_queues', {}).get('action', None)
|
||
model_inference_triggered = len(action_queue) == 0 if action_queue is not None else True
|
||
action = self.agent.select_action(observation)
|
||
return action, bool(model_inference_triggered)
|
||
|
||
|
||
class _RemotePolicyRunner:
|
||
def __init__(
|
||
self,
|
||
*,
|
||
worker_index: int,
|
||
server_index: int,
|
||
request_queue,
|
||
response_queue,
|
||
camera_names: list[str],
|
||
obs_horizon: int,
|
||
num_action_steps: int,
|
||
action_chunk_start: Optional[int] = None,
|
||
response_timeout_s: float = 30.0,
|
||
):
|
||
self.worker_index = int(worker_index)
|
||
self.server_index = int(server_index)
|
||
self.request_queue = request_queue
|
||
self.response_queue = response_queue
|
||
self.camera_names = list(camera_names)
|
||
self.obs_horizon = int(obs_horizon)
|
||
self.num_action_steps = int(num_action_steps)
|
||
self.action_chunk_start = None if action_chunk_start is None else int(action_chunk_start)
|
||
self.response_timeout_s = float(response_timeout_s)
|
||
self.local_queues = _new_local_policy_queues(self.obs_horizon)
|
||
self.uses_local_model = False
|
||
|
||
def reset(self):
|
||
self.local_queues = _new_local_policy_queues(self.obs_horizon)
|
||
|
||
def select_action(
|
||
self,
|
||
observation: Dict[str, torch.Tensor],
|
||
*,
|
||
episode_index: int,
|
||
timestep: int,
|
||
) -> tuple[torch.Tensor, bool]:
|
||
_populate_local_policy_queues(self.local_queues, observation)
|
||
model_inference_triggered = len(self.local_queues['action']) == 0
|
||
if model_inference_triggered:
|
||
batch = _prepare_local_policy_batch(
|
||
self.local_queues,
|
||
obs_horizon=self.obs_horizon,
|
||
camera_names=self.camera_names,
|
||
)
|
||
self.request_queue.put({
|
||
'type': 'predict_chunk',
|
||
'worker_index': self.worker_index,
|
||
'server_index': self.server_index,
|
||
'episode_index': int(episode_index),
|
||
'timestep': int(timestep),
|
||
'batch': _serialize_policy_batch(batch),
|
||
})
|
||
try:
|
||
response = self.response_queue.get(timeout=self.response_timeout_s)
|
||
except queue.Empty as exc:
|
||
raise RuntimeError(
|
||
f'worker {self.worker_index} timed out waiting for inference server {self.server_index}'
|
||
) from exc
|
||
if response.get('type') != 'predict_chunk_result':
|
||
raise RuntimeError(
|
||
f'worker {self.worker_index} received unexpected inference response: '
|
||
f'{response.get("type")}'
|
||
)
|
||
_enqueue_predicted_actions(
|
||
self.local_queues,
|
||
predicted_actions=response['actions'],
|
||
obs_horizon=self.obs_horizon,
|
||
num_action_steps=self.num_action_steps,
|
||
action_chunk_start=self.action_chunk_start,
|
||
)
|
||
|
||
if not self.local_queues['action']:
|
||
raise RuntimeError(f'worker {self.worker_index} received no executable action from server')
|
||
return self.local_queues['action'].popleft(), bool(model_inference_triggered)
|
||
|
||
|
||
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()
|
||
},
|
||
}
|
||
|
||
|
||
_TIMING_SAMPLE_KEYS = (
|
||
'obs_read_time_ms',
|
||
'preprocess_time_ms',
|
||
'inference_time_ms',
|
||
'env_step_time_ms',
|
||
'total_time_ms',
|
||
)
|
||
|
||
|
||
def _empty_merge_state() -> dict[str, list[float] | list[bool]]:
|
||
return {
|
||
**{key: [] for key in _TIMING_SAMPLE_KEYS},
|
||
'model_forward_flags': [],
|
||
}
|
||
|
||
|
||
def _normalize_num_workers(num_workers: int, num_episodes: int) -> int:
|
||
num_episodes = max(int(num_episodes), 0)
|
||
if num_episodes == 0:
|
||
return 0
|
||
return min(max(int(num_workers), 1), num_episodes)
|
||
|
||
|
||
def _split_episode_indices(num_episodes: int, num_workers: int) -> list[list[int]]:
|
||
active_workers = _normalize_num_workers(num_workers=num_workers, num_episodes=num_episodes)
|
||
if active_workers == 0:
|
||
return []
|
||
episode_indices = np.arange(int(num_episodes), dtype=np.int32)
|
||
return [
|
||
chunk.tolist()
|
||
for chunk in np.array_split(episode_indices, active_workers)
|
||
if len(chunk) > 0
|
||
]
|
||
|
||
|
||
def _plan_episode_box_poses(
|
||
num_episodes: int,
|
||
sampler=None,
|
||
) -> list[np.ndarray]:
|
||
if sampler is None:
|
||
sampler = sample_transfer_pose
|
||
return [
|
||
np.asarray(sampler(), dtype=np.float32).copy()
|
||
for _ in range(int(num_episodes))
|
||
]
|
||
|
||
|
||
def _merge_worker_summaries(
|
||
worker_summaries: list[dict[str, Any]],
|
||
artifact_paths: dict[str, Optional[str]],
|
||
) -> dict[str, Any]:
|
||
merged_episodes = []
|
||
merged_state = _empty_merge_state()
|
||
for worker_summary in worker_summaries:
|
||
merged_episodes.extend(worker_summary.get('episodes', []))
|
||
merge_state = worker_summary.get('_merge_state', {})
|
||
for key in _TIMING_SAMPLE_KEYS:
|
||
merged_state[key].extend(float(value) for value in merge_state.get(key, []))
|
||
merged_state['model_forward_flags'].extend(
|
||
bool(value) for value in merge_state.get('model_forward_flags', [])
|
||
)
|
||
|
||
merged_episodes = sorted(
|
||
merged_episodes,
|
||
key=lambda episode: int(episode.get('episode_index', 0)),
|
||
)
|
||
episode_rewards = [
|
||
float(episode.get('episode_reward', 0.0))
|
||
for episode in merged_episodes
|
||
]
|
||
episode_max_rewards = [
|
||
(
|
||
float(episode['episode_max_reward'])
|
||
if episode.get('episode_max_reward') is not None
|
||
else None
|
||
)
|
||
for episode in merged_episodes
|
||
]
|
||
valid_max_rewards = [
|
||
value for value in episode_max_rewards
|
||
if value is not None
|
||
]
|
||
summary = {
|
||
'num_episodes': len(merged_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(valid_max_rewards)) if valid_max_rewards else 0.0,
|
||
'episodes': merged_episodes,
|
||
'artifact_dir': artifact_paths.get('output_dir') if artifact_paths else None,
|
||
'artifacts': artifact_paths,
|
||
}
|
||
|
||
if merged_episodes:
|
||
summary.update({
|
||
'avg_inference_fps': float(np.mean([
|
||
float(episode.get('inference_fps', 0.0))
|
||
for episode in merged_episodes
|
||
])),
|
||
'avg_control_fps': float(np.mean([
|
||
float(episode.get('control_fps', 0.0))
|
||
for episode in merged_episodes
|
||
])),
|
||
'avg_obs_read_time_ms': _mean_or_zero(merged_state['obs_read_time_ms']),
|
||
'avg_preprocess_time_ms': _mean_or_zero(merged_state['preprocess_time_ms']),
|
||
'avg_inference_time_ms': _mean_or_zero(merged_state['inference_time_ms']),
|
||
'avg_env_step_time_ms': _mean_or_zero(merged_state['env_step_time_ms']),
|
||
'avg_total_time_ms': _mean_or_zero(merged_state['total_time_ms']),
|
||
'timing_summary': _summarize_timing_breakdown(
|
||
{
|
||
'obs_read': merged_state['obs_read_time_ms'],
|
||
'preprocess': merged_state['preprocess_time_ms'],
|
||
'inference': merged_state['inference_time_ms'],
|
||
'env_step': merged_state['env_step_time_ms'],
|
||
'loop_total': merged_state['total_time_ms'],
|
||
},
|
||
merged_state['model_forward_flags'],
|
||
),
|
||
})
|
||
|
||
return summary
|
||
|
||
|
||
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,
|
||
output_dir_override: Optional[str] = None,
|
||
) -> 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:
|
||
if output_dir_override:
|
||
output_dir = Path(str(output_dir_override)).expanduser().resolve()
|
||
else:
|
||
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)
|
||
|
||
|
||
def _public_summary(summary: dict[str, Any]) -> dict[str, Any]:
|
||
public_summary = dict(summary)
|
||
public_summary.pop('_merge_state', None)
|
||
return _json_friendly(public_summary)
|
||
|
||
|
||
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 _sample_task_reset_state(task_name: str):
|
||
if task_name == 'sim_air_insert_socket_peg':
|
||
return sample_air_insert_socket_peg_state()
|
||
if 'sim_transfer' in task_name:
|
||
return sample_transfer_pose()
|
||
raise NotImplementedError(f'Unsupported eval task reset sampling: {task_name}')
|
||
|
||
|
||
def _print_eval_config(cfg: DictConfig):
|
||
# 打印配置
|
||
print("=" * 80)
|
||
print("VLA 评估配置:")
|
||
print("=" * 80)
|
||
print(OmegaConf.to_yaml(cfg))
|
||
print("=" * 80)
|
||
|
||
|
||
def _build_episode_plans(
|
||
num_episodes: int,
|
||
box_poses: Optional[list[np.ndarray]] = None,
|
||
) -> list[dict[str, Any]]:
|
||
if box_poses is None:
|
||
return [
|
||
{
|
||
'episode_index': int(episode_index),
|
||
}
|
||
for episode_index in range(int(num_episodes))
|
||
]
|
||
return [
|
||
{
|
||
'episode_index': int(episode_index),
|
||
'box_pos': np.asarray(box_pos, dtype=np.float32).copy(),
|
||
}
|
||
for episode_index, box_pos in enumerate(box_poses)
|
||
]
|
||
|
||
|
||
def _run_eval_episode_plans(
|
||
cfg: DictConfig,
|
||
episode_plans: list[dict[str, Any]],
|
||
policy_runner,
|
||
worker_index: int = 0,
|
||
artifact_paths: Optional[dict[str, Optional[str]]] = None,
|
||
show_progress: bool = True,
|
||
) -> dict[str, Any]:
|
||
eval_cfg = cfg.eval
|
||
device = str(eval_cfg.device)
|
||
camera_names = list(eval_cfg.camera_names)
|
||
image_resize_shape = _resolve_eval_image_resize_shape(cfg)
|
||
task_description = eval_cfg.get('task_description', None)
|
||
artifact_paths = artifact_paths or _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()
|
||
merge_state = _empty_merge_state()
|
||
|
||
# 可选：动作平滑器
|
||
smoother = ActionSmoother(alpha=eval_cfg.smooth_alpha) if eval_cfg.use_smoothing else None
|
||
|
||
# =========================================================================
|
||
# 创建环境
|
||
# =========================================================================
|
||
env = None
|
||
try:
|
||
env = make_sim_env(eval_cfg.task_name, headless=eval_cfg.headless)
|
||
|
||
# =========================================================================
|
||
# 运行评估回合
|
||
# =========================================================================
|
||
all_stats = []
|
||
episode_rewards = []
|
||
episode_max_rewards = []
|
||
for plan in episode_plans:
|
||
episode_idx = int(plan['episode_index'])
|
||
task_state = plan.get('box_pos')
|
||
if task_state is None:
|
||
task_state = _sample_task_reset_state(str(eval_cfg.task_name))
|
||
elif isinstance(task_state, np.ndarray):
|
||
task_state = np.asarray(task_state, dtype=np.float32)
|
||
|
||
|
||
if show_progress:
|
||
print(f"\n{'='*60}")
|
||
print(f"回合 {episode_idx + 1}/{eval_cfg.num_episodes}")
|
||
print(f"{'='*60}\n")
|
||
|
||
env.reset(task_state)
|
||
|
||
# 为新回合重置 rollout policy 状态
|
||
policy_runner.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():
|
||
episode_iterator = range(eval_cfg.max_timesteps)
|
||
if show_progress:
|
||
episode_iterator = tqdm(episode_iterator, desc=f"回合 {episode_idx + 1}")
|
||
|
||
for t in episode_iterator:
|
||
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,
|
||
task_description=task_description,
|
||
)
|
||
end_preprocess = time.perf_counter()
|
||
|
||
# 选择动作（本地 agent 或远端 inference server）
|
||
start_inference = time.perf_counter()
|
||
action, model_inference_triggered = policy_runner.select_action(
|
||
observation,
|
||
episode_index=episode_idx,
|
||
timestep=t,
|
||
)
|
||
|
||
if (
|
||
getattr(policy_runner, 'uses_local_model', False)
|
||
and _is_cuda_device(device)
|
||
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))
|
||
merge_state['obs_read_time_ms'].append(step_timing_ms['obs_read_time_ms'])
|
||
merge_state['preprocess_time_ms'].append(step_timing_ms['preprocess_time_ms'])
|
||
merge_state['inference_time_ms'].append(step_timing_ms['inference_time_ms'])
|
||
merge_state['env_step_time_ms'].append(step_timing_ms['env_step_time_ms'])
|
||
merge_state['total_time_ms'].append(step_timing_ms['total_time_ms'])
|
||
merge_state['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 = {
|
||
'worker_index': int(worker_index),
|
||
'episode_index': int(episode_idx),
|
||
'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))
|
||
episode_max_rewards.append(
|
||
float(episode_max_reward) if episode_max_reward != float('-inf') else None
|
||
)
|
||
|
||
if show_progress:
|
||
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}")
|
||
|
||
# =========================================================================
|
||
# 总体统计
|
||
# =========================================================================
|
||
if show_progress:
|
||
print(f"\n{'='*60}")
|
||
print("评估完成!")
|
||
print(f"{'='*60}")
|
||
|
||
valid_max_rewards = [
|
||
reward for reward in episode_max_rewards
|
||
if reward is not None
|
||
]
|
||
|
||
summary = {
|
||
'num_episodes': len(episode_plans),
|
||
'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(valid_max_rewards)) if valid_max_rewards else 0.0,
|
||
'episodes': all_stats,
|
||
'artifact_dir': artifact_paths['output_dir'],
|
||
'artifacts': artifact_paths,
|
||
'_merge_state': merge_state,
|
||
}
|
||
|
||
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(merge_state['obs_read_time_ms'])
|
||
avg_preprocess_time = _mean_or_zero(merge_state['preprocess_time_ms'])
|
||
avg_inference_time = _mean_or_zero(merge_state['inference_time_ms'])
|
||
avg_env_step_time = _mean_or_zero(merge_state['env_step_time_ms'])
|
||
avg_total_time = _mean_or_zero(merge_state['total_time_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': merge_state['obs_read_time_ms'],
|
||
'preprocess': merge_state['preprocess_time_ms'],
|
||
'inference': merge_state['inference_time_ms'],
|
||
'env_step': merge_state['env_step_time_ms'],
|
||
'loop_total': merge_state['total_time_ms'],
|
||
},
|
||
merge_state['model_forward_flags'],
|
||
),
|
||
})
|
||
|
||
if show_progress:
|
||
print(f"\n总体统计 ({len(episode_plans)} 个回合):")
|
||
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)
|
||
public_summary = _public_summary(summary)
|
||
if artifact_paths['summary_json'] is not None:
|
||
_save_summary_json(artifact_paths['summary_json'], public_summary)
|
||
if artifact_paths['timing_json'] is not None:
|
||
_save_summary_json(artifact_paths['timing_json'], public_summary.get('timing_summary', {}))
|
||
if show_progress:
|
||
print()
|
||
return summary
|
||
finally:
|
||
video_recorder.close()
|
||
_close_env(env)
|
||
|
||
|
||
def _run_eval_worker(
|
||
cfg: DictConfig,
|
||
episode_plans: list[dict[str, Any]],
|
||
worker_index: int = 0,
|
||
artifact_paths: Optional[dict[str, Optional[str]]] = None,
|
||
show_progress: bool = True,
|
||
) -> dict[str, Any]:
|
||
eval_cfg = cfg.eval
|
||
|
||
log.info(f"🚀 从 {eval_cfg.ckpt_path} 加载模型...")
|
||
agent, _dataset_stats = load_checkpoint(
|
||
ckpt_path=eval_cfg.ckpt_path,
|
||
agent_cfg=cfg.agent,
|
||
device=str(eval_cfg.device),
|
||
)
|
||
policy_runner = _LocalPolicyRunner(agent)
|
||
return _run_eval_episode_plans(
|
||
cfg,
|
||
episode_plans=episode_plans,
|
||
policy_runner=policy_runner,
|
||
worker_index=worker_index,
|
||
artifact_paths=artifact_paths,
|
||
show_progress=show_progress,
|
||
)
|
||
|
||
|
||
def _run_remote_eval_worker(
|
||
cfg: DictConfig,
|
||
episode_plans: list[dict[str, Any]],
|
||
*,
|
||
worker_index: int,
|
||
server_index: int,
|
||
request_queue,
|
||
response_queue,
|
||
artifact_paths: Optional[dict[str, Optional[str]]] = None,
|
||
show_progress: bool = False,
|
||
) -> dict[str, Any]:
|
||
eval_cfg = cfg.eval
|
||
agent_cfg = cfg.agent
|
||
num_action_steps = int(agent_cfg.get('num_action_steps', eval_cfg.get('num_queries', 1)))
|
||
action_chunk_start = agent_cfg.get('action_chunk_start', None)
|
||
policy_runner = _RemotePolicyRunner(
|
||
worker_index=worker_index,
|
||
server_index=server_index,
|
||
request_queue=request_queue,
|
||
response_queue=response_queue,
|
||
camera_names=_resolve_policy_camera_names(cfg),
|
||
obs_horizon=int(agent_cfg.get('obs_horizon', eval_cfg.obs_horizon)),
|
||
num_action_steps=num_action_steps,
|
||
action_chunk_start=action_chunk_start,
|
||
response_timeout_s=float(eval_cfg.get('response_timeout_s', 300.0)),
|
||
)
|
||
return _run_eval_episode_plans(
|
||
cfg,
|
||
episode_plans=episode_plans,
|
||
policy_runner=policy_runner,
|
||
worker_index=worker_index,
|
||
artifact_paths=artifact_paths,
|
||
show_progress=show_progress,
|
||
)
|
||
|
||
|
||
def _run_eval_serial(cfg: DictConfig):
|
||
eval_cfg = cfg.eval
|
||
_configure_headless_mujoco_gl(eval_cfg)
|
||
artifact_paths = _resolve_artifact_paths(eval_cfg)
|
||
episode_plans = _build_episode_plans(eval_cfg.num_episodes)
|
||
summary = _run_eval_worker(
|
||
cfg,
|
||
episode_plans=episode_plans,
|
||
worker_index=0,
|
||
artifact_paths=artifact_paths,
|
||
show_progress=True,
|
||
)
|
||
return _public_summary(summary)
|
||
|
||
|
||
def _validate_parallel_eval_cfg(eval_cfg: DictConfig):
|
||
if not bool(eval_cfg.get('headless', False)):
|
||
raise ValueError('eval.num_workers > 1 requires eval.headless=true')
|
||
|
||
unsupported_exports = [
|
||
flag_name
|
||
for flag_name in (
|
||
'record_video',
|
||
'save_trajectory',
|
||
'save_trajectory_npz',
|
||
)
|
||
if bool(eval_cfg.get(flag_name, False))
|
||
]
|
||
if unsupported_exports:
|
||
joined_flags = ', '.join(unsupported_exports)
|
||
raise ValueError(
|
||
'eval.num_workers > 1 does not yet support parallel export for '
|
||
f'{joined_flags}'
|
||
)
|
||
|
||
|
||
def _is_cuda_device(device: Any) -> bool:
|
||
return str(device).lower().startswith('cuda')
|
||
|
||
|
||
def _resolve_cuda_devices(eval_cfg: DictConfig) -> list[int]:
|
||
if not _is_cuda_device(eval_cfg.get('device', 'cpu')):
|
||
return []
|
||
|
||
configured_devices = eval_cfg.get('cuda_devices', None)
|
||
if configured_devices is None:
|
||
return [0]
|
||
|
||
resolved_devices = [int(device_index) for device_index in configured_devices]
|
||
if not resolved_devices:
|
||
raise ValueError('eval.cuda_devices must not be empty when eval.device is CUDA')
|
||
if any(device_index < 0 for device_index in resolved_devices):
|
||
raise ValueError('eval.cuda_devices must contain non-negative logical CUDA device indices')
|
||
return resolved_devices
|
||
|
||
|
||
def _run_spawn_jobs(
|
||
payloads: list[dict[str, Any]],
|
||
max_workers: int,
|
||
worker_fn,
|
||
) -> list[Any]:
|
||
if not payloads:
|
||
return []
|
||
|
||
ctx = multiprocessing.get_context('spawn')
|
||
results = []
|
||
with concurrent.futures.ProcessPoolExecutor(
|
||
max_workers=int(max_workers),
|
||
mp_context=ctx,
|
||
) as executor:
|
||
future_to_payload = {
|
||
executor.submit(worker_fn, payload): payload
|
||
for payload in payloads
|
||
}
|
||
try:
|
||
for future in concurrent.futures.as_completed(future_to_payload):
|
||
results.append(future.result())
|
||
except Exception:
|
||
for future in future_to_payload:
|
||
future.cancel()
|
||
raise
|
||
return results
|
||
|
||
|
||
def _run_eval_worker_entry(payload: dict[str, Any]) -> dict[str, Any]:
|
||
if payload.get('_spawn_probe', False):
|
||
return {
|
||
'probe_value': int(payload['probe_value']),
|
||
'worker_index': int(payload.get('worker_index', -1)),
|
||
}
|
||
|
||
cfg = OmegaConf.create(payload['cfg'])
|
||
artifact_paths = _resolve_artifact_paths(
|
||
cfg.eval,
|
||
output_dir_override=payload.get('artifact_dir'),
|
||
)
|
||
return _run_eval_worker(
|
||
cfg,
|
||
episode_plans=list(payload.get('episode_plans', [])),
|
||
worker_index=int(payload.get('worker_index', 0)),
|
||
artifact_paths=artifact_paths,
|
||
show_progress=False,
|
||
)
|
||
|
||
|
||
def _build_parallel_worker_payloads(
|
||
cfg: DictConfig,
|
||
artifact_paths: dict[str, Optional[str]],
|
||
) -> tuple[list[dict[str, Any]], int]:
|
||
eval_cfg = cfg.eval
|
||
requested_workers = int(eval_cfg.get('num_workers', 1))
|
||
episode_splits = _split_episode_indices(
|
||
num_episodes=int(eval_cfg.num_episodes),
|
||
num_workers=requested_workers,
|
||
)
|
||
box_poses = _plan_episode_box_poses(int(eval_cfg.num_episodes))
|
||
resolved_cfg = OmegaConf.to_container(cfg, resolve=True)
|
||
payloads = []
|
||
workers_dir = None
|
||
if artifact_paths.get('output_dir') is not None:
|
||
workers_dir = Path(str(artifact_paths['output_dir'])) / 'workers'
|
||
workers_dir.mkdir(parents=True, exist_ok=True)
|
||
artifact_paths['workers_dir'] = str(workers_dir)
|
||
|
||
for worker_index, episode_indices in enumerate(episode_splits):
|
||
worker_artifact_dir = None
|
||
if workers_dir is not None:
|
||
worker_artifact_dir = workers_dir / f'worker_{worker_index:02d}'
|
||
worker_artifact_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
worker_cfg = json.loads(json.dumps(resolved_cfg))
|
||
worker_cfg['eval']['artifact_dir'] = (
|
||
str(worker_artifact_dir) if worker_artifact_dir is not None else None
|
||
)
|
||
payloads.append({
|
||
'cfg': worker_cfg,
|
||
'worker_index': int(worker_index),
|
||
'artifact_dir': str(worker_artifact_dir) if worker_artifact_dir is not None else None,
|
||
'episode_plans': [
|
||
{
|
||
'episode_index': int(episode_index),
|
||
'box_pos': np.asarray(box_poses[episode_index], dtype=np.float32).tolist(),
|
||
}
|
||
for episode_index in episode_indices
|
||
],
|
||
})
|
||
|
||
return payloads, len(episode_splits)
|
||
|
||
|
||
def _build_cuda_server_payloads(
|
||
cfg: DictConfig,
|
||
worker_payloads: list[dict[str, Any]],
|
||
cuda_devices: list[int],
|
||
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
||
resolved_cfg = OmegaConf.to_container(cfg, resolve=True)
|
||
server_payloads = [
|
||
{
|
||
'cfg': json.loads(json.dumps(resolved_cfg)),
|
||
'server_index': int(server_index),
|
||
'device_index': int(device_index),
|
||
'worker_indices': [],
|
||
}
|
||
for server_index, device_index in enumerate(cuda_devices)
|
||
]
|
||
|
||
assigned_workers = []
|
||
for worker_payload in worker_payloads:
|
||
server_index = int(worker_payload['worker_index']) % len(server_payloads)
|
||
assigned_payload = dict(worker_payload)
|
||
assigned_payload['server_index'] = server_index
|
||
assigned_payload['device_index'] = int(server_payloads[server_index]['device_index'])
|
||
assigned_workers.append(assigned_payload)
|
||
server_payloads[server_index]['worker_indices'].append(int(worker_payload['worker_index']))
|
||
|
||
return server_payloads, assigned_workers
|
||
|
||
|
||
def _inference_server_main(payload: dict[str, Any]) -> None:
|
||
request_queue = payload['request_queue']
|
||
response_queues = payload['response_queues']
|
||
server_index = int(payload.get('server_index', -1))
|
||
|
||
if payload.get('_spawn_probe', False):
|
||
while True:
|
||
message = request_queue.get()
|
||
if message.get('type') == 'shutdown_server':
|
||
return
|
||
if message.get('type') != 'predict_chunk':
|
||
continue
|
||
response_queues[int(message['worker_index'])].put({
|
||
'type': 'predict_chunk_result',
|
||
'server_index': server_index,
|
||
'actions': np.asarray([[[11.0], [22.0], [33.0]]], dtype=np.float32),
|
||
})
|
||
return
|
||
|
||
result_queue = payload.get('result_queue')
|
||
cfg = OmegaConf.create(payload['cfg'])
|
||
device = f'cuda:{int(payload["device_index"])}'
|
||
try:
|
||
agent, _dataset_stats = load_checkpoint(
|
||
ckpt_path=cfg.eval.ckpt_path,
|
||
agent_cfg=cfg.agent,
|
||
device=device,
|
||
)
|
||
if result_queue is not None:
|
||
result_queue.put({
|
||
'kind': 'server_ready',
|
||
'server_index': server_index,
|
||
})
|
||
while True:
|
||
message = request_queue.get()
|
||
message_type = message.get('type')
|
||
if message_type == 'shutdown_server':
|
||
return
|
||
if message_type != 'predict_chunk':
|
||
raise ValueError(f'unknown server message type: {message_type}')
|
||
|
||
batch = _deserialize_policy_batch(message['batch'], device=device)
|
||
with torch.inference_mode():
|
||
actions = agent.predict_action_chunk(batch)
|
||
if torch.cuda.is_available():
|
||
torch.cuda.synchronize()
|
||
response_queues[int(message['worker_index'])].put({
|
||
'type': 'predict_chunk_result',
|
||
'server_index': server_index,
|
||
'actions': actions.detach().cpu().numpy().astype(np.float32, copy=True),
|
||
})
|
||
except Exception as exc:
|
||
if result_queue is not None:
|
||
result_queue.put({
|
||
'kind': 'server_error',
|
||
'server_index': server_index,
|
||
'error': str(exc),
|
||
})
|
||
raise
|
||
|
||
|
||
def _env_worker_main(payload: dict[str, Any]) -> None:
|
||
worker_index = int(payload.get('worker_index', -1))
|
||
server_index = int(payload.get('server_index', -1))
|
||
request_queue = payload['request_queue']
|
||
response_queue = payload['response_queue']
|
||
result_queue = payload.get('result_queue')
|
||
|
||
if payload.get('_spawn_probe', False):
|
||
request_queue.put({
|
||
'type': 'predict_chunk',
|
||
'worker_index': worker_index,
|
||
'server_index': server_index,
|
||
'episode_index': 0,
|
||
'timestep': 0,
|
||
'batch': {
|
||
'qpos': np.zeros((1, 1, 1), dtype=np.float32),
|
||
'images': {},
|
||
},
|
||
})
|
||
response = response_queue.get(timeout=10.0)
|
||
result_queue.put({
|
||
'kind': 'worker_result',
|
||
'worker_index': worker_index,
|
||
'summary': {
|
||
'probe_worker_index': worker_index,
|
||
'probe_server_index': server_index,
|
||
'probe_actions': np.asarray(response['actions']).tolist(),
|
||
},
|
||
})
|
||
return
|
||
|
||
cfg = OmegaConf.create(payload['cfg'])
|
||
artifact_paths = _resolve_artifact_paths(
|
||
cfg.eval,
|
||
output_dir_override=payload.get('artifact_dir'),
|
||
)
|
||
|
||
try:
|
||
summary = _run_remote_eval_worker(
|
||
cfg,
|
||
episode_plans=list(payload.get('episode_plans', [])),
|
||
worker_index=worker_index,
|
||
server_index=server_index,
|
||
request_queue=request_queue,
|
||
response_queue=response_queue,
|
||
artifact_paths=artifact_paths,
|
||
show_progress=False,
|
||
)
|
||
except Exception as exc:
|
||
if result_queue is not None:
|
||
result_queue.put({
|
||
'kind': 'worker_error',
|
||
'worker_index': worker_index,
|
||
'server_index': server_index,
|
||
'error': str(exc),
|
||
})
|
||
raise
|
||
|
||
if result_queue is not None:
|
||
result_queue.put({
|
||
'kind': 'worker_result',
|
||
'worker_index': worker_index,
|
||
'server_index': server_index,
|
||
'summary': summary,
|
||
})
|
||
|
||
|
||
def _shutdown_processes(processes: list[tuple[Any, dict[str, Any]]], *, terminate: bool = False) -> None:
|
||
for process, _payload in processes:
|
||
if terminate and process.is_alive():
|
||
process.terminate()
|
||
for process, _payload in processes:
|
||
process.join(timeout=5.0)
|
||
if process.is_alive():
|
||
process.terminate()
|
||
process.join(timeout=5.0)
|
||
|
||
|
||
def _run_cuda_parallel_processes(
|
||
server_payloads: list[dict[str, Any]],
|
||
worker_payloads: list[dict[str, Any]],
|
||
) -> list[dict[str, Any]]:
|
||
if not worker_payloads:
|
||
return []
|
||
|
||
ctx = multiprocessing.get_context('spawn')
|
||
result_queue = ctx.Queue()
|
||
response_queues = [ctx.Queue() for _ in range(len(worker_payloads))]
|
||
request_queues = {}
|
||
server_processes: list[tuple[Any, dict[str, Any]]] = []
|
||
worker_processes: list[tuple[Any, dict[str, Any]]] = []
|
||
should_terminate = False
|
||
startup_timeout_s = float(
|
||
server_payloads[0]['cfg']['eval'].get('server_startup_timeout_s', 300.0)
|
||
) if server_payloads else 300.0
|
||
|
||
try:
|
||
for server_payload in server_payloads:
|
||
request_queue = ctx.Queue()
|
||
request_queues[int(server_payload['server_index'])] = request_queue
|
||
process_payload = dict(server_payload)
|
||
process_payload['request_queue'] = request_queue
|
||
process_payload['response_queues'] = response_queues
|
||
process_payload['result_queue'] = result_queue
|
||
process = ctx.Process(
|
||
target=_inference_server_main,
|
||
args=(process_payload,),
|
||
name=f'eval-inference-server-{int(server_payload["server_index"]):02d}',
|
||
)
|
||
process.start()
|
||
server_processes.append((process, server_payload))
|
||
|
||
pending_servers = {int(payload['server_index']) for _process, payload in server_processes}
|
||
startup_deadline = time.monotonic() + startup_timeout_s
|
||
while pending_servers:
|
||
remaining = startup_deadline - time.monotonic()
|
||
if remaining <= 0:
|
||
raise RuntimeError(
|
||
'Timed out waiting for CUDA inference servers to become ready'
|
||
)
|
||
try:
|
||
message = result_queue.get(timeout=min(0.2, remaining))
|
||
except queue.Empty:
|
||
for process, payload in server_processes:
|
||
if process.exitcode not in (None, 0):
|
||
raise RuntimeError(
|
||
f'CUDA inference server {int(payload["server_index"])} exited with code '
|
||
f'{process.exitcode}'
|
||
)
|
||
continue
|
||
|
||
message_kind = message.get('kind')
|
||
if message_kind == 'server_ready':
|
||
pending_servers.discard(int(message['server_index']))
|
||
continue
|
||
if message_kind == 'server_error':
|
||
raise RuntimeError(
|
||
f'CUDA inference server {int(message["server_index"])} failed: {message.get("error")}'
|
||
)
|
||
raise RuntimeError(f'Unexpected CUDA startup message: {message_kind}')
|
||
|
||
for worker_payload in worker_payloads:
|
||
worker_index = int(worker_payload['worker_index'])
|
||
process_payload = dict(worker_payload)
|
||
process_payload['request_queue'] = request_queues[int(worker_payload['server_index'])]
|
||
process_payload['response_queue'] = response_queues[worker_index]
|
||
process_payload['result_queue'] = result_queue
|
||
process = ctx.Process(
|
||
target=_env_worker_main,
|
||
args=(process_payload,),
|
||
name=f'eval-env-worker-{worker_index:02d}',
|
||
)
|
||
process.start()
|
||
worker_processes.append((process, worker_payload))
|
||
|
||
pending_workers = {int(payload['worker_index']) for payload in worker_payloads}
|
||
worker_summaries = {}
|
||
while pending_workers:
|
||
try:
|
||
message = result_queue.get(timeout=0.2)
|
||
except queue.Empty:
|
||
for process, payload in server_processes:
|
||
if process.exitcode not in (None, 0):
|
||
raise RuntimeError(
|
||
f'CUDA inference server {int(payload["server_index"])} exited with code '
|
||
f'{process.exitcode}'
|
||
)
|
||
for process, payload in worker_processes:
|
||
worker_index = int(payload['worker_index'])
|
||
if worker_index in pending_workers and process.exitcode not in (None, 0):
|
||
raise RuntimeError(
|
||
f'CUDA rollout worker {worker_index} exited with code {process.exitcode}'
|
||
)
|
||
continue
|
||
|
||
message_kind = message.get('kind')
|
||
if message_kind == 'worker_result':
|
||
worker_index = int(message['worker_index'])
|
||
worker_summaries[worker_index] = message['summary']
|
||
pending_workers.discard(worker_index)
|
||
continue
|
||
if message_kind == 'worker_error':
|
||
raise RuntimeError(
|
||
f'CUDA rollout worker {int(message["worker_index"])} failed: {message.get("error")}'
|
||
)
|
||
if message_kind == 'server_error':
|
||
raise RuntimeError(
|
||
f'CUDA inference server {int(message["server_index"])} failed: {message.get("error")}'
|
||
)
|
||
raise RuntimeError(f'Unexpected CUDA parallel message: {message_kind}')
|
||
|
||
return [
|
||
worker_summaries[worker_index]
|
||
for worker_index in sorted(worker_summaries)
|
||
]
|
||
except Exception:
|
||
should_terminate = True
|
||
raise
|
||
finally:
|
||
for request_queue in request_queues.values():
|
||
try:
|
||
request_queue.put({'type': 'shutdown_server'})
|
||
except Exception:
|
||
continue
|
||
_shutdown_processes(worker_processes, terminate=should_terminate)
|
||
_shutdown_processes(server_processes, terminate=should_terminate)
|
||
|
||
|
||
def _run_eval_parallel(cfg: DictConfig):
|
||
eval_cfg = cfg.eval
|
||
_configure_headless_mujoco_gl(eval_cfg)
|
||
_validate_parallel_eval_cfg(eval_cfg)
|
||
if _is_cuda_device(eval_cfg.get('device', 'cpu')):
|
||
_resolve_cuda_devices(eval_cfg)
|
||
return _run_eval_parallel_cuda(cfg)
|
||
return _run_eval_parallel_cpu(cfg)
|
||
|
||
|
||
def _run_eval_parallel_cpu(cfg: DictConfig):
|
||
eval_cfg = cfg.eval
|
||
artifact_paths = _resolve_artifact_paths(eval_cfg)
|
||
payloads, active_workers = _build_parallel_worker_payloads(cfg, artifact_paths)
|
||
|
||
try:
|
||
worker_summaries = _run_spawn_jobs(
|
||
payloads=payloads,
|
||
max_workers=max(active_workers, 1),
|
||
worker_fn=_run_eval_worker_entry,
|
||
)
|
||
except Exception as exc:
|
||
raise RuntimeError(f'Parallel rollout worker failed: {exc}') from exc
|
||
|
||
summary = _merge_worker_summaries(worker_summaries, artifact_paths)
|
||
public_summary = _public_summary(summary)
|
||
if artifact_paths.get('summary_json') is not None:
|
||
_save_summary_json(artifact_paths['summary_json'], public_summary)
|
||
if artifact_paths.get('timing_json') is not None:
|
||
_save_summary_json(artifact_paths['timing_json'], public_summary.get('timing_summary', {}))
|
||
return public_summary
|
||
|
||
|
||
def _run_eval_parallel_cuda(cfg: DictConfig):
|
||
eval_cfg = cfg.eval
|
||
_validate_parallel_eval_cfg(eval_cfg)
|
||
artifact_paths = _resolve_artifact_paths(eval_cfg)
|
||
worker_payloads, _active_workers = _build_parallel_worker_payloads(cfg, artifact_paths)
|
||
cuda_devices = _resolve_cuda_devices(eval_cfg)
|
||
server_payloads, assigned_worker_payloads = _build_cuda_server_payloads(
|
||
cfg,
|
||
worker_payloads=worker_payloads,
|
||
cuda_devices=cuda_devices,
|
||
)
|
||
|
||
try:
|
||
worker_summaries = _run_cuda_parallel_processes(
|
||
server_payloads=server_payloads,
|
||
worker_payloads=assigned_worker_payloads,
|
||
)
|
||
except Exception as exc:
|
||
raise RuntimeError(f'Parallel CUDA rollout failed: {exc}') from exc
|
||
|
||
summary = _merge_worker_summaries(worker_summaries, artifact_paths)
|
||
public_summary = _public_summary(summary)
|
||
if artifact_paths.get('summary_json') is not None:
|
||
_save_summary_json(artifact_paths['summary_json'], public_summary)
|
||
if artifact_paths.get('timing_json') is not None:
|
||
_save_summary_json(artifact_paths['timing_json'], public_summary.get('timing_summary', {}))
|
||
return public_summary
|
||
|
||
|
||
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_eval_config(cfg)
|
||
requested_workers = int(cfg.eval.get('num_workers', 1))
|
||
active_workers = _normalize_num_workers(
|
||
num_workers=requested_workers,
|
||
num_episodes=int(cfg.eval.get('num_episodes', 0)),
|
||
)
|
||
if active_workers <= 1:
|
||
return _run_eval_serial(cfg)
|
||
return _run_eval_parallel(cfg)
|
||
|
||
|
||
@hydra.main(version_base=None, config_path="../../vla/conf", config_name="config")
|
||
def main(cfg: DictConfig):
|
||
return _run_eval(cfg)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|