第五章：LeRobot策略训练与评估

在掌握了LeRobot的数据处理和模型使用基础之后，本章将深入探索LeRobot的核心能力：策略训练。策略训练将数据转化为智能行为，将算法理论转化为实际应用。LeRobot提供了完整的训练工具链，从简单的Python脚本训练到强大的命令行工具，为用户提供了灵活而强大的训练工具。

5.1 策略训练流程

理解训练流程的核心步骤与关键参数配置，是掌握LeRobot训练能力的基础。

5.1.1 训练流程解析

LeRobot官方提供的3_train_policy.py示例展示了完整的策略训练流程。这个示例使用Diffusion Policy在PushT数据集上进行训练，代码结构清晰，注释详细，是学习LeRobot训练系统的最佳起点。并且训练完成后可以使用2_evaluate_pretrained_policy.py进行评估。这种设计体现了LeRobot系统的模块化思想：训练和评估分离，便于独立使用和组合。

#!/usr/bin/env python

"""
基于LeRobot官方3_train_policy.py示例的完整训练流程演示。

这个脚本演示了如何使用LeRobot进行策略训练，包括：
1. 环境和依赖的正确导入
2. 训练配置的设置
3. 数据集的加载和预处理
4. 策略模型的初始化
5. 训练循环的实现
6. 训练过程的监控和日志记录
"""

from pathlib import Path
import torch

# LeRobot核心组件导入
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
from lerobot.common.datasets.utils import dataset_to_policy_features
from lerobot.common.policies.diffusion.configuration_diffusion import DiffusionConfig
from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy
from lerobot.configs.types import FeatureType

def main():
    """主训练函数"""
    
    print("=== LeRobot策略训练开始 ===")
    
    # 第一步：创建输出目录
    # 这是训练过程中保存模型检查点、日志和配置文件的地方
    output_directory = Path("outputs/train/example_pusht_diffusion")
    output_directory.mkdir(parents=True, exist_ok=True)
    print(f"输出目录: {output_directory}")
    
    # 第二步：设备选择和配置
    # LeRobot支持CUDA、CPU和MPS（Mac）设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    
    if torch.cuda.is_available():
        print(f"GPU信息: {torch.cuda.get_device_name()}")
        print(f"GPU内存: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
    
    # 第三步：训练参数设置
    # 这些参数控制训练的基本行为
    training_steps = 100  # 总训练步数，可根据需要调整
    log_freq = 1          # 日志记录频率，每步都记录
    
    print(f"训练步数: {training_steps}")
    print(f"日志频率: 每 {log_freq} 步")

#output
== LeRobot策略训练开始 ===
输出目录: outputs/train/example_pusht_diffusion
使用设备: cuda
GPU信息: NVIDIA RTX 4000 Ada Generation Laptop GPU
GPU内存: 12.5 GB
训练步数: 100
日志频率: 每 1 步

在这个初始化阶段，我们可以看到LeRobot训练系统的几个重要特点。首先是输出目录的管理，LeRobot会自动创建必要的目录结构来保存训练过程中的各种文件。其次是设备的自动检测和配置，系统会优先使用GPU，但也能够在CPU上运行，这为不同硬件环境的用户提供了灵活性。

接下来是数据集的加载和特征处理，这是训练流程中最关键的环节之一：

    # 第四步：数据集元数据加载
    # LeRobot使用元数据系统来管理数据集的结构信息
    print("\n=== 数据集配置 ===")
    dataset_metadata = LeRobotDatasetMetadata("lerobot/pusht")
    print(f"数据集: {dataset_metadata.repo_id}")
    print(f"总帧数: {dataset_metadata.total_frames}")
    print(f"总轨迹数: {dataset_metadata.total_episodes}")
    
    # 第五步：特征提取和分类
    # 将数据集特征转换为策略可以使用的格式
    features = dataset_to_policy_features(dataset_metadata.features)
    
    # 区分输入特征和输出特征
    # 输出特征通常是动作（ACTION类型）
    output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
    # 输入特征是除了动作之外的所有特征（观察、状态等）
    input_features = {key: ft for key, ft in features.items() if key not in output_features}
    
    print(f"输入特征: {list(input_features.keys())}")
    print(f"输出特征: {list(output_features.keys())}")
    
    # 显示特征的详细信息
    for key, feature in input_features.items():
        print(f"  {key}: {feature.shape}, dtype: {feature.type}")
    for key, feature in output_features.items():
        print(f"  {key}: {feature.shape}, dtype: {feature.type}")

#output
=== 数据集配置 ===
数据集: lerobot/pusht
总帧数: 25650
总轨迹数: 206
输入特征: ['observation.image', 'observation.state']
输出特征: ['action']
  observation.image: (3, 96, 96), dtype: VISUAL
  observation.state: (2,), dtype: STATE
  action: (2,), dtype: ACTION

这个特征处理过程展示了LeRobot数据系统的强大之处。系统能够自动识别数据集中的不同类型特征，并将它们正确地分类为输入和输出。这种自动化的特征管理大大简化了用户的工作，同时也确保了数据处理的一致性和正确性。

策略模型的初始化是训练流程的核心环节：

    # 第六步：策略配置和初始化
    print("\n=== 策略模型配置 ===")
    
    # 创建Diffusion Policy配置
    # DiffusionConfig包含了模型架构的所有参数
    cfg = DiffusionConfig(
        input_features=input_features, 
        output_features=output_features
    )
    
    print(f"策略类型: Diffusion Policy")
    print("输入特征:")
    for name, feature in cfg.input_features.items():
        print(f"  {name}: 形状={feature.shape}")
    print("输出特征:")
    for name, feature in cfg.output_features.items():
        print(f"  {name}: 形状={feature.shape}")
    
    # 初始化策略模型
    # dataset_metadata.stats包含了数据的统计信息，用于归一化
    policy = DiffusionPolicy(cfg, dataset_stats=dataset_metadata.stats)
    
    # 设置为训练模式并移动到指定设备
    policy.train()
    policy.to(device)
    
    # 显示模型参数信息
    total_params = sum(p.numel() for p in policy.parameters())
    trainable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad)
    print(f"总参数数量: {total_params:,}")
    print(f"可训练参数: {trainable_params:,}")

