Restore z-score normalization for pi0_fast_libero configs#971
Open
0xadvait wants to merge 1 commit into
Open
Conversation
The quantile-norm default for non-PI0 models silently changed normalization for the released pi0_fast_libero checkpoint, which was trained with z-score stats. Serving it with quantile normalization mis-scales the state input and the action unnormalization, which is why LIBERO evals with that checkpoint dropped to roughly 0% success. Adds an explicit per-config override and pins the two pi0_fast_libero configs back to z-score normalization; DROID FAST and pi05 configs are unchanged. Fixes Physical-Intelligence#849
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Evaluating the official
pi0_fast_liberocheckpoint on current main gives ~0% success on LIBERO (the arm spins in place). #849 has six independent reports, and one user bisected the last working commit to e458066.Root cause
At e458066,
LeRobotLiberoDataConfigbuilt itsDataConfigwithout settinguse_quantile_norm, so the pi0-FAST LIBERO configs trained and served with z-score normalization. Current main'sDataConfigFactory.create_base_configforce-setsuse_quantile_norm = model_type != PI0, which silently flippedpi0_fast_liberoandpi0_fast_libero_low_mem_finetuneto quantile normalization. The released checkpoint was trained with z-score stats, so the state input is mis-normalized going in and the action chunk is mis-unnormalized coming out. The scale mismatch on actions is what produces the spinning-arm behavior. DROID FAST configs are unaffected; they used quantile normalization both before and after that change.Evidence
Deterministic probe: the official checkpoint, a fixed synthetic LIBERO observation, greedy FAST decoding, run at e458066 and at current main. Preprocessed image hashes and transform sequences are identical at both commits, but:
With this PR applied to main, the same probe reproduces e458066's post-normalization state, token ids, and both action chunks exactly (max abs diff 0 on all four checks, two infer calls each).
Change
DataConfigFactorygainsuse_quantile_norm_override: bool | None(default None keeps the model-type default for every existing config).pi0_fast_liberoconfigs pin the override to False with a comment explaining the checkpoint compatibility constraint.Flag matrix after the change:
pi0_fast_liberoFalse,pi0_fast_libero_low_mem_finetuneFalse,pi0_fast_droidTrue,pi0_liberoFalse,pi05_liberoTrue.ruff checkandruff format --checkpass.data_loader_test.py: 4 passed, plustest_with_real_datasetwhich fails identically on clean main (pre-existing dataset-fetch issue, unrelated).Caveat
Anyone who finetuned pi0-FAST on LIBERO with current main trained under quantile normalization; for those checkpoints the override makes the choice explicit and per-config instead of a silent default flip.
Fixes #849