基于LeRobot与SO-101机械臂的ACT模型实战:从数据采集到实机部署

基于LeRobot与SO-101机械臂的ACT模型实战:从数据采集到实机部署

1. 项目概述:从零到一,让机械臂学会“看”和“抓”

如果你对机器人、具身智能或者强化学习感兴趣,并且手头恰好有一台像SO 101这样的桌面级机械臂,那么“训练一个ACT模型”这件事,听起来可能既酷炫又遥远。我最近花了大量时间,基于Meta开源的LeRobot框架,完整地走通了用SO 101机械臂训练ACT模型的全流程。这个过程远不止是跑通一个Demo,它涉及到硬件选型、环境搭建、数据采集、模型训练和部署验证等一系列环环相扣的步骤。今天,我就把这套从零开始的“保姆级”流程拆解开来,分享其中的核心思路、实操细节以及我踩过的那些坑。无论你是机器人领域的学生、爱好者,还是希望将前沿AI算法落地到实体机器人的工程师,这篇内容都能为你提供一份可直接参考的路线图。

ACT,全称Action Chunking with Transformers,是一种基于Transformer的模仿学习方法。它的核心思想是让模型像人一样,根据当前观察到的状态(比如摄像头看到的图像),预测未来一小段时间内连续的动作序列(比如机械臂末端执行器接下来10个时间步该怎么移动)。相比于传统的逐帧预测,这种“分块”预测的方式更符合人类行为的连续性,也让机械臂的动作看起来更平滑、更拟人。而LeRobot,则是Meta AI推出的一个旨在降低机器人学习门槛的PyTorch库,它集成了包括ACT在内的多种先进算法,并提供了从仿真到实机的一整套工具链。我们的目标,就是利用LeRobot,教会SO 101机械臂完成一个具体的任务,比如从桌面上抓取一个特定的物体。

2. 核心硬件与软件栈深度解析

2.1 SO-101机械臂:你的实体“学生”

在开始任何软件工作之前,我们必须充分了解我们的硬件“学生”——SO-101机械臂。这是一款开源的6自由度桌面级机械臂,其开源特性意味着我们可以获得完整的机械图纸、电路设计和固件代码,这对于深度定制和问题排查至关重要。

核心组件与参数调优要点:

  1. 关节与电机:SO-101通常使用舵机(如MG996R)或步进电机作为关节驱动器。这里有一个关键选择:位置控制 vs 扭矩控制。大多数开源套件默认是位置控制(发送目标角度),但更高级的模仿学习往往需要扭矩或速度控制,以实现更柔顺的交互。如果你的控制器(如STM32或树莓派+驱动板)支持,务必在固件层启用扭矩/速度控制接口。在LeRobot的数据采集中,我们记录的是关节的目标位置或扭矩指令,因此硬件控制模式必须与软件期望的模式对齐。

  2. 末端执行器:通常是二指夹爪。你需要校准其开合范围,并将其映射到一个标准化的数值区间(例如,0.0代表完全闭合,1.0代表完全张开)。这个映射关系需要在数据采集和模型推理时保持一致。

  3. 感知系统:这是ACT模型的“眼睛”。至少需要一个RGB摄像头,固定在工作区域上方(顶视)或机械臂本体上(眼在手)。我强烈推荐使用顶视摄像头,因为它能提供稳定、全局的视角,避免了机械臂运动带来的图像剧烈变化,简化了学习问题。摄像头分辨率无需过高,640x480足以,但帧率应稳定在30FPS以上。务必进行相机标定,获取内参和畸变系数,这对于后续如果需要做视觉伺服或更精确的坐标映射非常重要。

  4. 控制器:SO-101可能通过USB串口或ROS话题接收指令。你需要一个稳定的通信桥接器。我的方案是使用树莓派4B作为上位机,它运行LeRobot的数据采集和推理程序,并通过USB转TTL串口模块或GPIO(连接舵机控制板)与机械臂底层控制器通信。树莓派同时连接摄像头,完成图像采集和指令下发的一体化任务。