#output
=== 策略模型配置 ===
WARNING:root:Device 'None' is not available. Switching to 'cuda'.
策略类型: Diffusion Policy
输入特征:
  observation.image: 形状=(3, 96, 96)
  observation.state: 形状=(2,)
输出特征:
  action: 形状=(2,)
总参数数量: 262,709,044
可训练参数: 262,709,026

这里我们可以看到LeRobot策略系统的设计哲学。配置类（DiffusionConfig）将模型的所有参数集中管理，使得模型的创建和修改变得简单直观。同时，数据统计信息的自动集成确保了模型能够正确处理输入数据的归一化，这对训练的稳定性至关重要。

数据加载器的配置体现了LeRobot对时序数据处理的深度理解：

    # 第七步：时间戳配置
    # 这是LeRobot的一个重要特性：支持多时间步的观察和动作
    print("\n=== 数据加载配置 ===")
    
    delta_timestamps = {
        # 观察数据：使用当前和前一时刻的图像
        "observation.image": [-0.1, 0.0],
        # 状态数据：同样使用当前和前一时刻
        "observation.state": [-0.1, 0.0],
        # 动作数据：使用未来多个时间步的动作序列
        "action": [-0.1, 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4],
    }
    
    print("时间戳配置:")
    for key, timestamps in delta_timestamps.items():
        print(f"  {key}: {len(timestamps)} 个时间步 {timestamps}")
    
    # 实例化数据集
    dataset = LeRobotDataset("lerobot/pusht", delta_timestamps=delta_timestamps)
    print(f"数据集大小: {len(dataset)} 个样本")
    
    # 创建优化器
    optimizer = torch.optim.Adam(policy.parameters(), lr=1e-4)
    print(f"优化器: Adam, 学习率: 1e-4")
    
    # 创建数据加载器
    dataloader = torch.utils.data.DataLoader(
        dataset,
        num_workers=4,        # 多进程数据加载
        batch_size=64,        # 批次大小
        shuffle=True,         # 随机打乱数据
        pin_memory=device.type != "cpu",  # 内存固定（GPU加速）
        drop_last=True,       # 丢弃最后不完整的批次
    )
    
    print(f"数据加载器配置:")
    print(f"  批次大小: 64")
    print(f"  工作进程: 4")
    print(f"  内存固定: {device.type != 'cpu'}")

