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.
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
bash
1# fetch the final checkpoint (inference only — skipping train_state roughly halves the download)2hf download angkul07/pi05-abcego-sd \3 --include "11413/params/*""11413/assets/*""11413/_CHECKPOINT_METADATA"\4 --local-dir ./pi05-abcego-sd
python
1from openpi.policies import policy_config
2from openpi.training import config as _config
34cfg = _config.get_config("pi05_abcego_sd")5policy = policy_config.create_trained_policy(cfg,"./pi05-abcego-sd/11413")67action_chunk = policy.infer({8"observation/top_image": top_rgb,# HxWx3 uint89"observation/left_wrist_image": left_wrist_rgb,10"observation/right_wrist_image": right_wrist_rgb,11"observation/state": state_14d,# [L j0-5, L grip, R j0-5, R grip]12"prompt":"put the screwdriver in the bin",13})["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.
bash
1# ---- 0. environment ----2git clone https://github.com/angkul07/openpi.git /workspace/openpi
3cd /workspace/openpi
4git checkout fc939b753861b87a56cbeba821cc460009a7a56f
5GIT_LFS_SKIP_SMUDGE=1 uv sync67cat> vast_run/env.local <<'EOF'
8export WANDB_API_KEY=<your key>
9export ABCEGO_SD_ROOT=/workspace/abc-ego-lerobot
10EOF11chmod600 vast_run/env.local
1213# ---- 1. GPU health check (do not skip) ----14# A thermally throttled card halves data-parallel throughput while still reporting15# 100% utilization. Idle temps should be within ~10 C of each other.16nvidia-smi --query-gpu=index,clocks.sm,clocks.max.sm,temperature.gpu,power.draw,clocks_event_reasons.active --format=csv
1718# ---- 2. data ----19exportHF_HOME=/workspace/.hf_home HF_HUB_DISABLE_XET=120hf download angkul07/abc-ego-screwdriver --repo-type dataset \21 --local-dir /workspace/abc-ego-lerobot
22# must print v2.1 — the launcher refuses to migrate in place23python -c "import json;print(json.load(open('/workspace/abc-ego-lerobot/meta/info.json'))['codebase_version'])"2425# ---- 3. norm stats (~4 min; --skip-videos never touches pixels) ----26uv run scripts/compute_norm_stats.py --config-name pi05_abcego_sd \27 --max-frames 200000 --skip-videos
28# -> assets/pi05_abcego_sd/abcego_sd/norm_stats.json2930# ---- 4. train (~4 h 36 m) ----31tmux new -s train
32./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