RoboDyna — π0.5 multi-task policy (dual-arm UR5)

Evaluating this model and getting bad/lost behavior? Read EVAL_FIX.md and use robodyna_policy.py (both in this repo). The weights are verified good (reproduces training actions at correlation ~0.9997); "lost" behavior is an input-format bug — almost always images passed as HWC instead of CHW (3,H,W), or a missing repack transform.

A π0.5 (openpi) vision-language-action policy finetuned on the RoboDyna suite: 16 dynamic-manipulation tasks × 4 setups (base / opt1 / opt2 / opt1+2) = 3,200 demonstrations, dual-arm UR5 (two 6-DoF arms + grippers), 3 cameras.

• Base model: pi05_base (openpi), full finetune (all parameters).
• Checkpoint: step 22,000 (EMA params/). Training was still in progress when exported.
• Config name (openpi): pi05_robodyna_full.
• Action loss at export: ~0.006–0.009 (windowed mean; still descending).

⚠️ This is a training checkpoint, not a converged final model. It reproduces the training-data action distribution well but has not been validated on rollouts. Evaluate on-sim before trusting it.

---

Contents

params/                              # EMA model weights (orbax checkpoint) — this is the model
assets/robodyna_merged/norm_stats.json   # state/action normalization stats — REQUIRED for inference

train_state/ (optimizer state, resume-only) is intentionally not included.

---

The robot & I/O contract

Robot	dual-arm UR5 — two 6-DoF arms, each with a gripper
State	14-dim joint vector: [L_joint0..5, L_gripper, R_joint0..5, R_gripper]
Action	14-dim, same layout, action horizon = 50 (model predicts a 50-step chunk)
Cameras	head (overhead/base), left_wrist, right_wrist — RGB, 240×320 native
Control rate	16.67 Hz (dataset fps)

Robot vs. data format. The physical robot is a dual-arm UR5. The demonstrations are stored in the LeRobot aloha 14-dim layout and the model was trained through openpi's Aloha data pipeline (AlohaInputs/AlohaOutputs, cameras mapped to cam_high/cam_left_wrist/cam_right_wrist). So you will see "aloha" naming in the openpi config and transforms below — that is the data/adapter convention, not the arm hardware.

Action semantics (important): during training the 12 arm joints are delta actions (target − current joint) and the 2 grippers (indices 6 and 13) are absolute. The openpi policy's output transform inverts this for you — see below — so what you receive back is absolute joint targets ready to command. You do not delta-decode anything yourself.

---

Expected INPUT to the policy

Pass policy.infer(obs) a flat dict with these exact keys (this is the LeRobot column layout the model trained on):

obs = {
    # IMAGES ARE CHANNEL-FIRST (C,H,W) uint8 — NOT (H,W,C). Verified against openpi make_aloha_example.
    "observation.images.head":        np.uint8,  # (3, H, W), RGB, CHW
    "observation.images.left_wrist":  np.uint8,  # (3, H, W), RGB, CHW
    "observation.images.right_wrist": np.uint8,  # (3, H, W), RGB, CHW
    "observation.state":              np.float32,# (14,)  current joint positions (dataset convention)
    "prompt":                         str,       # the task instruction (see task list below)
}

• CHW, not HWC. Images must be (3, H, W). If your sim gives (H, W, 3), transpose: img.transpose(2, 0, 1). Feeding HWC either crashes the resize or silently transposes your image into noise — this was the #1 eval bug (a mislabeled earlier version of this card said HWC).
• RGB, not BGR. Trained on RGB. OpenCV BGR → convert first (img[::-1] on the channel axis).
• Any HxW is fine; resized to 224×224 internally (native 240×320). uint8 [0,255].
• Normalization, resize, tokenization, state/action padding to 32-dim, and delta-encoding are all applied inside the policy from norm_stats.json. Do not pre-normalize.
• Missing a camera? openpi substitutes a black image + masks it, but all 3 were present in training — provide all 3 for in-distribution behavior.

Expected OUTPUT from the policy

result = policy.infer(obs)
actions = result["actions"]   # np.float32, shape (50, 14)

• Shape (action_horizon=50, 14). Already un-normalized and converted delta → absolute, so each row is a directly-commandable 14-dim absolute joint target.
• Layout per row: [L_j0, L_j1, L_j2, L_j3, L_j4, L_j5, L_gripper, R_j0, R_j1, R_j2, R_j3, R_j4, R_j5, R_gripper].
• Grippers (indices 6, 13) are absolute positions in the dataset's gripper units.

How to consume the 50-step chunk in a sim loop (open-loop chunk with re-planning — the standard π0.5 pattern):