#output
=== 数据加载配置 ===
时间戳配置:
  observation.image: 2 个时间步 [-0.1, 0.0]
  observation.state: 2 个时间步 [-0.1, 0.0]
  action: 16 个时间步 [-0.1, 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4]
Resolving data files: 100%|████████████████████████████████████████████████████████████████████████████████████| 206/206 [00:00<00:00, 63391.54it/s]
数据集大小: 25650 个样本
优化器: Adam, 学习率: 1e-4
数据加载器配置:
  批次大小: 64
  工作进程: 4
  内存固定: True

时间戳配置是LeRobot的一个独特特性，它允许模型同时访问过去的观察和未来的动作序列。这种设计对于机器人控制任务特别重要，因为机器人需要基于历史信息预测未来的动作序列。动作序列的长度（16个时间步）体现了Diffusion Policy对长期规划的支持。

最后是训练循环的实现，这是整个训练流程的核心：

    # 第八步：训练主循环
    print("\n=== 开始训练 ===")
    
    step = 0
    done = False
    
    while not done:
        for batch in dataloader:
            # 将数据移动到指定设备
            batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) 
                    for k, v in batch.items()}
            
            # 前向传播：计算损失
            loss, _ = policy.forward(batch)
            
            # 反向传播
            loss.backward()
            
            # 优化器步骤
            optimizer.step()
            optimizer.zero_grad()
            
            # 日志记录
            if step % log_freq == 0:
                print(f"步骤: {step:5d} | 损失: {loss.item():.6f}")
            
            step += 1
            
            # 检查是否达到训练步数
            if step >= training_steps:
                done = True
                break
    
    print(f"\n=== 训练完成 ===")
    print(f"总训练步数: {step}")
    print(f"最终损失: {loss.item():.6f}")
    
    # 保存模型
    model_path = output_directory / "final_model.pth"
    torch.save(policy.state_dict(), model_path)
    print(f"模型已保存到: {model_path}")

if __name__ == "__main__":
    main()

#output
=== 开始训练 ===
步骤:     0 | 损失: 1.222766
步骤:     1 | 损失: 6.313533
……
步骤:    99 | 损失: 0.126722

=== 训练完成 ===
总训练步数: 100
最终损失: 0.126722
模型已保存到: outputs/train/example_pusht_diffusion/final_model.pth

这个训练循环的实现简洁而完整，体现了PyTorch训练的标准流程。值得注意的是，LeRobot的策略类已经封装了复杂的前向传播逻辑，用户只需要调用policy.forward(batch)就能获得损失值，这大大简化了训练代码的编写。

5.1.2 训练过程深度理解

训练过程中的关键概念需要进一步加深理解。

• 首先是时间戳配置的作用机制。LeRobot通过delta_timestamps参数控制模型能够访问的时间窗口。负值表示过去的时间步，正值表示未来的时间步。这种设计允许模型同时利用历史信息和未来规划，这对于复杂的机器人控制任务是必要的。
• 其次是批次处理的优化策略。LeRobot的数据加载器使用了多项优化技术：多进程数据加载减少了I/O等待时间，内存固定（pin_memory）加速了CPU到GPU的数据传输，数据打乱（shuffle）提高了训练的随机性和泛化能力。
• 最后是模型的设备管理。LeRobot自动处理模型和数据在不同设备间的移动，用户无需手动管理这些细节。这种自动化的设备管理既提高了代码的可移植性，也减少了出错的可能性。

5.2 命令行训练工具深度实践

在LeRobot框架中，命令行训练工具是一个强大而灵活的功能，它允许研究人员和开发者通过简单的命令行操作来配置、启动、监控和管理复杂的模型训练过程。本节将深入介绍LeRobot的命令行训练系统，帮助读者掌握如何高效地使用这一工具进行模型训练和实验管理。

