PI0 WARNING: Vision embedding key might need handling && WandB logging of key "loss_per_dim" was ignored

[Oooer8]

Ticket Type

🐛 Bug Report (Something isn't working)

Environment & System Info

lerobot=0.4.4

Description

Train PI0 as turorial.

Logs as follows:

WARNING 2026-03-08 14:47:13 ing_pi0.py:1121 Vision embedding key might need handling: paligemma_with_expert.paligemma.model.vision_tower.visi
on_model.embeddings.patch_embedding.bias                                                                                                     
WARNING 2026-03-08 14:47:13 ing_pi0.py:1121 Vision embedding key might need handling: paligemma_with_expert.paligemma.model.vision_tower.visi
on_model.embeddings.patch_embedding.weight                                                                                                   
Remapped 777 state dict keys                                                                                                                 
Warning: Could not remap state dict keys: Error(s) in loading state_dict for PI0Policy:                                                      
        Missing key(s) in state_dict: "model.paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight".                      

...

INFO 2026-03-08 14:59:15 ot_train.py:435 step:200 smpl:2K ep:5 epch:0.16 loss:0.375 grdn:15.012 lr:2.5e-06 updt_s:3.533 data_s:0.057
WARNING 2026-03-08 14:59:15 db_utils.py:181 WandB logging of key "loss_per_dim" was ignored as its type "<class 'list'>" is not handled by this wrapper.

There is the same question in PI05 #2307 (comment), but no one propose it in PI0 pipeline.

How soon it will be solved?

Context & Reproduction

No response

Relevant logs or stack trace

Checklist

• I have searched existing tickets to ensure this isn't a duplicate.
• I am using the latest version of the main branch.
• I have verified this is not an environment-specific problem.

Additional Info / Workarounds

No response

[he-yufeng]

Two separate bugs here:

1. embed_tokens.weight missing: the remap in _fix_pytorch_state_dict_keys only creates embed_tokens.weight from lm_head.weight (weight tying). If the checkpoint stores embed_tokens.weight directly, it passes through without the model. prefix mapping needed for that specific key path. Fix: explicitly handle embed_tokens.weight as a source key.

2. loss_per_dim dropped by WandB: the wrapper only handles scalar types. Fix: expand list/tuple metrics into per-index entries (loss_per_dim_0, loss_per_dim_1, etc.).

PR incoming.