# RoboPRO π₀.₅ (PyTorch) — step 30000 checkpoint

Fine-tuned **π₀.₅ (pi05)** VLA policy for the **Aloha-Agilex** bimanual robot, trained with [openpi](https://github.com/Physical-Intelligence/openpi) **PyTorch** trainer (`scripts/train_pytorch.py`, DDP) on the RoboPRO **top-cam** dataset (`roboreal_lerobot`). This folder holds the **eval weights** (`model.safetensors`); optimizer state is **not** included.

- **Base model:** `pi05_base` (Physical Intelligence), PyTorch port, ~3.6B params
- **Framework:** PyTorch, **safetensors** (`PI0Pytorch` keys) — this is **not** a JAX/orbax checkpoint
- **Precision:** bfloat16 (compute **and** stored master weights — see caveat below)
- **Training:** 30,000 steps, global batch 256 (3×H200, 85/GPU), cosine LR (peak 2.5e-5), ~2 epochs over 3.74M frames @ 25 Hz
- **Final train loss:** ~0.0043 (flow-matching)
- **JAX counterpart:** see [`jax_30000/`](../jax_30000) — same data/recipe, reached ~0.0021 (see caveat)

> ⚠️ **Precision caveat.** This PyTorch run stores the model weights themselves in **bf16** (not just compute). In the low-LR tail, weight updates fall below bf16's mantissa resolution and stall, so the final loss plateaus ~2× higher than the JAX run (0.0043 vs 0.0021), which keeps **fp32 master weights**. The fix (fp32 master + bf16 autocast + EMA) is documented in the training repo at `docs/pytorch_fp32master_ema_fix.md`. Treat this checkpoint as the **bf16 baseline**.

---

## Folder contents

```
model.safetensors                       # PI0Pytorch weights (load these)
assets/roboreal_lerobot/norm_stats.json # input/output normalization stats (REQUIRED)
metadata.pt                             # {global_step, config dict, timestamp}
train_config.py                         # the exact TrainConfig used for this run
```
> `optimizer.pt` (13 GB) is **not** uploaded — this checkpoint is for **inference/eval**, not for exact-resume training.

---

## Model configuration

| field | value | notes |
|---|---|---|
| `model_type` | **pi05** (`PI0Pytorch`, `pi05=True`) | flow-matching action expert |
| `paligemma_variant` | `gemma_2b` | vision-language backbone |
| `action_expert_variant` | `gemma_300m` | action expert |
| `action_dim` | **32** | 14 real dims, padded to 32 |
| `action_horizon` | **50** | timesteps per action chunk |
| `max_token_len` | **200** | prompt token budget |
| `discrete_state_input` | `True` | pi05 discretizes state into the prompt |
| `dtype` | `bfloat16` | compute + master weights |

**Training hyperparameters** (full dump in `train_config.py`):

| field | value |
|---|---|
| dataset | `roboreal_lerobot` (robopro_top_cam) — 15,999 eps / 3.74M frames @ 25 Hz |
| steps | 30,000 (≈2 epochs) |
| global batch | 256 (3×H200, 85/GPU) |
| LR schedule | cosine, warmup 1,000, peak `2.5e-5`, floor `2.5e-6` |
| optimizer | AdamW `b1=0.9 b2=0.95 eps=1e-8 wd=1e-10`, grad-clip `1.0` |
| precision | bfloat16 |
| EMA | none (not supported by the PyTorch trainer) |
| base weights | `pi05_base` (PyTorch), loaded via `pytorch_weight_path` |

---

## Inputs

Single-timestep observation dict: **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` (trained with the top-cam view in that slot, *not* a robot-head camera). Wrist cams map by side. In the raw dataset these are the `countertop` / `left` / `right` video keys respectively.
- 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.

### 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, then re-infer with the new observation.

### Delta vs. absolute

This config trains with `use_delta_joint_actions = True`, installing a paired transform around the model:

- **Training input** — `DeltaActions(mask)`: masked dims become **(target − current_state)** = deltas.
- **Inference output** — `AbsoluteActions(mask)`: masked dims become **(delta + current_state)** = absolute.

Mask `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` 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 current state yourself; the output transform already did.
- `AbsoluteActions` broadcasts the single observation `state` across all 50 timesteps.
- Feed the robot's true current joint positions as `state` (it's the base the arm deltas are added to).

---

## How to run inference (openpi, PyTorch)

Requires an openpi env with **PyTorch** (this project's `pi05_pt` conda env) and the `pi05_robopro_top_cam_pt` train config. The exact config is in this folder as **`train_config.py`** — its entry already exists in the training repo's `src/openpi/training/config.py`.

`create_trained_policy` **auto-detects PyTorch** by the presence of `model.safetensors` in the checkpoint dir (no code change vs. the JAX call):

```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_pt")

# checkpoint_dir must contain model.safetensors + assets/ (this folder after download)
policy = _policy_config.create_trained_policy(
    train_config,
    "/path/to/robopro_jax_30000/pytorch_30000",   # dir with model.safetensors + assets/
    robotwin_repo_id="roboreal_lerobot",           # picks assets/roboreal_lerobot/norm_stats.json
    pytorch_device="cuda",                         # or "cpu"
)

# 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 PyTorch because the checkpoint dir has `model.safetensors` (not `params/`).
- `norm_stats.json` **must** be present/loaded; without it actions are unnormalized and wrong.
- 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`.
- Weights are bf16; `create_trained_policy` casts selected params for inference automatically.