CHUNK_EXECUTE = 25          # execute the first K of the 50 predicted steps, then re-infer
while not done:
    obs = get_observation()                 # RGB images + 14-dim joint state + prompt
    actions = policy.infer(obs)["actions"]   # (50, 14) absolute joint targets
    for a in actions[:CHUNK_EXECUTE]:
        step_sim(a)                          # command a as absolute joint positions
        if done: break

Executing all 50 open-loop also works but re-planning every ~25 steps tracks dynamic objects better (these tasks involve moving balls/belts/targets). Tune CHUNK_EXECUTE to your control rate.

---

Minimal eval script

import numpy as np
from openpi.training import config as _config
from openpi.policies import policy_config
import openpi.transforms as T

cfg = _config.get_config("pi05_robodyna_full")

# REQUIRED: create_trained_policy defaults to an EMPTY repack, so you MUST pass one, and at
# inference it must NOT include the "actions" key (there are no future actions to look up).
infer_repack = T.Group(inputs=[T.RepackTransform({
    "images": {"cam_high": "observation.images.head",
               "cam_left_wrist": "observation.images.left_wrist",
               "cam_right_wrist": "observation.images.right_wrist"},
    "state":  "observation.state",
    "prompt": "prompt",
})])

policy = policy_config.create_trained_policy(
    cfg,
    checkpoint_dir="/path/to/this/download",   # dir containing params/ and assets/
    repack_transforms=infer_repack,            # <-- do not omit this
)

obs = {
    "observation.images.head":        head_rgb_chw_uint8,    # (3, H, W)  CHW, RGB
    "observation.images.left_wrist":  left_rgb_chw_uint8,    # (3, H, W)
    "observation.images.right_wrist": right_rgb_chw_uint8,   # (3, H, W)
    "observation.state":              joint_state_14.astype(np.float32),  # (14,)
    "prompt": "hold the cup, wait for a gap in the swaying curtain ...",
}
actions = policy.infer(obs)["actions"]   # (50, 14) absolute joint targets

create_trained_policy loads params/ and reads assets/robodyna_merged/norm_stats.json automatically. Verified: on real training frames in this exact format the policy reproduces the dataset actions with MAE ~0.01 and correlation ~0.9998 — if your eval looks lost, the mismatch is in how the obs is built (CHW vs HWC, missing repack, wrong/empty prompt, or BGR), not the weights.

Config note: pi05_robodyna_full maps head→cam_high, left_wrist→cam_left_wrist, right_wrist→cam_right_wrist, uses use_delta_joint_actions=True, prompt_from_task=True, and quantile normalization. If evaluating with stock openpi, register a TrainConfig with these fields (pi05=True, action_horizon=50) or reuse the one from the training fork.

---

Tasks (use the matching prompt string)

id	prompt (abbreviated)
0	press the button matching the target marble color … trapdoor
1	catch the red ball rolling off the ramp with the cup; ignore blue distractor
2	close the gripper around the rat/mouse when it pops up (catch_two_mice → both arms)
3	slide the bowl along the belt to catch the marble off the shelves
4	place a cup beyond the red line to catch the red ball off the valley ramp
5	pass the cup through a gap in the swaying curtain into the moving slot
6	place the square blocker in the green zone to keep the ball out of goal
7	attach the dart inside the yellow center of the moving target board
8	drop the ball into a wagon of the circling toy train (opt1: red wagon only)
9	press left/right to tilt the shelf so the marble lands in the bowl
10	set the tall block onto the moving conveyor without tipping
11	strike the red ball into the target pocket with the cue (no robot-ball contact)
12	press the matching red/green key per tile; skip black distractors
13	drop shaped blocks into matching holes on the rotating sorter cap
14	hold the ping-pong bat so the red ball hits the bat head
15	press each mole when it pops up; avoid rabbits (distractors)

Full prompt strings are in the dataset's meta/tasks.jsonl (Hoshipu/robodyna-lerobot-suite). Use the exact prompt string the model was trained on for best results.

---

Training details

• Robot: dual-arm UR5 (14-dim: [L_j0..5, L_gripper, R_j0..5, R_gripper]), aloha-format data pipeline.
• 4×H200, global batch 128 (32/GPU), 40k steps planned (this export = step 22k).
• Cosine LR: warmup 1k → peak 2.5e-5 → 2.5e-6, full finetune off pi05_base.
• Delta joint actions (mask: 12 arm joints delta, 2 grippers absolute), quantile norm.
• Dataset merged from 64 LeRobot v2.1 datasets into one (3,200 eps, 801,045 frames, 16 task prompts).