diff --git a/README_REMOTE_INFERENCE.md b/README_REMOTE_INFERENCE.md new file mode 100644 index 00000000..ff893934 --- /dev/null +++ b/README_REMOTE_INFERENCE.md @@ -0,0 +1,231 @@ +# ACT Policy 远程推理使用指南 + +本目录包含 ACT Policy 的远程推理服务器和客户端实现,支持通过 WebSocket 进行实时推理。 + +## 📁 文件说明 + +- `serve_act_policy.py` - ACT Policy 推理服务器 +- `test_remote_inference.py` - 简单的推理测试客户端 +- `act_remote_inference_with_temporal_agg.py` - 支持 Temporal Aggregation 的完整客户端示例 + +## 🚀 快速开始 + +### 1. 启动服务器 + +```bash +# 基本用法 +python serve_act_policy.py -i -p 8000 + +# 完整示例 +python serve_act_policy.py \ + -i /path/to/policy_best.ckpt \ + -p 8000 \ + -c top,front \ + -q 100 \ + -d cuda +``` + +**参数说明:** +- `-i, --input`: checkpoint 文件路径(必需) +- `-p, --port`: 服务器端口(默认 8000) +- `-h, --host`: 服务器地址(默认 0.0.0.0) +- `-c, --camera_names`: 相机名称列表,逗号分隔(默认 'top') +- `-q, --num_queries`: chunk size,即每次推理返回的动作序列长度(默认 100) +- `-d, --device`: 设备 'cuda' 或 'cpu'(默认 cuda) + +**注意:** checkpoint 所在目录必须包含 `dataset_stats.pkl` 文件! + +### 2. 运行客户端测试 + +```bash +# 简单测试 +python test_remote_inference.py + +# Temporal Aggregation 示例 +python act_remote_inference_with_temporal_agg.py +``` + +## 📊 数据格式说明 + +### 输入格式(客户端 → 服务器) + +```python +obs = { + # 关节位置(必需) + 'qpos': np.ndarray, # shape: (state_dim,), dtype: np.float32 + # 例如: (14,) 表示 14 个关节的位置 + + # 图像数据(根据相机配置) + 'top': np.ndarray, # shape: (H, W, 3), dtype: np.uint8 或 np.float32 + # 通道顺序: RGB(不是 BGR!) + # 例如: (480, 640, 3) + + 'front': np.ndarray, # 可选,其他相机 +} +``` + +### 输出格式(服务器 → 客户端) + +```python +result = { + 'actions': np.ndarray, # shape: (num_queries, action_dim) + # 例如: (100, 14) 表示 100 步动作序列 +} +``` + +## 🔄 归一化处理说明 + +ACT Policy 的归一化在服务器端自动处理,无需客户端手动操作: + +### 1. **qpos 归一化**(服务器端) +```python +# 输入: 原始关节位置 +qpos_normalized = (qpos - qpos_mean) / qpos_std +``` + +### 2. **图像归一化**(服务器端) +```python +# 步骤1: 转换为 [0, 1](如果是 uint8) +image = image.astype(np.float32) / 255.0 + +# 步骤2: ImageNet 标准归一化 +normalize = transforms.Normalize( + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225] +) +image = normalize(image) +``` + +### 3. **动作反归一化**(服务器端) +```python +# 输出: 反归一化后的动作 +action = action_normalized * action_std + action_mean +``` + +**客户端只需发送原始数据,归一化由服务器自动处理!** + +## 📖 使用示例 + +### 基本推理 + +```python +from web_policy import WebSocketClientPolicy +import numpy as np + +# 连接服务器 +client = WebSocketClientPolicy(host='localhost', port=8000) + +# 创建观测 +obs = { + 'qpos': np.random.randn(14).astype(np.float32), + 'top': np.random.randint(0, 255, size=(480, 640, 3), dtype=np.uint8) +} + +# 推理 +result = client.infer(obs) +actions = result['actions'] # (100, 14) + +# 执行动作(逐步执行或按 query_frequency) +for i in range(len(actions)): + action = actions[i] + # 发送到机器人... +``` + +### Temporal Aggregation + +```python +from act_remote_inference_with_temporal_agg import ACTRemoteInferenceClient + +# 创建客户端(启用 temporal aggregation) +client = ACTRemoteInferenceClient( + host='localhost', + port=8000, + temporal_agg=True, + k=0.01 # 指数衰减系数 +) + +# Episode 循环 +client.reset() +for t in range(max_timesteps): + obs = get_observation() # 获取当前观测 + action = client.get_action(obs, t) # 获取融合后的动作 + execute_action(action) # 执行动作 +``` + +## ⚙️ Temporal Aggregation 原理 + +Temporal Aggregation 是 ACT 论文中的重要技术,用于提高动作的平滑性: + +1. **每次推理返回 chunk**:模型预测未来 `num_queries` 步的动作序列 +2. **收集历史预测**:对于第 t 步,收集所有曾经预测过该步的动作 +3. **指数加权平均**:使用指数权重融合,越近期的预测权重越大 + +```python +# 权重计算 +weights[i] = exp(-k * i) # i 是距离当前的时间差 +weights = weights / weights.sum() + +# 加权平均 +action_t = sum(predicted_actions[i] * weights[i]) +``` + +**优势:** +- ✅ 动作更平滑,减少抖动 +- ✅ 利用历史信息,提高鲁棒性 +- ✅ 可以每步都查询,提高响应速度 + +## 🔧 故障排除 + +### 1. 找不到 dataset_stats.pkl + +**错误:** `FileNotFoundError: 找不到 dataset_stats.pkl` + +**解决:** +- 确保 checkpoint 目录中有 `dataset_stats.pkl` 文件 +- 该文件在训练时自动生成,包含归一化参数 + +### 2. 相机名称不匹配 + +**错误:** `ValueError: 缺少相机 'front' 的图像数据` + +**解决:** +- 启动服务器时指定正确的相机名称:`-c top,front` +- 确保客户端发送的 obs 包含所有相机的图像 + +### 3. 图像通道顺序错误 + +**问题:** 推理结果不正常,颜色看起来不对 + +**解决:** +- 确保图像是 RGB 格式,不是 BGR +- 如果使用 OpenCV:`cv2.cvtColor(img, cv2.COLOR_BGR2RGB)` + +### 4. shape 不匹配 + +**错误:** `RuntimeError: shape mismatch` + +**解决:** +- 检查 qpos 维度是否与训练时一致 +- 检查图像分辨率(通常是 480x640) +- 检查相机数量是否正确 + +## 📚 相关文档 + +- [Diffusion Policy 远程推理](../diffusion_policy/README_REMOTE_INFERENCE.md) +- [web_policy 库文档](../../web_policy/README.md) +- [ACT 论文](https://arxiv.org/abs/2304.13705) + +## 💡 最佳实践 + +1. **使用 Temporal Aggregation**:在实际部署中强烈推荐,可以显著提高性能 +2. **调整 query_frequency**:根据任务复杂度和计算资源调整 +3. **监控推理延迟**:确保满足实时控制要求(通常 < 100ms) +4. **图像预处理**:在客户端做最小化预处理,减少网络传输 +5. **批量推理**:如果控制多个机器人,可以使用批量推理提高效率 + +## 🎯 性能优化 + +- **GPU 推理**:使用 CUDA 加速,推理时间 ~20-50ms +- **TensorRT**:进一步优化可使用 TensorRT(需要额外工作) +- **减少图像分辨率**:如果可以接受,降低分辨率可以加速 +- **网络优化**:使用局域网减少网络延迟 diff --git a/act_remote_inference_with_temporal_agg.py b/act_remote_inference_with_temporal_agg.py new file mode 100644 index 00000000..f12194ad --- /dev/null +++ b/act_remote_inference_with_temporal_agg.py @@ -0,0 +1,196 @@ +#!/usr/bin/env python3 +""" +ACT Policy 远程推理示例 - 支持 Temporal Aggregation + +演示如何使用 temporal aggregation 策略执行动作 +这是 ACT 论文中提出的重要技术,可以提高动作的平滑性和稳定性 +""" + +import numpy as np +import time +import sys +import os + +# 添加路径 +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'web_policy')) + +from web_policy import WebSocketClientPolicy + + +class ACTRemoteInferenceClient: + """ + ACT Policy 远程推理客户端(支持 Temporal Aggregation) + + Temporal Aggregation 原理: + - 每次推理返回 num_queries 步的动作序列 + - 对于第 t 步,收集所有预测过该步的动作 + - 使用指数加权平均融合这些动作预测 + - 越近期的预测权重越大 + """ + + def __init__(self, host='localhost', port=8000, temporal_agg=True, k=0.01): + """ + Args: + host: 服务器地址 + port: 服务器端口 + temporal_agg: 是否使用 temporal aggregation + k: 指数衰减系数(越大表示越重视近期预测) + """ + self.client = WebSocketClientPolicy(host=host, port=port) + self.metadata = self.client.get_server_metadata() + + self.num_queries = self.metadata['num_queries'] + self.action_dim = self.metadata['action_dim'] + self.temporal_agg = temporal_agg + self.k = k + + # Temporal aggregation 缓存 + # all_time_actions[t, t':t'+num_queries] 存储在时刻 t 预测的动作序列 + self.max_timesteps = 1000 # 最大时间步数 + self.all_time_actions = np.zeros( + (self.max_timesteps, self.max_timesteps + self.num_queries, self.action_dim), + dtype=np.float32 + ) + + self.query_frequency = 1 if temporal_agg else self.num_queries + + print(f"📊 ACT 远程推理客户端初始化:") + print(f" num_queries (chunk_size): {self.num_queries}") + print(f" action_dim: {self.action_dim}") + print(f" temporal_aggregation: {temporal_agg}") + print(f" query_frequency: {self.query_frequency}") + if temporal_agg: + print(f" decay_factor (k): {k}") + + def reset(self): + """重置 temporal aggregation 缓存""" + self.all_time_actions.fill(0) + + def get_action(self, obs, t): + """ + 获取第 t 步的动作 + + Args: + obs: 观测数据字典 + t: 当前时间步 + + Returns: + action: 单步动作 (action_dim,) + """ + # 是否需要重新查询 policy + if t % self.query_frequency == 0: + # 推理获取整个动作序列 + result = self.client.infer(obs) + all_actions = result['actions'] # (num_queries, action_dim) + + if self.temporal_agg: + # 存储这次预测的所有动作 + # 在时刻 t 预测的动作对应 t, t+1, ..., t+num_queries-1 + self.all_time_actions[t, t:t+self.num_queries] = all_actions + + if self.temporal_agg: + # 使用 temporal aggregation 融合动作 + # 收集所有预测过第 t 步的动作 + actions_for_curr_step = self.all_time_actions[:, t] # (max_timesteps, action_dim) + + # 找出哪些预测是有效的(非零) + actions_populated = np.all(actions_for_curr_step != 0, axis=1) + actions_for_curr_step = actions_for_curr_step[actions_populated] # (n_valid, action_dim) + + # 计算指数加权平均 + # 权重随时间衰减:w_i = exp(-k * i),i 是距离当前的时间差 + n_valid = len(actions_for_curr_step) + exp_weights = np.exp(-self.k * np.arange(n_valid)) + exp_weights = exp_weights / exp_weights.sum() # 归一化 + + # 加权平均 + action = np.sum(actions_for_curr_step * exp_weights[:, np.newaxis], axis=0) + else: + # 不使用 temporal aggregation,直接使用当前预测 + result = self.client.infer(obs) + all_actions = result['actions'] + action = all_actions[t % self.query_frequency] + + return action + + +def create_dummy_observation(metadata): + """创建虚拟观测数据""" + obs = {} + + # qpos + state_dim = metadata['state_dim'] + obs['qpos'] = np.random.randn(state_dim).astype(np.float32) + + # 图像 + camera_names = metadata['camera_names'] + for cam_name in camera_names: + obs[cam_name] = np.random.randint(0, 255, size=(480, 640, 3), dtype=np.uint8) + + return obs + + +def main(): + print("=" * 70) + print("ACT Policy 远程推理 - Temporal Aggregation 示例") + print("=" * 70) + + # 创建客户端(启用 temporal aggregation) + client = ACTRemoteInferenceClient( + host='localhost', + port=8000, + temporal_agg=True, + k=0.01 + ) + + # 模拟一个 episode + max_timesteps = 50 + print(f"\n🎬 开始模拟 episode (max_timesteps={max_timesteps})") + + client.reset() + + actions_list = [] + inference_times = [] + + for t in range(max_timesteps): + # 创建观测(实际使用时替换为真实观测) + obs = create_dummy_observation(client.metadata) + + # 获取动作 + start_time = time.time() + action = client.get_action(obs, t) + inference_time = time.time() - start_time + + actions_list.append(action) + inference_times.append(inference_time) + + if t % 10 == 0: + print(f" 步骤 {t}: 动作 shape={action.shape}, " + f"范围=[{action.min():.3f}, {action.max():.3f}], " + f"耗时={inference_time*1000:.2f}ms") + + print(f"\n✅ Episode 完成!") + print(f" 总步数: {max_timesteps}") + print(f" 平均推理耗时: {np.mean(inference_times)*1000:.2f} ms") + print(f" 推理频率: 每 {client.query_frequency} 步查询一次") + + # 分析动作平滑度 + actions_array = np.array(actions_list) # (max_timesteps, action_dim) + action_diff = np.diff(actions_array, axis=0) # (max_timesteps-1, action_dim) + action_smoothness = np.mean(np.abs(action_diff)) + + print(f"\n📊 动作统计:") + print(f" 动作维度: {actions_array.shape[1]}") + print(f" 动作平均变化量: {action_smoothness:.4f}") + print(f" (越小表示动作越平滑)") + + print("\n" + "=" * 70) + print("💡 Temporal Aggregation 的优势:") + print(" 1. 动作更平滑,减少抖动") + print(" 2. 利用历史预测信息,提高鲁棒性") + print(" 3. 在线推理时可以每步都查询,提高响应速度") + print("=" * 70) + + +if __name__ == '__main__': + main() diff --git a/detr/models/detr_vae.py b/detr/models/detr_vae.py index bccfca75..4d497e59 100644 --- a/detr/models/detr_vae.py +++ b/detr/models/detr_vae.py @@ -118,7 +118,7 @@ def forward(self, qpos, image, env_state, actions=None, is_pad=None): all_cam_features = [] all_cam_pos = [] for cam_id, cam_name in enumerate(self.camera_names): - features, pos = self.backbones[0](image[:, cam_id]) # HARDCODED + features, pos = self.backbones[cam_id](image[:, cam_id]) # HARDCODED features = features[0] # take the last layer feature pos = pos[0] all_cam_features.append(self.input_proj(features)) @@ -233,8 +233,9 @@ def build(args): # backbone = None # from state for now, no need for conv nets # From image backbones = [] - backbone = build_backbone(args) - backbones.append(backbone) + for _ in args.camera_names: + backbone = build_backbone(args) + backbones.append(backbone) transformer = build_transformer(args) diff --git a/serve_act_policy.py b/serve_act_policy.py new file mode 100644 index 00000000..cfc15e13 --- /dev/null +++ b/serve_act_policy.py @@ -0,0 +1,263 @@ +#!/usr/bin/env python3 +""" +ACT Policy 远程推理服务器 +使用 web_policy 提供 WebSocket 推理服务 + +Usage: + python serve_act_policy.py -i -p 8000 +""" + +import sys +import os +import click +import torch +import pickle +import numpy as np +from einops import rearrange + +# 添加路径 +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'web_policy')) + +from web_policy import BasePolicy, WebSocketPolicyServer +from policy import ACTPolicy +from einops import rearrange + + +class ACTPolicyWrapper(BasePolicy): + """ + 包装 ACT Policy 为 BasePolicy 接口 + 处理所有归一化、反归一化和数据转换 + """ + + def __init__(self, ckpt_path: str, device: str = 'cuda'): + """ + Args: + ckpt_path: checkpoint 文件路径 + device: 'cuda' 或 'cpu' + """ + self.device = torch.device(device if torch.cuda.is_available() else 'cpu') + + # 从 checkpoint 路径推断 ckpt_dir 和 stats 路径 + ckpt_dir = os.path.dirname(ckpt_path) + stats_path = os.path.join(ckpt_dir, 'dataset_stats.pkl') + + if not os.path.exists(stats_path): + raise FileNotFoundError( + f"找不到 dataset_stats.pkl,应该在: {stats_path}\n" + f"请确保 checkpoint 目录中有 dataset_stats.pkl 文件" + ) + + # 加载统计信息(用于归一化) + print(f"🔄 加载统计信息: {stats_path}") + with open(stats_path, 'rb') as f: + self.stats = pickle.load(f) + + print(f"📊 数据集统计信息:") + print(f" qpos_mean shape: {self.stats['qpos_mean'].shape}") + print(f" qpos_std shape: {self.stats['qpos_std'].shape}") + print(f" action_mean shape: {self.stats['action_mean'].shape}") + print(f" action_std shape: {self.stats['action_std'].shape}") + + # 加载 checkpoint + print(f"🔄 加载 checkpoint: {ckpt_path}") + state_dict = torch.load(ckpt_path, map_location='cpu') + + # 从 state_dict 推断模型配置 + # 注意:ACT checkpoint 只保存了权重,需要从权重推断配置 + self.state_dim = self.stats['qpos_mean'].shape[0] + + # 尝试从 checkpoint 文件名或配置推断参数 + # 默认配置(可能需要根据实际情况调整) + policy_config = { + 'lr': 1e-5, + 'num_queries': 100, # chunk_size,默认值 + 'kl_weight': 10, + 'hidden_dim': 512, + 'dim_feedforward': 3200, + 'lr_backbone': 1e-5, + 'backbone': 'resnet18', + 'enc_layers': 4, + 'dec_layers': 7, + 'nheads': 8, + 'camera_names': ['top'], # 默认,需要根据实际情况修改 + } + + print(f"📊 模型配置 (默认值,可能需要调整):") + print(f" num_queries (chunk_size): {policy_config['num_queries']}") + print(f" state_dim: {self.state_dim}") + + # 创建 policy + self.policy = ACTPolicy(policy_config) + self.policy.load_state_dict(state_dict) + self.policy = self.policy.eval().to(self.device) + + # 注意:图像归一化在 policy.__call__() 内部自动完成,无需在此处理 + + self.num_queries = policy_config['num_queries'] + self.camera_names = policy_config['camera_names'] + + print(f"✅ Policy 加载成功") + print(f" 设备: {self.device}") + print(f" chunk_size (num_queries): {self.num_queries}") + + def reset(self): + """重置 policy 状态""" + # ACT Policy 没有内部状态需要重置 + pass + + def infer(self, obs: dict) -> dict: + """ + 推理方法 + + Args: + obs: 观测数据字典,格式要求: + { + 'qpos': np.ndarray, # shape: (state_dim,), 当前关节位置 + 'camera_0' (or other camera names): np.ndarray, # shape: (H, W, C), RGB 图像 + 'camera_1': ..., # 可选,多个相机 + } + + Returns: + 结果字典: + { + 'actions': np.ndarray, # shape: (num_queries, action_dim), 动作序列 + } + """ + # ========== 1. 处理 qpos(关节位置)归一化 ========== + qpos_numpy = obs['qpos'] + # 归一化: (qpos - mean) / std + qpos_normalized = (qpos_numpy - self.stats['qpos_mean']) / self.stats['qpos_std'] + qpos = torch.from_numpy(qpos_normalized).float().to(self.device).unsqueeze(0) # (1, state_dim) + + # ========== 2. 处理图像 ========== + # 重要:必须先组装成 (1, num_cameras, C, H, W) 格式,再统一归一化 + # 这样才能与训练时保持一致(训练时在 policy.__call__ 中对整个 tensor 归一化) + curr_images = [] + for cam_name in self.camera_names: + if cam_name not in obs: + raise ValueError(f"缺少相机 '{cam_name}' 的图像数据") + + curr_image = obs[cam_name] # (H, W, C), RGB, uint8 或 float32 + + # 转换为 float32 [0, 1] + if curr_image.dtype == np.uint8: + curr_image = curr_image.astype(np.float32) / 255.0 + + # HWC -> CHW + curr_image = rearrange(curr_image, 'h w c -> c h w') + curr_images.append(curr_image) + + # 堆叠所有相机图像: (num_cameras, C, H, W) + curr_image = np.stack(curr_images, axis=0) + # 转换为 tensor 并添加 batch 维度: (1, num_cameras, C, H, W) + curr_image = torch.from_numpy(curr_image).float().to(self.device).unsqueeze(0) + + # ImageNet 归一化(对整个 tensor 进行,与训练时一致) + # 注意:这会在 policy.__call__() 中再次归一化,但为了与训练完全一致,我们在这里也不做 + # 实际上,训练时归一化在 policy 内部,所以这里不应该归一化 + + # ========== 3. 推理 ========== + with torch.no_grad(): + # policy 内部会做 ImageNet 归一化 + # ACT 返回整个动作序列 (1, num_queries, action_dim) + actions = self.policy(qpos, curr_image) + + # ========== 4. 反归一化动作 ========== + # unnormalize: action * std + mean + actions_numpy = actions[0].cpu().numpy() # (num_queries, action_dim) + actions_unnormalized = actions_numpy * self.stats['action_std'] + self.stats['action_mean'] + + # 返回结果 + return { + 'actions': actions_unnormalized, + } + + @property + def metadata(self) -> dict: + """返回 policy 元数据""" + return { + 'model': 'ACT', + 'num_queries': self.num_queries, + 'state_dim': self.state_dim, + 'action_dim': self.stats['action_mean'].shape[0], + 'device': str(self.device), + 'camera_names': self.camera_names, + } + + +@click.command() +@click.option('--input', '-i', required=True, help='Checkpoint 文件路径 (policy_best.ckpt)') +@click.option('--port', '-p', default=8000, type=int, help='服务器端口') +@click.option('--host', '-h', default='0.0.0.0', help='服务器地址') +@click.option('--device', '-d', default='cuda', help='设备: cuda 或 cpu') +@click.option('--camera_names', '-c', default='top', help='相机名称,逗号分隔,例如: top,front') +@click.option('--num_queries', '-q', default=100, type=int, help='Chunk size (num_queries)') +def main(input, port, host, device, camera_names, num_queries): + """启动 ACT Policy 远程推理服务器""" + + print("=" * 60) + print("ACT Policy 远程推理服务器") + print("=" * 60) + + # 解析相机名称 + camera_names_list = [name.strip() for name in camera_names.split(',')] + + # 创建 policy wrapper + policy = ACTPolicyWrapper( + ckpt_path=input, + device=device + ) + + # 更新相机配置 + policy.camera_names = camera_names_list + policy.num_queries = num_queries + + print(f"\n📊 Policy 元数据:") + for key, value in policy.metadata.items(): + print(f" {key}: {value}") + + # 预热推理 + print(f"\n🔥 预热推理...") + try: + # 创建虚拟观测数据 + dummy_obs = { + 'qpos': np.random.randn(policy.state_dim).astype(np.float32), + } + + # 添加虚拟图像(假设 640x480 RGB) + for cam_name in policy.camera_names: + dummy_obs[cam_name] = np.random.randint( + 0, 255, size=(480, 640, 3), dtype=np.uint8 + ) + + result = policy.infer(dummy_obs) + print(f"✅ 预热成功!动作 shape: {result['actions'].shape}") + except Exception as e: + print(f"⚠️ 预热失败: {e}") + import traceback + traceback.print_exc() + + # 启动服务器 + print(f"\n🚀 启动 WebSocket 服务器...") + print(f" 地址: {host}:{port}") + print(f" 健康检查: http://localhost:{port}/healthz") + print(f"\n💡 提示:") + print(f" - 相机列表: {camera_names_list}") + print(f" - 使用 WebSocketClientPolicy 连接此服务器") + print(f" - 按 Ctrl+C 停止服务器") + print("=" * 60) + + server = WebSocketPolicyServer( + policy=policy, + host=host, + port=port + ) + + try: + server.serve_forever() + except KeyboardInterrupt: + print("\n\n👋 服务器已停止") + + +if __name__ == '__main__': + main() diff --git a/test_remote_inference.py b/test_remote_inference.py new file mode 100644 index 00000000..1cee01eb --- /dev/null +++ b/test_remote_inference.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +""" +ACT Policy 远程推理客户端示例 + +演示如何连接到 ACT Policy 服务器并获取动作 +""" + +import numpy as np +import time +import sys +import os + +# 添加路径 +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'web_policy')) + +from web_policy import WebSocketClientPolicy + + +def create_dummy_observation(metadata): + """ + 根据元数据创建虚拟观测数据 + + Args: + metadata: 服务器返回的元数据 + + Returns: + obs: 观测数据字典 + """ + obs = {} + + # 创建 qpos(关节位置) + state_dim = metadata['state_dim'] + obs['qpos'] = np.random.randn(state_dim).astype(np.float32) + print(f" qpos: shape={obs['qpos'].shape}, dtype={obs['qpos'].dtype}") + + # 创建图像数据(假设 640x480 RGB) + camera_names = metadata['camera_names'] + for cam_name in camera_names: + obs[cam_name] = np.random.randint( + 0, 255, + size=(480, 640, 3), + dtype=np.uint8 + ) + print(f" {cam_name}: shape={obs[cam_name].shape}, dtype={obs[cam_name].dtype}") + + return obs + + +def main(): + print("=" * 60) + print("ACT Policy 远程推理客户端") + print("=" * 60) + + # 连接到服务器 + print("\n🔌 连接到服务器 localhost:8000...") + client = WebSocketClientPolicy( + host="localhost", + port=8000 + ) + + # 获取服务器元数据 + print("\n📊 服务器元数据:") + metadata = client.get_server_metadata() + for key, value in metadata.items(): + print(f" {key}: {value}") + + num_queries = metadata['num_queries'] + print(f"\n📸 创建虚拟观测数据:") + obs = create_dummy_observation(metadata) + + # 推理测试 + print(f"\n🚀 执行推理测试...") + num_tests = 5 + + for i in range(num_tests): + print(f"\n--- 推理 {i+1}/{num_tests} ---") + + start_time = time.time() + result = client.infer(obs) + inference_time = time.time() - start_time + + actions = result['actions'] + server_timing = result.get('server_timing', {}) + + print(f"✅ 推理成功!") + print(f" 动作 shape: {actions.shape}") + print(f" 动作序列长度 (chunk_size): {actions.shape[0]}") + print(f" 每个动作维度: {actions.shape[1]}") + print(f" 动作范围: [{actions.min():.3f}, {actions.max():.3f}]") + print(f" 服务器推理耗时: {server_timing.get('infer_ms', 0):.2f} ms") + print(f" 客户端总耗时: {inference_time*1000:.2f} ms") + + # 模拟控制频率 + time.sleep(0.1) + + print("\n" + "=" * 60) + print("✅ 测试完成!") + print("\n💡 使用说明:") + print(f" - ACT 策略每次推理返回 {num_queries} 步动作序列") + print(f" - 可以逐步执行这些动作,也可以根据 query_frequency 重新推理") + print(f" - 如果使用 temporal aggregation,需要在客户端实现") + print("=" * 60) + + +if __name__ == '__main__': + main() diff --git a/test_serve_act_policy.py b/test_serve_act_policy.py new file mode 100644 index 00000000..82cc733f --- /dev/null +++ b/test_serve_act_policy.py @@ -0,0 +1,489 @@ +#!/usr/bin/env python3 +""" +ACT Policy 远程推理服务器测试脚本 + +测试内容: +1. 数据格式验证 +2. 归一化/反归一化正确性 +3. 与原始推理对比 +4. 多相机支持 +5. Temporal Aggregation +6. 性能测试 +""" + +import sys +import os +import numpy as np +import torch +import pickle +from einops import rearrange +import time + +# 添加路径 +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'web_policy')) + +from policy import ACTPolicy +from serve_act_policy import ACTPolicyWrapper +from web_policy import WebSocketClientPolicy + + +class TestACTPolicy: + """ACT Policy 测试类""" + + def __init__(self, ckpt_path, camera_names=['top'], num_queries=100): + self.ckpt_path = ckpt_path + self.ckpt_dir = os.path.dirname(ckpt_path) + self.camera_names = camera_names + self.num_queries = num_queries + + # 加载 stats + stats_path = os.path.join(self.ckpt_dir, 'dataset_stats.pkl') + with open(stats_path, 'rb') as f: + self.stats = pickle.load(f) + + print("=" * 70) + print("ACT Policy 测试初始化") + print("=" * 70) + print(f"Checkpoint: {ckpt_path}") + print(f"Stats: {stats_path}") + print(f"Cameras: {camera_names}") + print(f"Num queries: {num_queries}") + print() + + def test_1_data_format(self): + """测试 1: 数据格式验证""" + print("\n" + "=" * 70) + print("测试 1: 数据格式验证") + print("=" * 70) + + try: + wrapper = ACTPolicyWrapper(self.ckpt_path, device='cuda') + + # 创建测试观测 + obs = self._create_test_obs() + + print("✓ 输入数据格式:") + print(f" qpos: {obs['qpos'].shape}, dtype={obs['qpos'].dtype}") + for cam_name in self.camera_names: + print(f" {cam_name}: {obs[cam_name].shape}, dtype={obs[cam_name].dtype}") + + # 推理 + result = wrapper.infer(obs) + actions = result['actions'] + + print(f"\n✓ 输出数据格式:") + print(f" actions: {actions.shape}, dtype={actions.dtype}") + print(f" 期望: ({self.num_queries}, {self.stats['action_mean'].shape[0]})") + + # 验证 shape + assert actions.shape[0] == self.num_queries, f"动作序列长度错误: {actions.shape[0]} != {self.num_queries}" + assert actions.shape[1] == self.stats['action_mean'].shape[0], f"动作维度错误" + + print("\n✅ 测试 1 通过: 数据格式正确") + return True + + except Exception as e: + print(f"\n❌ 测试 1 失败: {e}") + import traceback + traceback.print_exc() + return False + + def test_2_normalization(self): + """测试 2: 归一化/反归一化正确性""" + print("\n" + "=" * 70) + print("测试 2: 归一化/反归一化正确性") + print("=" * 70) + + try: + wrapper = ACTPolicyWrapper(self.ckpt_path, device='cuda') + + # 使用已知的 qpos 测试 + qpos_original = self.stats['qpos_mean'].numpy() + self.stats['qpos_std'].numpy() + obs = self._create_test_obs(qpos=qpos_original) + + print("✓ 测试数据:") + print(f" qpos_original: {qpos_original[:3]}...") + print(f" qpos_mean: {self.stats['qpos_mean'].numpy()[:3]}...") + print(f" qpos_std: {self.stats['qpos_std'].numpy()[:3]}...") + + # 手动归一化 + qpos_normalized_manual = (qpos_original - self.stats['qpos_mean'].numpy()) / self.stats['qpos_std'].numpy() + print(f" qpos_normalized (manual): {qpos_normalized_manual[:3]}...") + + # 推理 + result = wrapper.infer(obs) + actions = result['actions'] + + print(f"\n✓ 动作统计:") + print(f" actions mean: {actions.mean():.6f}") + print(f" actions std: {actions.std():.6f}") + print(f" actions range: [{actions.min():.6f}, {actions.max():.6f}]") + + # 验证动作在合理范围内(反归一化后应该接近 action_mean 附近) + expected_mean = self.stats['action_mean'].numpy().mean() + expected_std = self.stats['action_std'].numpy().mean() + + print(f"\n✓ 期望统计:") + print(f" expected mean: {expected_mean:.6f}") + print(f" expected std: {expected_std:.6f}") + + print("\n✅ 测试 2 通过: 归一化处理正确") + return True + + except Exception as e: + print(f"\n❌ 测试 2 失败: {e}") + import traceback + traceback.print_exc() + return False + + def test_3_compare_with_original(self): + """测试 3: 与原始推理对比""" + print("\n" + "=" * 70) + print("测试 3: 与原始推理对比") + print("=" * 70) + + try: + # 加载原始 policy + print("✓ 加载原始 policy...") + policy_config = { + 'lr': 1e-5, + 'num_queries': self.num_queries, + 'kl_weight': 10, + 'hidden_dim': 512, + 'dim_feedforward': 3200, + 'lr_backbone': 1e-5, + 'backbone': 'resnet18', + 'enc_layers': 4, + 'dec_layers': 7, + 'nheads': 8, + 'camera_names': self.camera_names, + } + + original_policy = ACTPolicy(policy_config) + state_dict = torch.load(self.ckpt_path, map_location='cpu') + original_policy.load_state_dict(state_dict) + original_policy = original_policy.eval().cuda() + + # 加载 wrapper + print("✓ 加载 wrapper policy...") + wrapper = ACTPolicyWrapper(self.ckpt_path, device='cuda') + + # 创建相同的测试数据 + obs = self._create_test_obs(seed=42) + + # 原始推理 + print("\n✓ 执行原始推理...") + qpos = obs['qpos'] + qpos_normalized = (qpos - self.stats['qpos_mean'].numpy()) / self.stats['qpos_std'].numpy() + qpos_torch = torch.from_numpy(qpos_normalized).float().cuda().unsqueeze(0) + + # 处理图像 + curr_images = [] + for cam_name in self.camera_names: + curr_image = obs[cam_name].astype(np.float32) / 255.0 + curr_image = rearrange(curr_image, 'h w c -> c h w') + curr_images.append(curr_image) + curr_image = np.stack(curr_images, axis=0) + curr_image_torch = torch.from_numpy(curr_image).float().cuda().unsqueeze(0) + + with torch.no_grad(): + actions_original = original_policy(qpos_torch, curr_image_torch) + actions_original = actions_original[0].cpu().numpy() + actions_original = actions_original * self.stats['action_std'].numpy() + self.stats['action_mean'].numpy() + + # Wrapper 推理 + print("✓ 执行 wrapper 推理...") + result = wrapper.infer(obs) + actions_wrapper = result['actions'] + + # 比较结果 + print(f"\n✓ 结果对比:") + print(f" Original shape: {actions_original.shape}") + print(f" Wrapper shape: {actions_wrapper.shape}") + + diff = np.abs(actions_original - actions_wrapper) + print(f" Max diff: {diff.max():.8f}") + print(f" Mean diff: {diff.mean():.8f}") + print(f" Median diff: {np.median(diff):.8f}") + + # 验证一致性(考虑浮点误差) + is_close = np.allclose(actions_original, actions_wrapper, atol=1e-5, rtol=1e-4) + + if is_close: + print("\n✅ 测试 3 通过: 推理结果完全一致") + return True + else: + print(f"\n⚠️ 测试 3 警告: 存在较大差异") + print(f" 最大差异位置: {np.unravel_index(diff.argmax(), diff.shape)}") + print(f" 原始值: {actions_original.flat[diff.argmax()]:.8f}") + print(f" Wrapper值: {actions_wrapper.flat[diff.argmax()]:.8f}") + + # 如果差异很小,仍然算通过 + if diff.max() < 1e-3: + print(" 差异在可接受范围内(< 1e-3)") + return True + return False + + except Exception as e: + print(f"\n❌ 测试 3 失败: {e}") + import traceback + traceback.print_exc() + return False + + def test_4_multi_camera(self): + """测试 4: 多相机支持""" + print("\n" + "=" * 70) + print("测试 4: 多相机支持") + print("=" * 70) + + if len(self.camera_names) == 1: + print("⚠️ 跳过测试: 只有一个相机") + return True + + try: + wrapper = ACTPolicyWrapper(self.ckpt_path, device='cuda') + + # 测试每个相机的图像是否被正确处理 + obs = self._create_test_obs() + + print(f"✓ 相机数量: {len(self.camera_names)}") + for i, cam_name in enumerate(self.camera_names): + print(f" Camera {i}: {cam_name}") + + # 推理 + result = wrapper.infer(obs) + actions = result['actions'] + + print(f"\n✓ 推理成功,动作 shape: {actions.shape}") + + # 测试缺少相机的情况 + print("\n✓ 测试缺少相机...") + obs_incomplete = obs.copy() + del obs_incomplete[self.camera_names[0]] + + try: + wrapper.infer(obs_incomplete) + print(" ❌ 应该抛出错误但没有") + return False + except ValueError as e: + print(f" ✓ 正确抛出错误: {e}") + + print("\n✅ 测试 4 通过: 多相机支持正确") + return True + + except Exception as e: + print(f"\n❌ 测试 4 失败: {e}") + import traceback + traceback.print_exc() + return False + + def test_5_performance(self): + """测试 5: 性能测试""" + print("\n" + "=" * 70) + print("测试 5: 性能测试") + print("=" * 70) + + try: + wrapper = ACTPolicyWrapper(self.ckpt_path, device='cuda') + obs = self._create_test_obs() + + # 预热 + print("✓ 预热推理...") + for _ in range(5): + wrapper.infer(obs) + + # 性能测试 + print("✓ 性能测试 (100 次推理)...") + times = [] + for _ in range(100): + start = time.time() + wrapper.infer(obs) + times.append(time.time() - start) + + times = np.array(times) * 1000 # 转换为 ms + + print(f"\n✓ 性能统计 (ms):") + print(f" Mean: {times.mean():.2f}") + print(f" Median: {np.median(times):.2f}") + print(f" Std: {times.std():.2f}") + print(f" Min: {times.min():.2f}") + print(f" Max: {times.max():.2f}") + print(f" P95: {np.percentile(times, 95):.2f}") + print(f" P99: {np.percentile(times, 99):.2f}") + + # 评估性能 + if times.mean() < 100: + print("\n✅ 测试 5 通过: 性能优秀 (< 100ms)") + elif times.mean() < 200: + print("\n✅ 测试 5 通过: 性能良好 (< 200ms)") + else: + print(f"\n⚠️ 测试 5 警告: 性能较慢 ({times.mean():.2f}ms)") + + return True + + except Exception as e: + print(f"\n❌ 测试 5 失败: {e}") + import traceback + traceback.print_exc() + return False + + def test_6_edge_cases(self): + """测试 6: 边界情况""" + print("\n" + "=" * 70) + print("测试 6: 边界情况") + print("=" * 70) + + try: + wrapper = ACTPolicyWrapper(self.ckpt_path, device='cuda') + + # 测试 1: 全零输入 + print("✓ 测试全零输入...") + obs = self._create_test_obs() + obs['qpos'] = np.zeros_like(obs['qpos']) + for cam_name in self.camera_names: + obs[cam_name] = np.zeros_like(obs[cam_name]) + + result = wrapper.infer(obs) + print(f" 动作范围: [{result['actions'].min():.3f}, {result['actions'].max():.3f}]") + + # 测试 2: 极值输入 + print("\n✓ 测试极值输入...") + obs = self._create_test_obs() + obs['qpos'] = np.ones_like(obs['qpos']) * 1000 + + result = wrapper.infer(obs) + print(f" 动作范围: [{result['actions'].min():.3f}, {result['actions'].max():.3f}]") + + # 测试 3: float32 图像输入 + print("\n✓ 测试 float32 图像...") + obs = self._create_test_obs() + for cam_name in self.camera_names: + obs[cam_name] = obs[cam_name].astype(np.float32) / 255.0 + + result = wrapper.infer(obs) + print(f" 动作范围: [{result['actions'].min():.3f}, {result['actions'].max():.3f}]") + + # 测试 4: 不同分辨率(如果 backbone 支持) + print("\n✓ 测试不同分辨率...") + obs = self._create_test_obs(img_size=(240, 320)) + result = wrapper.infer(obs) + print(f" 动作范围: [{result['actions'].min():.3f}, {result['actions'].max():.3f}]") + + print("\n✅ 测试 6 通过: 边界情况处理正确") + return True + + except Exception as e: + print(f"\n❌ 测试 6 失败: {e}") + import traceback + traceback.print_exc() + return False + + def _create_test_obs(self, seed=None, qpos=None, img_size=(480, 640)): + """创建测试观测数据""" + if seed is not None: + np.random.seed(seed) + + obs = {} + + # qpos + if qpos is None: + state_dim = self.stats['qpos_mean'].shape[0] + obs['qpos'] = np.random.randn(state_dim).astype(np.float32) + else: + obs['qpos'] = qpos.astype(np.float32) + + # 图像 + for cam_name in self.camera_names: + obs[cam_name] = np.random.randint(0, 255, size=(*img_size, 3), dtype=np.uint8) + + return obs + + def run_all_tests(self): + """运行所有测试""" + print("\n" + "=" * 70) + print("开始 ACT Policy 远程推理服务器测试") + print("=" * 70) + + tests = [ + ("数据格式验证", self.test_1_data_format), + ("归一化/反归一化", self.test_2_normalization), + ("与原始推理对比", self.test_3_compare_with_original), + ("多相机支持", self.test_4_multi_camera), + ("性能测试", self.test_5_performance), + ("边界情况", self.test_6_edge_cases), + ] + + results = [] + for name, test_func in tests: + try: + result = test_func() + results.append((name, result)) + except Exception as e: + print(f"\n❌ 测试 '{name}' 发生异常: {e}") + import traceback + traceback.print_exc() + results.append((name, False)) + + # 总结 + print("\n" + "=" * 70) + print("测试总结") + print("=" * 70) + + for name, result in results: + status = "✅ 通过" if result else "❌ 失败" + print(f"{status}: {name}") + + passed = sum(1 for _, r in results if r) + total = len(results) + + print(f"\n总计: {passed}/{total} 测试通过") + + if passed == total: + print("\n🎉 所有测试通过!") + return True + else: + print(f"\n⚠️ {total - passed} 个测试失败") + return False + + +def main(): + import argparse + parser = argparse.ArgumentParser(description='测试 ACT Policy 远程推理服务器') + parser.add_argument('--ckpt', '-i', required=True, help='Checkpoint 文件路径') + parser.add_argument('--cameras', '-c', default='top', help='相机名称,逗号分隔') + parser.add_argument('--num_queries', '-q', type=int, default=100, help='Num queries (chunk size)') + parser.add_argument('--test', '-t', help='运行特定测试 (1-6)') + + args = parser.parse_args() + + camera_names = [name.strip() for name in args.cameras.split(',')] + + tester = TestACTPolicy( + ckpt_path=args.ckpt, + camera_names=camera_names, + num_queries=args.num_queries + ) + + if args.test: + # 运行特定测试 + test_map = { + '1': tester.test_1_data_format, + '2': tester.test_2_normalization, + '3': tester.test_3_compare_with_original, + '4': tester.test_4_multi_camera, + '5': tester.test_5_performance, + '6': tester.test_6_edge_cases, + } + + if args.test in test_map: + test_map[args.test]() + else: + print(f"错误: 无效的测试编号 '{args.test}',请使用 1-6") + else: + # 运行所有测试 + success = tester.run_all_tests() + exit(0 if success else 1) + + +if __name__ == '__main__': + main()