pi05-abcego-sd — π₀.₅ finetuned on 6.76 h of real YAM teleop, single task

Training checkpoints for the 100 %-teleop arm of the ABC-ego study: π₀.₅ (flow matching) finetuned from pi05_base on real bimanual YAM teleoperation of a single task — "put the screwdriver in the bin" — for a 14-DoF bimanual YAM robot.

This repo contains raw Orbax training checkpoints (params + train_state + norm stats), not a converted HF-format model. They are meant to be loaded by openpi.

⚠️ Read this first

Checkpoints here	5000, 10000, 11000, 11413 (final)
Evaluation	None. No held-out split, no success rate. Every number below is training loss.

The run trained the full 11,414 steps and completed cleanly (exit 0). 11413 is the final checkpoint — use it unless you specifically want an earlier one.

Two things to keep in mind:

• There is no evaluation, so "final" here means last, not best. Nothing in this repo establishes that 11413 outperforms 5000 on the actual task.
• This is exactly one epoch over the dataset. Every frame was seen once. Loss was still falling steeply at the end (see Results) — this model is almost certainly under-trained, not converged.

Checkpoint 9000 was also retained on disk by the rolling window but is not published here. 6000/7000/8000 were deleted during the run by max_to_keep=4.

Experiment

Field	Value
Config	pi05_abcego_sd (openpi TrainConfig)
Model	Pi0Config(pi05=True) — PaliGemma backbone + flow-matching action expert
PaliGemma variant	gemma_2b_lora
Init	gs://openpi-assets/checkpoints/pi05_base/params (12.5 GiB)
Action dim	32 (14 real, 18 zero-padded)
Action horizon	50
Max token len	200
Hardware	2× H100 80GB HBM3 (SXM), data-parallel, --fsdp-devices 1 (full model replica per GPU)
Wall clock	4 h 36 m 08 s for 11,414 steps (≈1.452 s/step)
Approx. compute cost	~$19 at $4.09/hr (rate not recorded on the box — adjust if yours differed)
Code revision	fc939b753861b87a56cbeba821cc460009a7a56f (angkul07/openpi fork)
W&B	project pi05-yam, run sd (7nf4dedz)

Hyperparameters

Field	Value
Batch size	64 (single source — all 64 from teleop)
Steps	11,414 = exactly one epoch (730,496 frames / 64)
Optimizer	AdamW, b1=0.9, b2=0.95, eps=1e-8, weight_decay=1e-10, clip_gradient_norm=1.0
LR schedule	Cosine, warmup 500 → peak 3.5e-5 → decay 3.5e-6 over 11,414 steps
EMA	off (ema_decay=None)
Dataloader workers	16
Checkpoint interval	every 1,000 steps, max_to_keep=4, keep_period=5000
Image augmentation	off — both image_augmentation=False and augment_config=None
Precision	fp32 params (param_norm tracked in fp32)

Two notes on the schedule:

• Warmup is 4.4 % of this run (500 of 11,414), against 2.1 % of the 23.6k-step pi05_yam7h_* arms. It was left at 500 for comparability, not because it was tuned here.
• The LR schedule is inherited, not tuned. openpi's own π₀.₅ recipes (pi05_libero, pi05_full_droid_finetune) use a constant 5e-5 with no cosine. The cosine schedule was held fixed across this whole study so the moving variable is the data, not the optimizer.

Data

A single source — no mixture, no oversampling. samples_per_batch=64 equals batch_size, so every batch is 64 teleop frames drawn from a reshuffled permutation (ordinary shuffled training).

Source	Episodes	Frames	Duration	Per batch	Gradient share	Epochs
angkul07/abc-ego-screwdriver	2,234	730,496	6.76 h @ 30 fps	64	100 %	1.00

Real YAM bimanual teleoperation, converted from MCAP by vast_run/mcap_to_lerobot.py. LeRobot v2.1, robot_type: yam, read as-is (the launcher asserts the version and refuses to migrate in place).

Observation space

