spikefly/pi05-arx-cot-math-step30k

π0.5 · ARX Lift2S 双臂把题目答案放到白板上(3 万步存档)

在 yaoxianze/arx_cot_math(89 局示教,单一任务:把题目的答案正确地放到白板上)上,从 lerobot/pi05_base 全参数微调 3 万步(全局批 32)的 π0.5,本仓是第 3 万步的存档(训练的最后一档)。 提示词逐字(含句末句点)Place the answers to the questions on the whiteboard properly.;每次推理出 70 步,执行其中第 1–35 步(和叠毛巾、挂衣架模型一样,和叠杯模型不同)。 本仓只有模型和部署工具,不含数据;复现检查和回放测试要用的数据按第 3 节从原数据仓下载、自己转换。来龙去脉和我们想要的反馈见 HANDOVER.md。

已经在用叠杯 / 叠毛巾 / 挂衣架模型的话,先看这几处不同

项	叠杯模型	叠毛巾模型	挂衣架模型	本仓(放答案)
提示词	stack three cups	fold the towel	三句之一(远 / 中 / 近段)	Place the answers to the questions on the whiteboard properly.(逐字,含句点)
每次推理出几步 / 执行第几步	50 / 第 1–25 步	70 / 第 1–35 步	70 / 第 1–35 步	70 / 第 1–35 步
夹爪读数范围(数据实测)	左 −3.34 … −1.34,右 −3.39 … −1.40	左 −3.32 … −0.04,右 −3.39 … −0.04	左 −3.331 … −0.467,右 −3.385 … −0.052	左 −3.328 … −1.545,右 −3.384 … −1.438(两只夹爪最多都只闭到约 −1.5)
归一化参数(在 policy_preprocessor* / policy_postprocessor* 里)	叠杯数据统计	毛巾数据统计	挂衣架数据统计	本数据统计;四个模型的处理器文件不能互换
复现检查要加的参数	不用加	--task-prompt "fold the towel" --repo-id local/arx_towel_pi05 --execute-horizon 35	--task-prompt from-dataset --repo-id local/arx_hanger_pi05 --execute-horizon 35	--task-prompt from-dataset --repo-id local/arx_math_pi05 --execute-horizon 35(只有一个任务串,from-dataset 和直接传那句话结果相同)

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-cot-math-step30k --local-dir pi05_math_step30k

约 9.4 GB。

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

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

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

• 下载后脚本会按目录数文件(meta 5、data 89、videos 267),和这个版本应有的数目对不上就停下。原数据仓里的深度图、标定、相机参数等目录不用下载。
• 校验规则:视频按字节比(7 个 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)从原数据仓 10aa998 重建过:视频和 parquet 全部和训练用的一致,在重建的数据上跑第 4 节的复现检查,数字和 deploy/expected/ 完全相同。

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

python deploy/fit_check_pi05.py --checkpoint . --dataset-root "$DATA" --output-dir fitcheck_bf16 \
    --task-prompt from-dataset --repo-id local/arx_math_pi05 --execute-horizon 35 \
    --episode-indices 0 11 22 33 44 55 66 77 88 --dtype bfloat16 --compare-to deploy/expected/action_fit_summary_bf16.json

• 9 局 = 在 89 局里每隔 11 局取一局(第 0、11、22 … 88 局);--task-prompt from-dataset 让每一局用它自己的任务串,和训练时一样(这份数据只有一个任务串)。
• 输出 fitcheck_bf16/action_fit_summary.json;脚本会读 deploy/tolerance.json(和 deploy/TOLERANCE.md 同一组数)逐项判,末尾打印 TOLERANCE_CHECK PASS/FAIL。必须过的几项:
  • 期望块执行段(第 1–35 步)最大差:全部 14 维 ≤ 0.1605,只看 12 个关节 ≤ 0.0452 rad(容差 = 3 × max(H200 跨精度差, 5090 实测差));
  • 拟合误差的相对差 ≤ 10%;块起点偏移的相对差 ≤ 13%;
  • client_path_test.pass 为 true(每帧图像差 ≤ 1/255)。
