spikefly/pi05-arx-stack-three-cups-step10k

π0.5 · ARX Lift2S 双臂叠三个杯子(1 万步存档)

v3(10-09):修正数据下载 / 安装 / 校验 / 客户端判据,加 5090 实测;权重没变,已经下载过的请重新拉取 README.md、HANDOVER.md、deploy/、tools/。

在 yaoxianze/arx_external_cam_cups(50 局示教)上,从 lerobot/pi05_base 全参数微调 2 万步(全局批 32)的 π0.5,本仓是第 1 万步的存档。同一次训练的另一档:spikefly/pi05-arx-stack-three-cups-step20k(2 万步)。 本仓只有模型和部署工具,不含数据;复现检查和回放测试要用的数据按第 3 节从原数据仓下载、自己转换。来龙去脉和我们想要的反馈见 HANDOVER.md。

先选哪一档

真机上两档都试。 下表是训练局上的离线误差(第 0 / 8 / 16 / 25 / 33 / 41 / 49 局,按部署方式每 25 步重推一次),只供参考:这些局都参与了训练,步数越多误差越低是预期的,分不出「学得更好」还是「记得更牢」,不能拿来选档。

存档	执行段关节拟合误差(平均)	块起点偏移 |â0 − s_t|(关节中位)
1 万步	0.0108 rad	0.0041 rad
2 万步	0.0064 rad	0.0034 rad

两档的块起点偏移都不到 1 帧的关节运动量。

1. 装环境

conda create -n pi05 python=3.12 -y && conda activate pi05
# RTX 5090(Blackwell)要用 CUDA 12.8 版的 torch,先装它:
pip install torch==2.7.1 torchvision==0.22.1 --index-url https://download.pytorch.org/whl/cu128
git clone https://github.com/huggingface/lerobot.git && cd lerobot
git checkout 89236ea0f4f81a81ca566081e20dd1ff5f823cbe
pip install -e ".[pi,dataset]"

clone 卡在 git-lfs 上时:最常见的是机器上只配置了 lfs 过滤器、没装 git-lfs,用 git -c filter.lfs.process= -c filter.lfs.smudge= -c filter.lfs.required=false clone https://github.com/huggingface/lerobot.git;装了 git-lfs 的可以用 GIT_LFS_SKIP_SMUDGE=1 git clone https://github.com/huggingface/lerobot.git。推理不需要 LFS 文件。

• 推理只要上游 lerobot 这个提交,不需要任何改动。
• [dataset] 带 PyAV、torchcodec、jsonlines、datasets、pandas、pyarrow;只装 [pi] 的话,数据转换会因缺 jsonlines 直接 ImportError,环境里也没有 PyAV。
• 看到 torchcodec 加载失败的提示可以忽略,复现检查固定用 pyav。
• 不需要登录 HF:本仓和原数据仓都公开;分词器在 tokenizer/ 里,不用另下 google/paligemma。
• 我们训练和出期望值用的是 Python 3.12.14、torch 2.7.1(CUDA 12.6)、torchvision 0.22.1、transformers 5.5.4、PyAV 15.1.0、NVIDIA H200。版本不同,只要复现检查在容差内就行。

2. 下载模型

hf download spikefly/pi05-arx-stack-three-cups-step10k --local-dir pi05_step10k

约 9.4 GB。

3. 准备数据(复现检查和回放测试要用;推理不需要)

从原数据仓下载固定版本 2d87a09(我们训练用的版本),用仓里的 tools/ 转成训练时用的 LeRobot v3.0 格式,并自动校验:

cd pi05_step10k
bash tools/prep_pi05_data.sh ~/arx_data            # 只下 meta/、data/、videos/(约 1.15 GB),转换后自动跑校验;已有同版本的原始数据,可以作第二个参数传入,就不再下载
DATA=~/arx_data/arx_cups_pi05                      # 你自己转换出来的目录,后面都用它
python tools/verify_dataset_v30.py "$DATA"         # 想单独再校验一次也可以