5.2.1 命令行训练系统概览

LeRobot的命令行训练系统以lerobot/scripts/train.py脚本为核心，该脚本提供了一个完整的训练流水线，包括配置加载、数据集准备、环境初始化、策略模型创建、训练循环执行以及结果保存等功能。通过这个脚本，用户可以轻松地训练各种机器人学习策略，而无需编写大量的代码。

训练脚本的核心功能：

1. 配置管理：加载和解析训练配置，支持命令行参数和配置文件
2. 数据集处理：实例化数据集并准备数据加载器
3. 环境初始化：（可选）创建对应的仿真环境用于评估
4. 策略模型创建：根据配置初始化策略模型
5. 训练循环执行：执行前向传播、反向传播、优化步骤等
6. 日志记录：记录训练过程中的各种指标
7. 模型评估：定期在环境中评估策略性能
8. 检查点保存：保存模型和训练状态

使用命令行训练工具的基本流程如下：

python lerobot/scripts/train.py [配置参数]

其中，[配置参数]可以是各种命令行选项，用于指定数据集、策略类型、环境设置等。这种方式使得用户可以在不修改代码的情况下，灵活地配置和控制训练过程。

5.2.2 配置系统深度解析

LeRobot的命令行训练工具基于Draccus配置系统，这是一个专为机器学习实验设计的配置管理工具。与传统的配置系统相比，Draccus提供了更严格的类型检查、更好的IDE支持和更灵活的参数覆盖机制。

配置层次结构：

训练脚本的主函数train接受一个TrainPipelineConfig对象作为参数：

@parser.wrap()
def train(cfg: TrainPipelineConfig):
    # 训练逻辑

TrainPipelineConfig是一个包含多个子配置的复合配置类，定义在lerobot/configs/train.py中：

@dataclass
class TrainPipelineConfig:
    dataset: DatasetConfig                  # 数据集配置
    env: envs.EnvConfig | None = None       # 环境配置（可选）
    policy: PreTrainedConfig | None = None  # 策略配置（可选）
    # 其他配置项...

这种层次化的配置结构使得用户可以通过点表示法来访问和设置深层次的配置项。例如，要设置数据集的仓库ID，可以使用--dataset.repo_id=lerobot/pusht。

配置方式：

LeRobot的命令行训练工具支持三种主要的配置方式：

1. 命令行参数：直接在命令行中指定配置项

   python lerobot/scripts/train.py --dataset.repo_id=lerobot/pusht --policy.type=diffusion
   

2. 配置文件：从JSON或YAML配置文件加载配置

   python lerobot/scripts/train.py --config_path=path/to/config.json
   

3. 混合方式：先从配置文件加载，然后用命令行参数覆盖特定项

   python lerobot/scripts/train.py --config_path=path/to/config.json --policy.learning_rate=1e-4
   

特殊配置参数：

1. 类型选择参数：使用.type后缀选择特定类型的组件

   python lerobot/scripts/train.py --policy.type=act --env.type=aloha
   

2. 路径参数：使用.path后缀指定预训练模型的路径

   python lerobot/scripts/train.py --policy.path=lerobot/act_aloha_sim_transfer_cube_human
   

3. 配置文件路径：使用--config_path指定配置文件位置

   python lerobot/scripts/train.py --config_path=outputs/train/my_experiment/checkpoints/last/pretrained_model/
   

5.2.3 实用训练命令示例

下面我们将通过一系列实用的命令示例，展示如何使用LeRobot的命令行训练工具进行不同类型的训练任务。

示例1：从头开始训练策略

python lerobot/scripts/train.py \
    --dataset.repo_id=lerobot/pusht \
    --policy.type=diffusion \
    --env.type=pusht \
    --output_dir=outputs/train/diffusion_pusht

这个命令将在PushT环境上训练一个Diffusion Policy，使用lerobot/pusht数据集，并将输出保存到outputs/train/diffusion_pusht目录。

示例2：在不同任务上训练ACT策略

