feat(pi05): auto-detect adaRMSNorm conditioning from checkpoint by pwolnows · Pull Request #858 · open-edge-platform/physical-ai-studio

pwolnows · 2026-07-27T19:16:19Z

Problem

OpenVINO/pi05-libero-fp16-ov on HuggingFace was exported from an older codebase that predates bf16 computation and adaRMS layer support in PaliGemmaWithExpertModel. Comparing it against the native lerobot/pi05_libero_finetuned_v044 checkpoint (which runs in bfloat16 with adaRMS conditioning) produces random-looking output.

Structural evidence from the OpenVINO IR:

Export	bf16 ops	fp32 ops	fp16 ops
`OpenVINO/pi05-libero-fp16-ov` (broken HF)	0	15877	3084
Current `main` branch export	10422	7619	2059
This PR's export	10422	7619	2059

The zero bf16 count in the HF model confirms it was traced from a model without bf16 computation. The current main branch (hardcoded use_adarms=[False, True]) already produces the correct graph with 10422 bf16 tensors — identical to this PR's export.

Root Cause

The HF model is stale. It was not re-exported after bf16/adaRMS support was added to PaliGemmaWithExpertModel. The model code on main was already correct.

What was missing was the config/detection layer: Pi05Config had no use_adarms or adarms_cond_dim fields, so:

Checkpoint introspection was impossible — there was no way to know programmatically whether a checkpoint was trained with adaRMS.
Pi05Config.use_adarms defaulted to False, meaning any code path that constructed the config without going through the hardcode in Pi05Model would build a model without adaRMS layers.

Changes

`config.py`

Added use_adarms: bool = True and adarms_cond_dim: int | None = None to Pi05Config. Default is True to match the model-level invariant: adaRMS is always enabled for the action expert in Pi05Model (matching upstream LeRobot Pi05). Removed the over-constrained __post_init__ guard that blocked Pi05Config() construction when adarms_cond_dim was None — the model derives the conditioning dim from action_expert_config.width, not from config.

`pretrained_utils.py`

Added detect_adarms_from_checkpoint(weights_file, hf_config):

Scans safetensors keys for *.input_layernorm.dense.weight to detect adaRMS weights
Injects use_adarms=True and adarms_cond_dim=<dim> into hf_config for checkpoints that predate the config fields

`policy.py`

Calls detect_adarms_from_checkpoint() in _from_hf() before Pi05Config.from_dict(), enabling automatic detection for legacy checkpoints.

`model.py`

Pi05Model.__init__ now accepts use_adarms: bool (default True) and passes [False, use_adarms] to PaliGemmaWithExpertModel. policy._initialize_model wires self.config.use_adarms through.

Validation

Validated with test_golden_action from openvinotoolkit/physicalai#208 — 30 samples from lerobot/libero_10_image, use_random_input_noise=False, on Intel Arc A770 GPU (openvino==2026.2.1). Results after commit 76c90d02 on this branch:

Export	Passing samples	Max L2	Mean L2	Threshold	Result
`OpenVINO/pi05-libero-fp16-ov` (stale HF)	1 / 30	14.32	7.40	2.0	❌ FAIL
Local re-export from this PR	30 / 30	1.87	0.17	2.0	✅ PASS

The stale HF model needs to be replaced with a re-export from current main (or this PR).

Testing

Unit tests in library/tests/unit/policies/test_pi05.py covering detect_adarms_from_checkpoint().

…ort parity - ACT.__init__ accepts pretrained_name_or_path (local dir or HF Hub id), resolving config.json + model.safetensors via _from_hf() and remapping LeRobot checkpoint keys to physicalai's native parameter names via _remap_lerobot_act_state_dict(). - preprocessor.py: fix image_resolution unpacking order in the resize/pad path (stored as (height, width) but _resize_with_ar_pad expects width first), which silently swapped dimensions for non-square configs. - Add scripts/compare_act_backends.py: loads a real ACT checkpoint, exports it to OpenVINO, and validates parity via (1) single-step predict_action_chunk numeric diff on real gym-aloha observations and (2) matching closed-loop AlohaTransferCube episodes on both backends.