• 下载后脚本会按目录数文件(meta 5、data 50、videos 200),和这个版本应有的数目对不上就停下。

• 校验规则:视频按字节比(8 个 mp4 的 sha256,和 tools/dataset_v30_SHA256SUMS.txt 一致);parquet 按内容比(行数、列名、列类型、逐列内容哈希,参考值在 tools/dataset_v30_reference.json)。不同 pyarrow 版本写出的 parquet 字节会不同(文件头记录的写入版本不一样),这不算错。meta/info.json、meta/stats.json 只作提示,对不上不影响推理和复现检查。最后一行是 VERIFY_DATASET_V30 PASS 就行。

• 我们用这套工具(上游 lerobot 89236ea)从原数据仓 2d87a09 重建过:视频和 parquet 全部和训练用的一致,在重建的数据上跑第 4 节的复现检查,数字和 deploy/expected/ 完全相同;parquet 的内容哈希在 pyarrow 25 和 26 下一致。

4. 复现检查(上真机前必须过)

python deploy/fit_check_pi05_cups.py --checkpoint . --dataset-root "$DATA" --output-dir fitcheck_bf16 \
    --episode-indices 0 25 --dtype bfloat16 --compare-to deploy/expected/action_fit_summary_bf16.json

• 输出 fitcheck_bf16/action_fit_summary.json;脚本会读 deploy/tolerance.json(和 deploy/TOLERANCE.md 同一组数)逐项判,末尾打印 TOLERANCE_CHECK PASS/FAIL:期望块最大差在容差内(执行段 ≤ 0.0509 rad、整块 ≤ 0.1216 rad;容差 = 3 × max(H200 跨精度差, 5090 实测差))、拟合误差和块起点偏移的相对差 ≤ 10%、client_path_test.pass 为 true(每帧图像差 ≤ 1/255)。
• 再用你自己机器人端的取图代码跑一遍:加 --client-fn 你的模块:函数 --client-frame-format bgr(或 rgb;模块要在 PYTHONPATH 上)。函数收相机原样交出的 uint8 HxWx3 帧,返回进模型的 float (3,H,W),取值 [0,1];模板 deploy/example_client_fn.py。BGR / RGB、缩放、裁剪接错,只有这样才抓得到。
• 客户端路径测试的关卡只看图像:每帧图像差 ≤ 1/255 才算 PASS。两条路径预测出的块差只作提示(client_path_test.chunk_advisory,超过执行段容差才打 WARN):部分显卡上两次前向即使输入完全相同也不逐位相同(我们在 5090 上看到最大约 0.03 rad),块差不为 0 是正常的。想逐位复现可以加 --deterministic(默认关;会变慢,不支持确定性的算子只告警)。
• 期望块用 CPU 上固定种子的噪声生成,所以换一块显卡也能逐块对比;5090 和 H200 之间有一点数值差,在容差内就正常(我们在 5090 上实测本档:执行段差 0.012–0.015 rad、整块差 0.037–0.041 rad)。

5. 回放测试(任何策略之前,与模型无关)

python deploy/replay_demo.py export --dataset-root "$DATA" --episode 0 --out replay_ep000_commands.csv
# 按第 0 局的布局摆好杯子,机械臂慢速移到第 0 行,再按 20 Hz 逐行下发,记录实测关节和视频
python deploy/replay_demo.py analyze --commands replay_ep000_commands.csv --measured replay_ep000_measured.jsonl

回答三件事:「拿实测位置当指令」能不能复现示教;夹爪停在约 −2.0 时夹不夹得住杯子(看视频);控制器跟踪滞后几帧(后面真机执行段的跟踪误差要和它比)。回放都做不成,先查硬件和指令语义。

6. 上真机的执行规则