python lerobot/scripts/train.py \
    --policy.type=act \
    --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
    --env.type=aloha \
    --env.task=AlohaTransferCube-v0 \
    --output_dir=outputs/train/act_aloha_transfer

这个命令将在ALOHA环境的立方体传输任务上训练一个ACT策略，使用lerobot/aloha_sim_transfer_cube_human数据集。

示例3：从配置文件加载训练设置

python lerobot/scripts/train.py \
    --config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
    --output_dir=outputs/train/act_aloha_transfer_2

这个命令将从之前训练运行的配置文件加载设置，并启动一个新的训练运行，输出保存到新的目录。

示例4：恢复中断的训练

python lerobot/scripts/train.py \
    --config_path=outputs/train/my_experiment/checkpoints/last/pretrained_model/ \
    --resume=true

这个命令将从最后一个检查点恢复之前中断的训练，继续训练过程。

示例5：延长训练步数

python lerobot/scripts/train.py \
    --config_path=outputs/train/my_experiment/checkpoints/last/pretrained_model/ \
    --resume=true \
    --steps=200000

这个命令将从最后一个检查点恢复训练，并将总训练步数增加到200,000步。

示例6：微调预训练策略

python lerobot/scripts/train.py \
    --policy.path=lerobot/act_aloha_sim_transfer_cube_human \
    --dataset.repo_id=lerobot/aloha_sim_insertion_human \
    --env.type=aloha \
    --env.task=AlohaInsertion-v0

这个命令将加载一个在ALOHA立方体传输任务上预训练的ACT策略，并在插入任务上进行微调。

5.2.4 训练输出与结果管理

LeRobot的命令行训练工具会生成多种输出，帮助用户跟踪训练进度、分析性能并保存训练结果。

训练日志：

训练过程中，脚本会在终端输出详细的日志信息，例如：

INFO 2025-06-13 13:35:12 ts/train.py:192 step:0 smpl:64 ep:1 epch:0.00 loss:1.112 grdn:15.387 lr:2.0e-07 updt_s:1.738 data_s:4.774

这些日志包含了丰富的训练信息，如当前步数、样本数、损失值、梯度范数、学习率等。评估日志则提供了策略在环境中的表现：

INFO 2025-06-13 13:38:45 ts/train.py:226 step:100 smpl:6K ep:52 epch:0.25 ∑rwrd:20.693 success:0.0% eval_s:120.266

日志指标解读：

• step: 当前训练步数
• smpl: 已处理的样本数
• ep: 已处理的数据集中的回合数
• epch: 已完成的训练轮数
• loss: 当前批次的损失值
• grdn: 梯度范数
• lr: 当前学习率
• updt_s: 网络参数更新时间（秒）
• data_s: 数据加载时间（秒）
• ∑rwrd: 评估回合的平均累积奖励
• success: 评估回合的成功率
• eval_s: 评估时间（秒）

检查点结构：

训练过程中，脚本会定期保存检查点到输出目录的checkpoints子目录中：

outputs/train/my_experiment/checkpoints/
├── 000100                          # 训练步数为100的检查点
│   ├── pretrained_model/           # 预训练模型目录
│   │   ├── config.json             # 策略配置
│   │   ├── model.safetensors       # 模型权重
│   │   └── train_config.json       # 训练配置
│   └── training_state/             # 训练状态目录
│       ├── optimizer_param_groups.json  # 优化器参数组
│       ├── optimizer_state.safetensors  # 优化器状态
│       ├── rng_state.safetensors        # 随机数生成器状态
│       ├── scheduler_state.json         # 学习率调度器状态
│       └── training_step.json           # 训练步数
├── 000200                          # 训练步数为200的检查点
└── last -> 000200                  # 指向最新检查点的符号链接

这种结构使得用户可以轻松地恢复训练、分析中间结果或使用特定检查点进行评估。

性能优化建议：

通过观察训练日志中的性能指标，用户可以识别并解决潜在的性能瓶颈：

1. 如果data_s（数据加载时间）过长，可以考虑：

   • 增加数据加载器的工作线程数（--dataloader.num_workers）
   • 调整批量大小（--train.batch_size）
   • 检查数据集存储位置，可能需要移至更快的存储设备

