Hugging Face's logo

cognaclee
/
Frank-Insert-Wipe

Robotics
openpi
pi0.5
franka
Model card Files
xet
 Community
Frank-Insert-Wipe / code /README.md
cognaclee's picture
cognaclee
Add files using upload-large-folder tool
c8d5563 verified
preview code
|
Raw History Blame Contribute Delete
16.1 kB

AdaSpeVLA 真机训练代码

整理日期：2026-11-10。 用途：在 insert-tube 和 wipe-whiteboard 两个任务上训练 AdaSpeVLA，并和各自的 π0.5 基线做对照。

不包含：检查点、数据集、模型权重、HF token / SSH key 等任何凭证。

目录

路径 内容
adaspe_train/adaspe/ plugin 本体：model.py（AdaSpePi05 模型）、bspline.py（B 样条、knot、timing 标签、位姿几何）、timing_data.py（把未来 13 帧实测 state 带进 batch）
adaspe_train/configs.py 训练 config（不改 OpenPI 源码，通过注册接入）。TASK 决定 config 名，如 bread_adaspe_action、bread_adaspe_timing_from20k
adaspe_train/train.py、adaspe_train/compute_norm_stats.py OpenPI 脚本的包装，先注册 config 再运行
adaspe_train/calibrate_timing_anchor.py 标定 knot / timing 标签用的 motion_scale 和 e_crit
adaspe_train/run_train.sh 全流程：预检 → Action 40k → 3 个 timing
adaspe_train/run_timing_one.sh 单独跑一个 timing，便于多卡并行
adaspe_train/strip_old_train_state.sh 磁盘不够时用：删掉旧检查点里的优化器状态，只留参数
adaspe_train/wait_gpus_and_launch.sh 等 GPU 空闲再启动
adaspe_train/tests/ 数值等价测试、冒烟测试、数据检查、标签检查
adaspe_train/assets/ bread 的归一化统计（示例），以及标签检查用的参考文件
env_patches/ 需要放进 venv site-packages 的补丁（见下文）
reference_adaspevla_source/ 原版 AdaSpeVLA PyTorch 代码，tests/test_bspline_parity.py 用它做对照

Recipe

  • OpenPI JAX，commit 215abfb，不改上游源码。
  • 骨干 PaliGemma 从官方 pi05_base 加载，全量微调；不加载 π0.5 action expert。plugin 头随机初始化，seed 0。
  • 数据配置：OpenPI LeRobotLiberoDataConfig，prompt_from_task=True；state 按 π0.5 默认离散化写进 prompt；图像 resize_with_pad 到 224，使用默认增强。
  • 归一化统计直接用基线模型仓库里的 norm_stats.json（分位数归一化），用前先重算一遍，确认一致。
  • action_horizon=13（1.2s @10Hz），三次 B 样条 K=8（端点固定，ridge λ=1e-4），查询间隔 0.2–0.8s。
  • flow 头：隐藏层 384，6 层，6 头，cond_dim 512，ODE 15 步。flow 的目标是归一化后的原始动作。
  • knot 密度目标和 timing 标签都按原版 AdaSpeVLA 计算：
    • 运动来自窗口内的实测位姿（相对窗口第一帧，位置差 + 旋转向量），按 motion_scale 归一化；
    • 夹爪用实测开度，以 0.5 为界判断切换；
    • timing 标签是第一次运动变化量超过 e_crit 或夹爪切换的时间，裁剪到 [0.2, 0.8]s。
  • 损失：Action 阶段 flow MSE + 0.5 × knot KL；timing 阶段 log-MSE。
  • global batch 32，AdamW（b1 0.9，b2 0.95，wd 1e-10，梯度裁剪 1.0），cosine（warmup 1k，peak 2.5e-5，decay 30k，end 2.5e-6），EMA 0.99，评测用 EMA 参数。
  • Action 训练 40k 步，保留 10k/20k/30k/40k。然后基于 20k/30k/40k 各训练一个 timing 头：10k 步，只训练 freq_head，其余全部冻结。
  • 每 1000 步存一次检查点，wandb 关闭，FSDP 设备数等于 GPU 数。

环境搭建

  1. 克隆 OpenPI 并切到 215abfb，uv sync --frozen --python 3.11。缺内核头文件时加 --no-install-package evdev。
  2. 把 env_patches/ 里的 datasets_list_shim.py 和 datasets_list_shim.pth 放进 external/openpi/.venv/lib/python3.11/site-packages/：让 datasets 3.6 认识 List 特征（训练、norm stats 和检查脚本都要用）。不要把 datasets 升级到 4.x。 仓库里的脚本自己会找到 adaspe_train/，不需要别的路径补丁。只有在仓库外的代码里 import configs（比如自己写的 policy server）时，才需要在仓库根目录运行一次： echo "$PWD/adaspe_train" > external/openpi/.venv/lib/python3.11/site-packages/adaspe_train_path.pth
  3. PaliGemma tokenizer 要预先下载到 HF cache。脚本默认离线运行（HF_HUB_OFFLINE=1）。

