PI0 policy suffers from weight/key mismatch on main when using transformers ≥ 4.52.0

[yongjincho]

Summary

The current main branch works only with transformers ≥ 4.52.0, but the pretrained PI0 weights no longer align with the model’s internal key names. As a result, training starts with an abnormally high loss (≈ 4.5 vs ≈ 0.5 on a healthy run).

Failing configuration

# 1. Check out the current main commit
git co 69901b9b6a2300914ca3de0ea14b6fa6e0203bd4

# 2. Install an older Transformers version
pip install transformers==4.51.3

# 3. Run a short training job
python lerobot/scripts/train.py --policy.path=lerobot/pi0 --policy.push_to_hub=false --dataset.repo_id=lerobot/aloha_sim_insertion_human --steps=2 --log_freq=1

Observed error

INFO 2025-06-30 18:19:27 ts/train.py:202 Start offline training on a fixed dataset
Traceback (most recent call last):
  ...
AttributeError: 'GemmaForCausalLM' object has no attribute 'embed_tokens'

current main with transformers == 4.52.0

git co 69901b9b6a2300914ca3de0ea14b6fa6e0203bd4

pip install transformers==4.52.0

python lerobot/scripts/train.py --policy.path=lerobot/pi0 --policy.push_to_hub=false --dataset.repo_id=lerobot/aloha_sim_insertion_human --steps=2 --log_freq=1

Observed output (low performance)

INFO 2025-06-30 18:24:46 ts/train.py:202 Start offline training on a fixed dataset
INFO 2025-06-30 18:24:51 ts/train.py:232 step:1 smpl:8 ep:0 epch:0.00 loss:4.314 grdn:246.739 lr:5.0e-08 updt_s:1.388 data_s:3.667
INFO 2025-06-30 18:24:52 ts/train.py:232 step:2 smpl:16 ep:0 epch:0.00 loss:4.609 grdn:235.238 lr:7.5e-08 updt_s:0.348 data_s:0.000
INFO 2025-06-30 18:24:52 ts/train.py:241 Checkpoint policy after step 2

Training runs, but the initial loss is > 4.0—much higher than expected.

Expected behaviour

# checkout just before "Fixing `PI0` Policy (#1297)"
git co 697c76f75e12a6e1a2ba09911c93e1e22c9b8f5c

pip install transformers==4.51.3

python lerobot/scripts/train.py --policy.path=lerobot/pi0 --dataset.repo_id=lerobot/aloha_sim_insertion_human --steps=2 --log_freq=1

Output

INFO 2025-06-30 18:30:17 ts/train.py:202 Start offline training on a fixed dataset
INFO 2025-06-30 18:30:25 ts/train.py:232 step:1 smpl:8 ep:0 epch:0.00 loss:0.521 grdn:7.218 lr:5.0e-08 updt_s:1.244 data_s:7.370
INFO 2025-06-30 18:30:26 ts/train.py:232 step:2 smpl:16 ep:0 epch:0.00 loss:0.385 grdn:7.831 lr:7.5e-08 updt_s:0.375 data_s:0.000
INFO 2025-06-30 18:30:26 ts/train.py:241 Checkpoint policy after step 2

The loss is much lower than before.

Diagnosis

The high loss appears to stem from a key mismatch between the pretrained PI0 checkpoint and the parameter names expected by transformers ≥ 4.52.0. Earlier commits (and transformers 4.51.x) still line up correctly.

[YushunXiang]

It is a good question.
I think the cause of this problem is LeRobot PR #1297 and huggingface/transformers#37033.

• Fixing PI0 Policy #1297: Bring

      def embed_language_tokens(self, tokens: torch.Tensor):
-        return self.paligemma.language_model.model.embed_tokens(tokens)
      def embed_language_tokens(self, tokens: torch.Tensor):
+        return self.paligemma.language_model.embed_tokens(tokens)

The language model is GemmaModel can use function embed_tokens directly. But I think it will not cause the loss problem and key mismatch?

[yongjincho]

Here are the actual safetensor keys from two checpoints.

• original_model_keys.txt
  • old version with transformers == 4.51.3