• 3 × RGB, stored at native 480×848, resized to 224×224 by ResizeImages: observation.images.top, observation.images.left_wrist, observation.images.right_wrist
• 14-D joint state (7 per arm: 6 joints + gripper), fed as a discrete state token (TokenizePrompt(discrete_state_input=True), the π₀.₅ convention)
• Language prompt from the LeRobot task string (prompt_from_task=True)

Action space

• 14-D, horizon 50, zero-padded to 32 by PadStatesAndActions(32)
• Layout [L j0-5, L grip, R j0-5, R grip]
• DeltaActions(make_bool_mask(6, -1, 6, -1)) — the 12 arm-joint dims are predicted relative to current state; the 2 gripper dims stay absolute. AbsoluteActions inverts this at inference.

Normalization — quantile (use_quantile_norm=True, q01/q99), asset abcego_sd, shipped in every checkpoint at <step>/assets/abcego_sd/norm_stats.json.

Why norm stats were recomputed rather than reused

The yam7h_* stats from the earlier arms were not copied across, even though the code path is identical. Different robot campaign, different joint distribution — rig A parks the left arm entirely while rigs B/C do not — so the q01/q99 quantiles move. Reusing them would have silently normalized against the wrong distribution. Computed fresh over 200k of the 730k frames with --skip-videos (norm stats never touch pixels), verified free of degenerate dims (min std 0.2167 state / 0.1333 actions).

What was deliberately not done

• No held-out split. exclude_episodes=() and holdout_fraction=0.0 — trains on 100 % of the data, by request. This is the reason there is no offline eval, and it was a deliberate trade, not an oversight.
• No augmentation. Both the data-side ImageAugmentConfig stack (ColorJitter + crop) and the model-side image_augmentation flag were switched off. The pi05_yam7h_* arms had them on; this run is the un-augmented reference.
• No gripper rescale. The rescale that the 7 h mixture needed exists to reconcile two sources that disagreed on the absolute value of "open". With one source there is nothing to reconcile, so the data is untouched.

Results

Training loss is the flow-matching objective, mean over dims of (v_θ(x_t, t) − u_t)² where u_t = noise − actions.

Step window	Loss (mean)	grad_norm (mean)	param_norm
0 – 100	0.04399	0.20910	1803.23
100 – 500	0.01930	0.07064	1803.25
500 – 1,000	0.01513	0.05425	1803.33
1,000 – 2,000	0.01326	0.04382	1803.48
2,000 – 4,000	0.01137	0.03966	1803.74
4,000 – 6,000	0.00949	0.03834	1804.00
6,000 – 8,000	0.00790	0.03768	1804.14
8,000 – 10,000	0.00670	0.03740	1804.19
10,000 – 11,414	0.00606	0.03777	1804.20

Per published checkpoint (200-step trailing mean):

Checkpoint	Loss	grad_norm
5000	0.00956	0.03740
10000	0.00636	0.03736
11000	0.00590	0.03795
11413 (final)	0.00581	0.03760

Other figures: first-step loss 0.08715; minimum single-step loss 0.00300 at step 9,786; last-500-step loss 0.00599 ± 0.00122; maximum gradient norm over the whole run 0.7204; total param_norm drift +0.97 (+0.054 %).

Horizon split at the end of training (last 200 steps):

	Loss
flow_loss_chunk_first (first action in the chunk)	0.00301
flow_loss_chunk_last (50th action in the chunk)	0.00940

Full per-step metrics for all 11,414 steps are in train_metrics.log.

Learnings and takeaways

1. The absolute loss value is not interpretable, and its smallness is not evidence of convergence. Loss reaches 0.044 in the first 100 steps and 0.019 by step 500. That looks alarming until you account for the objective's structure: 18 of the 32 action dims are zero padding, where x_t = t·noise makes the target u_t = noise = x_t / t a closed-form function of the input — free to fit. The remaining 14 dims are single-frame delta actions at 30 fps (so consecutive-frame deltas are small by construction), quantile-normalized, and the timestep is drawn from Beta(1.5, 1)·0.999 + 0.001 (mean t ≈ 0.6, skewed toward the easy noise-dominated end). Compare π₀.₅ runs to each other, never to an absolute threshold and never to a cross-entropy π₀-FAST number.