2. 如果GPU利用率低（可通过nvidia-smi命令查看），可以考虑：

   • 增加批量大小
   • 启用混合精度训练（--policy.use_amp=true）
   • 调整模型架构或优化器设置

3. 对于更详细的性能分析，可以使用PyTorch Profiler：

   python lerobot/scripts/train.py --profile=true [其他参数]
   

5.2.5 高级训练配置与优化

除了基本的训练配置外，LeRobot的命令行训练工具还提供了多种高级配置选项，用于优化训练过程和提高模型性能。

学习率调整：

python lerobot/scripts/train.py \
    --dataset.repo_id=lerobot/pusht \
    --policy.type=diffusion \
    --optim.lr=5e-5 \
    --optim.weight_decay=0.01

这个命令设置了初始学习率为5e-5，权重衰减为0.01。

学习率调度：

python lerobot/scripts/train.py \
    --dataset.repo_id=lerobot/pusht \
    --policy.type=diffusion \
    --optim.lr=5e-5 \
    --optim.scheduler.type=cosine \
    --optim.scheduler.warmup_steps=1000

这个命令启用了余弦学习率调度，并设置了1000步的预热期。

混合精度训练：

python lerobot/scripts/train.py \
    --dataset.repo_id=lerobot/pusht \
    --policy.type=diffusion \
    --policy.use_amp=true

这个命令启用了混合精度训练，可以显著提高训练速度，同时减少内存使用。

梯度裁剪：

python lerobot/scripts/train.py \
    --dataset.repo_id=lerobot/pusht \
    --policy.type=diffusion \
    --optim.grad_clip=1.0

这个命令设置了梯度裁剪阈值为1.0，有助于稳定训练过程。

Weights & Biases集成：

python lerobot/scripts/train.py \
    --dataset.repo_id=lerobot/pusht \
    --policy.type=diffusion \
    --wandb.enable=true \
    --wandb.project=my_lerobot_project \
    --wandb.name=diffusion_pusht_experiment

这个命令启用了Weights & Biases集成，将训练日志和指标发送到指定的项目和实验名称。

数据增强配置：

python lerobot/scripts/train.py \
    --dataset.repo_id=lerobot/pusht \
    --policy.type=diffusion \
    --policy.transforms.random_crop=true \
    --policy.transforms.random_flip=true

这个命令启用了随机裁剪和随机翻转数据增强。

通过这些高级配置选项，用户可以根据具体需求和实验目标，灵活地调整和优化训练过程，以获得更好的模型性能和训练效率。

5.2.6 命令行训练工具小结

LeRobot的命令行训练工具是一个功能强大、灵活易用的系统，它使得机器人学习策略的训练变得简单而高效。通过本节的学习，我们已经掌握了如何使用这一工具进行各种训练任务，包括从头训练、恢复训练、微调预训练模型等。

通过本章的学习，相信读者已经掌握了LeRobot策略训练的完整流程。LeRobot的策略训练系统体现了现代机器学习工程的最佳实践：模块化的设计、灵活的配置、完善的监控、科学的评估。掌握这些能力不仅有助于当前项目的成功，更为未来的研究和开发奠定了坚实的基础。

参考文献

[1] LeRobot官方文档. “Training Script Tutorial”. https://github.com/huggingface/lerobot/blob/main/examples/4_train_policy_with_script.md

[2] LeRobot官方示例. “Train Policy Example”. https://github.com/huggingface/lerobot/blob/main/examples/3_train_policy.py

[3] LeRobot官方示例. “Evaluate Pretrained Policy”. https://github.com/huggingface/lerobot/blob/main/examples/2_evaluate_pretrained_policy.py

[4] Draccus配置框架文档. “Configuration Management for ML”. https://github.com/dlwh/draccus

[5] LeRobot策略算法文档. “Policy Implementations”. https://github.com/huggingface/lerobot/tree/main/lerobot/common/policies

[6] Weights & Biases文档. “Experiment Tracking”. https://docs.wandb.ai/