lerobot

</div> <h3 align="center"> <p>用于现实世界机器人的最先进机器学习</p> </h3>

🤗 LeRobot旨在为PyTorch中的现实世界机器人提供模型、数据集和工具。其目标是降低机器人领域的准入门槛，使每个人都能贡献并从共享数据集和预训练模型中受益。

🤗 LeRobot包含已被证明可以迁移到现实世界的最先进方法，重点关注模仿学习和强化学习。

🤗 LeRobot已经提供了一系列预训练模型、包含人类收集的演示数据集，以及仿真环境，让你无需组装机器人就能开始使用。在接下来的几周内，我们计划为最实惠和最强大的真实机器人添加更多支持。

🤗 LeRobot在Hugging Face社区页面上托管预训练模型和数据集：huggingface.co/lerobot

仿真环境中预训练模型的示例

<table> <tr> <td><img src="https://yellow-cdn.veclightyear.com/835a84d5/6a8dd93e-d9ba-4b49-807f-0276af547b74.gif" width="100%" alt="ALOHA环境中的ACT策略"/></td> <td><img src="https://yellow-cdn.veclightyear.com/835a84d5/40ea11f5-a679-4ed9-98c5-694b3619a306.gif" width="100%" alt="SimXArm环境中的TDMPC策略"/></td> <td><img src="https://yellow-cdn.veclightyear.com/835a84d5/8c11d0e4-ca09-4728-ba5c-7efa34505fba.gif" width="100%" alt="PushT环境中的扩散策略"/></td> </tr> <tr> <td align="center">ALOHA环境中的ACT策略</td> <td align="center">SimXArm环境中的TDMPC策略</td> <td align="center">PushT环境中的扩散策略</td> </tr> </table>

致谢

感谢Tony Zaho、Zipeng Fu及其同事开源ACT策略、ALOHA环境和数据集。我们的版本改编自ALOHA和Mobile ALOHA。
感谢Cheng Chi、Zhenjia Xu及其同事开源扩散策略、Pusht环境和数据集，以及UMI数据集。我们的版本改编自扩散策略和UMI抓取器。
感谢Nicklas Hansen、Yunhai Feng及其同事开源TDMPC策略、Simxarm环境和数据集。我们的版本改编自TDMPC和FOWM。
感谢Antonio Loquercio和Ashish Kumar的早期支持。
感谢Seungjae (Jay) Lee、Mahi Shafiullah及其同事开源VQ-BeT策略并帮助我们将代码库适配到我们的仓库。该策略改编自VQ-BeT仓库。

安装

下载我们的源代码：

git clone https://github.com/huggingface/lerobot.git
cd lerobot

创建一个Python 3.10的虚拟环境并激活它，例如使用miniconda：

conda create -y -n lerobot python=3.10
conda activate lerobot

安装🤗 LeRobot：

pip install -e .

注意： 根据您的平台，如果在此步骤中遇到任何构建错误，您可能需要安装cmake和build-essential来构建我们的一些依赖项。在Linux上：sudo apt-get install cmake build-essential

对于仿真，🤗 LeRobot附带了可以作为额外功能安装的gymnasium环境：

例如，要安装带有aloha和pusht的🤗 LeRobot，请使用：

pip install -e ".[aloha, pusht]"

要使用Weights and Biases进行实验跟踪，请登录：

wandb login

（注意：您还需要在配置中启用WandB。请参见下文。）

目录结构

.
├── examples             # 包含演示示例，从这里开始了解LeRobot
|   └── advanced         # 包含更多高级示例，适合掌握基础知识的用户
├── lerobot
|   ├── configs          # 包含可在命令行中覆盖的所有选项的hydra yaml文件
|   |   ├── default.yaml   # 默认选择，加载pusht环境和扩散策略
|   |   ├── env            # 各种仿真环境及其数据集：aloha.yaml, pusht.yaml, xarm.yaml
|   |   └── policy         # 各种策略：act.yaml, diffusion.yaml, tdmpc.yaml
|   ├── common           # 包含类和实用工具
|   |   ├── datasets       # 各种人类演示数据集：aloha, pusht, xarm
|   |   ├── envs           # 各种仿真环境：aloha, pusht, xarm
|   |   ├── policies       # 各种策略：act, diffusion, tdmpc
|   |   ├── robot_devices  # 各种真实设备：dynamixel电机、opencv相机、koch机器人
|   |   └── utils          # 各种实用工具
|   └── scripts          # 包含通过命令行执行的函数
|       ├── eval.py                 # 加载策略并在环境中评估
|       ├── train.py                # 通过模仿学习和/或强化学习训练策略
|       ├── control_robot.py        # 遥控真实机器人、记录数据、运行策略
|       ├── push_dataset_to_hub.py  # 将数据集转换为LeRobot数据集格式并上传到Hugging Face hub
|       └── visualize_dataset.py    # 加载数据集并渲染其演示
├── outputs               # 包含脚本执行结果：日志、视频、模型检查点
└── tests                 # 包含用于持续集成的pytest实用工具

