Lakesenberg/pi0-rtc-cube-realtime-vla

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