注意:在采购或组装SO-101时,务必确认所有关节的零位(Home Position)并做好标记。不准确的零位会导致运动学模型计算出的正向运动学(末端位置)与实际位置偏差巨大,使得基于坐标的训练几乎不可能成功。我建议先用一个简单的脚本控制每个关节单独运动到极限位置,记录下电机读数,并以此机械定义零位。

2.2 软件生态:LeRobot及其“朋友们”

软件栈的搭建是项目的地基,一个清晰、隔离的环境能避免无数依赖冲突的噩梦。

1. 操作系统与基础环境:我选择Ubuntu 22.04 LTS作为树莓派和开发机的操作系统。其长期的软件支持和庞大的ROS/机器人社区资源是关键。首先,为树莓派安装64位的Ubuntu Server镜像。

2. 核心框架:LeRobot安装与配置LeRobot是项目的核心。我强烈建议使用Conda或Venv创建独立的Python环境。

# 创建并激活conda环境 conda create -n lerobot python=3.10 conda activate lerobot # 安装PyTorch(根据你的CUDA版本选择,树莓派上安装CPU版本) # 对于开发机(有NVIDIA GPU): pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 对于树莓派(ARM CPU): pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu # 安装LeRobot核心库 pip install lerobot

安装后,运行python -c “import lerobot; print(lerobot.__version__)”验证是否成功。

3. 通信中间件:ROS 2 Humble的选择LeRobot与实体机器人的交互可以通过其自带的lerobot.scripts工具完成,这些工具底层通常封装了PyRobot或直接的硬件接口。但对于SO-101这种自定义程度高的机械臂,我们往往需要自己编写通信节点。ROS 2是一个更通用、更强大的选择。我选用ROS 2 Humble,因为它与Ubuntu 22.04兼容性好,且是LTS版本。

在树莓派上安装ROS 2 Humble:

# 设置locale和软件源 sudo apt update && sudo apt install locales sudo locale-gen en_US en_US.UTF-8 sudo update-locale LC_ALL=en_US.UTF-8 LANG=en_US.UTF-8 export LANG=en_US.UTF-8 # 添加ROS 2仓库并安装 sudo apt install software-properties-common sudo add-apt-repository universe sudo apt update && sudo apt install curl -y sudo curl -sSL https://raw.githubusercontent.com/ros/rosdistro/master/ros.key -o /usr/share/keyrings/ros-archive-keyring.gpg echo “deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/ros-archive-keyring.gpg] http://packages.ros.org/ros2/ubuntu $(. /etc/os-release && echo $UBUNTU_CODENAME) main” | sudo tee /etc/apt/sources.list.d/ros2.list > /dev/null sudo apt update sudo apt install ros-humble-desktop python3-colcon-common-extensions

安装完成后,在~/.bashrc中 source ROS 2环境:source /opt/ros/humble/setup.bash

4. 相机驱动与视觉处理:使用cv_camerausb_camROS 2包来驱动USB摄像头。

sudo apt install ros-humble-cv-camera ros-humble-image-transport-plugins ros-humble-cv-bridge

同时,你需要安装OpenCV (pip install opencv-python) 用于图像预处理,如裁剪、归一化等,这些预处理步骤将在数据采集时直接应用,以减少训练时的计算负担。

5. 自定义ROS 2接口包:你需要创建一个ROS 2包,用于定义机械臂的状态(关节角度、夹爪开合度)和动作(目标关节角度或速度)的消息类型,并编写一个简单的节点,订阅LeRobot推理出的动作话题,将其转换为串口或PWM指令发送给SO-101的底层控制器。反之,该节点也需要发布从控制器读取的关节状态和从摄像头获取的图像到ROS话题上,供LeRobot的数据采集器订阅。

3. 数据采集:高质量数据是成功的一半

ACT是一个模仿学习模型,其性能上限几乎由示范数据的质量决定。低质量、不一致的数据会导致模型学不会或行为怪异。

3.1 设计示范任务