可视化数据集

查看示例1，了解如何使用我们的数据集类，该类可以自动从Hugging Face hub下载数据。

您还可以通过从命令行执行我们的脚本来本地可视化hub上数据集中的片段：

python lerobot/scripts/visualize_dataset.py \
    --repo-id lerobot/pusht \
    --episode-index 0

或者从本地文件夹中的数据集，使用根DATA_DIR环境变量（在以下情况下，将在./my_local_data_dir/lerobot/pusht中搜索数据集）

DATA_DIR='./my_local_data_dir' python lerobot/scripts/visualize_dataset.py \
    --repo-id lerobot/pusht \
    --episode-index 0

它将打开rerun.io并显示相机流、机器人状态和动作，如下所示：我们的脚本还可以可视化存储在远程服务器上的数据集。有关更多说明，请参阅 python lerobot/scripts/visualize_dataset.py --help。

`LeRobotDataset` 格式

LeRobotDataset 格式的数据集使用起来非常简单。它可以从 Hugging Face hub 上的仓库或本地文件夹中加载，例如 dataset = LeRobotDataset("lerobot/aloha_static_coffee")，并且可以像任何 Hugging Face 和 PyTorch 数据集一样进行索引。例如，dataset[0] 将从数据集中检索单个时间帧，其中包含作为 PyTorch 张量的观察和动作，可以直接输入到模型中。

LeRobotDataset 的一个特点是，除了通过索引检索单个帧外，我们还可以通过设置 delta_timestamps 为相对于索引帧的相对时间列表，来检索多个与索引帧具有时间关系的帧。例如，使用 delta_timestamps = {"observation.image": [-1, -0.5, -0.2, 0]}，可以为给定索引检索 4 个帧：3 个"之前"的帧（分别在索引帧之前 1 秒、0.5 秒和 0.2 秒），以及索引帧本身（对应于 0 条目）。有关 delta_timestamps 的更多详细信息，请参阅示例 1_load_lerobot_dataset.py。

在底层，LeRobotDataset 格式使用了几种序列化数据的方法，如果您计划更深入地使用这种格式，了解这些方法可能会有所帮助。我们尝试创建一种灵活yet简单的数据集格式，以涵盖强化学习和机器人领域中存在的大多数类型的特征和特性，包括仿真和现实世界中的情况，重点关注相机和机器人状态，但也可以轻松扩展到其他类型的感官输入，只要它们可以用张量表示。

以下是使用 dataset = LeRobotDataset("lerobot/aloha_static_coffee") 实例化的典型 LeRobotDataset 的重要细节和内部结构组织。具体特征会因数据集而异，但主要方面保持不变：

数据集属性：
  ├ hf_dataset：一个 Hugging Face 数据集（由 Arrow/parquet 支持）。典型特征示例：
  │  ├ observation.images.cam_high (VideoFrame)：
  │  │   VideoFrame = {'path': mp4 视频路径, 'timestamp' (float32): 视频中的时间戳}
  │  ├ observation.state (float32 列表)：例如机械臂关节位置
  │  ... (更多观察)
  │  ├ action (float32 列表)：例如机械臂关节目标位置
  │  ├ episode_index (int64)：此样本的剧集索引
  │  ├ frame_index (int64)：此样本在剧集中的帧索引；每个剧集从 0 开始
  │  ├ timestamp (float32)：剧集中的时间戳
  │  ├ next.done (bool)：表示剧集结束；每个剧集的最后一帧为 True
  │  └ index (int64)：整个数据集中的通用索引
  ├ episode_data_index：包含每个剧集的起始和结束索引的两个张量
  │  ├ from (1D int64 张量)：每个剧集的第一帧索引 — 形状 (num episodes,) 从 0 开始
  │  └ to：(1D int64 张量)：每个剧集的最后一帧索引 — 形状 (num episodes,)
  ├ stats：数据集中每个特征的统计信息（最大值、平均值、最小值、标准差）字典，例如
  │  ├ observation.images.cam_high：{'max': 具有相同维度数的张量（例如图像为 `(c, 1, 1)`，状态为 `(c,)`），等}
  │  ...
  ├ info：数据集元数据字典
  │  ├ codebase_version (str)：用于跟踪创建数据集的代码库版本
  │  ├ fps (float)：数据集记录/同步的每秒帧数
  │  ├ video (bool)：指示帧是否编码为 mp4 视频文件以节省空间或存储为 png 文件
  │  └ encoding (dict)：如果是视频，这记录了用于编码视频的 ffmpeg 主要选项
  ├ videos_dir (Path)：存储/访问 mp4 视频或 png 图像的位置
  └ camera_keys (字符串列表)：在数据集返回的项目中访问相机特征的键（例如 `["observation.images.cam_high", ...]`）

