Enhance pi0 model inference

[xuaner233]

This is a simple change for pi0 model inference, along with minor fix for loss_dict in training part.

1. Pass task into observation for VLA model(pi0)
2. Update loss_dict stats data format.

What this does

Add inference support for pi0 model.

How it was tested

• first train pi0 with dataset, e.g.

python lerobot/scripts/train.py \
  --steps=40000 \
  --policy.type=pi0 \
  --dataset.repo_id=xuaner233/so100_grasp_place_20250313 \
  --wandb.enable=true \
  --wandb.disable_artifact=true

• then inference with trained pi0 model, set the control.single_task as pi0's text prompt for task:

HF_USER=xuaner233
REPO_ID="${HF_USER}/eval_pi0_so100_test"

python lerobot/scripts/control_robot.py \
  --robot.type=so100 \
  --control.type=record \
  --control.fps=30 \
  --control.single_task="Grasp a white cube and put it in the bin." \
  --control.repo_id=${REPO_ID} \
  --control.tags='["pi0"]' \
  --control.warmup_time_s=5 \
  --control.episode_time_s=300 \
  --control.reset_time_s=10 \
  --control.num_episodes=1 \
  --control.push_to_hub=false \
  --control.policy.device=cuda \
  --control.policy.path=outputs/train/2025-03-14_pi0/checkpoints/last/pretrained_model

[imstevenpmwork]

I left a comment in code

[imstevenpmwork]

What is the reasoning of getting the mean here?

Wouldn't it be better to use .detach() in here instead of clone()?

[xuaner233]

Hi Steven, thanks for the review!

mean() is to align the value of losses in loss_dict, which means these value need to align with the loss_dict["l2_loss"] aka the single mean loss value. Otherwise wandb would complain as below:

WARNING 2025-03-26 09:52:43 db_utils.py:116 WandB logging of key "losses_after_forward" was ignored as its type is not handled by this wrapper.
WARNING 2025-03-26 09:52:43 db_utils.py:116 WandB logging of key "losses_after_rm_padding" was ignored as its type is not handled by this wrapper.

As for the removal of clone() is because: instead of adding the actual loss "data", just calcuate mean() seems doesn't need the clone() or detach() for this?