路径

代码里没有写死的路径：所有脚本都从自己所在的位置推出仓库根目录（Python 用 __file__，shell 用 BASH_SOURCE），仓库放在哪里都可以，不用改路径。

根目录（即本仓库）下应有的布局，每一项都可以是软链接（比如指到更大的磁盘）：

  • external/openpi：OpenPI 和 venv
  • adaspe_train：这份代码
  • reference_adaspevla_source：原版 AdaSpeVLA 代码（数值等价测试用；也可以用环境变量 ADASPEVLA_REF 指定别的位置）
  • data/lerobot/<repo_id>：LeRobot 数据集
  • checkpoints/：训练输出
  • test/franka/<实验名>/：日志
  • .cache/：HF、JAX 等缓存

pi05_base 参数默认在 .cache/openpi/openpi-assets/checkpoints/pi05_base/params（configs.py 的 PI05_BASE）。

换成新任务要改的地方

两个任务的格式和 bread 一致：两路 128×128 图像，19 维 state，7 维动作，10Hz，各只有一条 task 文本。

任务 数据集（revision） 基线模型（revision） 基线 norm_stats 在仓库里的路径
insert-tube dddddsds31341/insert-tube-lerobot（70bcd1f86ed35fc93a0422a564472355242a51ab），102 条 / 16,403 帧 dddddsds31341/Franka-Insert-Tube-pi05-step-20000（e97d90f5bed133b24e8de4189eb25254641e69cf） assets/bofang/insert-tube-lerobot/norm_stats.json
wipe-whiteboard dddddsds31341/wipe-whiteboard-lerobot（2c8cfa8ddef1f83cb09ca7881a0944cd637cd6d0），112 条 / 29,795 帧 dddddsds31341/Franka-Wipe-Whiteboard-pi05-step-20000（628df08d8cea3965d8a5ece6f57a8f60dd242c2e） assets/bofang/wipe-whiteboard-lerobot/norm_stats.json

  1. configs.py

    • REPO_ID 改成新数据集。
    • 把基线的 norm_stats.json 放到 adaspe_train/assets/baseline/<REPO_ID>/norm_stats.json。注意基线仓库里的目录名是 bofang/...，要按我们的 REPO_ID 放。
    • TASK 改成新任务名（如 insert_tube），config 名会跟着变。检查点放在 checkpoints/adaspe/<TASK>_adaspe_*，不会和别的任务混在一起。

  2. 重新标定，不能沿用 bread 的数（只读数据，CPU 上运行）：

    • 改 calibrate_timing_anchor.py 里的 parquet 路径，然后运行。
    • 把输出的 motion_scale 和 e_crit 填进 configs.py 的 TIMING_MOTION_SCALE 和 TIMING_E_CRIT。
    • 记录标签分布，即 0.2s、0.8s 各占多少。

  3. 核对 state 含义。代码假设：

    • state[0] 是夹爪开度（开 1，合 0）；
    • state[4:7] 是位置；
    • state[7:10] 是 Euler xyz（scipy "xyz"）。

    先用 tests/check_data.py 和几条 episode 确认新数据也是这样。wipe-whiteboard 可能全程拿着板擦，夹爪不开合；这不影响训练，只是 timing 事件全部来自运动。

  4. run 脚本：运行时设置环境变量 TASK，要和 configs.py 里的 TASK 一致。日志默认写到 test/franka/<TASK>_adaspe/。

  5. tests/check_timing_labels.py 依赖 assets/timing_labels/<TASK>_H<H>.npy。先用新任务的数据运行 calibrate_timing_anchor.py 重新生成这个文件，再跑检查。

运行顺序

  1. 重算 norm stats，和基线逐位比较：python adaspe_train/compute_norm_stats.py --config-name pi05_<TASK>_franka_ref
  2. 标定（见上文）。
  3. CPU 测试：
    • tests/test_bspline_parity.py：和原版的误差 ≤1e-12；
    • tests/smoke_model.py；
    • tests/check_data.py；
    • tests/check_timing_labels.py：0 处不一致。
  4. 预检和正式训练：CUDA_VISIBLE_DEVICES=<卡> TASK=<任务名> bash adaspe_train/run_train.sh。可以先设 PREFLIGHT_ONLY=1，只跑预检：训练 30 步，再从断点续训，确认 loss 有限、续训正常。
  5. timing 想并行的话，Action 跑完后每个 timing 用一张卡：CUDA_VISIBLE_DEVICES=4 TASK=<任务名> bash adaspe_train/run_timing_one.sh 20，30、40 同理。