Pi05 models fine-tuned with adaRMSNorm (adaRMS) conditioning have dense projection layers in every action-expert LayerNorm block (`input_layernorm.dense.weight/bias`, `post_attention_layernorm.dense.weight/bias`). These layers were absent from Pi05Config so the export pipeline silently omitted them, producing an OpenVINO IR that was missing the entire conditioning branch. Comparing such an export against the native bfloat16 model yielded L2 ≈ 7.7, cosine-sim ≈ 0.51 — effectively random output. After this fix the L2 drops to ≈ 0.047 (cosine-sim 0.9999), consistent with the expected bfloat16→float32 numerical noise. Changes ------- library/src/physicalai/policies/pi05/config.py • Added two new optional fields to Pi05Config: use_adarms: bool = False adarms_cond_dim: int | None = None These mirror the per-expert use_adarms / adarms_cond_dim parameters already consumed by Pi05Model but were never persisted in the config. library/src/physicalai/policies/pi05/pretrained_utils.py • Added detect_adarms_from_checkpoint(weights_file, hf_config): Opens the safetensors checkpoint, scans for any key matching `*.input_layernorm.dense.weight`, and if found injects use_adarms=True and adarms_cond_dim=<projection_width> into the hf_config dict in-place before Pi05Config is constructed. This makes adaRMS support fully automatic — no manual config edits. library/src/physicalai/policies/pi05/policy.py • Calls detect_adarms_from_checkpoint() inside _from_hf() immediately after normalization_mode detection and before Pi05Config.from_dict(), so any downstream code (training, export, inference) gets a correctly populated config.

Copilot

Pull request overview

This PR addresses incorrect Pi05 exports when adaRMSNorm conditioning layers exist in a checkpoint but are not represented in Pi05Config, by adding config fields and checkpoint-key-based auto-detection prior to config parsing. It also includes ACT-related changes (HF/LeRobot checkpoint loading + an OpenVINO parity comparison script), which broadens the PR scope beyond the stated Pi05 focus.

Changes:

Add use_adarms / adarms_cond_dim to Pi05Config and auto-detect these from safetensors keys during pretrained loading.
Wire adaRMS auto-detection into Pi05 _from_hf() so downstream export/tracing uses the correct graph.
Add ACT LeRobot checkpoint loading/remapping support and a new script to compare native vs OpenVINO ACT parity (plus a small ACT preprocessor fix).

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
scripts/compare_act_backends.py	New script to compare native PyTorch ACT vs OpenVINO export numerically and in closed-loop rollouts.
library/src/physicalai/policies/pi05/pretrained_utils.py	Adds adaRMS detection by scanning safetensors keys and injecting config fields.
library/src/physicalai/policies/pi05/policy.py	Invokes adaRMS auto-detection during HF checkpoint loading before `Pi05Config.from_dict()`.
library/src/physicalai/policies/pi05/config.py	Adds adaRMS configuration fields to the Pi05 config dataclass.
library/src/physicalai/policies/act/preprocessor.py	Fixes argument order when resizing images (height/width vs width/height).
library/src/physicalai/policies/act/policy.py	Adds `pretrained_name_or_path` loading, LeRobot state-dict remapping, and weight loading during initialization.

- pretrained_utils: only skip adaRMS auto-detection when both use_adarms and adarms_cond_dim are already set; log the exception when checkpoint key scan fails - Pi05Config: validate that adarms_cond_dim is set when use_adarms is True - ACT: document that pretrained_name_or_path overrides other __init__ args - remove scripts/compare_act_backends.py, mistakenly pushed in PR open-edge-platform#858

sovrasov · 2026-07-28T12:00:43Z

Nice catch, thanks Piotr! Interesting thing is that the model still demonstrate good perf numbers on libero even with incorrect weights loading

Covers the reviewer request on PR open-edge-platform#858: mocked safetensors checkpoints exercising dense-layernorm detection, the no-op case, the already-set short-circuit, re-detection when adarms_cond_dim is missing, and the scan-error logging path.

sovrasov · 2026-07-28T15:53:16Z

@pwolnows could you share a script / snippet how pi05 was created and exported? There are some parameters that can affect exact outputs matching, but preserve results in envs, because the policy performs denoising internall

eugene123tw · 2026-07-28T15:57:06Z

