Views
No views yet
| 文件 | 说明 |
|---|---|
converted_checkpoint.pkl | 转换后的模型权重（pickle + bfloat16） |
pi0_infer.py | Triton 加速推理引擎 |
norm_stats.json | 归一化统计量（state/action 的 mean/std） |
infer_example.py | 推理示例代码 |
pip install torch triton numpy1import pickle
2import numpy as np
3import torch
4from pi0_infer import Pi0Inference
5
6# 加载模型
7checkpoint = pickle.load(open('converted_checkpoint.pkl', 'rb'))
8infer = Pi0Inference(checkpoint, num_views=2, chunk_size=50)
9
10# 输入（需要归一化）
11observation_images = torch.randn(2, 224, 224, 3, dtype=torch.bfloat16, device="cuda") # [-1, 1]
12observation_state = torch.randn(32, dtype=torch.bfloat16, device="cuda") # mean/std 归一化
13diffusion_noise = torch.randn(50, 32, dtype=torch.bfloat16, device="cuda")
14
15# 推理
16actions = infer.forward(observation_images, observation_state, diffusion_noise)
17# 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 - mean) / (std + 1e-6)，然后 pad 到 32 维action * (std + 1e-6) + mean，然后取前 6 维image / 255.0 * 2.0 - 1.0，resize 到 224x224norm_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 |