• 第一次上机:先悬空或低速,只打印一块 50 步动作,核对关节范围、顺序、夹爪方向都对了再正式执行;确认急停可用。
• 输入:三路相机 640×480 RGB——头 camera_h、左腕 camera_l、右腕 camera_r(外部相机不用);14 维关节 observation.state,顺序 [左 j0…j5, 左夹爪, 右 j0…j5, 右夹爪],关节为弧度,夹爪为原始读数(全开约 −3.34 / −3.39(左 / 右)、抓住杯子时约 −2.0,都是数据实测;按数据集约定 0 为全闭,请以你们机器人的实际读数为准);提示词逐字 stack three cups。
• 输出:一块 50 步 × 14 维绝对关节目标(已反归一化),顺序同 state。
• 执行:20 Hz;第 0 步约等于「现在的位置」,不执行;执行第 1–25 步,然后重新推理。
• 块起点渐入:默认开——前 5 步(k 可调)从实测关节线性过渡到预测轨迹,缓解每次重推的起点跳变;开 / 关在真机上对比。
• 夹爪:直接下发模型输出值。
• 资源(bf16):RTX 5090 上我们实测一次推理(出一整块 50 步)中位 204 ms、最大 206 ms,显存峰值约 8.9 GB(nvidia-smi 看到 10.6 GB,含 CUDA 上下文),同步循环(推理 + 执行 25 步)每块约 1.45 s;H200 上一次推理约 0.28 s。
• 详细契约见 deploy/IO_CONTRACT.md。最小推理代码:

import numpy as np
import torch
from lerobot.configs.policies import PreTrainedConfig
from lerobot.policies.factory import make_pre_post_processors
from lerobot.policies.pi05 import PI05Policy

ckpt = "pi05_step10k"  # hf download 下来的目录
cfg = PreTrainedConfig.from_pretrained(ckpt)
cfg.device, cfg.dtype = "cuda", "bfloat16"
policy = PI05Policy.from_pretrained(ckpt, config=cfg).eval()
pre, post = make_pre_post_processors(
    policy_cfg=cfg, pretrained_path=ckpt,
    preprocessor_overrides={"device_processor": {"device": "cuda"}},
    postprocessor_overrides={"device_processor": {"device": "cpu"}},
)

def to_tensor(rgb_uint8):  # 相机帧 uint8 HxWx3,已转成 RGB;不要自己缩放、不要自己归一化
    return torch.from_numpy(np.ascontiguousarray(rgb_uint8)).permute(2, 0, 1).float() / 255.0

batch = {
    "task": ["stack three cups"],
    "observation.state": torch.tensor(state14, dtype=torch.float32).reshape(1, -1),
    "observation.images.camera_h": to_tensor(head_rgb).unsqueeze(0),
    "observation.images.camera_l": to_tensor(left_rgb).unsqueeze(0),
    "observation.images.camera_r": to_tensor(right_rgb).unsqueeze(0),
}
with torch.no_grad():
    chunk = post(policy.predict_action_chunk(pre(batch)))[0]  # (50, 14);执行 chunk[1:26]

7. 真机上查训推不一致

import sys; sys.path.insert(0, "deploy")               # 在仓根目录下
from rollout_consistency import RolloutLogger
log = RolloutLogger("run_001.jsonl")                    # 一局一个文件
log.inference(step=t, state=state_t, chunk=chunk, exec_from=1, exec_n=25)
log.state(step=t + k, state=measured_joints)            # 每个控制步记一次实测关节
log.close()

跑完执行:

python deploy/rollout_consistency.py run_*.jsonl --reference deploy/expected/action_fit_summary_bf16.json

• 块起点偏移 |â0 − s_t|:模型估计的「手臂现在在哪」和实测关节之差,不需要标注,直接反映模型看到的东西和训练时是否一致。离线基准见 deploy/TOLERANCE.md;真机上超过离线 3 倍会给出排错提示(只是提示,不是关卡)。
• 执行段跟踪 |â_k − s_{t+k}|:含控制器跟踪滞后,离线没有这一项,要和回放测试测出的跟踪误差比。

