Noctis12138/pi0-rokae-6task

pi0 Rokae 6 任务模型（30000 步 EMA，真机测试用）

和 bhz 的 pi0.5 6 任务模型（maple0617/pi05-robot32-30k，robot_only）是同一套配方，只把模型换成 pi0： 同样的数据（rokae_bimanual_6task rev 3ef2361）、同样的动作表示、同样的归一化统计（norm_stats.json 与 pi0.5 版数值完全一致）、 同样的相机输入、同样的 batch 32 / 30000 步 / 学习率 / EMA 0.999 / 腕部相机 dropout。 如果真机上已经能跑 pi0.5 那一版，输入输出的适配代码可以原样复用，只改模型配置这一处（见下文）。

和 pi0.5 版部署时唯一要改的地方

from openpi.models.pi0_config import Pi0Config
model_config = Pi0Config(pi05=False, action_horizon=50)   # max_token_len 默认 48；pi0.5 版是 pi05=True, max_token_len=200
params = openpi.models.model.restore_params("30000/params", dtype=jnp.bfloat16)
model = model_config.load(params)
actions = model.sample_actions(rng, observation, num_steps=10)   # [B, 50, 32]，前 14 维有效

prompt 用 PaligemmaTokenizer(max_len=48).tokenize(prompt, state=None) 编码（pi0.5 版是 max_len=200）。6 个 prompt 都远短于 48 个 token。

输入（与 pi0.5 版相同）

键	内容
base_0_rgb	头部相机 observation.images.head
left_wrist_0_rgb	左腕相机 observation.images.left_wrist
right_wrist_0_rgb	右腕相机 observation.images.right_wrist

• 图像：resize_with_pad 到 224×224，uint8 → x/255*2-1（范围 [-1, 1]），三路的 image_mask 都为 True。
• state：不使用实测状态，输入全零（32 维）。
• prompt：任务原文，见下表。

输出动作（与 pi0.5 版相同）

模型输出 50 步 × 32 维，只有前 14 维有效：左手 7 维 + 右手 7 维，每只手为 [dx, dy, dz, rx, ry, rz, grip]。

1. 反归一化（分位数）：a = (y + 1) / 2 * (q99 - q01 + 1e-6) + q01，q01/q99 取 30000/assets/rokae_bimanual_6task/norm_stats.json。
2. 还原绝对位姿：位移和旋转都相对于该 chunk 第 0 步动作的位姿 (p0, R0)（训练时用的是 action[0]，不是实测状态；部署时用当前位姿作为锚点）： p_k = p0 + R0 @ d_k，R_k = R0 @ exp(r_k)（r_k 为旋转向量）。
3. 夹爪：第 7 维是绝对值，不是相对量。

任务	prompt（必须一字不差）
block_cup_placement	place the block into the cup with the left arm, then place the cup onto the plate with the right arm
cross_layer_swap	Use the left hand to place the red block on the higher shelf, and use the right hand to place the orange block on the lower shelf.
cup_stack_and_place	stack one cup into the other, then place the stacked cups onto the plate
drawer_block_place	open the drawer, put the pink block into the drawer, then close the drawer
toothbrush_in_cup_to_yellow_shelf	place the cup on the yellow shelf with the left arm, then put the toothbrush into the cup with the right arm
pomegranate_orange_to_purple_box	put the orange and the pomegranate into the purple box

仓库内容

路径	内容
30000/params/	EMA 推理参数，Orbax/JAX 原始格式（约 12 GB）
30000/assets/rokae_bimanual_6task/norm_stats.json	动作归一化统计（14 维，q01/q99/mean/std）
recipe_assets/robot/	训练时使用的同一份统计及其 manifest
recipe.json	模型、优化器、学习率、数据版本等完整配方
metrics.jsonl	训练曲线（每 20 步一条）

不含优化器状态、训练数据和中间 checkpoint，不能用本仓库续训。

已做的检查

• 训练正常跑完 30000 步，最终 checkpoint 写盘完成（Saved checkpoint 30000）。
• 用上面的部署代码（Pi0Config(pi05=False, action_horizon=50) + restore_params + sample_actions，10 步去噪）加载本仓库的 30000/params， 在 4 个随机训练帧上采样：输出形状 [4, 50, 32]，全部为有限值；有效 14 维的归一化误差平均 0.045（中位数 0.011）， 对比"全零预测"的 0.297。分项：左手位移 0.066 / 旋转 0.054 / 夹爪 0.014，右手位移 0.037 / 旋转 0.037 / 夹爪 0.031。 这只说明权重和输入输出格式是对的，不代表真机成功率（训练帧上的误差不等于泛化）。
• norm_stats.json 与训练时使用的统计逐项一致，也与 pi0.5 版的统计一致。
• 上传后已与训练服务器上的文件逐个核对 sha256。

尚未上真机。

真机测试建议

为了和 TurboVLA（ActionMap / ACT）、pi0.5 可比：场景摆放和相机不动，每个任务至少 10 次； 三路视频（头部 + 两个腕部）都录下来，连同控制日志一起回传。

许可

本模型由 OpenPI 的 pi0_base 全量微调得到。OpenPI 代码按 Apache 2.0（见 LICENSE）； 模型中源自 Gemma 的部分受 Gemma Terms of Use 约束（见 LICENSE_GEMMA.txt、NOTICE）， 包括其中的 Gemma Prohibited Use Policy。训练数据不包含在本仓库中。