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

0ffeea4 verified about 1 month ago

8.82 kB

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

Fine-tuned π₀.₅ (pi05) VLA policy for the Aloha-Agilex bimanual robot, trained with 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/ — 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

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):

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.