LeRobot训练推理实操：25分钟搞定数据集增强

今天，和大家聊一聊如何使用 π₀.₅ (pi05) 策略在自定义数据集上进行训练和推理的完整操作。如果你手头上有至少600行的实验数据集，这份教程会帮助你轻松入门。

准备工作

1. 安装 LeRobot 和 π₀.₅ 依赖包

想要顺利进行训练，首先需要安装好相关工具和依赖：

# 安装 LeRobot
pip install lerobot

# 安装 pi05 所需的额外依赖
pip install -e ".[pi]"

# 若您的版本是 lerobot 0.4.0，用以下方式安装：
pip install "lerobot[pi]@git+https://github.com/huggingface/lerobot.git"

2. 设置 Hugging Face Hub 访问权限

Hugging Face Hub 是一个功能强大的平台，用于数据集和模型的上传与存储。在使用前，需要您先登录：

# 登录 Hugging Face Hub
huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential

# 获取你的 Hugging Face 用户名
HF_USER=$(hf auth whoami | head -n 1)
echo $HF_USER

3. 确认数据集格式

LeRobot 的数据集需要符合 v3.0 版本的标准格式，主要包括以下内容：

• 观察数据（Observations）：状态向量或摄像头图像；
• 动作数据（Actions）：机器人行为的矢量；
• 任务描述（Task Descriptions）：每个实验步骤的文本描述；
• 元数据（Metadata）：实验边界和统计信息。

特别注意：

• 数据集不能少于 600 行（frame）；
• 必须明确划分每段实验的边界；
• 不同实验的数据特征维度需要保持一致。

数据集准备

在开始训练之前，确保您的数据集包含训练所需的统计特征。π₀.₅ 策略依赖 分位数统计（quantile statistics）来进行标准化。如果您的数据集中缺少这个特征，需要手动补充。

第一步：对数据集进行统计增强

通过以下脚本，为数据集添加分位数统计（如 q01、q99）：

python src/lerobot/datasets/v30/augment_dataset_quantile_stats.py \
  --repo-id=${HF_USER}/your_dataset_name

这个工具会自动完成以下操作：

• 加载你的数据集；
• 计算状态特征和动作特征的分位数（如 q01、q10、q50、q90、q99）；
• 更新统计特征到元数据中。

如果不想使用分位数统计，还可以选择均值/标准差（MEAN_STD）标准化方法，我们将在训练过程中再详细讲解。

第二步：验证数据集统计信息

验证数据是否已添加分位数统计：

from lerobot.datasets.lerobot_dataset import LeRobotDataset

dataset = LeRobotDataset("${HF_USER}/your_dataset_name")
print(dataset.meta.stats)  # 应打印出 q01 和 q99 等统计信息

策略训练

完成数据集准备后，就可以开始策略的训练了。

第一步：基础训练指令

以下是一个用 π₀.₅ 策略对自定义数据集进行微调的完整指令：

lerobot-train \
  --dataset.repo_id=${HF_USER}/your_dataset_name \
  --policy.type=pi05 \
  --output_dir=./outputs/pi05_training \
  --job_name=pi05_training \
  --policy.repo_id=${HF_USER}/my_pi05_policy \
  --policy.pretrained_path=lerobot/pi05_base \
  --policy.compile_model=true \
  --policy.gradient_checkpointing=true \
  --wandb.enable=true \
  --policy.dtype=bfloat16 \
  --steps=3000 \
  --policy.device=cuda \
  --batch_size=32

第二步：关键参数解释

新媒网跨境认为，要想顺利完成策略训练，理解关键参数尤为重要。以下是一些参数的详细介绍：

• --dataset.repo_id：指向您的 Hugging Face 数据集的位置，例如 ${HF_USER}/your_dataset_name；
• --policy.type：策略类型，训练 π₀.₅ 时必须指定为 pi05；
• --policy.pretrained_path：预训练模型，可以选择 lerobot/pi05_base；
• --policy.compile_model：检测是否开启 PyTorch 模型编译以优化训练速度；
• --policy.gradient_checkpointing：是否开启梯度检查点来减少显存占用；
• --policy.dtype：训练时数据类型，推荐使用 bfloat16；
• --batch_size：批量训练大小，需根据显存调整；
• --steps：总训练步数，依据数据集具体规模灵活调整。

第三步：使用均值/标准差标准化的方案

如果您的数据集未添加分位数统计，可以在训练时改用 MEAN_STD 标准化：

lerobot-train \
  --dataset.repo_id=${HF_USER}/your_dataset_name \
  --policy.type=pi05 \
  --policy.normalization_mapping='{"ACTION": "MEAN_STD", "STATE": "MEAN_STD", "VISUAL": "IDENTITY"}' \
  --output_dir=./outputs/pi05_training \
  --job_name=pi05_training \
  --policy.repo_id=${HF_USER}/my_pi05_policy

推理与评估

完成训练后，下一步就是测试你的策略，看看它的实际表现如何。

第一种方式：在真实机器人上推理

通过以下指令，在物理机器人上运行您的策略：

lerobot-record \
  --robot.type=your_robot_type \
  --robot.port=/dev/ttyACM1 \
  --robot.cameras="{front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
  --robot.id=my_robot_id \
  --display_data=false \
  --dataset.repo_id=${HF_USER}/eval_pi05 \
  --dataset.single_task="Your task description" \
  --dataset.num_episodes=10 \
  --policy.path=${HF_USER}/my_pi05_policy

指令中的几个关键参数：

• --policy.path：您训练完成的策略路径，可以是本地路径，也可以是 Hugging Face 上的 repo ID；
• --dataset.repo_id：推理数据存储位置；
• --dataset.single_task：具体任务描述；
• --dataset.num_episodes：推理时需完成的实验次数。

第二种方式：通过 Python API 运行推理

如果您更倾向于用 Python 代码调用训练好的策略，可以参考以下代码：

from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.policies.pi05.modeling_pi05 import PI05Policy
from lerobot.robots.your_robot import YourRobot, YourRobotConfig

# 加载策略模型
policy = PI05Policy.from_pretrained("${HF_USER}/my_pi05_policy")

# 初始化机器人
robot_config = YourRobotConfig(port="/dev/ttyACM1", id="my_robot_id")
robot = YourRobot(robot_config)

# 开始推理
robot.connect()
policy.run(robot)  # 运行策略
robot.disconnect()

风险与常见问题

新媒网跨境预测，在使用 π₀.₅ 策略时，以下问题或许会偶尔碰到：

问题1：训练中提示 "QUANTILES normalization mode requires q01 and q99 stats"

解决方法：运行以下补充分位数统计的脚本：

python src/lerobot/datasets/v30/augment_dataset_quantile_stats.py \
  --repo-id=${HF_USER}/your_dataset_name

或者，直接在训练过程中选用 MEAN_STD 标准化方式（见上文）。

问题2：训练时内存不足（OOM）

解决方法：

1. 减少 batch size，比如降低到 --batch_size=16 或更低；
2. 确保开启 --policy.gradient_checkpointing=true；
3. 使用混合精度训练（--policy.dtype=bfloat16）。

总结

希望通过这篇教程，让大家对如何使用 π₀.₅ 策略在自定义数据集上进行训练和推理有更清晰的认知。新媒网跨境建议，在探索新技术的同时，多站在实际应用的角度考虑如何进一步优化操作流程。

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/lerobot-train-in-25-min-with-dataset.html