# RoboPRO π₀.₅ (JAX) — step 30000 checkpoint Fine-tuned **π₀.₅ (pi05)** VLA policy for the **Aloha-Agilex** bimanual robot, trained with [openpi](https://github.com/Physical-Intelligence/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 ```python 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`. ```python 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.