pi05-yam7h-ea β€” Ο€β‚€.β‚… finetuned on a 7-hour YAM bimanual mixture

Training checkpoints for arm E-A of the YAM 7-hour mixture study: Ο€β‚€.β‚… (flow matching) finetuned from pi05_base on a 50/50 mixture of real teleop and retargeted human video, 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 15000, 20000, 23000, 23599 (final)
Evaluation None. No held-out split, no success rate. Every number below is training loss.

The run trained the full 23,600 steps and completed cleanly. 23599 is the final checkpoint β€” use it unless you specifically want an earlier one. The intermediate 21000 and 22000 checkpoints were not retained (max_to_keep=4).

Because there is no evaluation, "final" here means last, not best. Nothing in this repo establishes that 23599 outperforms 15000 on any task.


Experiment

Field Value
Config pi05_yam7h_ea (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 SXM (NVLink), data-parallel, --fsdp-devices 1 (full model replica per GPU)
Wall clock 9 h 06 m 42 s for 23,600 steps (β‰ˆ1.39 s/step avg; 1.4–2.2 s/it observed)
Approx. compute cost ~$40 at $4.45/hr
Code revision ced0204ec233c6f636584c35d632ec28c78d3bda (angkul07/openpi fork)
W&B project pi05-yam, run ea

Hyperparameters

Field Value
Batch size 64 (32 teleop + 32 ego, fixed per batch)
Steps 23,600
Optimizer AdamW, clip_gradient_norm=1.0
LR schedule Cosine, warmup 500 β†’ peak 3.5e-5 β†’ decay 3.5e-6 over 23,600 steps
EMA off (ema_decay=None)
Dataloader workers 16
Checkpoint interval every 1,000 steps, max_to_keep=4
Precision fp32 params (param_norm tracked in fp32)

Note on the LR: openpi's own Ο€β‚€.β‚… recipes (pi05_libero, pi05_full_droid_finetune) use 5e-5 held constant (no cosine). The cosine schedule here was held fixed across all three arms of this study (E-A/E-B/E-C) so that the only moving variable is the mixture ratio and run length β€” not because cosine was found to be better.


Data

A 7-hour mixture of two sources, sampled at a fixed 32/32 per batch (a StratifiedBatchSampler, not probabilistic sampling β€” every batch has exactly 32 of each).

Source Frames Per batch Gradient share Epochs over the run
angkul07/abc-teleop (7h subset) 251,703 32 50.0% 3.00
angkul07/EgoDex-PickPlace-YAM-14dof-multiview 505,837 32 50.0% 1.49
Total 757,540 64 1,510,400 frames seen

Because the two sources differ in size, the 50/50 gradient split means teleop is oversampled 1.50Γ— relative to its storage share and ego is undersampled 0.75Γ—.

Observation space

  • 3 Γ— RGB at 224Γ—224: base_0_rgb (top), left_wrist_0_rgb, right_wrist_0_rgb
  • 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
  • DeltaActions(mask=(TrueΓ—6, False, TrueΓ—6, False)) β€” the 12 arm-joint dims are predicted relative to current state; the 2 gripper dims are absolute

Normalization β€” quantile (use_quantile_norm=True, q01/q99), asset yam7h_p50, shipped in each checkpoint at <step>/assets/yam7h_p50/norm_stats.json.

Train-only augmentation β€” ImageAugment(brightness=0.15, contrast=0.15, saturation=0.1, hue=0.03, crop_scale=(0.95, 1.0)), geometric crop applied to base_0_rgb only (wrist cameras keep their framing).

Both sources are LeRobot v2.1, 30 fps, 224Γ—224, 14-D state and action β€” schema-identical by construction, so they can be mixed without any per-source adapter.


Data preparation β€” what we changed

Neither source is off-the-shelf. This section documents every transformation applied between the raw data and what the model saw, because several of them materially affect how the numbers above should be read.

1. Teleop subset (abc-teleop β†’ the 7 h arm)

Real YAM bimanual teleoperation, single task: "put the screwdriver in the bin".

  • Held-out episodes physically removed. The 249 episodes published as angkul07/abc-teleop-holdout were excluded at selection time, leaving 761 episodes / 251,703 frames / 2.33 h, renumbered 0..760.
  • Because the exclusion is physical rather than an index filter, the training config sets exclude_episodes=() β€” passing the original holdout indices here would double-exclude the wrong episodes, since the indices no longer refer to the same clips.
  • Otherwise unmodified: no rescaling, no filtering, no augmentation at rest.

2. Ego source (EgoDex human video β†’ YAM robot data)

This source did not exist as robot data. It was produced by a three-stage pipeline.

(a) Retargeting β€” human hands β†’ YAM joints. EgoDex ships 3D hand pose (SE3), so hand estimation is skipped and the poses feed the retargeting core directly:

  • End-effector derivation from hand keypoints: a virtual finger k_vf = 0.7Β·index_tip + 0.3Β·middle_tip, with the grasp point at 0.5Β·(thumb_tip + k_vf). Gripper aperture comes from the thumb-to-index distance.
  • Base placement chosen per clip by a BPP grid search over torso poses.
  • IK: mink + daqp, joint-limit-safe, solving to YAM's 14 DoF (2 Γ— [6 joints + gripper]) at 30 fps. Orientation is tracked at low weight (IK_ORIENTATION_COST = 0.4) rather than as a hard 6-D target β€” position dominates, so the arm retains nullspace redundancy.

(b) Multi-view synthesis (Method Aβ€²). EgoDex has one egocentric camera; teleop has three. To make the schemas identical, three views were synthesized from that single camera:

View How
top square center-crop β†’ 224Γ—224
left_wrist / right_wrist 620 px crop tracking that hand's projected grasp point β†’ 224Γ—224

Each view gets a fixed virtual lens (barrel distortion, tilt, colour jitter) and is encoded to match teleop exactly: 224Γ—224, h264, yuv420p, 30 fps. Two locked conventions:

  • Left/right are anatomical, not screen-relative. The "crossed" look β€” the right hand appearing on frame-left in an egocentric view β€” is real, not a mirror bug.
  • Wrist views are never blank. The crop centre is edge-clamped so the window always lies fully in frame. When a hand leaves the FOV the view rides the nearest edge showing hand-adjacent workspace, rather than going black.

(c) Episode selection. Objects are not a metadata field; they were parsed out of the task sentence ("Pick up a ⟨object⟩ from the ⟨source⟩ and place it …"). Objects with fewer than 30 episodes were dropped, then the surviving pool was subsampled proportionally per object (seed 0) to hit the 4.67 h target β€” preserving the object mix rather than truncating it. Result: 111 objects, 4,174 episodes, 505,837 frames.

3. Gripper rescale β€” the one change made specifically for Ο€β‚€.β‚…

This is the most consequential data change in this run, and it was not needed for the Ο€β‚€-FAST arms.

Grippers stay absolute through DeltaActions, then get quantile-normalized over the mixture β€” so both sources' gripper values land on one shared [βˆ’1, 1] axis. Measured on the 7 h mixture, they disagreed about what "open" means:

Dim teleop open β†’ normalized ego open β†’ normalized
R_grip 0.992 β†’ +0.984 0.443 β†’ βˆ’0.114
L_grip 0.993 β†’ +0.986 0.392 β†’ βˆ’0.216

Ego's "open" sat closer to teleop's closed than to teleop's open. Under Ο€β‚€-FAST this is survivable β€” it samples one discrete token, i.e. one mode. Under flow matching with few denoising steps it is not: the model mode-averages between widely separated modes, and the midpoint of a bimodal open-gripper target is a half-open hand that grasps nothing.

Ego's absolute gripper scale is a retargeting artifact (it derives from human thumb-index distance), while teleop's is the real YAM gripper β€” so teleop is the target range and only the ego columns were touched.

Key choices:

  • Anchored on robust percentiles (p1/p99), not mode-to-mode. Teleop's gripper is near-binary (p50 = 0.986, 13% of frames closed); ego's is a continuous aperture spanning [0.02, 0.86]. Mapping mode-to-mode implies a gain of 5.0 on R_grip, which would saturate **50%** of ego frames at the top of the range. The percentile anchoring is monotone, preserves ego's within-source aperture ordering, and clips 0.21%.
  • observation.state and action are mapped separately. Teleop's closed gripper reads 0.270 in state but 0.054 in action β€” that gap is real gripper-servo lag, and norm_stats carries independent entries for each. This does break ego's action[t] == state[t+1] identity on dims 6 and 13 only; nothing in the training path depends on it.
  • Reversible. Original columns were dumped to an .npz sidecar before any write, and a sentinel blocks double-application.

Applied result (piecewise-linear knots, ego β†’ teleop range):

Column:dim ego open before ego open after teleop open saturated
state:6 (R_grip) 0.441 0.852 0.992 0.21%
state:13 (L_grip) 0.401 0.908 0.992 0.21%
action:6 (R_grip) 0.441 0.855 0.991 0.21%
action:13 (L_grip) 0.402 0.916 0.993 0.21%

Reproducibility note. The rescale was applied in place to the local copy after download, not upstream β€” so pulling angkul07/EgoDex-PickPlace-YAM-14dof-multiview fresh from the Hub gives you the unrescaled data. Check for meta/gripper_rescale.json (the sentinel, which records the full knot plan) to tell which version you have, and re-run vast_run/pi05/rescale_ego_gripper.py --apply followed by scripts/compute_norm_stats.py if it is absent.

4. Norm stats recomputed after the rescale

The rescale invalidates any earlier statistics, so assets/pi05_yam7h_ea/yam7h_p50 was recomputed over the post-rescale mixture (quantile, 200k frames, --skip-videos). The resulting stats are shipped inside every checkpoint. Recomputation was verified bit-identical across two independent machines (max |Ξ”| = 0.00000).

What was deliberately not changed

  • Inter-arm collisions in the ego source were kept, not filtered. Roughly 2/3 of retargeted clips have some inter-arm penetration (~55 mm median depth, ~41% of frames). This is an embodiment mismatch β€” two YAM arms cannot occupy the volume two human hands did β€” and filtering would discard most of the data for a defect no retargeter setting removes. Recorded as metadata; accept-and-flag.
  • Teleop was not rescaled, resampled, or augmented at rest. It is the reference range.
  • No holdout was carved from either source for this run (holdout_fraction=0.0).

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.05178 0.18415 1803.24
100 – 500 0.02858 0.06547 1803.25
500 – 1,000 0.02427 0.05503 1803.33
1,000 – 2,000 0.02205 0.05205 1803.50
2,000 – 4,000 0.01990 0.05018 1803.81
4,000 – 6,000 0.01792 0.04964 1804.20
6,000 – 10,000 0.01581 0.04927 1804.68
10,000 – 15,000 0.01323 0.04957 1805.13
15,000 – 20,000 0.01086 0.04969 1805.32
20,000 – 23,600 0.00969 0.04974 1805.37

Per published checkpoint (200-step trailing mean):

Checkpoint Loss grad_norm
15000 0.01197 0.04878
20000 0.01012 0.05033
23000 0.00957 0.04768
23599 (final) 0.00932 0.04941

Other figures: first step loss 0.1007; minimum single-step loss 0.00554 at step 20,432; last-500-step loss 0.00940 Β± 0.00142; total param_norm drift over the whole run +2.14 (+0.12%).

Full per-step metrics are in train_metrics_full.log (all 23,600 steps). train_metrics.log is the copy that was snapshotted alongside the checkpoints and stops at step 23,500.


Learnings and takeaways

1. The absolute loss value is not interpretable, and its smallness is not evidence of convergence. Loss reaches 0.05 by step 100 and 0.02 by step 500. This 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 in magnitude 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 or to a cross-entropy Ο€β‚€-FAST number.

2. The run was not length-limited. Loss fell monotonically in every window with no plateau β€” 0.05178 β†’ 0.00969, a 5.3Γ— reduction, still descending at 23,600. Whether the extra steps buy real capability is exactly what arm E-C (47,200 steps, same ratio) is meant to answer; this run cannot.

3. Optimization was uneventful. grad_norm settled at ~0.049 by step 500 and stayed flat for the remaining 23k steps β€” no spikes, no instability, and clip_gradient_norm=1.0 never bound (max observed gradient norm is ~20Γ— below the clip). The LR schedule and warmup can be considered validated for this setup.

4. LoRA finetuning barely moves the base model. param_norm drifted +0.12% across the entire run. With gemma_2b_lora and a 3.5e-5 peak LR, this is a light touch on pi05_base β€” most of the capability is inherited, not learned here. That is the intent, but it also means the run is unlikely to have overfit the 7h mixture in 3 epochs.

5. Aggregate loss cannot answer the question this experiment was designed to ask. A single scalar over a 50/50 mixture cannot tell you whether the model is learning from the teleop half, the ego half, or fitting the zero-padding. Per-source (loss_teleop / loss_ego) and real-vs-padded-dim loss splits were designed for this run but were not deployed β€” the metrics here are the aggregate only. This is the single highest-value change for the next run, and it is cheap: split jnp.square(v_t βˆ’ u_t) before the final mean, and carry a source_index through MixtureDataset into the train step.

6. Mixing sources on an absolute channel needs an explicit range alignment, and the architecture decides how badly. The gripper rescale above is the concrete instance: two sources agreed on the semantics of "open" but not on its value, and quantile normalization over the mixture happily mapped ego-open to βˆ’0.11 while teleop-open sat at +0.98. Ο€β‚€-FAST tolerated this because it samples a discrete token; Ο€β‚€.β‚…'s flow matching mode-averages, so the same data defect turns into a half-open gripper that grasps nothing. The general lesson: any channel left absolute through DeltaActions is a channel where cross-source scale disagreement survives normalization β€” audit those explicitly before mixing, and re-audit when changing architecture. The corollary on method: anchor such alignments on robust percentiles, not on modes; mode-matching here would have saturated half the ego frames instead of 0.21%.

7. Infrastructure notes worth reusing.

  • 2Γ— H100 SXM at ~1.39 s/step is ~2.2Γ— the throughput of 2Γ— A100 SXM4 for this config, at ~2.1Γ— the price β€” roughly cost-neutral, but it halves the wall clock, which matters for iteration speed.
  • Bulk-downloading a many-file LeRobot dataset from HF hits a 1,000-API-requests / 5-minutes rate limit. snapshot_download issues ~1 request per file and will report success while silently incomplete. git clone + git-lfs batches ~100 objects per request and is the reliable path for datasets with >10k small files.
  • Ο€β‚€.β‚… checkpoints are 13 GB each (25–30% larger than Ο€β‚€-FAST, from the +311M expert params and their optimizer state). At max_to_keep=4 budget ~52 GB of live checkpoint storage.

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 on both sources β€” there is no validation split, no success rate, no hardware rollout. Loss curves only.
  • Language conditioning comes almost entirely from the ego source. The teleop set has 1 distinct task string; the ego set has 471. Any instruction-following the model shows is learned from retargeted human video, not from real robot demonstrations.
  • Half the gradient comes from synthesized data. The ego source's two wrist views are keypoint-guided synthesis from a single egocentric camera, and its actions are IK retargeting of human hand poses. The sim-to-real gap of that half is unmeasured.
  • The gripper rescale is unvalidated in the only way that counts. Its motivation is sound and its statistics are checked (0.21% saturation, monotone, ordering preserved), but no rollout has confirmed that the aligned gripper actually grasps better than the unaligned one. It is a well-argued correction, not a measured improvement. The .npz sidecar makes the A/B cheap for whoever runs it.
  • Ego wrist views are synthesized from a camera that never saw the wrist. When a hand leaves the egocentric FOV, the "wrist" view rides the frame edge showing adjacent workspace. That is deliberate β€” better than a black frame β€” but it means some fraction of ego wrist observations are not views of the gripper at all.
  • The ego retargeting residual is reported as exactly 0.000. That is suspicious β€” worth confirming it is genuinely near-zero tracking rather than a residual that is never written. Unverified.
  • Inter-arm collisions are inherent to the YAM embodiment in the retargeted data and were accepted as a known limitation, not fixed.
  • Single run, single seed. No ablation isolates the mixture ratio; arms E-B (62.5% teleop) and E-C (2Γ— length) are the intended comparisons and are separate runs.

Repo layout

15000/                       # and 20000/, 23000/, 23599/
  _CHECKPOINT_METADATA
  params/                    # Orbax OCDBT β€” inference weights (~7.2 GB)
  train_state/               # optimizer state β€” for resuming only (~6.2 GB)
  assets/yam7h_p50/norm_stats.json
train_metrics.log            # snapshot copy, steps 0–23,500
train_metrics_full.log       # complete, steps 0–23,599

Each checkpoint is ~13.4 GB; the repo is ~50 GiB total. For inference you only need params/ and assets/ β€” train_state/ is ~46% of the bytes and is needed only to resume training.

Usage

# fetch the final checkpoint (inference only β€” skipping train_state halves the download)
hf download angkul07/pi05-yam7h-ea \
  --include "23599/params/*" "23599/assets/*" "23599/_CHECKPOINT_METADATA" \
  --local-dir ./pi05-yam7h-ea
from openpi.policies import policy_config
from openpi.training import config as _config

cfg = _config.get_config("pi05_yam7h_ea")
policy = policy_config.create_trained_policy(cfg, "./pi05-yam7h-ea/23599")

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,
    "prompt":                        "pick up the block and place it in the bin",
})["actions"]                                          # (50, 14)

To resume training from a checkpoint, point openpi at the parent directory:

uv run scripts/train.py pi05_yam7h_ea --exp-name ea --fsdp-devices 1 --resume

License

Apache-2.0, inherited from openpi and pi05_base. The underlying datasets carry their own terms; the EgoDex-derived source inherits EgoDex's license.

Downloads last month

-

Downloads are not tracked for this model. How to track
Video Preview
loading

Datasets used to train angkul07/pi05-yam7h-ea