8. 常见错误怎么判

现象	多半是	怎么查
复现检查期望块差远超容差,数值接近 TOLERANCE.md 里的「送 BGR 时」	颜色通道接反(OpenCV 取图是 BGR)	用 --client-fn 跑客户端路径测试
复现检查不过,但客户端路径测试正常	数据没按第 3 节转换,或转换结果和我们的不一致	看 python tools/verify_dataset_v30.py "$DATA" 的结果
真机上块起点偏移是离线的好几倍	相机机位 / 曝光 / 光照变了,或三路相机接错位置	逐路存一帧图,和你转换出来的数据里的帧对比
动作关节错位、方向反	state 顺序或单位不对	打印 state,和你转换出来的数据里同一姿态的一帧对比
动作幅度明显不对	客户端自己做了归一化(处理器里已经有)	只送原始读数和 [0,1] 图像
动作乱、不像在叠杯子	提示词不是逐字 stack three cups
每块开头一跳	执行了第 0 步,或没开渐入	见第 6 节
5090 上和期望值有小差	显卡 / 内核数值差	bf16 下在容差内即正常
客户端路径测试 PASS,但打了块差 WARN	显卡内核非确定性(两次前向不逐位相同)	图像差为 0 就不是取图的问题;想排除可加 --deterministic 再跑
CUDA error: no kernel image is available	torch 不是 CUDA 12.8 版	见第 1 节
跟踪误差大、块起点偏移正常	控制器滞后	和回放测试的跟踪误差比

训练细节

• 数据:yaoxianze/arx_external_cam_cups 版本 2d87a09(50 局、24,003 帧、20 fps),转成 LeRobot v3.0(本仓不含数据;转换脚本和逐文件校验和在 tools/)。
• 输入:头 + 两腕三路相机(外部相机不用),14 维关节(补零到 32 维),提示词 stack three cups。
• 动作:14 维绝对关节,块长 50;这份数据的动作列等于同一帧的关节读数,所以模型输出的第 0 步约等于当前位置。
• 归一化:关节和动作用分位数(q01 / q99),图像不归一化;参数已存进 policy_preprocessor* / policy_postprocessor*。
• 优化:2 卡 × 每卡 16 = 全局批 32,2 万步;AdamW(β = 0.9 / 0.95,权重衰减 0.01,梯度裁剪 1.0),学习率 2.5e-5,预热 1000 步后余弦衰减到 2.5e-6;bf16,梯度检查点;不做图像增强,不用 EMA;视觉编码器不冻结。
• 底座:lerobot/pi05_base(本次用的 model.safetensors 与其主分支逐字节一致)。
• 代码:上游 lerobot 提交 89236ea;训练时另有两个只管存档的本地开关,不影响推理。

文件

路径	内容
model.safetensors、config.json	π0.5 权重与配置
policy_preprocessor*、policy_postprocessor*	相机改名、分位数归一化 / 反归一化参数
tokenizer/	分词器
deploy/	部署契约、容差(TOLERANCE.md 和机器可读的 tolerance.json)、复现检查、客户端模板、回放测试、训推一致性工具、期望结果
tools/	从源数据集重建 v3.0 数据的脚本(复现检查 / 回放测试要用);校验脚本 verify_dataset_v30.py 和参考值(dataset_v30_reference.json、dataset_v30_SHA256SUMS.txt)
README.md、HANDOVER.md	用法;来龙去脉与反馈
SHA256SUMS	全部文件的校验和(sha256sum -c SHA256SUMS)

许可

权重派生自 lerobot/pi05_base(其模型卡的 license 字段为 gemma;底层为 PaliGemma),使用须遵守 Gemma 使用条款。本仓不分发数据,数据见 yaoxianze/arx_external_cam_cups。