jax_30000/model_params.md

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