斯坦福UMI代码解析：刷盘机器人Universal Manipulation Interface代码的整体解读

v_JULY_v

已于 2025-02-02 22:35:32 修改

阅读量3.3k

点赞数 28

分类专栏： ALOHA/umi/dexcap中的模仿学习与动作预测文章标签：斯坦福UMI 通用操作界面 UMI代码解析手持夹爪收集数据

于 2025-02-02 21:32:39 首次发布

本文链接：https://blog.csdn.net/v_JULY_v/article/details/145422746

版权

ALOHA/umi/dexcap中的模仿学习与动作预测专栏收录该内容

16 篇文章

订阅专栏

前言

本文一开始是属于此文《UMI——斯坦福刷盘机器人：通过手持夹爪革新数据收集方式，且使用Diffusion Policy预测动作》的第四部分，但为了把原理部分和代码解析更好的解耦

加之，近期春节期间时间相对充裕正在盘点之前的一系列文章，盘到UMI时，觉得可以把代码解析抽取出来，独立成文

本解读基本来自我司「七月在线」具身智能开发组的远根同学，之前之所以分享其中的部分出来，当时是想招纳可以针对UMI共同做二次开发的朋友

第一部分示例数据、运行SLAM管道、数据集转换

1.1 Download example data 下载示例数据

(umi)$ wget --recursive --no-parent --no-host-directories --cut-dirs=2 --relative --reject="index.html*" https://real.stanford.edu/umi/data/example_demo_session/

网络问题，可以访问https链接，直接浏览器下载，然后统一放到当前目录example_demo_session文件夹内

数据集包含以下三个部分

1.1.1 一个夹爪校准视频

GX010214.mp4

1.1.2 几个收集的任务演示数据视频

GX010215.mp4

GX010232.mp4

GX010234.mp4

GX010236.mp4

GX010238.mp4

1.1.3 一个SLAM地图视频

mapping.mp4

1.2 Run SLAM pipeline 运行 SLAM 管道

(umi)$ python run_slam_pipeline.py example_demo_session

这里 example_demo_session 如果路径报错，请换成绝对路径

整个pipeline用于提取收集到的gopro数据，进行数据处理
最终目的：使用GoPro数据恢复机器人动作

具体分为以下步骤：

1.2.0 00_process_videos.py

上一个下载的所有示例数据，把目录结构进行整理。

每个视频都要单独处理，用ExifTool 工具包，提取每个视频的相机序列号+拍摄时间，作为文件夹的名称，每个视频放入各自的文件夹内

1.2.1 01_extract_gopro_imu.py：提取gopro惯性测量单元数据(imu)

这步用于提取gopro惯性测量单元数据(imu)，提取方式是拉的docker镜像，直接使用的外部仓库：GitHub - urbste/OpenImuCameraCalibrator: Camera calibration tool
且是C++写的，直接看提取结果，保存在imu_data.json文件中，总共提取了6种数据：GoPro Tags，这六种数据分别如下所示

1.2.1.1 ACCL (Accelerometer)加速度计

加速度计测量物体在三个方向上的加速度，通常分别是 X 轴、Y 轴和 Z 轴。这些数据用于检测物体的运动和方向变化，共有四个值

value：代表了三个轴的加速度值：x 轴、y 轴和 z 轴。单位：m/s2
cts: 采样的时间戳，可能表示这个数据点是在第x个采样时钟周期中采集的
data: 采样的日期和时间
temperature: 采样的温度

{
"value":[8.37410071942446,0.5875299760191847,4.9352517985611515],
"cts":78.018,
"date":"2024-01-10T18:54:47.109Z",
"temperature [°C]":51.544921875
}

1.2.1.2 GYRO (Gyroscope)陀螺仪

陀螺仪测量物体在三个方向上的角速度，即物体围绕每个轴旋转的速度。陀螺仪用于确定物体的姿态和运动状态，对于检测旋转和倾斜非常有效

value：代表了三个轴的角速度：x 轴、y 轴和 z 轴，单位：rad/s
cts: 采样的时间戳
data: 采样的日期和时间
temperature: 采样的温度

{
"value":[0.06496272630457935,0.0724174653887114,-0.027689030883919063],
"cts":78.018,
"date":"2024-01-10T18:54:47.109Z",
"temperature [°C]":51.544921875
}

1.2.1.3 GPS5 (Global Positioning System)全球定位系统

GPS 传感器提供位置数据，包括经度、纬度、高度以及速度。GPS 数据用于定位和导航

