From f0b969ae48ba763c0cf321023938bf35337476d1 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Sat, 13 Sep 2025 11:23:58 +0200 Subject: [PATCH] do not rename normalization layers --- src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py | 6 +++++- src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py b/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py index 9ff71152a..f53432cd3 100644 --- a/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py +++ b/src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py @@ -915,11 +915,15 @@ class PI05OpenPIPolicy(PreTrainedPolicy): fixed_state_dict = model._fix_pytorch_state_dict_keys(original_state_dict, model.config) # Then add "model." prefix for all keys that don't already have it + # Exclude normalization layers as they are direct attributes of the policy, not part of self.model remapped_state_dict = {} remap_count = 0 for key, value in fixed_state_dict.items(): - if not key.startswith("model."): + if not key.startswith("model.") and not any( + key.startswith(prefix) + for prefix in ["normalize_inputs.", "normalize_targets.", "unnormalize_outputs."] + ): new_key = f"model.{key}" remapped_state_dict[new_key] = value remap_count += 1 diff --git a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py index 549dc0a9b..f3c73d6f8 100644 --- a/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py +++ b/src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py @@ -934,11 +934,15 @@ class PI0OpenPIPolicy(PreTrainedPolicy): fixed_state_dict = model._fix_pytorch_state_dict_keys(original_state_dict, model.config) # Then add "model." prefix for all keys that don't already have it + # Exclude normalization layers as they are direct attributes of the policy, not part of self.model remapped_state_dict = {} remap_count = 0 for key, value in fixed_state_dict.items(): - if not key.startswith("model."): + if not key.startswith("model.") and not any( + key.startswith(prefix) + for prefix in ["normalize_inputs.", "normalize_targets.", "unnormalize_outputs."] + ): new_key = f"model.{key}" remapped_state_dict[new_key] = value remap_count += 1