从一个简单的任务开始,比如“从固定位置抓取一个方块并移动到另一个固定位置”。这个任务看似简单,但包含了接近、抓取、提升、移动、放置等多个子阶段,是测试流程的完美起点。

场景搭建:

  • 工作台:使用纯色、无纹理的桌面(如绿色或蓝色),这与方块物体形成高对比度,后期如果需要可以简化背景分割。
  • 物体:使用一个颜色鲜艳、形状规则(如立方体)的物体。避免反光或透明物体。
  • 标记:可以在桌面上用胶带贴出物体的“初始位置”和“目标位置”,方便演示时保持一致性。

3.2 采集工具与流程

LeRobot提供了命令行工具lerobot.scripts.dataset来辅助采集,但它更适用于标准接口的机器人。对于我们的自定义设置,我选择编写一个灵活的Python采集脚本。核心逻辑如下:

  1. 初始化:连接ROS 2节点,订阅/joint_states(机械臂状态) 和/camera/image_raw(图像),并创建发布器用于发送动作指令(在“手动控制”模式下)。
  2. 控制模式切换
    • 手动示教模式:通过键盘(如ikjl控制末端上下左右,uo控制夹爪)或游戏手柄实时控制机械臂。这是采集专家示范的主要方式。
    • 自动复位模式:编写一个函数,让机械臂自动运动到一个预设的“起始位置”,确保每次演示的开始状态一致。
  3. 数据记录
    • 以固定的频率(如10Hz)同步记录以下数据:
      • 观察 (observation):当前RGB图像(经过下采样和归一化到[0,1])。
      • 状态 (state):当前所有关节的角度(弧度制)和夹爪开合度。
      • 动作 (action):在手动模式下,记录操作者当前发送给机械臂的动作指令(关节角度增量或目标角度)。
      • 结束标志 (done):当前演示是否结束(例如任务成功完成或中途失败)。
    • 将每一次记录保存为一个字典,并添加到一个列表里。
  4. 单次演示 (Episode) 管理
    • 按下‘r’键开始一次新的演示录制,机械臂自动复位。
    • 操作者手动控制完成抓取-放置任务。
    • 按下‘s’键停止本次录制,并将本次演示的所有数据(一个列表的字典)保存为一个独立的.pkl.hdf5文件,同时清空缓存,准备下一次录制。

关键参数与格式:

  • 动作维度:如果你的SO-101是6自由度+夹爪,那么动作空间就是7维。决定你的动作是绝对位置相对位置增量还是速度。对于ACT,通常使用相对位置增量(delta position),因为它更容易学习且能产生更平滑的动作。例如,action = [delta_joint1, delta_joint2, ..., delta_gripper]
  • 图像预处理:立即将图像从BGR转换为RGB,并调整大小为[3, 240, 320](C, H, W)。归一化到[0, 1]。这个处理必须在存入数据集之前完成,保证训练和推理时的一致性。
  • 数据存储:我使用h5py库将数据存储为HDF5格式。它的好处是能高效存储大量数组数据,并且LeRobot内置了对HDF5数据集的支持。每个episode作为一个group,里面包含observationsactionsstates等dataset。

3.3 采集实操心得与数据量评估

  • “20-50法则”:对于这个简单的抓取任务,我采集了约50次成功的演示。经验表明,少于20次,模型很难学到稳健的策略;50-100次通常能取得不错的效果。关键在于质量而非绝对数量。确保每次演示都是流畅、成功的。
  • 引入多样性:在物体位置和姿态上加入微小变化(在初始和目标位置附近随机扰动几厘米),可以极大地提升模型的泛化能力。
  • 同步是魔鬼:务必确保图像、状态、动作的时间戳是严格对齐的。ROS的message_filters库的ApproximateTimeSynchronizer可以帮上大忙。不同步的数据会导致模型学习到错误的因果关系。
  • 记录所有元数据:在数据集中保存相机内参、关节限位、控制模式(位置/速度)等。这些信息在后续的仿真复现或模型迁移时至关重要。