value：通常代表了 GPS 接收器测量的位置和速度信息。第一个和第二个元素（37.3996899, -122.1021419）通常代表纬度和经度,第三个元素（-43.125）可能表示海拔高度,第四个和第五个元素（0.469, 0.48）可能是 GPS 接收器的卫星状态信息
cts: 采样的时间戳
data: 采样的日期和时间
fix: 这个值通常表示 GPS 接收器的位置 fix 质量。在这个例子中，fix 值为 3，这通常意味着接收器具有较好的位置 fix，通常是一个卫星定位 fix，其中 1 是最差，5 是最佳。
precision: 这个值表示 GPS 位置测量的精度。在这个例子中，precision 值为 576，这个数值可能表示接收器测量位置的精度范围。
altitude system: 这个字段指定了用于表示海拔高度的参考系统。在这个例子中，“MSLV” 可能是指 mean sea level variation（平均海平面变化），这是一种相对于某个平均海平面的高度参考系统

{
"value":[37.3996899,-122.1021419,-43.125,0.469,0.48],
"cts":78.856,
"date":"2024-01-10T18:54:47.109Z",
"fix":3,
"precision":576,
"altitude system":"MSLV"
}

1.2.1.4 CORI( Camera Orientation)相机姿态

加速度计测量物体在三个方向上的加速度，通常分别是 X 轴、Y 轴和 Z 轴。这些数据用于检测物体的运动和方向变化，共有4个值

value：代表了三个轴的加速度值：x 轴、y 轴和 z 轴。
cts: 采样的时间戳
data: 采样的日期和时间
temperature: 采样的温度

{
"value":[0.999969481490524,0.002044740134891812,0.0016174810022278512,-0.0003662221137119663],
"cts":80.255,
"date":"2024-01-10T18:54:47.109Z"
}

1.2.1.5 IORI（Image Orientation）图像姿态

通常用于图像处理和计算机视觉领域，用于描述图像在三维空间中的方向和位置

// 待更

1.2.1.6 GRAV (Gravity Vector)重力向量

重力向量通常包括描述重力加速度的一组数值

// 待更

1.2.2 02_create_map.py

处理mapping地图视频数据，并生成地图。主要是用Docker来运行外部库ORB_SLAM3(Simultaneous Localization and Mapping，即同时定位与地图构建)系统

输入：上文《4.2.1 01_extract_gopro_imu.py：提取gopro惯性测量单元数据(imu)》中的的imu_data.json 和原MP4视频

输出：

mapping_camera_trajectory.csv
这是SLAM系统生成的相机轨迹文件，通常包含了相机在空间中的位置和方向信息
1) frame_idx：帧索引，表示该帧在视频中的顺序编号
2) timestamp: 时间戳，表示该帧的拍摄时间，通常是以秒或毫秒为单位
3) state：状态字段，通常用于表示SLAM系统的当前状态，例如是否初始化、是否丢失、是否关键帧等
4) is_lost：是否丢失标志，表示该帧是否丢失或无法追踪
5) is_keyframe：是否关键帧标志，表示该帧是否被选为关键帧，关键帧是SLAM中用于地图构建的重要帧
6) x、y、z：相机在三维空间中的平移坐标
7) q_x、q_y、q_z、q_w：相机姿态的四元数表示，用于描述相机的旋转。其中， q_x、q_y、q_z 分别表示轴向量的xyz分量，q_w 是四元数的第一个元素，表示旋转的轴向量
```
 frame_idx,timestamp,state,is_lost,is_keyframe,x,y,z,q_x,q_y,q_z,q_w
0,0.000000,1,true,false,0,0,0,0,0,0,0
1,0.016683,2,true,false,0,0,0,0,0,0,0
2,0.033367,2,true,false,0,0,0,0,0,0,0
3,0.050050,2,true,false,0,0,0,0,0,0,0
4,0.066733,2,true,false,0,0,0,0,0,0,0
```
map_atlas.osa
这个文检是个二进制文件，无法打开，一般是map_atlas.osa 文件是一个专用的 ORB_SLAM3 地图文件格式，它不是公开的标准格式，所以没有具体的公开文档说明其内部结构
ORB_SLAM3 是用c++写的，所以，要看懂，需要掌握slam和c++编程

不过，通常这类 SLAM系统的地图文件会包含以下类型的数据：
a) 关键帧数据：关键帧是 SLAM 系统中用于定位和地图构建的参考框架。它们通常包含了位置、姿态信息以及与之相关的地图点的标识符。
b) 地图点数据：地图点是 SLAM 系统中在环境中固定的特征点，如角点或边缘。它们在多个关键帧中被观测和记录，以帮助建立一个连续的位置估计。
c) 连接关系：关键帧之间的连接关系表示了时间上的连续性，这对于递归定位和构建地图至关重要。
d) 其他参数和信息：这可能包括与特定地图创建过程相关的参数，如滤波器的配置、重定位参数等