注意：Action 开始后不要改卡数。检查点只能在保存时的同样卡数上续训；global batch 32 要能被卡数整除，所以用 1、2 或 4 张卡。

交付时请给我

  • 最终的 Action（20k/30k/40k）和 timing EMA 参数，以及所有文件的 sha256；
  • loss 曲线；
  • 实际用的 config 导出、标定出的 motion_scale / e_crit、标签分布、git commit；
  • 和本文档不一致的地方，逐条列出。

Jean Zay 上的配置（2026-11-10）

只在 Jean Zay 上用。别的机器不需要 adaspe_train/slurm/，直接用 run_train.sh / run_timing_one.sh。

  • WORK 的 inode 配额紧张，所以 external/（OpenPI 215abfb + uv venv）、data/、checkpoints/、.cache/ 都是指向 $SCRATCH/adaspevla_train/ 的软链接。SCRATCH 上 30 天不访问的文件会被清掉，交付用的参数要拷走。
  • 任务由环境变量 TASK 选择（bread / insert_tube / wipe_whiteboard），各任务的 REPO_ID、motion_scale、e_crit 写在 adaspe_train/tasks.py；configs.py、标定脚本和测试都从这里读，不用再手改 configs.py。还没标定的任务只注册 pi05_<TASK>_franka_ref。
  • 标定结果：assets/calibration/<TASK>_H<H>.json；标签文件：assets/timing_labels/<TASK>_H<H>.npy（原 bread 的 timing_labels_measured.npy 保留不动，是 H=13 旧 recipe 的）。标定命令：TASK=<任务> python adaspe_train/calibrate_timing_anchor.py，结果要抄进 tasks.py（标定和 check 用的都是 adaspe/labels.py 里和训练相同的标签函数）。
  • pi05_base：.cache/openpi/openpi-assets/checkpoints/pi05_base/params（官方 GCS，与 bucket 列表逐文件核对过大小）。tokenizer 在 .cache/openpi/big_vision/paligemma_tokenizer.model（OpenPI 从 OPENPI_DATA_HOME 读，不是 HF cache）。
  • 计算节点不联网，所有数据/权重都在本地；脚本里 unset PYTHONPATH LD_LIBRARY_PATH、PYTHONNOUSERSITE=1，避免模块和 ~/.local 混进 venv。
  • SLURM（adaspe_train/slurm/）：都要在仓库根目录提交（SLURM 运行的是脚本的副本，根目录和日志路径都取提交目录），账号用 -A <账号> 给，脚本里不写账号。
    • prep_checks.slurm：CPU 上跑 norm stats 重算+比较、parity、smoke、check_data、check_timing_labels（两个任务）。
    • preflight.slurm：H100 dev，4 卡单独跑 preflight，作业名必须是 adaspe-pf-<TASK>，写 preflight_passed（可选：dev QoS 全集群共用 128 卡，常排不上；train_task.slurm 第一节自己会先跑 preflight，失败就停链）。
    • train_task.slurm：H100 t3 20 h 一节的链式作业（singleton，自动排下一节，最后一节或真失败时取消）。4 卡 FSDP 4 跑 Action 40k，然后 3 个 timing 各占一张卡并行。 sbatch -A <账号> --job-name=adaspe-insert_tube --export=ALL,TASK=insert_tube adaspe_train/slurm/train_task.slurm（后面各节自动用同一个账号）
    • 日志：test/franka/<作业名>-<id>.out 和 test/franka/<TASK>_adaspe/。

Recipe 修订（2026-11-10，和上面 "Recipe" 不同的地方）

开关都在 adaspe_train/tasks.py，方括号里是原值；改回原值并重新标定即可复现原 recipe（adaspe/labels.py 在原设置下和旧标签逐个一致，0 处不一致）。

项 现在 原来 原因（两个任务的数据）
ACTION_HORIZON 16（1.5 s） 13（1.2 s） 13 没有必须的理由；16 拟合误差配 K=11 反而更小
NUM_CONTROL_POINTS 11 8 归一化动作的位姿 RMSE 0.024/0.021（原 0.029/0.025）
query interval 0.2–1.5 s（上限 = horizon） 0.2–0.8 s 原上限 72–76% 的标签被顶住，并且强制截掉 chunk 后 1/3
夹爪事件（timing 标签、knot 密度） 夹爪命令的开/合状态 实测开度过 0.5 实测开度比命令晚中位 19 帧（~1.9 s，前 ~1.4 s 完全不动，机械臂静止），超过 chunk 长度
knot 密度 1 + 10·tanh(mv / e_crit) 1 + 10·tanh(mv) 原公式 tanh 在 68–73% 的帧间隔饱和，62–63% 的窗口 knot 目标近似均匀；现在 3–4%
推理 sample_spline 返回 execute_steps 和 dense_actions：前缀至少执行到预测的夹爪切换 只有 query interval 静止时下一帧观测不变，切换若落在前缀之后会被一直推后

