A1astor/pi05_ur5_gripper

π0.5 LoRA fine-tune — UR5 + gripper "catch the cube" (openpi JAX, A1astor)

Official openpi (JAX) LoRA fine-tune of pi05_base on 44 real UR5 demonstrations collected with gamepad teleoperation. Trained on a single A100-40GB (multi-GPU NCCL corruption workaround, openpi issue #822).

Robot & task

• Robot: UR5 + gripper, 2 cameras: observation.images.global (external) + observation.images.wrist
• Task: catch the cube
• Action space: 7-dim = 6 absolute joint positions (rad) + gripper (0-1), 30 Hz (trained in delta space via DeltaActions, gripper absolute; inference outputs absolute joint targets)
• Dataset: A1astor/ur5_gripper (private) — 44 episodes / 26,637 frames, LeRobot v3.0

Training recipe (config pi05_ur5_gripper, exp gripper_v4_singlegpu)

Item	Value
Base	pi05_base (gs://openpi-assets), bfloat16 frozen backbone
PEFT	LoRA (gemma_2b_lora + gemma_300m_lora), EMA off
Batch / steps	32 × 20000 steps (~19.5h on 1× A100-40GB)
Optimizer	AdamW, clip_gradient_norm=1.0, cosine 2.5e-5 → 2.5e-6, warmup 1000
Final train loss	0.0097 (grad_norm 0.093)
Data hygiene	action clamp ±6 (episode-boundary chunks), wrist_2 dead-dim stats widened

Repo layout: params/ (orbax checkpoint, inference weights), assets/ (norm stats). Optimizer state (train_state/) kept locally, not needed for inference.

Usage (openpi local inference)

from openpi.training.config import get_config
from openpi.policies import policy_config

config = get_config("pi05_ur5_gripper")
policy = policy_config.create_trained_policy(config, "<path-to-this-dir>")
obs = {
    "base_rgb": global_cam.read(),   # uint8 HWC
    "wrist_rgb": wrist_cam.read(),
    "joints": rtde_r.getActualQ(),   # 7-dim: 6 joints + gripper
    "prompt": "catch the cube",
}
result = policy.infer(obs)           # -> (action_horizon, 7) absolute targets

Also servable via scripts/serve_policy.py policy:checkpoint --policy.config=pi05_ur5_gripper.

Known fixes carried in the training code

• checkpoints.py: host-side snapshot before async save (race wrote NaN into float32 trainable params in earlier checkpoints).
• Single-GPU training only on this host: 4-GPU FSDP with NCCL_P2P_DISABLE=1 silently corrupted allreduces (NaN within 100-600 steps, batch-dependent).