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-holdoutwere excluded at selection time, leaving 761 episodes / 251,703 frames / 2.33 h, renumbered0..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 at0.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 on50%** 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%.R_grip, which would saturate ** observation.stateandactionare mapped separately. Teleop's closed gripper reads 0.270 in state but 0.054 in action β that gap is real gripper-servo lag, andnorm_statscarries independent entries for each. This does break ego'saction[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
.npzsidecar 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-multiviewfresh from the Hub gives you the unrescaled data. Check formeta/gripper_rescale.json(the sentinel, which records the full knot plan) to tell which version you have, and re-runvast_run/pi05/rescale_ego_gripper.py --applyfollowed byscripts/compute_norm_stats.pyif 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_downloadissues ~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). Atmax_to_keep=4budget ~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.0on 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
.npzsidecar 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.