4. ACT模型训练:参数背后的逻辑

有了高质量的数据集,我们就可以开始训练ACT模型了。LeRobot让启动训练变得非常简单,但理解其背后的超参数意义,才能有效地调优。

4.1 数据集准备与加载

首先,将采集的HDF5数据组织成LeRobot要求的格式。假设你的数据存储在~/so101_dataset/目录下,每个episode是一个episode_xxxx.hdf5文件。你需要创建一个dataset.yaml配置文件:

# dataset.yaml name: so101_pick_and_place local_path: /home/pi/so101_dataset fps: 10 # 你的采集频率 image_obs_keys: - observation state_obs_keys: - state action_keys: - action

然后,使用LeRobot的工具将其转换为标准格式:

lerobot convert_dataset --data.dir ~/so101_dataset --dataset.id so101_pick_and_place

4.2 训练配置详解

LeRobot使用Hydra进行配置管理。我们创建一个训练配置文件train_act.yaml

# train_act.yaml defaults: - base - _self_ policy: _target_: lerobot.models.act.act_policy.ACTPolicy encoder: “vip” # 使用VIP预训练的视觉编码器,能大幅提升学习效率 num_queries: 100 # Transformer解码器查询的数量,影响动作序列的预测粒度 chunk_size: 10 # 预测的动作序列长度(时间步) weight_decay: 0.0001 lr: 0.0001 # 学习率,对于微调VIP编码器,这个值通常较小 warmup_steps: 1000 dataset: name: so101_pick_and_place root_path: /home/pi/so101_dataset training: batch_size: 8 # 根据GPU内存调整,树莓派上可能只能用1或2 num_epochs: 100 # 训练轮数 eval_every: 5 # 每5轮评估一次 save_every: 10 # 每10轮保存一次检查点 device: cuda # 开发机用cuda,树莓派上改为cpu environment: eval_episodes: 10 # 评估时运行的episode数量

关键参数解析:

  • encoder: “vip”:这是最大的技巧。VIP是Meta开源的视觉表示模型,在大量机器人视频上进行了预训练。使用它作为图像编码器,相当于让模型拥有了强大的“视觉常识”,我们只需要微调(fine-tune)它即可,这比从头训练一个CNN编码器要快得多、效果好得多。
  • chunk_size: 10:这是ACT的核心。模型一次预测未来10个时间步的动作。在推理时,我们执行这10个动作,然后获取新的观测,再预测下一个10步,如此循环。这个值越大,模型规划得越远,但训练也越难。10是一个不错的起点。
  • num_queries: 100:Transformer解码器的查询数,可以理解为模型内部用于规划动作的“思维向量”的数量。更多的查询可能带来更精细的控制,但也会增加计算量。
  • lr: 0.0001:由于使用了预训练的VIP,学习率要设置得较小,避免破坏已有的良好视觉特征。

4.3 启动训练与监控

在开发机(有GPU)上运行训练:

lerobot train_policy policy=act dataset=so101_pick_and_place training.device=cuda

训练开始后,重点关注TensorBoard日志中的以下曲线:

  • train/total_loss:总损失值,应该随着训练轮数下降并逐渐平稳。
  • train/action_loss:动作预测的损失,这是主要优化目标。
  • eval/success_rate:在验证集上的成功率。这是衡量模型性能的黄金指标。如果这个指标迟迟不上升,可能是数据有问题或模型容量不足。

训练时间参考:在单张RTX 3060 GPU上,使用50条episode的数据,训练100个epoch大约需要2-4小时。

5. 模型部署与实机推理:从数字到物理

训练好的模型只是一个.ckpt文件。让它能在真实的SO-101上运行,是最后也是最激动人心的一步。

5.1 导出与优化模型

首先,将训练好的检查点导出为TorchScript或ONNX格式,以提升推理速度并简化部署。LeRobot可能没有直接导出功能,我们需要写一个小脚本:

