twanghcmut/pi05-SO101-Multitask

π₀.₅ SO101 Multitask (openpi, PyTorch)

π₀.₅ finetuned on hungho77/so101-multitask using the official openpi PyTorch trainer — not a third-party port. One checkpoint covers all three tasks; the language prompt selects the behaviour.

Tasks

Use the exact wording; the model has not been trained on paraphrases.

#	Instruction	Episodes
0	Pick up the banana and place it in the bot, then close the lid	50
1	Pick blue cube and place on red cube	49
2	Pick all cubes and place into cup	44

Robot

SO101 follower arm, 5 joints plus gripper, 30 fps.

State / action dim	6 — shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll + gripper
Cameras	overhead → base_0_rgb, wrist → left_wrist_0_rgb; right_wrist_0_rgb padded and masked
Action horizon	50
Action representation	5 arm joints as deltas from the chunk's first state; gripper absolute

The policy's output transform converts the deltas back, so infer() returns absolute joint targets.

Files

file	purpose
model.safetensors	trained weights, 812 tensors (7.0 GB)
metadata.pt	openpi checkpoint metadata
assets/hungho77/so101-multitask/norm_stats.json	normalization statistics — required
so101_policy.py	input/output transforms, drop into src/openpi/policies/
train_config_snippet.py	the LeRobotSO101DataConfig + TrainConfig to add to src/openpi/training/config.py

optimizer.pt (13 GB) is deliberately not included: it only exists to resume training.

Usage

This is an openpi checkpoint, so it loads through openpi's config system rather than from_pretrained. Add so101_policy.py and the config snippet to your openpi checkout, then:

# serve
uv run scripts/serve_policy.py policy:checkpoint \
  --policy.config=pi05_so101 --policy.dir=/path/to/this/checkpoint

# or in-process
import openpi.policies.policy_config as policy_config
import openpi.training.config as config

policy = policy_config.create_trained_policy(
    config.get_config("pi05_so101"), "/path/to/this/checkpoint"
)
action = policy.infer({
    "observation/image": top_rgb,          # (H, W, 3) uint8
    "observation/wrist_image": wrist_rgb,  # (H, W, 3) uint8
    "observation/state": joints,           # (6,) float32
    "prompt": "Pick blue cube and place on red cube",
})["actions"]                              # (50, 6) absolute joint targets

Training

Hyperparameters follow openpi's own pi05_aloha_pen_uncap example — its defaults for learning-rate schedule, optimizer and action horizon were kept rather than tuned.

Base	gs://openpi-assets/checkpoints/pi05_base, converted to PyTorch bf16
Steps	6,000
Batch size	128 (768k samples = 11.4 epochs over 67,496 frames)
LR schedule	openpi default cosine: warmup 1,000 → peak 2.5e-5, decay_steps=30_000
Optimizer	openpi default AdamW (b1 0.9, b2 0.95, wd 1e-10, clip 1.0)
Precision	bfloat16
EMA	off — openpi's default is 0.99 but its PyTorch trainer does not support EMA
Hardware	1× H100 80GB, peak 71 GB, 100% utilization
Runtime	13.8 h at 8.28 s/step
Final loss	0.00302 (mean over steps 5,751–6,000)

The LR barely decayed. decay_steps=30_000 is openpi's default but this run is only 6,000 steps, so the learning rate ended at 2.35e-5 — 94% of peak. The loss was still falling when training stopped, so this checkpoint is not a converged minimum. Set decay_steps equal to num_train_steps if you want a clean decay, or train the full 30,000 steps.

Open-loop evaluation

Two trajectories per task, walked in chunk-length strides, predictions compared against the ground-truth action chunk. Errors are in the dataset's own action units.

traj	task	MAE (h=16)	MAE (h=50)
0	banana	1.126	1.249
25	banana	1.074	1.374
60	blue/red cube	0.941	1.474
80	blue/red cube	1.048	1.641
110	cubes→cup	1.048	1.868
130	cubes→cup	0.795	1.641
	average	1.005	1.541

Dividing by the absolute action std (mean 22.33 across dims) gives 4.50% of action scale at horizon 16, and 6.90% at the full horizon 50. Error grows along a chunk, so a horizon-50 number must never be compared against a horizon-16 one from another model.

Note the checkpoint's own norm_stats.json holds delta std for the arm joints (mean 9.71). That is a different quantity from the absolute-unit errors above; dividing one by the other gives a meaningless figure.

For reference, GR00T N1.7 trained on the same data with the same batch size and step count scores 5.73% at horizon 16, so π₀.₅ is ~21% lower error on 5 of 6 trajectories. Both numbers carry the same caveat below.

Limitations

No held-out split. All 143 episodes were used for training, so the numbers above are measured on training data. They show the policy fits its data and is not degenerate; they do not measure generalization and they are not a success rate. Treat real-robot performance as unmeasured.

One episode (index 49) is a 5-frame aborted recording, 0.007% of frames, left in.

wrist_roll barely moves in this dataset (absolute std 0.95 against 6–36 for the other joints), so expect little controllability on that joint.

Deployment note

The data is 30 fps. A client interpolating between model steps must derive its sub-step count from that: smooth_step = control_hz / 30. Reusing a value from a 15 fps robot stretches every trajectory by 2× and the arm creeps without finishing the task.