Xiaoyan97/so100-pi0fast

Model Description

so100-pi0fast is the finetuned result of Hangzhou Theseus Yihai Rong Silicon Technology Co., LTD. based on the lerobot/pi0fast-base model and the anthonyav/so100-lego-v3 dataset

How to use

from lerobot.common.policies.pi0fast.modeling_pi0fast import PI0FASTPolicy

policy = PI0FASTPolicy.from_pretrained("Xiaoyan97/so100-pi0fast")
policy.eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
policy.to(device)
while True:
    start_time = time.perf_counter()
    obs = robot.capture_observation()
    for key in obs:
        if "image" in key:
            img = obs[key].type(torch.float32) / 255
            obs[key] = img.permute(2, 0, 1).unsqueeze(0)
        else:
            obs[key] = obs[key].unsqueeze(0)
            obs[key] = obs[key]
        obs["task"] = [prompt]
    with torch.no_grad():
        action = policy.select_action(obs).squeeze(0).cpu().numpy()
    robot.send_action(torch.tensor(action))
    dt = time.perf_counter() - start_time
    busy_wait(1 / 30 - dt)
    print(f"Frame processed in {dt:.4f} seconds")