import torch from lerobot.models.act.act_policy import ACTPolicy # 加载配置和模型权重 config = … # 加载你的训练配置 model = ACTPolicy.from_pretrained(“/path/to/your/checkpoint.ckpt”, config=config) model.eval() # 创建一个示例输入 example_obs = { “observation”: torch.randn(1, 3, 240, 320), # 批大小, 通道, 高, 宽 “state”: torch.randn(1, 7) # 批大小, 状态维度 } # 导出为TorchScript traced_script = torch.jit.trace(model, example_obs, strict=False) traced_script.save(“act_model_so101_traced.pt”)

5.2 编写实机推理节点

在树莓派上,我们编写一个ROS 2节点act_inference_node.py,其工作流程如下:

  1. 初始化:加载TorchScript模型,连接ROS 2,订阅/camera/image_raw/joint_states
  2. 预处理:对收到的图像进行与训练时完全一致的预处理( resize, RGB转换,归一化)。将关节状态转换为Tensor。
  3. 推理循环
    • 将当前观测(图像和状态)输入模型。
    • 模型输出未来chunk_size步的动作序列。
    • 我们只取序列中的第一个动作(即时动作)执行。这是模型预测的、基于当前状态的最优立即动作。
    • 将动作(关节增量)通过ROS话题或串口发送给机械臂控制器。
    • 等待一个控制周期(例如0.1秒,对应10Hz),然后获取新的观测,重复此过程。
  4. 终止条件:当模型输出一个特殊的“终止动作”(如果训练时有定义),或机械臂到达了预设的目标区域(可通过视觉或关节位置判断),则停止本次任务。

关键代码片段(伪代码):

import torch import rclpy from rclpy.node import Node from sensor_msgs.msg import Image from your_custom_msgs.msg import JointState, ActionCommand class ACTInferenceNode(Node): def __init__(self): super().__init__(‘act_inference_node’) # 加载模型 self.model = torch.jit.load(‘act_model_so101_traced.pt’) self.model.eval() # 订阅与发布 self.image_sub = self.create_subscription(Image, ‘/camera/image_raw’, self.image_callback, 10) self.js_sub = self.create_subscription(JointState, ‘/joint_states’, self.js_callback, 10) self.cmd_pub = self.create_publisher(ActionCommand, ‘/action_commands’, 10) self.current_obs = {‘image’: None, ‘state’: None} self.timer = self.create_timer(0.1, self.inference_step) # 10Hz控制频率 def inference_step(self): if self.current_obs[‘image’] is not None and self.current_obs[‘state’] is not None: with torch.no_grad(): # 预处理并转换为Tensor image_tensor = preprocess(self.current_obs[‘image’]).unsqueeze(0) # 增加批次维度 state_tensor = torch.tensor(self.current_obs[‘state’]).unsqueeze(0).float() # 推理 action_pred = self.model({‘observation’: image_tensor, ‘state’: state_tensor}) # 取第一个动作 immediate_action = action_pred[0, 0].cpu().numpy() # 假设action_pred形状为 [batch, chunk, dim] # 发布动作指令 cmd_msg = ActionCommand() cmd_msg.joint_deltas = immediate_action.tolist() self.cmd_pub.publish(cmd_msg)

5.3 实机运行与调试

  1. 安全第一:首次运行时,务必用手放在急停开关旁。将机械臂的动作幅度限制在较小范围(通过缩放immediate_action),逐步增加。
  2. 状态同步:确保推理节点接收到的关节状态是最新的。有时通信延迟会导致状态滞后,使模型基于“过去”的状态做决策,引发不稳定。可以在机械臂控制器端为状态信息添加高精度时间戳。
  3. 视觉一致性:实机环境的光照、背景必须与采集数据时尽可能一致。强烈的变化会导致模型性能下降。可以考虑在预处理中加入简单的自动白平衡或直方图均衡化来增强鲁棒性。
  4. 失败处理:实机运行中,机械臂可能会卡住或做出意外动作。在你的推理节点中加入安全监控:例如,检测关节是否接近限位、末端是否与桌面发生碰撞(可通过电流突变或简单的位置判断),一旦触发,立即停止当前动作序列,并回退到安全位置。

