1. 环境准备与项目理解大家好我是老张一个在机器人行业摸爬滚打了十来年的工程师。最近PI团队开源的π0模型在圈子里火得不行它号称能让机械臂像人一样通过看和语言指令来学习完成任务。我手头正好有一台Aubo i5机械臂寻思着能不能把这“最强大脑”给它装上让它变得更聪明。折腾了小半个月从一脸懵到成功让机械臂完成抓取踩了不少坑也总结了一套行之有效的部署流程。今天我就把这份从零开始的实战指南分享给你目标是让你也能亲手把π0模型部署到自己的Aubo机械臂上完成一个简单的视觉抓取任务。首先咱们得搞清楚π0是个啥。简单来说π0是一个多模态的机器人基础模型。它和我们以前用的那些需要精确编程轨迹的模型不同它更像是一个“实习生”——你给它看摄像头画面视觉告诉它要干什么语言指令它就能自己学习并输出机械臂该怎么动的动作序列。这对于需要灵活应对不同物体、不同摆放位置的场景比如无序抓取、分拣简直是神器。而Aubo i5是一款国产的轻型协作机械臂性价比高在实验室和中小型工厂里很常见。我们的目标就是让π0的“大脑”驱动Aubo的“身体”。在开始动手之前你需要准备好以下几样东西硬件一台Aubo i5机械臂或其他型号但关节数建议为6或7以匹配模型输出维度至少两个RGB摄像头一个顶置俯瞰工作区域一个安装在机械臂末端腕部以及一台性能尚可的电脑最好有NVIDIA GPU能大大加速训练。软件基础你的电脑上需要安装好Ubuntu 20.04或22.04系统这是机器人开发最友好的环境熟悉基本的Linux命令行操作以及Python编程。如果你对Docker有些了解那后续的环境隔离会轻松很多。关键资源访问GitHub的能力用于获取π0的开源代码。原始文章里提到了项目地址这里我们直接使用官方仓库。你需要对git命令有基本了解。整个项目的核心逻辑是这样的我们先要收集Aubo机械臂执行任务的数据比如移动、抓取然后把这些数据转换成π0模型能理解的格式接着用这些数据对预训练的π0模型进行微调最后把训练好的模型部署上去让它根据实时摄像头画面来控制机械臂。听起来步骤不少但别怕我们一步一步来我会把每个环节的细节和可能遇到的“坑”都讲清楚。2. 搭建开发环境与获取代码万事开头难环境配置往往是第一道坎。π0项目依赖一个比较新的Python包管理工具uv以及JAX等深度学习框架。为了不把你本机的Python环境搞乱我强烈建议使用conda或venv创建一个独立的虚拟环境。2.1 创建并激活虚拟环境打开你的终端我们首先创建一个名为pi0_aubo的Python 3.10环境经测试3.10版本兼容性最好。# 如果你使用conda conda create -n pi0_aubo python3.10 -y conda activate pi0_aubo # 如果你使用venv python3.10 -m venv pi0_aubo_env source pi0_aubo_env/bin/activate激活后你的命令行提示符前面应该会出现(pi0_aubo)的字样这表示你已经在这个独立的环境里了。2.2 安装UV和项目依赖接下来安装uv它是一个快速的Python包安装器和解析器比传统的pip快很多。curl -LsSf https://astral.sh/uv/install.sh | sh安装完成后可能需要重启终端或者运行source ~/.bashrc来让uv命令生效。现在我们去获取π0的源代码。# 克隆开源仓库这里使用原始文章参考的仓库 git clone https://github.com/c237814486/OA-Stereo.git cd OA-Stereo进入项目目录后你会看到很多文件和文件夹。我们需要安装项目依赖。项目通常会提供一个pyproject.toml或requirements.txt文件。使用uv来同步依赖是最稳妥的方式。uv sync这个过程会下载安装JAX、Flax、PyTorch用于数据处理等一系列依赖包时间可能会有点长取决于你的网络。这里有个大坑JAX的安装需要对应你的CUDA版本。如果你有GPU并且希望使用GPU加速训练请务必在安装后根据JAX官方文档的指示安装对应CUDA版本的JAX。例如对于CUDA 12.1pip install --upgrade jax[cuda12_pip]0.4.28 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html没有GPU的话安装CPU版本的JAX即可但训练速度会慢很多。环境装好后你可以运行一个简单的Python语句import jax来测试是否成功。3. 数据采集与格式转换模型训练离不开数据。π0期望的数据格式是一种名为LeRobot Dataset V2的结构化格式。但我们的Aubo机械臂原始数据很可能是一堆*.hdf5文件里面记录了每个时间步的关节角度、图像和动作指令。所以我们需要一个“翻译官”把Aubo的数据“翻译”成π0能看懂的格式。3.1 理解数据格式要求π0模型在训练时每个数据样本一帧需要包含以下信息observation.images.top: 一个形状为(3, 480, 640)的数组表示顶置摄像头的RGB图像通道、高、宽。observation.images.wrist: 腕部摄像头的图像格式同上。observation.state: 一个长度为7的数组表示Aubo机械臂的7个关节角度6个关节1个夹爪。action: 一个长度为7的数组表示下一时刻机械臂应该执行的7个关节角度/夹爪状态。你的原始Aubo数据可能存储在不同的地方比如通过ROS的rosbag记录或者你自己写脚本保存的数组。关键是要把它们整理成一个个的episode_*.hdf5文件每个文件代表一次完整的任务尝试例如一次从A点抓取物体放到B点的全过程。3.2 编写数据转换脚本原始文章里提供了一个非常关键的脚本convert_aubo_data_to_lerobot.py。这个脚本就是我们的“翻译官”。你需要根据自己数据的实际情况来修改它。我来带你过一遍核心修改点。首先找到脚本中定义机器人关节和摄像头的地方。对于Aubo i5通常是6个旋转关节加1个气动或电动夹爪。# 在 create_empty_dataset 函数内部或附近修改 motors 列表 motors [ joint1, joint2, joint3, joint4, joint5, joint6, gripper ] cameras [top, wrist] # 确保你的摄像头命名与此一致其次检查你的HDF5文件内部的数据路径。脚本中默认从/observations/qpos读取状态从/action读取动作从/observations/images/top和/observations/images/wrist读取图像。你必须用h5py库打开你的一个数据文件查看里面的确切数据集路径并相应修改load_raw_episode_data函数中的键名。例如如果你的数据里状态存储在/robot_state那么就需要修改state torch.from_numpy(ep[/robot_state][:])最后运行转换脚本。假设你把所有的原始*.hdf5文件都放在了/home/yourname/aubo_raw_data目录下。uv run examples/hhy_test/convert_aubo_data_to_lerobot.py --raw-dir /home/yourname/aubo_raw_data --repo-id your-name/aubo-dataset这个命令会在本地创建一个LeRobot格式的数据集。--repo-id参数可以随意命名这是本地目录名的一部分。如果一切顺利你会在~/.lerobot/datasets/your-name/aubo-dataset/下看到转换好的数据。这里有个经验之谈第一次运行时建议先处理一两个episode试试水加上--episodes 0 1这样的参数快速验证脚本是否正确避免处理大量数据后才发现格式错误。4. 配置训练任务与模型微调数据准备好之后我们就要告诉π0模型“这是你的新课本请根据这些Aubo的数据来学习。” 这就需要配置训练任务。4.1 创建自定义训练配置在项目路径src/openpi/training/config.py中我们需要添加一个自己的训练配置。直接复制原始文章里的TrainConfig是一个很好的起点但一定要理解每个参数的含义。# 在 config.py 的 TRAIN_CONFIGS 列表里添加你的配置 TrainConfig( namepi0_aubo_pick_and_place, # 给你的任务起个名字 modelpi0.Pi0Config( paligemma_variantgemma_2b_lora, # 视觉-语言骨干网络使用LoRA微调 action_expert_variantgemma_300m_lora, # 动作预测网络使用LoRA微调 ), dataLeRobotAuboDataConfig( # 注意这里换成了我们自定义的Aubo数据配置类 repo_idyour-name/aubo-dataset, # 与转换脚本中的repo-id一致 asset_idyour-name/aubo-dataset, default_promptPick the block and place it into the box., # 你的任务指令 ), batch_size1, # 根据你的GPU内存调整一开始可以设为1 weight_loaderweight_loaders.CheckpointWeightLoader(gs://openpi-assets/checkpoints/pi0_base/params), # 加载预训练权重 num_train_steps20_000, # 训练步数对于微调几千到几万步可能就够了 freeze_filterpi0.Pi0Config( paligemma_variantgemma_2b_lora, action_expert_variantgemma_300m_lora ).get_freeze_filter(), # 冻结大部分参数只微调LoRA层节省资源且防止过拟合 ),关键修改1数据配置类。你可能需要创建一个LeRobotAuboDataConfig类或者直接修改现有的LeRobotAlohaDataConfig。核心是确保robot_type等参数与你的Aubo机械臂匹配。如果不想麻烦可以暂时沿用LeRobotAlohaDataConfig但心里要明白我们的机器人是Aubo。关键修改2提示词Prompt。default_prompt非常重要它定义了任务的自然语言描述。在训练和推理时模型都会看到这个指令。请用简洁、清晰的英文描述你的任务例如“Pick the red screw and place it on the blue marker.”4.2 创建Aubo专用的策略输入模块π0模型原版是针对Aloha双机机械臂设计的其输入处理模块AlohaInputs期望特定的摄像头名称和状态维度。我们的Aubo是单臂状态维度是76关节1夹爪而原模型可能期望14维双臂。所以我们需要创建一个新的输入模块。按照原始文章的指引复制src/openpi/policies/aloha_policy.py为aubo_policy.py并进行关键修改class AuboInputs(transforms.DataTransformFn): action_dim: int 7 # 最重要的修改将动作维度改为7匹配Aubo EXPECTED_CAMERAS: ClassVar[tuple[str, ...]] (top, wrist) # 我们有两个摄像头 def __call__(self, data: dict) - dict: # 这里可能不需要像Aloha那样进行复杂的空间转换 # 直接获取状态和图像 state data[state] # 假设已经是7维 in_images data[images] # 确保状态是7维如果不是可能需要填充或截断 state transforms.pad_to_dim(state, self.action_dim) # 填充到7维 # 组织图像输入模型期望的键名是固定的 base_image in_images[top] images { base_0_rgb: base_image, left_wrist_0_rgb: in_images.get(wrist, np.zeros_like(base_image)), } # ... 其余部分参考原文件主要处理图像掩码和动作同时记得在src/openpi/models/model.py中确认IMAGE_KEYS元组与上面定义的图像键名一致。通常只需要base_0_rgb和left_wrist_0_rgb。4.3 启动模型训练在开始训练前我们需要计算数据集的归一化统计量均值和标准差这对稳定训练很重要。uv run scripts/compute_norm_stats.py --config-name pi0_aubo_pick_and_place然后就可以开始激动人心的训练了使用以下命令XLA_PYTHON_CLIENT_MEM_FRACTION0.99 uv run scripts/train.py pi0_aubo_pick_and_place --exp_name my_first_aubo_run --overwrite解释一下参数XLA_PYTHON_CLIENT_MEM_FRACTION0.99这是为JAX/XLA设置GPU内存分配0.99表示几乎使用全部GPU内存避免内存碎片。pi0_aubo_pick_and_place就是我们之前在config.py里定义的配置名。--exp_name给这次实验起个名字日志和检查点会保存在以它命名的目录下。--overwrite如果同名实验已存在则覆盖。训练开始后你会在终端看到损失值loss逐渐下降。这个过程可能需要几个小时甚至更久取决于你的数据量、步数和GPU性能。你可以定期查看logs/目录下的TensorBoard日志来监控训练进度。5. 模型部署与实时推理验证训练完成后我们得到了一个微调好的模型检查点checkpoint。接下来就是最有趣的环节让模型“活”起来实时控制机械臂。5.1 编写实时推理脚本我们需要一个脚本它能够实时读取两个摄像头的画面。获取当前机械臂的关节状态或者在开环测试时使用上一次发送的指令作为当前状态。将图像和状态组装成模型需要的观测observation格式。调用加载好的策略policy进行推理得到未来一系列动作。将第一个动作或前几个动作发送给Aubo机械臂执行。循环执行1-5步。原始文章中的eval.py就是一个很好的起点。它使用了pyrealsense2库来读取Intel RealSense摄像头。如果你用的是其他摄像头如USB摄像头可能需要换成OpenCV的VideoCapture。脚本的核心修改点在于与Aubo机械臂的通信。Aubo通常通过TCP/IP或ROS提供控制接口。你需要替换掉示例中socket通信的部分改为调用Aubo SDK如aubo_robot_sdk中的函数来发送关节角度指令。例如假设你有这样一个Aubo控制函数import aubo_robot_sdk as aubo def send_joint_angles_to_aubo(joint_angles): joint_angles: 长度为7的列表前6个是关节角度弧度第7个是夹爪开合度0-1000 # 连接到机械臂假设IP是192.168.1.100 robot aubo.Robot() robot.connect(192.168.1.100, 8899) # 设置运动模式为关节空间运动 robot.set_move_mode(aubo.MoveMode.JOINT) # 发送前6个关节角度 robot.move_joint(joint_angles[:6]) # 控制夹爪假设是气动通过数字输出控制 gripper_value joint_angles[6] robot.set_digital_output(0, gripper_value 500) # 简单阈值控制 robot.disconnect()然后在你的主循环中将policy.infer(observation)得到的actions数组的第一行即下一个时刻的动作取出调用这个函数发送给机械臂。5.2 处理动作空间与安全模型输出的动作空间是它学习时使用的空间可能与Aubo机械臂的真实关节空间有差异。在AuboInputs类中我们可能省略了_decode_aloha和_encode_actions_inv这类空间转换函数因为Aubo和Aloha的关节定义可能不同。这是一个至关重要的安全点。在将模型动作发送给真实机械臂之前务必进行仿真测试或严格限幅。你可以先让机械臂在“拖拽示教”模式或零力模式下让模型控制虚拟关节观察其运动轨迹是否合理。同时一定要在代码中加入关节角度限幅确保发送的指令不会超出Aubo每个关节的运动范围防止损坏设备或造成危险。def safe_clip_joint_angles(angles): limits [(-3.14, 3.14), (-2.0, 2.0), ...] # Aubo各关节的弧度制限位 clipped [] for i, (low, high) in enumerate(limits): clipped.append(np.clip(angles[i], low, high)) return clipped5.3 运行与调试确保摄像头连接正常机械臂上电并处于可远程控制的安全状态。运行你的推理脚本python eval_aubo.py如果一切配置正确你应该能看到摄像头窗口弹出并且机械臂开始根据模型的计算结果缓慢移动。第一次运行时动作可能很怪异或幅度很小这是正常的。你需要观察模型是否对物体和指令有反应例如当物体移动时机械臂的预期目标点是否也移动动作是否平滑有无剧烈抖动夹爪的开合时机是否正确调试是一个迭代的过程。你可能需要回头检查数据质量标注是否准确、训练是否充分损失是否已收敛、或者模型输入输出处理是否有误。耐心和细致的日志记录是你最好的帮手。把关键数据如图像、状态、模型输出的动作都保存下来方便离线分析。当看到机械臂颤颤巍巍但最终成功抓起目标物体时那种成就感是无与伦比的。整个部署过程虽然繁琐但每一步都加深了对机器人视觉控制系统的理解从数据管道、模型训练到最后的实时控制闭环这本身就是一个非常宝贵的全栈经验。