LeRobotDataset 使用几种广泛使用的文件格式来序列化其各个部分，即：

hf_dataset 使用 Hugging Face 数据集库序列化为 parquet
视频以 mp4 格式存储以节省空间，或存储为 png 文件
episode_data_index 使用 safetensor 张量序列化格式保存
stats 使用 safetensor 张量序列化格式保存
info 使用 JSON 保存

数据集可以无缝地上传/下载到 HuggingFace hub。要在本地数据集上工作，您可以将 DATA_DIR 环境变量设置为您的根数据集文件夹，如上述数据集可视化部分所示。

评估预训练策略

查看示例 2，该示例说明了如何从 Hugging Face hub 下载预训练策略，并在其对应的环境中运行评估。

我们还提供了一个更强大的脚本，可以在同一次推理中并行评估多个环境。以下是使用托管在 lerobot/diffusion_pusht 上的预训练模型的示例：

python lerobot/scripts/eval.py \
    -p lerobot/diffusion_pusht \
    eval.n_episodes=10 \
    eval.batch_size=10

注意：训练完自己的策略后，您可以使用以下命令重新评估检查点：

python lerobot/scripts/eval.py -p {OUTPUT_DIR}/checkpoints/last/pretrained_model

有关更多说明，请参阅 python lerobot/scripts/eval.py --help。

训练您自己的策略

查看示例 3，该示例说明了如何在 Python 中使用我们的核心库训练模型，以及示例 4，该示例展示了如何从命令行使用我们的训练脚本。

通常，您可以使用我们的训练脚本轻松训练任何策略。以下是在 Aloha 仿真环境中训练 ACT 策略的示例，该策略使用人类收集的插入任务轨迹：

python lerobot/scripts/train.py \
    policy=act \
    env=aloha \
    env.task=AlohaInsertion-v0 \
    dataset_repo_id=lerobot/aloha_sim_insertion_human \

实验目录会自动生成，并在终端中以黄色显示。它看起来像 outputs/train/2024-05-05/20-21-12_aloha_act_default。您可以通过在 train.py Python 命令中添加以下参数来手动指定实验目录：

    hydra.run.dir=your/new/experiment/dir

在实验目录中，会有一个名为 checkpoints 的文件夹，其结构如下：

checkpoints
├── 000250  # 训练步骤 250 的检查点目录
│   ├── pretrained_model  # Hugging Face 预训练模型目录
│   │   ├── config.json  # Hugging Face 预训练模型配置
│   │   ├── config.yaml  # 合并的 Hydra 配置
│   │   ├── model.safetensors  # 模型权重
│   │   └── README.md  # Hugging Face 模型卡片
│   └── training_state.pth  # 优化器/调度器/随机数生成器状态和训练步骤

要使用 wandb 记录训练和评估曲线，请确保您已经运行了 wandb login 作为一次性设置步骤。然后，在运行上述训练命令时，通过添加以下内容来启用配置中的 WandB：

    wandb.enable=true

运行的 wandb 日志链接也会在终端中以黄色显示。以下是它们在浏览器中的样子示例：

注意：为了提高效率，在训练过程中，每个检查点都会在少量剧集上进行评估。您可以使用 eval.n_episodes=500 来评估比默认更多的剧集。或者，在训练之后，您可能希望在更多剧集上重新评估最佳检查点或更改评估设置。有关更多说明，请参阅 python lerobot/scripts/eval.py --help。

重现最先进的技术（SOTA）

我们已经组织了配置文件（位于 lerobot/configs 下），以便它们能够重现各自原始工作中给定模型变体的 SOTA 结果。只需运行：

python lerobot/scripts/train.py policy=diffusion env=pusht

即可重现 Diffusion Policy 在 PushT 任务上的 SOTA 结果。