6. 常见问题与排查实录

在这一路上,我遇到了无数问题。下面这个表格总结了一些最典型的“坑”及其解决方案:

问题现象可能原因排查步骤与解决方案
训练时loss不下降或震荡剧烈1. 学习率设置不当。
2. 数据质量差(噪声大、不一致)。
3. 动作量纲问题(如角度与位置混用)。
4. VIP编码器未正确加载或冻结。
1. 尝试降低学习率(如1e-5),或使用学习率热身(warmup)。
2. 可视化你的动作数据,检查是否有异常值。确保演示是成功的。
3. 统一动作空间,确保所有维度都在相近的数值范围(如[-1, 1]),可进行归一化。
4. 检查模型参数,确认VIP部分的参数是否在训练初期被冻结(通常应冻结一部分层)。
实机运行时动作抽搐或不稳定1. 推理频率与控制频率不匹配。
2. 状态反馈延迟。
3. 模型在训练时未见过当前状态(分布外)。
4. 关节执行器有回差或精度不足。
1. 确保推理节点定时器的周期与训练数据频率(FPS)一致。
2. 在ROS 2中使用rqt_graph检查话题通信延迟,优化节点或使用更轻量的消息类型。
3. 在实机运行时,将观测到的状态输入模型,同时记录下模型预测的动作和实际状态,形成一个“在线数据集”。用这个数据对模型进行少量额外的微调(fine-tuning),能显著提升在实机上的适应性。
4. 检查机械臂硬件,对于舵机,可以考虑在指令输出后加入小的死区或低通滤波。
机械臂无法准确到达目标位置1. 运动学模型不准确(DH参数错误)。
2. 视觉定位误差。
3. 模仿学习本身的精度限制。
1.这是最常见的原因!重新校准SO-101的DH参数和关节零位。用一个已知的末端位置,反推关节角度,与实际读数对比。
2. 进行精确的相机标定,并实现手眼标定(如果使用眼在手配置)。对于顶视相机,可以做一个简单的像素到世界坐标的线性映射。
3. ACT等模仿学习方法旨在学习策略而非精确的轨迹。如果需要毫米级精度,应在模仿学习输出的动作基础上,结合一个基于位置的反馈控制器(如PID)。
抓取成功率低1. 夹爪控制不精确。
2. 训练数据中抓取瞬间的接触状态信息不足。
3. 物体在抓取时滑动。
1. 精细校准夹爪的开合度与指令值的关系。在抓取动作的最后阶段,改用“力控”模式,让夹爪以恒定的力闭合,而不是固定的位置。
2. 在状态观测中引入触觉或电流反馈(如果夹爪有压力传感器或电机电流检测)。即使没有,也可以将夹爪电机的负载电流作为一个额外的状态特征,这能隐式地反映抓取接触情况。
3. 在夹爪内侧粘贴增加摩擦力的材料(如橡胶片)。
模型在仿真中表现好,实机差仿真与现实之间的“现实鸿沟”。1. 在仿真中加入噪声:给关节控制添加延迟、给观测图像添加噪声和模糊、随机化摩擦力和物体质量。
2. 使用域随机化:在训练时,随机改变仿真中的纹理、光照、物体颜色等,让模型学会关注任务本质而非视觉表象。
3. 直接从实机采集少量数据对仿真训练的模型进行微调。

我个人最深刻的体会是:硬件校准的精度,直接决定了算法性能的上限。在软件和算法上花费数天调试可能都解决不了的问题,往往通过重新标定一次相机或调整一下机械臂的零位就迎刃而解。机器人是软硬件深度结合的系统,必须对两者都抱有同等的敬畏和耐心。

最后,这个项目只是一个起点。ACT模型成功运行后,你可以尝试更复杂的任务、多任务学习、加入语言指令(与CLIP等模型结合),甚至向强化学习进阶。SO-101和LeRobot为你提供了一个绝佳的实验平台,剩下的,就是你的想象力和动手能力了。