• changed_model_keys.txt
  • current main with transformers == 4.52.0

and diff is following:

[YushunXiang]

I strongly suggest the LeRobot repo add kwargs["strict"] = True in Line 168, lerobot/common/policies/factory.py

lerobot/lerobot/common/policies/factory.py

Lines 164 to 171 in 483be9a

	if cfg.pretrained_path:
	# Load a pretrained policy and override the config if needed (for example, if there are inference-time
	# hyperparameters that we want to vary).
	kwargs["pretrained_name_or_path"] = cfg.pretrained_path
	policy = policy_cls.from_pretrained(**kwargs)
	else:
	# Make a fresh policy.
	policy = policy_cls(**kwargs)

like:

        if cfg.pretrained_path:
            # Load a pretrained policy and override the config if needed (for example, if there are inference-time
            # hyperparameters that we want to vary).
            kwargs["pretrained_name_or_path"] = cfg.pretrained_path
+           kwargs["strict"] = True
            policy = policy_cls.from_pretrained(**kwargs)
        else:
            # Make a fresh policy.
            policy = policy_cls(**kwargs)

And I get the missing keys and unexpected keys:

mismatch-keys.json

[YushunXiang]

I found that in line381-388
huggingface/transformers#37033,

+ class PaliGemmaForConditionalGeneration(PaliGemmaPreTrainedModel, GenerationMixin):
+     _checkpoint_conversion_mapping = {
+         "^language_model.model": "model.language_model",
+         "^vision_tower": "model.vision_tower",
+         "^multi_modal_projector": "model.multi_modal_projector",
+         "^language_model.lm_head": "lm_head",
+     }
+     _tied_weights_keys = ["lm_head.weight"]

caused this bug. It will convert "^language_model.model" to "model.language_model" and so on.

[YushunXiang]

You can see #1415, just for your reference. @yongjincho

[kumarhans]

@YushunXiang I tried your PR and I see high loss still:

Command run:
python lerobot/scripts/train.py --policy.type=pi0 --dataset.repo_id=lerobot/aloha_sim_insertion_human --policy.repo_id=kumarhans/test_pi0 --steps=2 --log_freq=1

INFO 2025-07-02 21:32:04 ts/train.py:232 step:1 smpl:8 ep:0 epch:0.00 loss:2.561 grdn:90.310 lr:5.0e-08 updt_s:1.109 data_s:0.414
INFO 2025-07-02 21:32:04 ts/train.py:232 step:2 smpl:16 ep:0 epch:0.00 loss:2.265 grdn:82.791 lr:7.5e-08 updt_s:0.177 data_s:0.000
INFO 2025-07-02 21:32:04 ts/train.py:241 Checkpoint policy after step 2
INFO 2025-07-02 21:32:28 ts/train.py:283 End of training

[YushunXiang]

@YushunXiang I tried your PR and I see high loss still:

Command run:
python lerobot/scripts/train.py --policy.type=pi0 --dataset.repo_id=lerobot/aloha_sim_insertion_human --policy.repo_id=kumarhans/test_pi0 --steps=2 --log_freq=1

INFO 2025-07-02 21:32:04 ts/train.py:232 step:1 smpl:8 ep:0 epch:0.00 loss:2.561 grdn:90.310 lr:5.0e-08 updt_s:1.109 data_s:0.414
INFO 2025-07-02 21:32:04 ts/train.py:232 step:2 smpl:16 ep:0 epch:0.00 loss:2.265 grdn:82.791 lr:7.5e-08 updt_s:0.177 data_s:0.000
INFO 2025-07-02 21:32:04 ts/train.py:241 Checkpoint policy after step 2
INFO 2025-07-02 21:32:28 ts/train.py:283 End of training

Can you give me the match information between the model and state dict.

kwargs["strict"] = True in Line 168, lerobot/common/policies/factory.py

[YushunXiang]

@YushunXiang I tried your PR and I see high loss still:

Command run:
python lerobot/scripts/train.py --policy.type=pi0 --dataset.repo_id=lerobot/aloha_sim_insertion_human --policy.repo_id=kumarhans/test_pi0 --steps=2 --log_freq=1

