Can we fine-tune pi0 with a dataset of different robot configurations? I believe we can. Otherwise, pi0 is not cross-embodiment? But how?!

[dalishi]

Hi, I was trying to fine-tune the lerobot pi0_base model with lerobot/aloha_sim_transfer_cube_human dataset. The pi0 requires:
3 cameras: observation.images.base_0_rgb/left_wrist_0_rgb/right_wrist_0_rgb, type VISUAL, shape [224,224,3], and observation.state and action are shape 32.

However, lerobot/aloha_sim_transfer_cube_human dataset has only one single camera observation.images.top, and it is video (not image) of shape [480,640,3], and observation.state and action are of shape 14.

Following is the script and arguments I used to train (offline training):

lerobot-train \
  --policy.type=pi0 \
  --policy.pretrained_path=/workspace/hf_cache/hub/models--lerobot--pi0_base/snapshots/e3a5218ef7a5903445baf2cb656912fc35dc8712 \
  --policy.repo_id=lerobot/pi0_base \
  --policy.device=cuda \
  --policy.dtype=bfloat16 \
  --policy.n_obs_steps=1 \
  --policy.image_resolution='[224, 224]' \
  --policy.empty_cameras=2 \
  --policy.input_features='{"observation.images.top": {"shape":[3,224,224], "type":"VISUAL"}, "observation.state": {"shape":[14], "type":"STATE"}}' \
  --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
  --dataset.root=/workspace/hf_cache/hub/datasets--lerobot--aloha_sim_transfer_cube_human/snapshots/f20d5ec6357efe737b6441017b219f3f7c46766a \
  --batch_size=8 \
  --steps=2000 \
  --log_freq=10 \
  --save_freq=500 \
  --output_dir=outputs/pi0_finetune_aloha \
  --policy.push_to_hub=false`

The training process was able to run to the end except for one warning:

WARNING 2025-10-23 10:15:57 ing_pi0.py:1046 Vision embedding key might need handling: paligemma_with_expert.paligemma.model.vision_tower.vision_model.embeddings.patch_embedding.bias
WARNING 2025-10-23 10:15:57 ing_pi0.py:1046 Vision embedding key might need handling: paligemma_with_expert.paligemma.model.vision_tower.vision_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".

Finally, loss is around 0.165. I don't know whether what I have done is right or wrong. I just don't feel right. I attached the whole log for information. Could experts here advise me on the issues in my training? Thanks so much!

pi0 train log 3.txt

[robotopodia]

Hey, were you able to solve the problem