• 为什么全维和关节分两项:这一档 bf16 和 fp32 之间,执行段全维最大差 0.0535 落在夹爪维(夹爪开合在几步之内走完几个原始单位,两种精度把开合时刻挪一点点,在最陡那一步就差得多),关节只差 0.0126 rad。只看全维的话,关节上的接线错可能被夹爪的余量盖住,所以关节单独一项,两项都要过。整块(第 0–69 步)的两项只作提示:第 36 步以后不执行,两次采样在这段本来就分叉得最多。
• 再用你自己机器人端的取图代码跑一遍:加 --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):部分显卡上两次前向即使输入完全相同也不逐位相同,块差不为 0 是正常的。想逐位复现可以加 --deterministic(默认关;会变慢,不支持确定性的算子只告警)。
• 期望块用 CPU 上固定种子的噪声生成,所以换一块显卡也能逐块对比。这一档我们还没在 5090 上测过;容差里的 5090 部分借用叠杯 1 万步档在 5090 上的实测(执行段 0.015、整块 0.041;这部分差主要来自显卡内核的非确定性)。本档执行段和整块的全维两项由 H200 自己的跨精度差决定;只看关节的两项(执行段 0.0452、整块 0.1217 rad,后者是提示项)由借来的数决定:本档 H200 上关节的跨精度差(执行段 0.0126、整块 0.0214 rad)都比借来的数小。

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

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

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

6. 上真机的执行规则

• 第一次上机:先悬空或低速,只打印一块 70 步动作,核对关节范围、顺序、夹爪方向都对了再正式执行;确认急停可用。
• 输入:三路相机 640×480 RGB——头 camera_h、左腕 camera_l、右腕 camera_r;14 维关节 observation.state,顺序 [左 j0…j5, 左夹爪, 右 j0…j5, 右夹爪],关节为弧度,夹爪为原始读数(数据实测范围见上表;按数据集约定 0 为全闭,请以你们机器人的实际读数为准);提示词逐字 Place the answers to the questions on the whiteboard properly.。
• 输出:一块 70 步 × 14 维绝对关节目标(已反归一化),顺序同 state。
• 执行:20 Hz;第 0 步约等于「现在的位置」,不执行;执行第 1–35 步,然后重新推理。
• 块起点渐入:默认开——前 5 步(k 可调)从实测关节线性过渡到预测轨迹,缓解每次重推的起点跳变;开 / 关在真机上对比。
• 夹爪:直接下发模型输出值。
• 资源(bf16):H200 上一次推理(出一整块 70 步)的耗时看节点:我们在两台 H200 节点上分别测到中位 271 ms(最大 309 ms)和中位 283 ms(最大 317 ms),同步循环(推理 + 执行 35 步)每块约 2.02–2.03 s;显存峰值约 8.9 GB;测的时候两台节点上都还有别人的作业在跑。结构和块长相同的挂衣架模型之前测到的是 215 ms 和 300 ms。这一档还没在 5090 上测;参考叠杯模型在 5090 上中位 204 ms(块长 50)。
• 详细契约见 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_math_step30k"  # 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"}},
)

TASK = "Place the answers to the questions on the whiteboard properly."  # 逐字,含句末句点

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": [TASK],
    "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]  # (70, 14);执行 chunk[1:36]

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=35)
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(这一档训练局上关节中位 0.0042 rad,约 0.58 帧的单关节运动量;每次取偏移最大的关节、按每帧最快关节的运动量算约 1.13 帧);真机上超过离线 3 倍会给出排错提示(只是提示,不是关卡)。
• 执行段跟踪 |â_k − s_{t+k}|:含控制器跟踪滞后,离线没有这一项,要和回放测试测出的跟踪误差比。

8. 常见错误怎么判

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

训练细节

• 数据:yaoxianze/arx_cot_math 版本 10aa998(89 局、33,572 帧、20 fps,每局 281–530 帧,中位 359),只有一个任务「Place the answers to the questions on the whiteboard properly.」;全部 89 局都参与训练,没有留出局。转成 LeRobot v3.0(本仓不含数据;转换脚本和逐文件校验和在 tools/)。
• 输入:头 + 两腕三路相机,14 维关节(补零到 32 维),提示词 = 上面那句任务串。
• 动作:14 维绝对关节,块长 70,部署时每次执行 35 步;这份数据的动作列等于同一帧的关节读数,所以模型输出的第 0 步约等于当前位置。
• 归一化:关节和动作用分位数(q01 / q99),图像不归一化;参数已存进 policy_preprocessor* / policy_postprocessor*。
• 优化:2 卡 × 每卡 16 = 全局批 32,3 万步(约 28.6 遍数据);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_cot_math。