PI0.5 Insert Carrot — Epoch 3

这是使用 LeRobot 对 PI0.5 全参数微调得到的 epoch 3 checkpoint(step 2355),任务为 insert carrot into the hole

兼容版本

建议使用训练时的官方 LeRobot commit:

git clone https://github.com/huggingface/lerobot.git
cd lerobot
git checkout e40b58a8dfa9e7b86918c374791599d070518d11
uv sync --extra pi --extra feetech

模型为 BF16,建议在支持 BF16 的 NVIDIA GPU 上运行。仓库中的 preprocessor 和 postprocessor 是模型的一部分,必须一起加载;不要只加载 model.safetensors

在 SO-100/SO-101 上推理

下面示例假设机械臂是 SO-100、两个 OpenCV 相机分别对应 frontside。请把串口和视频设备改成实际值;若使用 SO-101,将 so100_follower 改为 so101_follower

uv run lerobot-rollout \
  --strategy.type=base \
  --inference.type=rtc \
  --inference.rtc.execution_horizon=10 \
  --inference.rtc.max_guidance_weight=10.0 \
  --policy.path=Aikwed/pi0.5-insert-carrot-epoch-3 \
  --robot.type=so100_follower \
  --robot.port=/dev/ttyACM0 \
  --robot.cameras="{front: {type: opencv, index_or_path: /dev/video0, width: 640, height: 480, fps: 30}, side: {type: opencv, index_or_path: /dev/video2, width: 640, height: 480, fps: 30}}" \
  --task="insert carrot into the hole" \
  --fps=30 \
  --duration=60 \
  --device=cuda

必须使用 --inference.type=rtc。该 checkpoint 使用相对动作,而此 LeRobot 版本的同步 inference engine 不支持相对动作策略;RTC 会正确处理 action chunk 及相对动作的重新锚定。

输入输出约定

输入 observation:

  • observation.state: 6 维 float32,顺序为 shoulder_pan.pos, shoulder_lift.pos, elbow_flex.pos, wrist_flex.pos, wrist_roll.pos, gripper.pos
  • observation.images.front: 640×480 RGB。
  • observation.images.side: 640×480 RGB。
  • 文本任务:insert carrot into the hole

训练数据中前五个关节使用角度制,gripper 使用 0–100 标度。推理时必须沿用训练数据的关节顺序、标度、机械臂标定和相机视角。

输出是 6 维绝对关节目标。模型内部预测前五个关节的相对变化量,postprocessor 会自动加回当前 state;gripper 始终是绝对值。不要在机器人端再次手动累加前五维。

Python 单个 action chunk 示例

下面示例用于检查模型加载和输出。它一次性对整个 50-step chunk 做后处理,因此相对动作只锚定到本次输入 state。实时部署建议使用上面的 RTC rollout。

import numpy as np
import torch

from lerobot.configs import PreTrainedConfig
from lerobot.policies import (
    get_policy_class,
    make_pre_post_processors,
    prepare_observation_for_inference,
)

model_id = "Aikwed/pi0.5-insert-carrot-epoch-3"
device = torch.device("cuda")

cfg = PreTrainedConfig.from_pretrained(model_id)
cfg.device = str(device)
policy = get_policy_class(cfg.type).from_pretrained(model_id, config=cfg).to(device).eval()
preprocessor, postprocessor = make_pre_post_processors(
    policy_cfg=cfg,
    pretrained_path=model_id,
    preprocessor_overrides={"device_processor": {"device": str(device)}},
)

# 每个 episode 开始时重置缓存。
policy.reset()
preprocessor.reset()
postprocessor.reset()

# front_rgb 和 side_rgb: np.uint8, shape (480, 640, 3)
# joint_state: np.float32, shape (6,)
raw_observation = {
    "observation.state": joint_state.astype(np.float32),
    "observation.images.front": front_rgb.astype(np.uint8),
    "observation.images.side": side_rgb.astype(np.uint8),
}

observation = prepare_observation_for_inference(
    raw_observation,
    device=device,
    task="insert carrot into the hole",
    robot_type="so_follower",
)
observation = preprocessor(observation)

with torch.inference_mode():
    relative_chunk = policy.predict_action_chunk(observation)  # (1, 50, 6)
    absolute_chunk = postprocessor(relative_chunk).squeeze(0)  # (50, 6), CPU

first_absolute_action = absolute_chunk[0]
print(first_absolute_action)

首次运行还会从 Hugging Face 下载 PaliGemma tokenizer。真实机械臂测试前请确认急停、关节限位和相机名称,并先以低风险姿态短时运行。

Downloads last month
13
Safetensors
Model size
4B params
Tensor type
F32
·
BF16
·
Video Preview
loading