RoboDyna/pi05-baseline-bs256

pi05-robodyna-baseline-bs256

π0.5 fine-tune on the RoboDyna benchmark (dual UR5 + WSG grippers, 35 tasks).

Checkpoint

Field	Value
Base	openpi-assets/checkpoints/pi05_base/params (π0.5)
Training config	pi05_robodyna (openpi TrainConfig, see training_config/train_config.py)
Exp name	baseline_bs256_30k
Ckpt uploaded	ckpt-25000 (step 25000 of a planned 30000)
Loss at step 25000	~0.0046 (grad_norm ~0.05, param_norm ~1806)
Why not 30k	Slurm 12h time limit hit at step ~25.8k; last saved ckpt is 25000.

Training setup (as-run)

Sbatch launched training with these overrides on top of the openpi TrainConfig:

python scripts/train.py pi05_robodyna \
    --exp-name baseline_bs256_30k \
    --overwrite \
    --fsdp-devices 8 \
    --batch-size 256 \
    --num-train-steps 30000 \
    --log-interval 100 \
    --save-interval 5000 \
    --keep-period 10000

Config-file defaults (see training_config/train_config.py):

• Model: Pi0Config(pi05=True, action_horizon=50)
• Data: LeRobotAlohaDataConfig(repo_id="robodyna", adapt_to_pi=False, use_delta_joint_actions=True)
• Repack: images = head → cam_high, left_wrist → cam_left_wrist, right_wrist → cam_right_wrist
• LR: CosineDecay, warmup=1000, peak=2.5e-5, decay=2.5e-6 over 30000 steps
• Optimizer: AdamW, clip_gradient_norm=1.0, ema_decay=0.999
• BS 256 at launch (overrides the file's default BS 32) on 8× H200, FSDP

Hardware / env: 1× gpu-h200-106, 8× H200, XLA mem 0.95, NCCL_NVLS_ENABLE=0.

Dataset (RoboDyna)

Field	Value
Robot	dual UR5 + WSG grippers (14-D state/action)
Total episodes	4050
Total frames	1,374,883
Total tasks	35
FPS	~16.67
Cameras (3)	head, left_wrist, right_wrist
LeRobot version	v2.1
Local path	/work/markhsp/datasets/robodyna (7.5 GB)
Build script	training_config/build_robodyna_lerobot.py
norm_stats	Baked into ckpt-25000/assets/robodyna/norm_stats.json

Layout

ckpt-25000/
├── _CHECKPOINT_METADATA
├── assets/
│   └── robodyna/
│       └── norm_stats.json      # normalization stats (baked in)
└── params/                       # π0.5 weights (Orbax checkpoint)
training_config/
├── train_config.py              # excerpt of openpi TrainConfig for pi05_robodyna
├── train_30k.sbatch             # slurm launch script (as-run)
├── pi05_robodyna_config.patch   # adds the TrainConfig to openpi
├── pi05_robodyna_workers16.patch  # num_workers=16 adjustment
├── build_robodyna_lerobot.py    # dataset builder (source→LeRobot v2.1)
└── compute_norm_stats_fast.py   # norm_stats generator

train_state/ (~31 GB Orbax optimizer/rng state) is intentionally NOT included. To resume, re-train from the base and load params/.