angkul07/pi0-fast-yam

pi0-FAST Fine-Tuning — YAM Bimanual

Fine-tuning of pi0-FAST (PaliGemma-3B VLA + FAST action tokenizer) on a bimanual YAM robot teleoperation dataset, trained with the openpi framework (JAX/Flax).

⚠️ Format note: These are openpi checkpoints (Orbax/JAX) — not HuggingFace PEFT adapters. Load them with openpi (create_trained_policy), not PeftModel.from_pretrained. See Load Example.

Training Setup

Base model	pi0-FAST (PaliGemma 3B + FAST action tokenizer)
Framework	openpi (JAX / Flax nnx)
Config	pi0_fast_yam_low_mem_finetune
Variant	gemma_2b_lora (LoRA fine-tune)
Robot	YAM bimanual (14-dim state/action = 6 joints + 1 gripper per arm)
Dataset	Kavin60606/yam_pi0fast_train
Cameras	3 (top + left_wrist + right_wrist)
Hardware	2× A100-80GB SXM4, FSDP
Precision	bfloat16

Hyperparameters

Param	Value
Action dim	14
Action horizon	50
Max token len	300
PEFT method	LoRA (gemma_2b_lora)
Optimizer	Adam
Scheduler	Cosine decay with warmup
Warmup steps	1,000
Peak LR	3.5e-5
Min LR (cosine end)	3.5e-6
Decay steps	7,000
Batch size	64
Total steps	7,000
EMA	disabled
FSDP devices	1

Action representation

Actions use delta joint encoding — arm joints are trained as deltas from the current state, grippers as absolute values:

delta mask = make_bool_mask(6, -1, 6, -1)
# [6 arm joints → delta] [1 gripper → absolute]  × 2 arms  = 14 dims
# gripper indices: 6 (left) and 13 (right)

Normalization statistics (norm_stats.json) are bundled in each checkpoint's assets/ directory.

Checkpoints

Each step subdir is a full openpi checkpoint containing params/, train_state/ (optimizer moments + step, for resuming), assets/ (norm stats), and _CHECKPOINT_METADATA.

Step	Notes
2000/	intermediate (first run)
2500/	intermediate (first run)
5000/	intermediate (second run)
6999/	final model — full 7,000-step training

The model was trained in two phases: an initial run through step 2,500, then resumed from that checkpoint and trained to completion at step 6,999. Older intermediate checkpoints were pruned (max_to_keep=1 + keep_period=5000).

Load Example (openpi)

from openpi.training import config as _config
from openpi.policies import policy_config

# 1. Load the training config used for this model
config = _config.get_config("pi0_fast_yam_low_mem_finetune")

# 2. Point at a downloaded checkpoint step directory (e.g. the final 6999/)
checkpoint_dir = "path/to/pi0-fast-yam/6999"

# 3. Build the inference policy (loads params + norm stats from assets/)
policy = policy_config.create_trained_policy(config, checkpoint_dir)

# 4. Run inference
#    obs must contain the 3 camera views + observation.state + prompt
action_chunk = policy.infer(obs)["actions"]

Download a checkpoint with the HF CLI:

huggingface-cli download angkul07/pi0-fast-yam --include "6999/*" --local-dir ./pi0-fast-yam

Reproducibility

Camera / key mapping (dataset → policy), applied via openpi's RepackTransform:

{
  "observation.images.top": "base_0_rgb",
  "observation.images.left_wrist": "left_wrist_0_rgb",
  "observation.images.right_wrist": "right_wrist_0_rgb",
  "observation.state": "state",
  "action": "actions"
}

Full config is defined by pi0_fast_yam_low_mem_finetune in the openpi fork: angkul07/openpi → src/openpi/training/config.py.

To resume training from any checkpoint (uses the bundled train_state/):

uv run scripts/train.py pi0_fast_yam_low_mem_finetune \
  --exp-name yam_run --fsdp-devices 1 --resume