预训练策略及其重现详细信息可在 https://huggingface.co/lerobot 的"Models"部分找到。

贡献

如果您想为 🤗 LeRobot 做出贡献，请查看我们的贡献指南。

添加新数据集

要将数据集添加到 hub，您需要使用具有写入权限的令牌登录，该令牌可以从 Hugging Face 设置生成：

huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential

然后指向您的原始数据集文件夹（例如 data/aloha_static_pingpong_test_raw），并使用以下命令将数据集推送到 hub：

python lerobot/scripts/push_dataset_to_hub.py \
--raw-dir data/aloha_static_pingpong_test_raw \
--out-dir data \
--repo-id lerobot/aloha_static_pingpong_test \
--raw-format aloha_hdf5

有关更多说明，请参阅 python lerobot/scripts/push_dataset_to_hub.py --help。如果您的数据集格式不受支持，请在 lerobot/common/datasets/push_dataset_to_hub/${raw_format}_format.py 中实现您自己的格式，可以参考以下示例：pusht_zarr、umi_zarr、aloha_hdf5 或 xarm_pkl。

添加预训练策略

一旦您训练了一个策略，您可以使用类似 ${hf_user}/${repo_name} 的 hub ID（例如 lerobot/diffusion_pusht）将其上传到 Hugging Face hub。

首先，您需要在实验目录中找到检查点文件夹（例如 outputs/train/2024-05-05/20-21-12_aloha_act_default/checkpoints/002500）。其中应包含一个 pretrained_model 目录，里面应该有：

config.json：策略配置的序列化版本（遵循策略的数据类配置）。
model.safetensors：一组 torch.nn.Module 参数，以 Hugging Face Safetensors 格式保存。
config.yaml：一个综合的 Hydra 训练配置，包含策略、环境和数据集配置。策略配置应与 config.json 完全匹配。环境配置对于想要评估您的策略的人来说很有用。数据集配置仅作为可重现性的记录。

要将这些上传到 hub，请运行以下命令：

huggingface-cli upload ${hf_user}/${repo_name} path/to/pretrained_model

查看 eval.py 以了解其他人如何使用您的策略。

通过性能分析改进您的代码

以下是一个用于分析策略评估的代码片段示例：

from torch.profiler import profile, record_function, ProfilerActivity

def trace_handler(prof):
    prof.export_chrome_trace(f"tmp/trace_schedule_{prof.step_num}.json")

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(
        wait=2,
        warmup=2,
        active=3,
    ),
    on_trace_ready=trace_handler
) as prof:
    with record_function("eval_policy"):
        for i in range(num_episodes):
            prof.step()
            # 插入要分析的代码，可能是 eval_policy 函数的整个主体

引用

如果您愿意，可以使用以下方式引用本工作：

@misc{cadene2024lerobot,
    author = {Cadene, Remi and Alibert, Simon and Soare, Alexander and Gallouedec, Quentin and Zouitine, Adil and Wolf, Thomas},
    title = {LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch},
    howpublished = "\url{https://github.com/huggingface/lerobot}",
    year = {2024}
}

此外，如果您使用了任何特定的策略架构、预训练模型或数据集，建议引用原作者的工作，如下所示：

Diffusion Policy

@article{chi2024diffusionpolicy,
	author = {Cheng Chi and Zhenjia Xu and Siyuan Feng and Eric Cousineau and Yilun Du and Benjamin Burchfiel and Russ Tedrake and Shuran Song},
	title ={Diffusion Policy: Visuomotor Policy Learning via Action Diffusion},
	journal = {The International Journal of Robotics Research},
	year = {2024},
}

ACT or ALOHA

@article{zhao2023learning,
  title={Learning fine-grained bimanual manipulation with low-cost hardware},
  author={Zhao, Tony Z and Kumar, Vikash and Levine, Sergey and Finn, Chelsea},
  journal={arXiv preprint arXiv:2304.13705},
  year={2023}
}

TDMPC

@inproceedings{Hansen2022tdmpc,
	title={Temporal Difference Learning for Model Predictive Control},
	author={Nicklas Hansen and Xiaolong Wang and Hao Su},
	booktitle={ICML},
	year={2022}
}

VQ-BeT

@article{lee2024behavior,
  title={Behavior generation with latent actions},
  author={Lee, Seungjae and Wang, Yibin and Etukuru, Haritheja and Kim, H Jin and Shafiullah, Nur Muhammad Mahi and Pinto, Lerrel},
  journal={arXiv preprint arXiv:2403.03181},
  year={2024}
}