Z ↑
|
| Y (left)
| ↗
| /
|/
└──────────→ X (forward)
$$T_{0\to6} = \prod_{i=1}^{6} R_z(\theta_i + o_i) \cdot T_z(d_i) \cdot T_x(a_i) \cdot R_x(\alpha_i)$$
1import torch
2from prismatic.models import load_vla
3from prismatic.vla.action_tokenizer import ACTION_TOKENIZERS
4from PIL import Image
5
6# 加载模型
7vla = load_vla("checkpoints/minivla_libero90/checkpoints/step-122500-epoch-55-loss=0.0743.pt",
8 hf_token="dummy", load_for_training=False)
9# 加载微调权重
10ckpt = torch.load("epoch1-minmax-smooth.pt", map_location="cpu")
11vla.load_state_dict(ckpt["model"], strict=False)
12vla = vla.cuda().eval()
13
14# 设置归一化参数
15vla.norm_stats = {"mycobot_task2": stats} # stats from dataset_statistics.json
16vla.action_tokenizer = ACTION_TOKENIZERS["extra_action_tokenizer"](vla.llm_backbone.tokenizer)
17
18# 推理
19image = Image.open("camera_view.jpg").convert("RGB")
20action = vla.predict_action(image,
21 "put the yellow block on the shelf onto the black table on the right",
22 unnorm_key="mycobot_task2")
23
24# action = [dx, dy, dz, droll, dpitch, dyaw, gripper] (7维)
25# dx,dy,dz: mm | droll,dpitch,dyaw: ° | gripper: [0,1]