标定（H=16）：insert-tube e_crit 11.98，标签 0.2 s 12.0% / 1.5 s 54.6%；wipe-whiteboard e_crit 8.96，0.2 s 10.2% / 1.5 s 60.0%。 在开爪命令处，样条在命令那一帧的值中位数 +0.93 / +0.98（均匀 knot 时 +0.66）。

另外修了一个原代码里的 bug：configs.py 里 import openpi.models.model as _model 被后面的 def _model(stage) 覆盖，两个 weight loader 调 _model.restore_params 会在训练开始时报 AttributeError。现在模块别名是 _openpi_model，config 函数叫 model_config。

部署（不在本仓库）：客户端调用 sample_spline，执行 dense_actions[:, :execute_steps]（每步一个 0.1 s 的增量命令），然后重新观测、推理。

变体：pi0.5 action expert 生成控制点（2026-11-13，VARIANT=expert）

动机：真机上 π0.5 基线的精度明显好于 AdaSpeVLA。两者的最大差别在生成动作的部分。AdaSpeVLA 用 20.4 M 的 flow 头，从零训练，只看 backbone 最后一层投影后的 token；π0.5 的 Gemma-300M expert 有 311 M transformer 参数，另有 117 M adaRMS 时间调制层，在 pi05_base 里预训练过，每一层都 attend 到 backbone 的 KV。原插件加载时把它丢掉了。样条本身不是瓶颈：用本 recipe 的样条（H=16、K=11）拟合 demo，9 步前缀累计误差的 p95 ≤ 一个典型步长的 5%（慢速阶段也一样）。

做法（adaspe/model.py 的 generator="pi05_expert"；默认的 "flow_head" 代码和参数初始化都不变，旧检查点照常用）：

  • 网络是完整的 π0.5（PaliGemma + action expert），pi05_base 的全部权重都加载。PluginWeightLoader 会检查 expert 和 action_in_proj / action_out_proj / time_mlp_* 是否齐全、形状是否一致，缺一个就报错。
  • expert 的输入是 K=11 个控制点 token，π0.5 原本是 action token。token = action_in_proj(带噪控制点，按 OpenPI 补零到 32 维) + cp_knot_proj(该控制点 Greville 位置的 sin-cos 编码、knot 向量、log horizon)。cp_knot_proj 初始化为 0，所以第 0 步 expert 收到的输入和 π0.5 完全一样。
  • flow 按 π0.5 的约定：t=1 是噪声，时间采样 Beta(1.5, 1)，ODE 10 步（timing 阶段取特征也用 10 步）。adaRMS 时间调制、attention mask、位置编码和 KV cache 的用法都照搬 Pi0。
  • knot 要从 prefix 预测，所以先单独跑一遍 prefix 拿 KV cache，再跑 expert（和 π0.5 推理时一样）。梯度穿过 KV cache，所以训练等价于 π0.5 的 prefix+suffix 联合前向（tests/test_expert_parity.py 有对比）。
  • 不变的部分：knot 头、timing 头、标签、样条拟合、损失（flow MSE + 0.5 × knot KL）、数据、norm stats、40k 步、保留 10k/20k/30k/40k、timing 阶段。

运行：和原变体一样，--export 里加 VARIANT=expert，作业名前缀用 adaspe-expert-。原变体是 adaspe-，两个变体的作业链互不干扰：

sbatch -A <账号> --qos=qos_gpu_h100-dev --time=02:00:00 --job-name=adaspe-expert-dev-insert_tube \
    --export=ALL,TASK=insert_tube,VARIANT=expert adaspe_train/slurm/train_task.slurm

config：<TASK>_adaspe_expert_action、<TASK>_adaspe_expert_timing_from{20,30,40}k。日志在 test/franka/<TASK>_adaspe_expert/，检查点在 checkpoints/adaspe/<TASK>_adaspe_expert_*。名字统一由 adaspe_train/names.sh 生成。 测试（CPU）：tests/test_expert_parity.py、tests/smoke_model.py（两种 generator 都测）、tests/check_config_names.py，都在 slurm/prep_checks.slurm 里。