Views
No views yet
1from lerobot.common.policies.pi0fast.modeling_pi0fast import PI0FASTPolicy
2
3policy = PI0FASTPolicy.from_pretrained("Xiaoyan97/so100-pi0fast")
4policy.eval()
5device = "cuda" if torch.cuda.is_available() else "cpu"
6policy.to(device)
7while True:
8 start_time = time.perf_counter()
9 obs = robot.capture_observation()
10 for key in obs:
11 if "image" in key:
12 img = obs[key].type(torch.float32) / 255
13 obs[key] = img.permute(2, 0, 1).unsqueeze(0)
14 else:
15 obs[key] = obs[key].unsqueeze(0)
16 obs[key] = obs[key]
17 obs["task"] = [prompt]
18 with torch.no_grad():
19 action = policy.select_action(obs).squeeze(0).cpu().numpy()
20 robot.send_action(torch.tensor(action))
21 dt = time.perf_counter() - start_time
22 busy_wait(1 / 30 - dt)
23 print(f"Frame processed in {dt:.4f} seconds")