lerobot/pi0 · The correct script to load this dataset

Apr 17

from lerobot.common.policies.pi0.modeling_pi0 import PI0Policy
import torch

# Load the pretrained policy from HF
policy = PI0Policy.from_pretrained("lerobot/pi0")

# Prepare a sample batch (obs, goal) for testing
batch = {
    "image": torch.randn(1, 3, 224, 224),  # Replace with real image
    "goal": ["scrub the plate"]
    }

# Forward pass
action = policy.select_action(batch)

print("Predicted joint actions:", action)

Note the uppercase 'I' in PIOPolicy , in the model card the case is Pi0Policy which is incorrect.

amir06

Apr 18

hello i am using the xArm environnement but i dont know what i have to pus in batch can you help me ?

amir06

Apr 18

File "C:\Users\SOUICI amir\lerobot\lerobot\common\policies\pi0\modeling_pi0.py", line 349, in prepare_images
raise ValueError(
ValueError: All image features are missing from the batch. At least one expected. (batch: dict_keys(['image', 'goal'])) (image_features:{})