INFO 2025-07-02 21:32:04 ts/train.py:232 step:1 smpl:8 ep:0 epch:0.00 loss:2.561 grdn:90.310 lr:5.0e-08 updt_s:1.109 data_s:0.414
INFO 2025-07-02 21:32:04 ts/train.py:232 step:2 smpl:16 ep:0 epch:0.00 loss:2.265 grdn:82.791 lr:7.5e-08 updt_s:0.177 data_s:0.000
INFO 2025-07-02 21:32:04 ts/train.py:241 Checkpoint policy after step 2
INFO 2025-07-02 21:32:28 ts/train.py:283 End of training

Can you give me the match information between the model and state dict?

kwargs["strict"] = True in Line 168, lerobot/common/policies/factory.py

[kumarhans]

@YushunXiang I tried your PR and I see high loss still:
Command run:
python lerobot/scripts/train.py --policy.type=pi0 --dataset.repo_id=lerobot/aloha_sim_insertion_human --policy.repo_id=kumarhans/test_pi0 --steps=2 --log_freq=1

INFO 2025-07-02 21:32:04 ts/train.py:232 step:1 smpl:8 ep:0 epch:0.00 loss:2.561 grdn:90.310 lr:5.0e-08 updt_s:1.109 data_s:0.414
INFO 2025-07-02 21:32:04 ts/train.py:232 step:2 smpl:16 ep:0 epch:0.00 loss:2.265 grdn:82.791 lr:7.5e-08 updt_s:0.177 data_s:0.000
INFO 2025-07-02 21:32:04 ts/train.py:241 Checkpoint policy after step 2
INFO 2025-07-02 21:32:28 ts/train.py:283 End of training

Can you give me the match information between the model and state dict?

kwargs["strict"] = True in Line 168, lerobot/common/policies/factory.py

I added this code on top of your PR:

        # Load a pretrained policy and override the config if needed (for example, if there are inference-time
        # hyperparameters that we want to vary).
        kwargs["pretrained_name_or_path"] = cfg.pretrained_path
        kwargs["strict"] = True
        policy = policy_cls.from_pretrained(**kwargs)
    else:
        # Make a fresh policy.
        policy = policy_cls(**kwargs)

Where can I get the key mismatches?

[YushunXiang]

@kumarhans in line 313, src/lerobot/policies/pi0/modeling_pi0.py:

msg = model.load_state_dict(transformed_state_dict, strict=strict)
print(msg)

What's more, I think you may forget to add --policy.path /path/pi0 in your running command.

[OpenJarvisAI]

has anybody occured that pi0 result very bad ( not work at all ) issue?

[YushunXiang]

has anybody occured that pi0 result very bad ( not work at all ) issue?

Yes, after I solved this issue with #1415, I got a huge boost in the success rate of my pick and place task (10%->80%).

[OpenJarvisAI]

Hi, your shape mismatch why effect succes rate? In my opionion, when your shape is mismatch, you can't train at all.

Do u mean your model loaded and trained, but actually some weights not loaded? (shouldn't happen, I trained succesfully even with transformer 4.52.0)

[YushunXiang]

@OpenJarvisAI
My model loaded and trained, but actually some weights not loaded.
The program runs successfully because the strict parameter in safetensors.torch.load_model(model, model_file, strict=strict, device=map_location) defaults to False, which doesn't affect the program's running and doesn't output any warning or error information

lerobot/src/lerobot/policies/pretrained.py

Lines 129 to 143 in 63ddfef

	@classmethod
	def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
	if packaging.version.parse(safetensors.__version__) < packaging.version.parse("0.4.3"):
	load_model_as_safetensor(model, model_file, strict=strict)
	if map_location != "cpu":
	logging.warning(
	"Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
	" This means that the model is loaded on 'cpu' first and then copied to the device."
	" This leads to a slower loading time."
	" Please update safetensors to version 0.4.3 or above for improved performance."
	)
	model.to(map_location)
	else:
	safetensors.torch.load_model(model, model_file, strict=strict, device=map_location)
	return model

You can set strict=True and try again.