1.2.3 03_batch_slam.py

这个和上一步一样，上一步是mapping地图视频生成的轨迹信息，这步是批量生成任务演示数据的轨迹信息，

输入：原始mp4视频、上一步生成的map_atlas.osa、imu_data.json
输出：相机轨迹信息camera_trajectory.csv

和mapping_camera_trajectory.csv一样的内容

注意：生成轨迹不包括夹爪校准视频，也不生成map_atlas.osa

1.2.4 04_detect_aruco.py

处理所有视频文件中的 ArUco 标记检测。ArUco 是一种在计算机视觉中用于标定和校准相机，以及进行图像和视频处理任务的可视化标记。

这里有[OpenCV] aruco Markers识别原理：

[OpenCV] aruco Markers识别 - 一郎哥哥 - 博客园

输入：原视频MP4

输出：tag_detection.pkl

frame_idx: 图像帧的索引或序号。
time: 检测到ArUco标记的时间戳（如果有的话）。
tag_dict: 一个包含ArUco标记详细信息的字典
- rvec: 标记在相机坐标系中的旋转矢量。这是一个表示标记相对于上一帧位置旋转的向量。
- tvec: 标记在相机坐标系中的平移矢量。这表示标记自上一帧以来在三维空间中的平移
- corners: 标记在图像中的四个角落的坐标。这些坐标通常是以像素为单位的

总共有很多图像帧，以下是第一帧的内容。0,1,13代表其中3个标记的id

[{
  'frame_idx':0,
  'time':0,
  'tag_dict':
    { 
      0: {
      'rvec': [ 1.61436202,  0.00735438, -0.06256209], 
      'tvec': [-0.06100964,  0.06824645,  0.07121038]), 
      'corners': [[ 869.44006, 1493.346  ],
                 [ 958.62274, 1519.135  ],
                 [ 907.301  , 1584.2545 ],
                 [ 814.1008 , 1547.9397 ]]}, 
      1: {
      'rvec': [1.60104593, 0.02810863, 0.00837912], 
      'tvec': [0.06514053, 0.0669615 , 0.0683511 ]), 
      'corners': [[1787.3639, 1514.471 ],
                 [1876.1868, 1488.2197],
                 [1929.7797, 1542.3733],
                 [1836.7034, 1578.136 ]]
          }, 
      13: {
      'rvec': [ 2.14044012, -0.083175  , -0.02542734], 
      'tvec': [-0.03199416,  0.11906246,  0.6503338 ]), 
      'corners': [[1227.5304, 1100.9562],
                 [1402.6522, 1089.1622],
                 [1423.2228, 1219.2614],
                 [1212.5001, 1232.1893]]
          }
    }
},
...
]

1.2.5 05_run_calibrations.py

自动化两个不同的校准过程：首先是通过SLAM标签校准，其次是夹爪范围的校准

1.2.5.1 SLAM标签校准

用到上一步生成的地图视频tag数据：tag_detection.pkl

见上面视频的3个ArUco标签

方法：通过比较摄像机姿态和标签检测数据，计算摄像机姿态与标签检测之间的变换矩阵。
输入：tag_detection.pkl（标签检测数据）、camera_trajectory.csv（摄像机姿态数据）
输出：校准后的标签变换矩阵tx_slam_tag.json

具体原理涉及到相机位姿的原理：

Python 相机位姿变换_python extrinsic-CSDN博客

一个4x4变换矩阵可以表示为：

[ R | t ]

[ 0 | 1 ]

这个矩阵可以分解为：

旋转矩阵R：前3x3个元素构成了旋转矩阵。这个旋转矩阵将摄像机坐标系中的点转换到标签坐标系中。
平移向量t：最后一行的前3个元素是平移向量，它表示标签在摄像机坐标系中的位置。
单位变换[ 0 | 1 ]：最后一行最后一个元素是1，表示矩阵的变换是单位变换，即保持标签的尺度不变。

在这个例子中，旋转矩阵R的元素表明标签在摄像机坐标系中有一个明显的旋转。平移向量t的元素表明标签相对于摄像机的位置。最后一个元素1表示变换保持标签的尺度不变。

每个变换矩阵都是这样的形式，它们共同描述了摄像机姿态与标签检测之间的变换关系

{
  "tx_slam_tag": [
  0.9746958492903609,  -0.13639169126733092,  -0.1771025087665814,  0.03227576843478422
  -0.22336783914501668, -0.5636180628698587,  -0.7952619048442,     0.09451740467473949
  0.00864894182569355,  0.8146974937377948,    -0.5798217172392903,  0.6434902619256598
  0.0,                  0.0,                    0.0,                   1.0
]
}

