LeRobot入门教程 10:LeRobot 策略推理

目录:00_LeRobot入门教程_目录.md

1. 从「随机乱动」到「会干活」

上一册机械臂在乱戳,因为动作是随机采样的。这一册换成真正训练过的策略模型。

用的模型是 lerobot/pi0_libero_finetuned——π0 是 Physical Intelligence 提出的视觉-语言-动作(VLA)模型,把「看画面 + 读指令」直接映射成「下一步动作」。Hub 上这个权重是在 LIBERO 数据上微调过的版本。

VLA 和普通图像分类模型的区别:

普通模型VLA 策略
输入一张图图像 + 机器人状态 + 语言指令
输出类别动作序列(action chunk)
使用方式单次前向闭环:观察 → 出动作 → 执行 → 再观察

2. 三段式:模型、策略侧处理器、环境侧适配器

LeRobot 把「模型」和「数据搬运」拆开了。一次推理要串起三个东西:

组件创建方式负责什么
策略模型make_policy(cfg, env_cfg=...)真正的神经网络,输入张量、输出动作
策略侧预处理/后处理make_pre_post_processors(cfg, ...)归一化、图像缩放、动作反归一化(un-normalize)
环境侧适配器make_env_pre_post_processors(...)把仿真环境给的数据格式变成策略要的格式,再把动作变回环境能吃的格式

3. 加载模型和环境

import torch
from lerobot.envs import make_env, make_env_pre_post_processors, preprocess_observation
from lerobot.envs.configs import LiberoEnv
from lerobot.policies import make_policy, make_pre_post_processors
from lerobot.configs.policies import PreTrainedConfig
from lerobot.utils.constants import ACTION

MODEL = "lerobot/pi0_libero_finetuned"
env_cfg = LiberoEnv(task="libero_spatial", task_ids=[0], observation_height=256, observation_width=256)
env = list(list(make_env(env_cfg, n_envs=1).values())[0].values())[0]
cfg = PreTrainedConfig.from_pretrained(MODEL)
cfg.pretrained_path = MODEL
cfg.device = "cuda"
policy = make_policy(cfg, env_cfg=env_cfg).eval()
pre, post = make_pre_post_processors(cfg, pretrained_path=MODEL,
                                     preprocessor_overrides={"device_processor": {"device": "cuda"}})
env_pre, env_post = make_env_pre_post_processors(env_cfg=env_cfg, policy_cfg=cfg)
Creating LIBERO envs | suites=['libero_spatial'] | n_envs(per task)=1 | init_states=True
Restricting to task_ids=[0]
[info] using task orders [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
Built vec env | suite=libero_spatial | task_id=0 | n_envs=1
The PI0 model is a direct port of the OpenPI implementation. 
This implementation follows the original OpenPI structure for compatibility. 
Original implementation: https://github.com/Physical-Intelligence/openpi
Loading model from: lerobot/pi0_libero_finetuned

stderr:

WARNING:root:Vision embedding key might need handling: model.paligemma_with_expert.paligemma.model.vision_tower.vision_model.embeddings.patch_embedding.bias

stderr:

WARNING:root:Vision embedding key might need handling: model.paligemma_with_expert.paligemma.model.vision_tower.vision_model.embeddings.patch_embedding.weight
✓ Loaded state dict from model.safetensors
All keys loaded successfully!

几个关键行:

  • PreTrainedConfig.from_pretrained(MODEL) 读的是 Hub 上的 config.json,它决定模型结构(层数、动作维度、图像尺寸),结构必须和权重匹配;
  • cfg.pretrained_path = MODEL 告诉后续加载去这里找权重文件;
  • cfg.device = "cuda",同时 preprocessor_overrides={"device_processor": {"device": "cuda"}} 让预处理也把张量搬到 GPU——两处都要设,只设一处会导致张量在 CPU/GPU 之间来回拷;
  • .eval() 切到推理模式,关掉 dropout 之类的训练期行为;
  • make_env_pre_post_processors(env_cfg=env_cfg, policy_cfg=cfg) 需要两边的信息:环境提供字段名和形状,策略提供归一化统计量。

输出里那两行 WARNING: Vision embedding key might need handling 是加载视觉主干时的提示,末尾的 All keys loaded successfully! 才是结论——权重全部加载成功。

4. 闭环推理:一次观察换一个动作

obs, info = env.reset()
frames = []
for _ in range(20):
    o = preprocess_observation(obs)
    o["task"] = list(env.call("task_description"))
    o = pre(env_pre(o))
    with torch.inference_mode():
        a = post(policy.select_action(o))
    obs, r, *_ = env.step(env_post({ACTION: a})[ACTION].cpu().numpy())
    frames.append(obs["pixels"]["image"][0][::-1].copy())
show(frames)
env.close()

10_LeRobot策略推理_0.gif

5. 把数据流拆开看

这一格是全套教程里信息密度最高的一段,逐行拆:

env.reset()                    → obs(仿真原始格式)
   ↓ preprocess_observation(obs)   → 统一的键名/布局
   ↓ o["task"] = ...               → 补上语言指令(策略的输入之一)
   ↓ env_pre(o)                    → 环境侧适配(形状、归一化)
   ↓ pre(o)                        → 策略侧预处理(图像缩放、设备搬运)
   ↓ policy.select_action(o)       → 神经网络前向
   ↓ post(a)                       → 动作反归一化
   ↓ env_post({ACTION: a})[ACTION] → 变回环境能执行的格式
   ↓ env.step(...)                 → 机械臂动一步

几个容易忽略的点:

细节说明
o["task"] = list(env.call("task_description"))指令必须喂给策略。少了这一行,模型不知道要干什么,动作会退化成无意义抖动
torch.inference_mode()比 no_grad() 更彻底地关掉梯度记录,省显存也更快。推理一律用它
post(policy.select_action(o))模型原始输出是归一化空间里的动作,post 负责映射回真实控制量。漏掉 post 机械臂会疯动
env_post(...)[ACTION].cpu().numpy()仿真只吃 numpy,张量留在 GPU 上会报错,所以 .cpu()
for _ in range(20)只跑 20 步是为了出 GIF;真实评测要跑到 episode 结束(见 11 册《评测》)

跑完会看到机械臂朝杯子方向移动,明显比上一册的随机乱动「有目的」——但 20 步通常做不完整个任务,所以画面不会出现「任务完成」。

6. 常见问题

现象原因处理
CUDA out of memoryπ0 较大,显存吃紧;或没关环境先 env.close(),再确认没有别的进程占卡;必要时用小 batch
机械臂动作剧烈抖动漏了 post() 反归一化检查是否写成 policy.select_action(o) 直接 env.step
模型完全不响应指令漏了 o["task"]补上 task_description
加载很慢/卡住没走加速通道回 07 册《环境与网络》
shape mismatch环境分辨率和模型训练分辨率不一致保持 observation_height=256