2. One epoch was not enough, and the loss curve says so plainly. Loss fell monotonically in every window with no plateau — 0.04399 → 0.00606, a 7.3× reduction, still descending at the last checkpoint. num_train_steps was set to exactly one epoch as a principled starting point, not because one epoch was known to be sufficient. Nothing here suggests the run had stopped learning, and the cheapest next experiment is simply 2–3× the steps.

3. The far end of the action chunk is 3.1× harder than the near end. chunk_first ends at 0.00301 while chunk_last ends at 0.00940. This is expected — predicting 50 steps ahead is genuinely harder than predicting 1 — but it is worth stating because the aggregate loss hides it, and because it is the number to watch if you shorten or lengthen action_horizon. It also means the aggregate is dominated by the hard tail of the chunk, not by the part of the prediction the controller actually executes first.

4. LoRA finetuning barely moves the base model. param_norm drifted +0.054 % across the entire run — less than half the drift of the 23.6k-step 7 h arms (+0.12 %), consistent with this run being 2× shorter. With gemma_2b_lora and a 3.5e-5 peak LR this is a very light touch on pi05_base: most of the capability is inherited, not learned here. It also means overfitting is implausible after a single epoch — which is a reason to distrust the low loss as evidence of task competence, not a reason to trust it.

5. Optimization was completely uneventful. grad_norm settled to ~0.038 by step 1,000 and stayed flat for the remaining 10k steps — no spikes, no instability. The maximum gradient norm observed over the entire run was 0.7204, so clip_gradient_norm=1.0 never once bound. The warmup and schedule can be considered validated for this setup; if anything, the clip threshold is loose enough to be doing nothing at all.

6. Language is doing nothing in this run. The dataset has exactly one task string, so prompt carries no discriminative signal — every sample gets the same conditioning. This is expected for a single-task finetune, but it means: (a) nothing here tests instruction following, and (b) any comparison against a multi-task arm is confounded by this, not just by the data volume.

7. A single-source config still has to be expressed as a "mixture of one." Not an ML result, but it cost real time. create_torch_dataset() hardcodes root=None on the non-mixture path, so a plain single-source config requires the dataset to live at $HF_LEROBOT_HOME/<repo_id> and be symlinked into place; and the launcher asserts data.mixture is non-empty. Using MixtureSource with samples_per_batch == batch_size sidesteps both and changes nothing about sampling — StratifiedBatchSampler with one source draws all 64 indices from a reshuffled permutation, i.e. ordinary shuffled training.

8. Aggregate loss cannot answer what this run was set up to ask. A single scalar tells you the model fits its own training distribution. It cannot tell you whether the policy grasps a screwdriver. Per-dim splits (real vs. zero-padded) and the per-horizon split above are cheap and already partly wired; a real held-out split is the missing piece, and it is the single highest-value change for the next run.

Caveats and open questions

These are known and unresolved. They bound what the numbers above can support.

• No evaluation of any kind. holdout_fraction=0.0, exclude_episodes=() on the only source — no validation split, no success rate, no hardware rollout. Loss curves only. Every claim about this model's capability is currently unsupported.
• Exactly one epoch, still descending. See takeaway 2. Treat this as an early checkpoint of an unfinished training curve, not a finished model.
• Single task, single seed. One task string, one run, no ablation. There is nothing here to isolate the effect of any choice made above.
• No augmentation, so generalization is untested. The model saw each frame exactly once, unperturbed. Robustness to lighting, camera shift, or object pose is entirely unmeasured and there is reason to expect it is weak.
• Images are trained at 224×224, downsampled from 480×848 non-square source. The aspect ratio is not preserved by ResizeImages; the model sees horizontally squashed frames. This is consistent between train and inference so it is not a bug, but it does mean the effective spatial resolution is lower than the stored data suggests.
• Gripper dims stay absolute through DeltaActions. With one source that is safe. If this checkpoint is ever mixed or compared with another data source, that channel is exactly where cross-source scale disagreement survives normalization — audit it before mixing.

Repo layout

5000/                        # and 10000/, 11000/, 11413/
  _CHECKPOINT_METADATA
  params/                    # Orbax OCDBT — inference weights (6.68 GiB)
  train_state/               # optimizer state — for resuming only (5.81 GiB)
  assets/abcego_sd/norm_stats.json
