RoboPRO π₀.₅ (JAX) — step 30000 checkpoint

Fine-tuned π₀.₅ (pi05) VLA policy for the Aloha-Agilex bimanual robot, trained with openpi (JAX/Flax) on the RoboPRO top-cam dataset (roboreal_lerobot). This repo holds the eval weights only (no optimizer state).

Base model: pi05_base (Physical Intelligence), ~3.6B params
Framework: JAX / Flax, orbax checkpoint (this is not a PyTorch/safetensors checkpoint)
Precision: bfloat16
Training: 30,000 steps, global batch 192, cosine LR (peak 2.5e-5), ~1.5 epochs over 3.74M frames @ 25 Hz
Final train loss: ~0.0021 (flow-matching)

Repo contents

params/                         # orbax model weights (load these)
assets/roboreal_lerobot/
    norm_stats.json             # input/output normalization stats (REQUIRED)
_CHECKPOINT_METADATA

⚠️ train_state/ (optimizer) is not included — this checkpoint is for inference/eval only, not for resuming training.

Inputs

The policy consumes a single-timestep observation dict with 3 camera images + a 14-D robot state + a language prompt.

1. Cameras (3× RGB)

policy key	physical view	shape	dtype
`cam_high`	overhead / countertop camera (looking down at the table)	`[3, H, W]` (CHW)	`uint8`, 0–255
`cam_left_wrist`	left-arm wrist camera	`[3, H, W]`	`uint8`, 0–255
`cam_right_wrist`	right-arm wrist camera	`[3, H, W]`	`uint8`, 0–255

RGB, channel-first [3, H, W]. Images are internally resized to 224×224, so any input resolution works (training used 240×320).
Camera mapping is critical: feed your countertop/overhead view as cam_high (the model was trained with the top-cam view in that slot, not a robot-head camera). Wrist cams map by side.
All three cameras are required.

2. State — `state`

float32[14], raw joint positions (radians) + gripper, absolute, in Aloha convention.

Order (same for state and action):

0 left_waist    1 left_shoulder  2 left_elbow      3 left_forearm_roll
4 left_wrist_angle  5 left_wrist_rotate  6 left_gripper
7 right_waist   8 right_shoulder 9 right_elbow    10 right_forearm_roll
11 right_wrist_angle 12 right_wrist_rotate 13 right_gripper

Feed raw physical values — normalization (quantile, from norm_stats.json) and the Aloha→pi convention conversion happen inside the policy.

3. Prompt — `prompt`

Natural-language task instruction, e.g. "put the mouse on the pad". Trained on 1,622 instruction variants across 80 tasks.

Observation dict shape

observation = {
    "state": np.ndarray,          # float32 [14]
    "images": {
        "cam_high":       np.ndarray,  # uint8 [3, H, W]  (countertop)
        "cam_left_wrist": np.ndarray,  # uint8 [3, H, W]
        "cam_right_wrist":np.ndarray,  # uint8 [3, H, W]
    },
    "prompt": str,
}

Output

policy.infer(observation)["actions"] returns an action chunk:

Shape [50, 14] — 50 future timesteps (action_horizon=50), 14-D per step.
Absolute joint-position targets in Aloha convention, same 14-D order as state.
De-normalized to physical units (you feed raw, you get raw).
At 25 Hz, the 50-step chunk ≈ 2 s of motion. Typical control: execute the first k actions (e.g. pi0_step steps), then re-infer with the new observation.

Why the output is absolute (delta vs. absolute)

This config trains with use_delta_joint_actions = True, which installs a paired transform around the model:

Training input — DeltaActions(mask): actions[:, :dims] -= where(mask, state, 0) → masked dims become (target − current_state) = deltas.
Inference output — AbsoluteActions(mask): actions[:, :dims] += where(mask, state, 0) → masked dims become (delta + current_state) = absolute.

The mask is make_bool_mask(6, -1, 6, -1) = [True×6, False, True×6, False]:

dims	joints	mask	model learns	returned
0–5, 7–12	6 arm joints per arm	`True`	delta	absolute (state re-added on output)
6, 13	grippers	`False`	absolute	absolute

So the network internally predicts arm-joint deltas, but AbsoluteActions runs on the output and adds back the observation's state, so the policy returns absolute joint-position targets. Grippers are absolute throughout.

Practical implications for eval:

Send the returned actions directly as target joint positions — do not add the current state yourself; the output transform already did.
AbsoluteActions broadcasts the single observation state across all 50 timesteps, so every action in the chunk is absolute relative to the state you passed at that inference call (standard openpi behavior).
The state you feed therefore affects the arm outputs (it's the base the deltas are added to); feed the robot's true current joint positions.

How to run inference (openpi, JAX)

Requires an openpi env with JAX (this project's pi05 conda env) and the pi05_robopro_top_cam_jax train config (defines the repack + Aloha transforms + norm stats binding). The exact config is included in this repo as train_config.py — paste its TrainConfig(...) entry into the _CONFIGS list in your openpi src/openpi/training/config.py.

from openpi.policies import policy_config as _policy_config
from openpi.training import config as _config

train_config = _config.get_config("pi05_robopro_top_cam_jax")

# checkpoint_dir must contain params/ and assets/ (this repo's root after download)
policy = _policy_config.create_trained_policy(
    train_config,
    "/path/to/robopro_jax_30000",     # dir with params/ + assets/
    robotwin_repo_id="roboreal_lerobot",  # picks assets/roboreal_lerobot/norm_stats.json
)

# Build the observation (feed COUNTERTOP cam as cam_high; images CHW uint8)
obs = {
    "state": state_14,                       # float32[14], absolute joints
    "images": {
        "cam_high":        countertop_chw,   # uint8[3,H,W]
        "cam_left_wrist":  left_chw,
        "cam_right_wrist": right_chw,
    },
    "prompt": "put the mouse on the pad",
}

actions = policy.infer(obs)["actions"]   # np.ndarray [50, 14], absolute joint targets
# execute actions[:k] on the robot, then re-infer

Notes:

Loading is auto-detected as JAX because the checkpoint has params/ (not model.safetensors).
If your runtime provides differently-named observation keys, apply a repack so images land under cam_high / cam_left_wrist / cam_right_wrist, state under state, and set prompt.
norm_stats.json must be present/loaded; without it actions are unnormalized and wrong.