adaRMS is intentionally hardcoded in Pi05Model, not a configuration option:

In model.py, the adaRMS setting is hardcoded:

self.paligemma_with_expert = PaliGemmaWithExpertModel(
    paligemma_config,
    action_expert_config,
    use_adarms=[False, True],  # VLM: disabled, Action expert: ENABLED
    precision=dtype,
    image_size=self._image_resolution[0],
    freeze_vision_encoder=freeze_vision_encoder,
    train_expert_only=train_expert_only,
)

This design is intentional and aligned with upstream LeRobot Pi05, adaRMS is always enabled for the action expert ([False, True]), never disabled.

Adding use_adarms and adarms_cond_dim to Pi05Config would not change the model behavior (it's hardcoded).

Pi05Model was hardcoding use_adarms=[False, True] regardless of the Pi05Config.use_adarms field, making the detect_adarms_from_checkpoint() auto-detection added in this PR completely ineffective. Changes: - Pi05Model.__init__: add use_adarms: bool = True parameter - Pi05Model.__init__: pass [False, use_adarms] to PaliGemmaWithExpertModel - policy._initialize_model: pass use_adarms=self.config.use_adarms This ensures checkpoints without adaRMS weights (use_adarms=False) build a model without dense layernorm layers, so load_state_dict does not silently discard keys with strict=False and the model runs correctly. Validated with test_golden_action.py against a local re-export: max L2=1.39, mean L2=0.50 (threshold=2.0) — all 3 samples PASS (vs L2≈6-7 before the fix)

…design adaRMS is always enabled for the action expert in Pi05Model (use_adarms=True is the default in Pi05Model.__init__). Pi05Config previously defaulted to False, which would silently disable adaRMS when loading checkpoints via any path that bypassed detect_adarms_from_checkpoint. Change the default to True to match the invariant enforced upstream (LeRobot Pi05 always enables adaRMS for the action expert). Remove the __post_init__ guard that required adarms_cond_dim to be set when use_adarms=True: the model derives adarms_cond_dim from action_expert_config.width directly; Pi05Config. adarms_cond_dim is only an optional detection hint in pretrained_utils.py.

pwolnows added 2 commits July 27, 2026 20:50

Copilot AI review requested due to automatic review settings July 27, 2026 19:16

pwolnows requested a review from a team as a code owner July 27, 2026 19:16

github-actions Bot added the Library label Jul 27, 2026

Copilot started reviewing on behalf of pwolnows July 27, 2026 19:17 View session

Copilot AI reviewed Jul 27, 2026

View reviewed changes

pwolnows added 4 commits July 28, 2026 08:01

Rever update of policy.py

31148f9

Revert change in policy.py

81e5699

Revert change in preprocessor.py

775236e

pwolnows mentioned this pull request Jul 28, 2026

feat(inference): golden-action integration test suite with quantile-stats cache and GPU auto-detect openvinotoolkit/physicalai#208

Draft

sovrasov reviewed Jul 28, 2026

View reviewed changes

Comment thread library/src/physicalai/policies/pi05/pretrained_utils.py

pwolnows added 2 commits July 28, 2026 21:14

pwolnows closed this Jul 29, 2026

Navigation Menu

feat(pi05): auto-detect adaRMSNorm conditioning from checkpoint - #858

feat(pi05): auto-detect adaRMSNorm conditioning from checkpoint#858
pwolnows wants to merge 9 commits into
open-edge-platform:mainfrom
pwolnows:feat/pi05-export-adarms-conditioning

pwolnows commented Jul 27, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sovrasov commented Jul 28, 2026

Uh oh!

sovrasov commented Jul 28, 2026

Uh oh!

eugene123tw commented Jul 28, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

pwolnows commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Root Cause

Changes

config.py

pretrained_utils.py

policy.py

model.py

Validation

Testing

Related

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sovrasov commented Jul 28, 2026

Uh oh!

sovrasov commented Jul 28, 2026

Uh oh!

eugene123tw commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

pwolnows commented Jul 27, 2026 •

edited

Loading

`config.py`

`pretrained_utils.py`

`policy.py`

`model.py`

eugene123tw commented Jul 28, 2026 •

edited

Loading