[Vertax42]

把 feature/rlt-jax 同步到最新 main，并按 RLT 原论文（rlt.pdf）对齐在线 RL 的几处实现。

建议按 commit 逐个 review：第一个是合并 main，diff 里大部分来自 main（PR #24 等），真正的 RLT 改动集中在后面 5 个 commit。

Commits

1. 0e1121f 合并 main
   • 冲突：cylinder_stack_1009 和 strap_tape_1010 的 VLA 配置采用 main 的版本，strap_tape 用 Xense/strap-tape-taccap-1015。
   • PR feat: add Elite/Rebot robot support, extend XTac-UMI training, and clean up JAX attention #24 删除 cuDNN attention 后要跟着改的地方：rlt/vla.py 删掉对 use_cudnn_attention 的引用（否则加载 VLA 时报 AttributeError）；insert_ethernet_0928 配置删掉 cuDNN 字段，以及启动说明里会把 FSDP 梯度算错的 XLA 参数。
2. 5468a47 采集和训练同时进行（论文 pp. 5-6）
   • 新增 BackgroundTrainer 线程，机器人一边采集一边训练，每个动作块都用最新权重。以前每轮结束后机器人要停下来等训练。
   • 权重、经验回放、随机数和计数器由一把锁保护；梯度计算在锁外进行，取 actor 输出和提交数据不会被训练卡住。
   • 每轮数据仍在轮末一次性提交，UTD 规则不变。后台训练出错时，下一个动作块会直接报错停下。
   • 开关是 rl.async_training，默认 true，resume 时可以改。
3. 3eba8fe 动作块 C 改为 10（论文附录 B）：默认值和 3 个任务配置都改了，30Hz 下一个动作块从 0.67s 缩短到 0.33s。
4. 2739370 经验回放容量从 6000 改为 10 万条（约 1.85 小时的标注数据）
   • 6000 条只够存约 7 分钟，warmup 和早期的人工干预数据很快就被覆盖了。论文把整次运行的数据都留在同一个经验回放里。
   • 实测每条 25.6 KB，存满约占 2.6 GB 内存，每个 RL checkpoint 也会大这么多。
   • 新增统计：开始覆盖旧数据时日志会提示，W&B 记录 round/replay_evicted。
5. 44f637e RLT 测试统一用完整 float32 精度做矩阵乘法
   • 修复几个等价性测试时过时不过的问题，test_remat_matches 之前每次都失败。
   • 精度是全局设置的，因为 jax.default_matmul_precision 上下文只对当前线程生效，管不到后台训练线程。
6. 5d6eca9 RL token 的输入去掉任务指令 token（论文 §IV-A 脚注 1）
   • 只去掉 Task: ..., 部分，保留 pi05 prompt 里离散化的 State: ...。
   • 生成 prefix 缓存和在线提取特征两处都应用，开关是 model.drop_task_tokens，默认 true。

⚠️ 兼容性

• C 从 20 改为 10：actor 和 critic 输出层的大小随之改变，用 C=20 训练的 RL checkpoint 不能在新配置下 resume，加载时契约校验会拒绝。
• drop_task_tokens 默认开启：
  • 已有的 prefix 缓存会被 check_identity 拒绝。需要用 precompute_prefix.py --overwrite 重新生成，再重训 token；或者设 model.drop_task_tokens: false 继续用旧缓存。
  • 旧的 token checkpoint 在在线阶段仍然可用，会按完整 prompt 处理。

没做的

• VLA 和 RL token 一起微调：论文中 token 的重建损失对 VLA 做了 stop gradient（"VLA considered frozen with respect to L_ro"），两边梯度互不影响，所以和我们先微调 VLA、再在冻结的 VLA 上训 token 等价。
• 本体状态加速度：我们的观测里没有速度。
• 部署时自动切换到 RL 策略：暂缓。

测试

• RLT 全部测试（src/openpi/rlt、rlt_policy_test、rlt_mode_test）连续两次跑出 82 passed, 9 skipped，test_remat_matches 也包含在内。新增测试：
  • learner_test.py：后台训练与同步训练结果完全一致；训练进行时仍能取 actor 输出、提交数据；中途打断再 resume 不丢待训练的更新；后台出错能报出来；经验回放的覆盖统计。
  • collector_test.py：后台训练运行时的多轮采集集成测试。
  • prompt_tokens_test.py：用真实 PaliGemma tokenizer 验证，只丢掉指令、保留 State，包括指令本身含 "State" 的情况。
  • features_test.py：在线特征和缓存路径在开/关两种设置下一致。
• 还没在真机上验证。第一次上机时请留意 chunk/feature_ms（VLA 推理和后台训练抢同一块 GPU）和经验回放占用的内存。

🤖 Generated with Claude Code