train_metrics.log            # complete, all 11,414 steps

Each checkpoint is 12.5 GiB; the repo is **50 GiB** total. For inference you only need params/ and assets/ — train_state/ is ~47 % of the bytes and is needed only to resume training.

Usage

# fetch the final checkpoint (inference only — skipping train_state roughly halves the download)
hf download angkul07/pi05-abcego-sd \
  --include "11413/params/*" "11413/assets/*" "11413/_CHECKPOINT_METADATA" \
  --local-dir ./pi05-abcego-sd

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

cfg = _config.get_config("pi05_abcego_sd")
policy = policy_config.create_trained_policy(cfg, "./pi05-abcego-sd/11413")

action_chunk = policy.infer({
    "observation/top_image":         top_rgb,          # HxWx3 uint8
    "observation/left_wrist_image":  left_wrist_rgb,
    "observation/right_wrist_image": right_wrist_rgb,
    "observation/state":             state_14d,        # [L j0-5, L grip, R j0-5, R grip]
    "prompt":                        "put the screwdriver in the bin",
})["actions"]                                          # (50, 14) absolute joint targets

AbsoluteActions has already added the state back, so the returned 14-D actions are absolute joint targets — feed them to the controller directly.

To resume training, point openpi at the parent directory:

uv run scripts/train.py pi05_abcego_sd --exp-name sd --fsdp-devices 1 --resume

Training recipe

Full reproduction from a bare 2× H100 80GB box.

# ---- 0. environment ----
git clone https://github.com/angkul07/openpi.git /workspace/openpi
cd /workspace/openpi
git checkout fc939b753861b87a56cbeba821cc460009a7a56f
GIT_LFS_SKIP_SMUDGE=1 uv sync

cat > vast_run/env.local <<'EOF'
export WANDB_API_KEY=<your key>
export ABCEGO_SD_ROOT=/workspace/abc-ego-lerobot
EOF
chmod 600 vast_run/env.local

# ---- 1. GPU health check (do not skip) ----
# A thermally throttled card halves data-parallel throughput while still reporting
# 100% utilization. Idle temps should be within ~10 C of each other.
nvidia-smi --query-gpu=index,clocks.sm,clocks.max.sm,temperature.gpu,power.draw,clocks_event_reasons.active --format=csv

# ---- 2. data ----
export HF_HOME=/workspace/.hf_home HF_HUB_DISABLE_XET=1
hf download angkul07/abc-ego-screwdriver --repo-type dataset \
  --local-dir /workspace/abc-ego-lerobot
# must print v2.1 — the launcher refuses to migrate in place
python -c "import json;print(json.load(open('/workspace/abc-ego-lerobot/meta/info.json'))['codebase_version'])"

# ---- 3. norm stats (~4 min; --skip-videos never touches pixels) ----
uv run scripts/compute_norm_stats.py --config-name pi05_abcego_sd \
  --max-frames 200000 --skip-videos
# -> assets/pi05_abcego_sd/abcego_sd/norm_stats.json

# ---- 4. train (~4 h 36 m) ----
tmux new -s train
./vast_run/run_yam.sh pi05_abcego_sd

run_yam.sh fails fast on a bad config before allocating a GPU, skips norm stats if they are already present, tees to /workspace/logs/pi05_abcego_sd/train.log, and auto-resumes if the checkpoint directory already holds numbered steps (FRESH=1 forces a clean restart).

If you change batch_size, recompute num_train_steps and decay_steps together — one epoch is 730,496 / batch. At other batch sizes: 32 → 22,828 | 48 → 15,219 | 96 → 7,610. Changing one without the other silently changes the epoch count and leaves the cosine schedule misaligned with the run length.

Budget ~50 GiB of live checkpoint storage (max_to_keep=4 at ~12.5 GiB each) plus 12.5 GiB for the pi05_base init and ~9 GiB for the dataset.

Related

• angkul07/abc-ego-screwdriver — the training dataset (LeRobot v2.1)
• angkul07/pi05-yam7h-ea — π₀.₅ on the 7 h 50/50 teleop+ego mixture, 23.6k steps