| # 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. |
|
|
