Add model_params.md (inputs/outputs/config) for pytorch_30000

pytorch_30000/model_params.md

@@ -0,0 +1,171 @@
+# 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.