PyTorch 加速器

LeRobot 支持多种硬件加速选项，适用于训练和推理。

这些选项包括：

• CPU：CPU 执行所有计算，不使用专用加速器
• CUDA：通过 NVIDIA GPU 进行加速，并通过 ROCm 支持 AMD GPU（参见 AMD GPU）
• MPS：使用 Apple Silicon GPU 进行加速
• XPU：使用英特尔集成和独立 GPU 进行加速

入门

要使用特定的加速器，应安装相应版本的 PyTorch。

对于 CPU、CUDA 和 MPS 后端，请按照 PyTorch 安装页面 提供的说明进行操作。 对于 XPU 后端，请遵循 PyTorch 文档 中的说明。

AMD GPU（ROCm）

PyTorch 的 ROCm 构建版本通过 torch.cuda 命名空间暴露 HIP 运行时，因此 --policy.device=cuda 即可选择 AMD GPU，无需在 LeRobot 侧进行任何更改。在选择特定显卡时，优先使用 HIP_VISIBLE_DEVICES 而非 CUDA_VISIBLE_DEVICES。

在现有的 ROCm PyTorch 之上安装 LeRobot 需要格外小心：LeRobot 对 torch、torchvision 和 numpy 的版本约束可能与 ROCm 镜像中预装的版本不同，而常规安装会从 PyPI 拉取 torch，其 Linux 轮子包是 CUDA 构建版本，因此不加约束的安装会替换掉 ROCm 构建版本，导致 GPU 无法使用。请使用 uv pip install --override 固定这些包的版本，以便依赖树的其余部分正常解析。docker/Dockerfile.rocm 已实现此操作，是获取可用环境的最快方式；有关 ROCm 所需的运行参数，请参阅 docker/README.md。

验证安装

安装完成后，可以通过运行以下命令来验证加速器的可用性

import torch
print(torch.<backend_name>.is_available())  # <backend_name> is cuda, mps, or xpu

如何运行训练或评估

要选择所需的加速器，请在运行 lerobot-train 或 lerobot-eval 时使用 --policy.device 标志。例如，要在 Apple Silicon 上使用 MPS，请运行：

lerobot-train \
    --policy.device=mps ...

lerobot-eval \
    --policy.device=mps ...

不过，在大多数情况下，系统会自动检测到加速器的存在，因此可以在 CLI 命令中省略 policy.device 参数。

混合精度

训练精度由 --accelerator.mixed_precision 控制，它接受 no（默认值）、bf16 和 fp16：

lerobot-train \
    --policy.type=act \
    --accelerator.mixed_precision=bf16 ...

fp16 需要加速器：在 --policy.device=cpu 下会被拒绝，因为 accelerate 不会构建损失缩放器，运行会静默回退到全精度。bf16 在 CPU 上也能自动转换，因此不受此限制。

bf16 需要硬件支持（NVIDIA GPU 需为 Ampere 架构或更新版本）。在 V100 或 T4 等较旧的加速器上，请使用 fp16。

fp16 和损失缩放器

fp16 的指数范围远小于 bf16，因此在反向传播之前会对梯度进行放大，并在优化器更新之前将其缩小。LeRobot 将此任务委托给 accelerate 的 GradScaler，它还能检测溢出梯度并跳过相应的优化器更新。在分片训练下，梯度是 DTensor，PyTorch 会在设备网格上归约溢出标志，因此任一 rank 上的溢出都会导致所有 rank 跳过更新。

跳过更新属于正常现象，尤其是在初始阶段，当 scale 从初始值向下校准时。在日志中，这表现为 scale 指标下降。scale 会随每个检查点一起保存（training_state/scaler_state.json），并在恢复训练时还原，因此恢复后的运行无需重新校准。

如果默认值不适合某个策略，可以调整缩放器：

lerobot-train \
    --policy.type=act \
    --accelerator.mixed_precision=fp16 \
    --accelerator.grad_scaler.init_scale=1024 \
    --accelerator.grad_scaler.growth_interval=2000 ...

TF32

在 Ampere 及更新版本的 NVIDIA GPU 上，LeRobot 默认使 float32 矩阵乘法和 cuDNN 卷积使用 TF32（--tf32=true）。使用 --tf32=false 可保持完整的 float32 精度，而 --tf32=null 则保留 PyTorch 的默认设置。