Pi0 RTC Cube - realtime-vla 推理部署

基于 realtime-vla 的 Triton 加速推理，在 RTX 4090 上实现 ~27ms 延迟（2 视角）。

文件说明

文件	说明
`converted_checkpoint.pkl`	转换后的模型权重（pickle + bfloat16）
`pi0_infer.py`	Triton 加速推理引擎
`norm_stats.json`	归一化统计量（state/action 的 mean/std）
`infer_example.py`	推理示例代码

环境要求

RTX 4090（或支持 bfloat16 的 GPU）
CUDA 12.x
Python 3.10+

pip install torch triton numpy

快速开始

import pickle
import numpy as np
import torch
from pi0_infer import Pi0Inference

# 加载模型
checkpoint = pickle.load(open('converted_checkpoint.pkl', 'rb'))
infer = Pi0Inference(checkpoint, num_views=2, chunk_size=50)

# 输入（需要归一化）
observation_images = torch.randn(2, 224, 224, 3, dtype=torch.bfloat16, device="cuda")  # [-1, 1]
observation_state = torch.randn(32, dtype=torch.bfloat16, device="cuda")                # mean/std 归一化
diffusion_noise = torch.randn(50, 32, dtype=torch.bfloat16, device="cuda")

# 推理
actions = infer.forward(observation_images, observation_state, diffusion_noise)
# actions shape: (50, 32), 前 6 维有效

输入格式

输入	Shape	说明
`observation_images`	`(2, 224, 224, 3)` bfloat16	2 个视角，像素值归一化到 [-1, 1]
`observation_state`	`(32,)` bfloat16	state 用 mean/std 归一化后 pad 到 32 维
`diffusion_noise`	`(50, 32)` bfloat16	随机高斯噪声

输出格式

输出	Shape	说明
`actions`	`(50, 32)` bfloat16	50 步 action chunk，前 6 维有效（5 关节 + 1 夹爪）

归一化/反归一化

推理时需要手动处理归一化：

State 归一化: (state - mean) / (std + 1e-6)，然后 pad 到 32 维
Action 反归一化: action * (std + 1e-6) + mean，然后取前 6 维
Image 归一化: image / 255.0 * 2.0 - 1.0，resize 到 224x224

归一化参数在 norm_stats.json 中。

模型信息

参数	值
训练数据集	`Lakesenberg/RTC_red_cube_on_blue_cube_retry`
任务 Prompt	`leader follower teleop`
视角数量	2（camera1 主视角 + camera2 辅助视角）
Action Chunk	50 步
有效 Action 维度	6（5 关节 + 1 夹爪）
训练步数	5000
最终 Loss	~0.017