1.2.5.2 夹爪范围的校准

用到上一步生成的夹爪校准视频tag数据：tag_detection.pkl

同样见上面视频夹爪上的左右两个个ArUco标签，这里只用了两个标签，因为是通过夹爪的开和关来校准，桌面上tag没有作用。

输入：tag_detection.pkl

输出：gripper_range.json

gripper_id：有时候可以有两个夹爪
tag_id：左右标签的id
width：夹爪极限开和关夹爪的距离尺寸

{
  "gripper_id": 0,
  "left_finger_tag_id": 0,
  "right_finger_tag_id": 1,
  "max_width": 0.1290424054542236,
  "min_width": 0.04346810704159828
}

1.2.6 06_generate_dataset_plan.py

1.2.6.1 相机坐标系到夹爪尖坐标系的转换

1.2.6.2 加载机器人夹爪的校准数据

1.2.6.3 提取每个视频的元数据

1.2.6.4 视频文件进行分组

1.2.6.5 识别每个视频中的夹爪硬件ID

1.2.6.6 确定每个摄像头在演示中的左右位置

1.2.6.7 准备模型训练的数据集

1.3 数据集转成模型训练的格式

1.3.1 数据存储在ReplayBuffer对象中

1.3.2 视频帧处理

1.3.2.1 转为RGB

1.3.2.2 标记检测

1.3.2.3 遮罩

1.3.2.4 鱼眼镜头

1.3.2.5 镜像处理

1.3.2.6 压缩图像

1.3.3 最终数据结构

第二部分 Training Diffusion Policy 训练模型

2.1 入口 train.py

import sys
# use line-buffering for both stdout and stderr
sys.stdout = open(sys.stdout.fileno(), mode='w', buffering=1)
sys.stderr = open(sys.stderr.fileno(), mode='w', buffering=1)
import hydra
from omegaconf import OmegaConf
import pathlib
from diffusion_policy.workspace.base_workspace import BaseWorkspace

# allows arbitrary python code execution in configs using the ${eval:''} resolver
# 注册一个自定义解析器 “eval”，该解析器允许在配置文件中使用 Python 表达式，可以查看yaml文件
OmegaConf.register_new_resolver("eval", eval, replace=True)

#使用 Hydra 的 @hydra.main 装饰器定义主函数 main，函数接受一个 OmegaConf 配置对象作为参数。
@hydra.main(
    version_base=None, # 表示不使用版本控制。
    config_path=str(pathlib.Path(__file__).parent.joinpath(
        'diffusion_policy','config')) # 指定了配置文件所在的路径。
)

def main(cfg: OmegaConf):
    # resolve immediately so all the ${now:} resolvers
    # will use the same time.
    OmegaConf.resolve(cfg) # 解析配置文件中的所有变量
    cls = hydra.utils.get_class(cfg._target_)
    workspace: BaseWorkspace = cls(cfg)
    workspace.run()
if __name__ == "__main__":
    main()

--config-name是 Hydra 模块专用的，用于指定要使用的配置文件的基本名称。

task.dataset_path 是配置文件中的一个键，配置文件位于diffusion_policy/config/task/umi.yml中

这个配置文件通过--config-name=train_diffusion_unet_timm_umi_workspace.yaml 中继承：

# defaults 部分指定了嵌套的配置文件
defaults:
  - _self_ # 表示当前文件
  - task: umi # 表示另一个名为 umi 的配置文件

2.2 Model模型架构

作者做了很多实验，配置文件可以配置不同的模型，当前使用CLIP预训练的ViT-B/16视觉编码器来训练扩散策略

使用：vit_base_patch16_clip_224.openai（视觉编码器） + ConditionalUnet1D（扩散策略）

先用vit_base_patch16_clip_224.openai（视觉编码器）编码特征，在用这些特征给到ConditionalUnet1D（扩散策略）模型生成预测的动作

2.3 Dataset 训练集格式

2.3.1 先采样动作数据

2.3.2 再转换为适合神经网络输入的格式

2.4 Batch数据格式

2.5 输出

第三部分真实机械臂上部署

3.1 UR5机械臂和电脑的通讯方式

3.2 WSG50 夹具和电脑的通讯方式

3.3 ABB机械臂和电脑的通讯方式

3.4 其他国产替代夹具的通讯方式

3.5 改造代码

// 待更，当然如果着急的话，欢迎加入本文开头所说的「大模型机器人二次开发线下营」