admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-24 10:16:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

do not rename normalization layers

Browse Source
This commit is contained in:
Pepijn
2025-09-13 11:23:58 +02:00
parent a9d54cbddb
commit f0b969ae48
2 changed files with 10 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py
+5 -1
View File
@@ -915,11 +915,15 @@ class PI05OpenPIPolicy(PreTrainedPolicy):
fixed_state_dict = model._fix_pytorch_state_dict_keys(original_state_dict, model.config) 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 # 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 = {} remapped_state_dict = {}
remap_count = 0 remap_count = 0
for key, value in fixed_state_dict.items(): 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}" new_key = f"model.{key}"
remapped_state_dict[new_key] = value remapped_state_dict[new_key] = value
remap_count += 1 remap_count += 1
src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py
+5 -1
View File
@@ -934,11 +934,15 @@ class PI0OpenPIPolicy(PreTrainedPolicy):
fixed_state_dict = model._fix_pytorch_state_dict_keys(original_state_dict, model.config) 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 # 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 = {} remapped_state_dict = {}
remap_count = 0 remap_count = 0
for key, value in fixed_state_dict.items(): 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}" new_key = f"model.{key}"
remapped_state_dict[new_key] = value remapped_state_dict[new_key] = value
remap_count += 1 remap_count += 1