Skip to the content.

第13章 视觉模仿学习:从手持夹爪到跨具身策略迁移

原课程: Robots That Learn, UC Berkeley CS294-291, Spring 2026 本章基于 Lecture 8B: Visual Imitation Learning 核心论文: Cheng Chi, Zhenjia Xu, Chuer Pan, Eric Cousineau, Benjamin Burchfiel, Siyuan Feng, Russ Tedrake, Shuran Song. “Universal Manipulation Interface: In-The-Wild Robot Teaching Without In-The-Wild Robots,” RSS 2024. 项目主页: https://umi-gripper.github.io / arXiv:2402.10329


一、本章概要

核心问题: 假如你想让机器人学会洗碗——拧开水龙头、拿起脏盘子、用海绵擦去番茄酱、再把干净盘子放回沥水架——你该如何向机器人”演示”这个任务?传统方法主要有两条路:一是在实验室中用遥操作(teleoperation)控制真实机器人采集数据,二是利用互联网上的被动人类视频。前者需要昂贵的硬件设备和专业操作员,且受限于实验室环境;后者则面临巨大的具身差距(embodiment gap)——人类的手和机器人的夹爪在运动学上截然不同,且被动视频缺乏显式的动作标注。

UMI 的回答: Cheng Chi 等人在 RSS 2024 发表的 Universal Manipulation Interface (UMI) 提供了一条颠覆性的中间路径:用一个搭载 GoPro 相机的 3D 打印手持夹爪作为数据采集设备,在任何环境中以接近零准备时间采集人类演示数据,然后将学到的策略零样本(zero-shot)部署到不同的机器人平台上(UR5 六轴工业臂、Franka FR2 七轴协作臂)。其核心洞察是:通过精心设计的硬件接口(鱼眼镜头 + 侧镜隐性立体视觉 + IMU 感知 SLAM 跟踪)和策略接口(推理时延匹配 + 相对轨迹动作表示 + 扩散策略),可以同时实现动作多样性(动态投掷、双臂折叠、长程清洗)和跨具身可迁移性。

本章的核心命题: 视觉模仿学习与传统行为克隆的本质区别在于,前者必须解决三个在实验室遥操作中不存在或不显著的问题:

  1. 从单目视频中恢复精确的 6-DoF 末端执行器位姿。手持夹爪没有编码器,没有外部动捕系统,只有一个 GoPro 相机的 RGB 视频和内置 IMU 数据。如何恢复度量尺度的(metric-scale)全局位姿?
  2. 构建包含深度信息的 3D 视觉表征。单目相机没有直接深度感知。UMI 如何通过鱼眼镜头 + 物理侧镜的巧妙组合,在单张 RGB 图像中编码隐性立体信息?
  3. 跨具身策略映射。手持夹爪的运动学与目标机器人不同,数据采集环境的坐标系统与部署环境不同。如何设计一个与具体硬件平台无关的策略接口?

这三个问题构成了视觉模仿学习的理论支柱,也是本章从数学第一性原理展开论述的主线。

与前后章节的关系: 第12章 (行为克隆) 介绍了模仿学习的基本范式——Diffusion Policy 作为动作生成模型。本章在此基础上引入视觉感知的维度:第12章的策略假设观测和动作都定义在机器人本体坐标系中,而本章的策略必须从原始像素中推断世界状态。第8章 (机器人动力学) 提供了运动学约束的数学工具,本章用它们来做演示数据的可行性过滤。第4章 (扩散模型入门) 的 DDPM/DDIM 理论基础被 UMI 直接应用于动作生成。


二、核心概念与数学基础

2.1 SLAM 姿态计算:从像素到 SE(3) 位姿

UMI 的姿态恢复系统是整个框架的基石。没有精确的 6-DoF 末端执行器位姿,就没有可用的动作标签;没有度量尺度,就无法实现双臂协调和数据过滤。

2.1.1 视觉-惯性 SLAM 的问题形式化

给定一个图像序列 \(\{\mathbf{I}_1, \mathbf{I}_2, \ldots, \mathbf{I}_T\}\) 和同步的 IMU 测量(加速度计 \(\tilde{\mathbf{a}}_t\) 和陀螺仪 \(\tilde{\boldsymbol{\omega}}_t\)),SLAM 的目标是联合估计:

  1. 相机位姿轨迹: \(\mathbf{T}_t \in \text{SE}(3)\),即从世界坐标系到相机坐标系的刚体变换,参数化为旋转矩阵 \(\mathbf{R}_t \in \text{SO}(3)\) 和平移向量 \(\mathbf{t}_t \in \mathbb{R}^3\):
\[\mathbf{T}_t = \begin{bmatrix} \mathbf{R}_t & \mathbf{t}_t \\ \mathbf{0}^\top & 1 \end{bmatrix} \in \text{SE}(3)\]
  1. 3D 地图点: \(\{\mathbf{X}_j \in \mathbb{R}^3\}_{j=1}^{M}\),即场景中可被多帧观测到的静态特征点在世界坐标系中的位置。
  2. IMU 偏置和尺度: 加速度计偏置 \(\mathbf{b}^a\)、陀螺仪偏置 \(\mathbf{b}^g\) 以及重力方向 \(\mathbf{g}\)。

这构成了一个非线性最小二乘问题(来自 ORB-SLAM3 的视觉-惯性 Bundle Adjustment [Campos et al., 2021]):

\[\boxed{\min_{\mathbf{T}, \mathbf{X}, \mathbf{b}} \underbrace{\sum_{t} \sum_{j \in \mathcal{O}(t)} \rho\left(\|\mathbf{x}_{t,j} - \pi(\mathbf{T}_t, \mathbf{X}_j)\|_{\Sigma_v}^2\right)}_{\text{视觉重投影误差}} + \underbrace{\sum_{t} \|\mathbf{e}_{\text{IMU}}(\mathbf{T}_t, \mathbf{T}_{t+1}, \mathbf{b})\|_{\Sigma_I}^2}_{\text{IMU 预积分误差}}}\]

其中:

2.1.2 IMU 预积分的数学推导

IMU 测量的物理模型为(局部 IMU 坐标系 \(\mathcal{B}\)):

\[\begin{aligned} \tilde{\mathbf{a}}_t &= \mathbf{R}_t^\top (\mathbf{a}_t - \mathbf{g}) + \mathbf{b}^a_t + \boldsymbol{\eta}^a \\ \tilde{\boldsymbol{\omega}}_t &= \boldsymbol{\omega}_t + \mathbf{b}^g_t + \boldsymbol{\eta}^g \end{aligned}\]

其中 \(\mathbf{a}_t\) 是世界坐标系下的真实加速度,\(\boldsymbol{\omega}_t\) 是真实角速度,\(\boldsymbol{\eta}^a \sim \mathcal{N}(0, \sigma_a^2 \mathbf{I})\) 和 \(\boldsymbol{\eta}^g \sim \mathcal{N}(0, \sigma_g^2 \mathbf{I})\) 是高斯白噪声,偏置 \(\mathbf{b}\) 建模为随机游走:\(\dot{\mathbf{b}} = \boldsymbol{\eta}^b\)。

在两个关键帧之间(从时刻 \(i\) 到 \(j\)),IMU 的运动学方程为:

\[\begin{aligned} \mathbf{R}_j &= \mathbf{R}_i \prod_{k=i}^{j-1} \text{Exp}\left((\tilde{\boldsymbol{\omega}}_k - \mathbf{b}^g_k - \boldsymbol{\eta}^g_k) \Delta t\right) \\ \mathbf{v}_j &= \mathbf{v}_i + \mathbf{g} \Delta t_{ij} + \sum_{k=i}^{j-1} \mathbf{R}_k (\tilde{\mathbf{a}}_k - \mathbf{b}^a_k - \boldsymbol{\eta}^a_k) \Delta t \\ \mathbf{t}_j &= \mathbf{t}_i + \sum_{k=i}^{j-1} \left[\mathbf{v}_k \Delta t + \frac{1}{2} \mathbf{g} \Delta t^2 + \frac{1}{2} \mathbf{R}_k (\tilde{\mathbf{a}}_k - \mathbf{b}^a_k - \boldsymbol{\eta}^a_k) \Delta t^2\right] \end{aligned}\]

关键观察:上述方程中 \(\mathbf{R}_i, \mathbf{v}_i, \mathbf{t}_i\) 是变量,而 IMU 测量部分可以预积分——独立于位姿估计计算。定义相对运动增量 \(\Delta \mathbf{R}_{ij}, \Delta \mathbf{v}_{ij}, \Delta \mathbf{t}_{ij}\):

\[\boxed{\begin{aligned} \Delta \mathbf{R}_{ij} &\triangleq \prod_{k=i}^{j-1} \text{Exp}\left((\tilde{\boldsymbol{\omega}}_k - \mathbf{b}^g_k) \Delta t\right) \\ \Delta \mathbf{v}_{ij} &\triangleq \sum_{k=i}^{j-1} \Delta \mathbf{R}_{ik} (\tilde{\mathbf{a}}_k - \mathbf{b}^a_k) \Delta t \\ \Delta \mathbf{t}_{ij} &\triangleq \sum_{k=i}^{j-1} \left[\Delta \mathbf{v}_{ik} \Delta t + \frac{1}{2} \Delta \mathbf{R}_{ik} (\tilde{\mathbf{a}}_k - \mathbf{b}^a_k) \Delta t^2\right] \end{aligned}}\]

预积分的美妙之处在于:这些量仅依赖 IMU 测量和偏置,与位姿变量 \(\mathbf{T}_i, \mathbf{T}_j\) 无关。在优化过程中,只需通过 \(\delta \mathbf{b}\) 的一阶泰勒展开修正预积分值,无需重新计算:

\[\Delta \mathbf{R}_{ij}(\mathbf{b}^g + \delta \mathbf{b}^g) \approx \Delta \mathbf{R}_{ij}(\mathbf{b}^g) \cdot \text{Exp}\left(\frac{\partial \Delta \mathbf{R}_{ij}}{\partial \mathbf{b}^g} \delta \mathbf{b}^g\right)\]

2.1.3 UMI 对 ORB-SLAM3 的关键改进

作者发现原始的 ORB-SLAM3 单目-惯性 SLAM 系统在手持快速运动和弱纹理场景中不够鲁棒。UMI 实现了两个关键改进:

改进 1: 地图即初始化 (Map as Initialization): 原始 ORB-SLAM3 的定位模式在重定位到一个已有地图后不会再创建新的地图点或执行全局优化。对于 UMI 的操控任务——场景在演示过程中被动态改变(物体被移动、盘子被拿起)——这种静态假设导致跟踪失败。UMI 修改 ORB-SLAM3,使其在加载已有地图后继续正常的 SLAM 操作,将已有地图点作为初始优化种子,同时允许创建新的地图点和关键帧:

\(\text{原始: } \mathcal{M}_{\text{localize}} = \mathcal{M}_{\text{map}} \quad \text{(只读)}\) \(\text{UMI 改进: } \mathcal{M}_{\text{init}} = \mathcal{M}_{\text{map}},\; \mathcal{M}(t+1) = \mathcal{M}(t) \cup \{\text{新观测特征点}\}\)

改进 2: 标记增强初始化 (Marker-enhanced Initialization): 在室外环境或重复纹理(树木、图案墙面)场景中,单目 SLAM 的初始化存在尺度歧义和特征误匹配问题。UMI 利用已知尺寸的基准标记(fiducial markers, ArUco 标记 [Garrido-Jurado et al., 2014])来消除可能的歧义解释。注意:标记仅在建图阶段使用,采集演示数据时不放置标记——这意味着策略学习到的观测分布与部署时更一致。

2.1.4 UMI 的数据采集流程:建图-定位-演示

UMI 的 SLAM 系统采用一种先建图后定位(map-then-localize)的数据采集范式:

步骤 1 — 建图: 在一个新场景中,演示者缓慢移动手持夹爪,首先覆盖场景的各个区域,然后沿着预期任务动作的轨迹做更密集的扫描。此过程通常只需约 1 分钟。可选地,在桌面/地面放置一个 fiducial 标记以增强建图鲁棒性。

步骤 2 — 演示: 演示者执行目标任务,每次都记录独立的 MP4 视频文件。所有同一场景内的演示视频都被重定位到步骤 1 建立的同一地图中,共享同一世界坐标系。这是 UMI 实现双臂协调的关键基础设施:虽然两个夹爪的视频分别被重定位,但由于它们在同一个地图坐标系中,任意时刻两个夹爪之间的相对位姿可以直接计算:

\[\mathbf{T}_{\text{gripper}_1 \to \text{gripper}_2}^{(t)} = \mathbf{T}_{\text{world} \to \text{gripper}_2}^{(t)} \cdot (\mathbf{T}_{\text{world} \to \text{gripper}_1}^{(t)})^{-1}\]

SLAM 精度: 在动捕(OptiTrack)基准数据集上,UMI 的 SLAM 系统达到了位置平均绝对轨迹误差(ATE)6.1mm 和旋转 3.5° 的精度。双夹爪之间的相对位姿误差(RPE)为位置 10.1mm 和旋转 0.8°——这对于双臂协调任务来说是高度实用的精度水平。


2.2 3D 视觉表征:从单目 RGB 到隐性深度

UMI 的硬件设计核心是一个反直觉的决策:只用腕部安装的单目相机作为唯一的感知传感器——没有外部相机、没有深度传感器、没有机载计算单元。这带来了极大的可移植性和低成本的代价是:如何从单目 RGB 中提取足够的视觉上下文和深度信息?

2.2.1 鱼眼投影模型

UMI 使用 155° 视场角(FoV)的鱼眼镜头附件。与标准针孔相机模型不同,鱼眼镜头的投影不能用 \(3 \times 3\) 的内参矩阵 \(\mathbf{K}\) 完整描述,需要使用非线性投影模型。常用的 Kannala-Brandt 模型 [Kannala & Brandt, 2006] 将入射角 \(\theta\) 映射为图像半径 \(r\):

\[r(\theta) = k_1 \theta + k_2 \theta^3 + k_3 \theta^5 + k_4 \theta^7\]

鱼眼镜头的关键优势在机器人操控中体现在两方面:(1) 视觉上下文: 大 FoV 覆盖了更多场景信息,使得策略在决策时能感知到物体间的空间关系,对标实验显示去除鱼眼镜头后成功率从 100% 降至 55%;(2) SLAM 鲁棒性: 更大的 FoV 意味着相邻帧之间的视觉特征重叠更多,提高了视觉跟踪的稳定性 [Zhang et al., 2016]。

为什么不使用矫正后的针孔图像? 直接使用原始鱼眼图像作为策略输入是一个重要的设计选择。将 155° FoV 的鱼眼图像矫正到针孔模型会导致严重的边缘拉伸(stretching)——矫正图像的外围区域(鱼眼中被压缩的有用信息)被拉伸放大而失真,而中心区域(包含最重要的操作物体信息)被压缩到一个极小的区域中。UMI 的实验显示策略直接处理原始鱼眼图像效果更好,因为:(a)鱼眼效应保留了中心的高分辨率;(b)卷积神经网络可以自然地适应非线性畸变模式;(c)运动中的相机与鱼眼效应的结合起到了类似随机裁剪(random cropping)的数据增强效果,使策略学会关注任务相关物体而非背景结构。

2.2.2 侧镜的几何学:隐性立体视觉

UMI 最具创造性的硬件设计是在夹爪两侧放置物理镜子,使镜面内容出现在鱼眼图像的边缘区域。这实际上构建了一个无额外成本和重量的隐性立体视觉系统。

几何等价性: 一面镜子 \(\mathcal{M}\) 产生的视图等价于一个反射相机——该虚拟相机的光心是真实相机光心关于镜面的反射。形式上,设镜面的平面方程为 \(\mathbf{n}^\top \mathbf{x} + d = 0\)(\(\mathbf{n}\) 是镜面法向量),真实相机位姿为 \((\mathbf{R}_c, \mathbf{t}_c)\),则虚拟相机的位姿为:

\[\begin{aligned} \mathbf{R}_v &= \mathbf{R}_c (\mathbf{I} - 2\mathbf{n}\mathbf{n}^\top) \\ \mathbf{t}_v &= \mathbf{t}_c - 2(\mathbf{n}^\top \mathbf{t}_c + d) \mathbf{n} \end{aligned}\]

关键结果是:虚拟相机的光心与真实相机的光心不同。这意味着两幅镜中视图和主视图一起构成了一个三目立体视觉系统,三个有效光心提供了三角测量所需的不同视角(如图 4(a) 所示)。当物体在主相机视图中被遮挡时(例如,放在盘子上的番茄酱被盘子边缘挡住),它仍可能在某一个镜中有清晰可见的投影。

数字反射(Digital Reflection): 实验中发现了一个关键的预处理步骤:直接将镜中原始图像送入策略会导致性能下降(从 90% 降至 85%)。这是因为物体在镜中的所有方向(orientation)与原物相反——例如杯柄在镜中是水平镜像的。如果在策略输入中同时出现原物和其镜像,视觉编码器会接收到矛盾的方向信号。解决方案是将镜中内容数字反射回来(在裁剪后做水平翻转),使得三个视图中的物体朝向一致。加上这一步骤后,性能恢复到 100%。

形式化描述: 令原始鱼眼图像为 \(\mathbf{I} \in \mathbb{R}^{H \times W \times 3}\)。定义两个裁剪区域 \(\mathcal{R}_L, \mathcal{R}_R\) 分别对应左右镜面内容。数字反射预处理后,策略的视觉输入为三个通道的拼接:

\[\mathbf{I}_{\text{policy}} = [\mathbf{I}_{\text{main}} ; \text{flip}(\mathbf{I}[\mathcal{R}_L]) ; \text{flip}(\mathbf{I}[\mathcal{R}_R])]\]

其中 \(\mathbf{I}_{\text{main}}\) 是主视图区域,\(\text{flip}(\cdot)\) 是水平翻转操作。

2.2.3 连续夹爪宽度控制

UMI 区别于许多前驱工作的另一个关键设计是将夹爪的二元(开/合)动作替换为连续宽度控制。连续宽度控制通过跟踪夹爪手指上的 fiducial 标记 [Garrido-Jurado et al., 2014] 实现,记录每个时间步的夹爪指间宽度 \(w_t \in [w_{\text{min}}, w_{\text{max}}]\)。

连续控制的物理基础是串联弹性末端执行器(Series Elastic End-Effector, SEED)原理 [Suh et al., 2022]:通过调节 TPU 软手指(95A 硬度,带肋状纹理)的形变程度,UMI 可以隐式地记录和控制夹持力——更大的指间宽度压缩意味着更大的夹持力。这个设计使得像拧水龙头(需要力顺从性)和舀起薄衣服(需要精确的力控制)这样的接触密集型操作成为可能。


2.3 跨具身策略映射:从手持夹爪到任意机械臂

UMI 的核心设计目标之一是使学到的策略与具体机器人平台无关。也就是说,同一个训练好的策略 checkpoint 可以直接部署到 UR5 工业和 Franka FR2 协作臂上,而无需任何微调。这要求策略接口在观察空间和动作空间两个层面都做到硬件无关。

2.3.1 相对末端执行器轨迹

传统的绝对位姿动作表示需要一个全局参考坐标系——在实验室中可以通过机器人基座标定获得,但在野外环境中采集数据时,夹爪的绝对世界坐标对策略来说是无意义且有害的信息。UMI 完全规避了这个问题,将所有的位姿表示为相对于当前末端执行器位姿的相对量。

定义(相对轨迹,Relative Trajectory): 对于从时刻 \(t_0\) 开始的长度为 \(H\) 的动作序列,第 \(k\) 步(\(k = 0, 1, \ldots, H-1\))的动作为:

\[\boxed{\mathbf{a}_k = \mathbf{T}_{t_0}^{-1} \cdot \mathbf{T}_{t_0 + k} \in \text{SE}(3)}\]

即期望在时刻 \(t_0 + k\) 时,末端执行器位姿相对于当前位姿的刚体变换。相比两种替代方案:

相对轨迹的优势是:每一步都相对于同一个参考点 \(t_0\),没有累积误差;同时又不需要绝对坐标系。UMI 的对标实验证实了这一点:相对轨迹实现 100% 成功率,优于增量动作的 80% 和绝对动作的 25%。

SE(3) 上的相对轨迹计算: 在 SLAM 输出中,第 \(t\) 帧的相机位姿为 \(\mathbf{T}_t^{\text{cam}}\)。通过已知的相机-夹爪外参 \(\mathbf{T}_{\text{cam} \to \text{ee}}\)(通过 CAD 模型确定),末端执行器位姿为:

\[\mathbf{T}_t^{\text{ee}} = \mathbf{T}_t^{\text{cam}} \cdot \mathbf{T}_{\text{cam} \to \text{ee}}\]

第 \(t_0\) 帧开始的相对轨迹动作序列为:

\[\mathbf{a}_k = (\mathbf{T}_{t_0}^{\text{ee}})^{-1} \cdot \mathbf{T}_{t_0 + k}^{\text{ee}}, \quad k = 0, 1, \ldots, H-1\]

对于实践中的网络输入,\(\text{SE}(3)\) 变换被参数化为 9 维向量:位置 \(\mathbf{p} \in \mathbb{R}^3\) 和 6D 连续旋转表示 [Zhou et al., 2019] \(\mathbf{r} \in \mathbb{R}^6\)(用旋转矩阵的前两列以避免不连续性)。加上夹爪宽度,完整的动作向量为 \(\mathbb{R}^{10}\)。

2.3.2 本体感知的相对表示

观察空间的相对表示同样重要。历史末端执行器位姿作为本体感知(proprioception) 输入,同样表示为相对轨迹的形式:

\[\mathbf{o}_k^{\text{proprio}} = (\mathbf{T}_{t-\text{obs\_horizon}}^{\text{ee}})^{-1} \cdot \mathbf{T}_{t-\text{obs\_horizon}+k}^{\text{ee}}, \quad k = 0, 1, \ldots, \text{obs\_horizon}-1\]

当观测窗口(observation horizon)设置为 2 时,相对姿态序列本质上提供了速度信息——\(\mathbf{T}_{t-2}^{-1} \cdot \mathbf{T}_{t-1}\) 刻画了末端执行器在最近时间步中的运动方向和速率。这个设计与腕部相机观测一起,使得 UMI 系统完全免标定:在策略执行过程中移动机器人基座不会影响任务表现(实验证实,Fig. 10(a)),只要操作物体仍在可达范围内。

2.3.3 推理时延匹配(Inference-time Latency Matching)

时延是机器人系统中极易被忽视但至关重要的实际问题。在手持夹爪的数据采集阶段,所有信息流(RGB 图像、IMU 数据、夹爪宽度视觉估计)都记录在同一个 GoPro MP4 文件中——它们之间没有相对时延。然而在策略部署阶段,各信息流通过分布式的微控制器采集,具有不同的传输和处理延迟。

观测时延匹配: 对于每个观测流(相机、末端执行器位姿、夹爪宽度),UMI 分别物理测量其端到端延迟(相机通过 QR 码时间戳方法,本体感知通过机器人时间戳或 ICMP round-trip time)。推理时,所有观测流在时间上对齐到延迟最大的流(通常是相机):

\[\text{对于图像时间戳 } t_{\text{obs}}, \quad \mathbf{o}_{\text{proprio}} \leftarrow \text{Interpolate}(\text{proprio\_stream}, t_{\text{obs}} - l_{\text{camera}} + l_{\text{proprio}})\]

动作时延匹配: UMI 策略输出的动作序列假设从最后一个观测帧开始即时执行。但实际机器人系统的执行延迟 \(l_{\text{exec}}\) 意味着指令到达末端执行器的时间比策略预期的要晚。UMI 通过提前发送指令来补偿:丢弃已过时的前几个动作,只执行时间戳在 \(t_{\text{now}} + l_{\text{exec}}\) 之后的动作:

\[\mathbf{a}_{\text{execute}} = \{\mathbf{a}_k \mid t_{\text{desired}}(k) \geq t_{\text{now}} + l_{\text{exec}}\}\]

时延匹配对动态任务的影响极为显著。在动态投掷任务中,禁用推理时延匹配(将所有测量的延迟设为零)导致成功率从 87.5% 降至 57.5%——时延不匹配破坏了机器人达到期望投掷速度所需的手眼协调和精确时序。

2.3.4 基于运动学的数据过滤

由于数据采集时下游部署机器人的运动学约束是未知的,部分演示动作可能在目标机器人上不可行(超出关节限位、达到奇异性、速度超出限制等)。UMI 采用简单但有效的运动学可行性过滤(Kinematic Feasibility Filtering):

给定机器人运动学模型(已知基座位姿和连杆参数),对每个演示轨迹的动作序列进行逆运动学(IK)求解。如果轨迹中任意时间步的期望位姿无有效 IK 解,或者某关节的角度增量超出限制,该演示轨迹被剔除出训练集:

\[\mathcal{D}_{\text{filtered}} = \{\tau \in \mathcal{D} \mid \forall \mathbf{T} \in \tau, \exists \mathbf{q} \in \mathcal{Q}_{\text{feasible}}: \text{FK}(\mathbf{q}) = \mathbf{T}\}\]

过滤后的数据集确保策略学到的动作分布与目标机器人的运动学可行域一致。实验表明,这一简单策略在实践中效果良好,尽管作者在 Limitations 中指出未来可以探索能自动适配不同运动学约束的具身感知(embodiment-aware)策略学习框架。


三、论文精读:UMI

3.1 问题的精确表述

Chi et al. (2024) 将视觉模仿学习问题形式化为以下映射的学习:

\[\pi: \mathcal{O}_{t-H_o:t} \to \mathcal{A}_{t:t+H_a}\]

其中:

3.2 论文的方法论架构

UMI 的方法论分为两个相互耦合的部分:演示接口设计(Demonstration Interface, HD1-HD6)和策略接口设计(Policy Interface, PD1-PD3)。这种”硬件-策略协同设计”的哲学贯穿全文——硬件设计的选择直接影响策略能学到什么,而策略接口的抽象决定了跨具身迁移的可行性。

演示接口设计 (HD1-HD6):

设计要素 关键决策 解决的问题
HD1: 腕部相机 GoPro 仅作为传感器和记录设备,无外部相机 最小化观测具身差距;便携性(整套设备仅 780g)
HD2: 鱼眼镜头 155° FoV,原始鱼眼图像直接输入策略 提供充足视觉上下文;SLAM 鲁棒性;自然数据增强
HD3: 侧镜隐性立体 物理侧镜 + 数字反射预处理 无额外重量的多视角立体深度感知
HD4: IMU 感知跟踪 视觉-惯性联合 SLAM(修改的 ORB-SLAM3) 度量尺度恢复;快速运动下的鲁棒跟踪
HD5: 连续夹爪控制 fiducial 标记跟踪 + SEED 原理 精确释放时序(投掷);隐式力控制(拧水龙头)
HD6: 运动学数据过滤 IK 可行性检查 确保数据与目标机器人运动学兼容

策略接口设计 (PD1-PD3):

设计要素 关键决策 解决的问题
PD1: 推理时延匹配 观测流同步 + 提前发送动作指令 消除训练/推理时域不匹配
PD2: 相对位姿表示 相对轨迹动作 + 相对本体感知 免除全局标定;跨具身迁移
PD3: 相对夹爪间本体感知 双臂协调信号(只对双臂任务) 双臂操作的时空同步

3.3 扩散策略在 UMI 中的应用

UMI 将扩散策略(Diffusion Policy, Chi et al., 2023)作为其核心的动作生成模型。与原始 Diffusion Policy 的直接应用相比,UMI 做了以下关键适配:

动作表示: 动作序列 \(\mathbf{A} = [\mathbf{a}_0, \mathbf{a}_1, \ldots, \mathbf{a}_{H_a-1}] \in \mathbb{R}^{H_a \times 10}\)(每步 10 维:位置 3 + 旋转 6 + 夹爪宽度 1)作为相对轨迹。训练时对干净动作 \(\mathbf{A}^0\) 进行 \(k\) 步前向扩散:

\[\mathbf{A}^k = \sqrt{\bar{\alpha}_k} \mathbf{A}^0 + \sqrt{1 - \bar{\alpha}_k} \boldsymbol{\epsilon}, \quad \boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I})\]

其中 \(\bar{\alpha}_k = \prod_{s=1}^{k} \alpha_s\),\(\alpha_k = 1 - \beta_k\),\(\beta_k\) 是 noise schedule。

视觉条件: 视觉编码器(ViT-B/16 或 ViT-L/14,CLIP 预训练)将原始鱼眼图像(含数字反射后的镜面内容)处理为视觉特征 Token \(\mathbf{z}_{\text{vis}} \in \mathbb{R}^{N_{\text{tokens}} \times d_{\text{vit}}}\)。对于大多数任务,观测窗口 \(H_o = 2\)——使用 2 帧历史图像提供运动线索。对于大规模 in-the-wild 实验,使用 1 帧以加速 ViT-L 的训练和推理。

条件去噪: 噪声预测网络 \(\boldsymbol{\epsilon}_\theta(\mathbf{A}^k, k, \mathbf{z}_{\text{vis}}, \mathbf{z}_{\text{proprio}})\) 以视觉特征和本体感知(历史相对位姿)为条件,预测添加的噪声。训练目标为标准的 DDPM 损失:

\[\mathcal{L}(\theta) = \mathbb{E}_{k \sim \mathcal{U}(1,K), \mathbf{A}^0, \boldsymbol{\epsilon}} \left[\|\boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta(\mathbf{A}^k, k, \mathbf{z}_{\text{vis}}, \mathbf{z}_{\text{proprio}})\|^2\right]\]

推理时使用 DDIM 采样器 [Song et al., 2021],从纯噪声开始通过 \(K_{\text{infer}} = 16\) 步去噪生成动作(训练使用 \(K_{\text{train}} = 50\) 步),获得干净的预测动作 \(\hat{\mathbf{A}}^0\)。

训练超参数汇总(来自论文 Table A1):

任务 观测窗口 动作窗口 频率 ViT 架构 预训练 扩散学习率 ViT学习率 Epochs GPU
杯子整理 2 6 10Hz ViT-B/16 CLIP 3e-4 3e-5 250 4xA10g
动态投掷 2 6 20Hz ResNet-34 无 3e-4 3e-4 350 4xA10g
双臂叠衣 2 6 10Hz ResNet-34 无 3e-4 3e-4 100 4xA10g
洗碗 2 6 10Hz ViT-B/16 CLIP 3e-4 3e-5 90 8xA10g
杯子(in-the-wild) 1 6 10Hz ViT-L/14 CLIP 3e-4 3e-5 50 8xA100

3.4 关键实验分析

3.4.1 能力实验(Capability Experiments)

杯子整理 — 测试具有多模态动作分布(顺时针和逆时针旋转把手的两种等价解)的精确操作。UMI 达到 20/20 成功。关键消融:

动态投掷 — 测试 UMI 捕获和迁移快速流畅的人类动作的能力。成功率 105/120 (87.5%)。禁用时延匹配后降至 69/120 (57.5%):观-执不同步破坏了手眼协调,夹爪-机械臂动作时序错位导致投掷速度不足。

双臂叠衣 — 测试双臂时空协调。成功率 14/20 (70%)。禁用夹爪间相对位姿(PD2.3)后降至 6/20 (30%)——两个机械臂在同时抓起衬衫底边时失去同步,经常一个臂错过抓取。

洗碗 — 超长程任务(7 个顺序依赖步骤),涉及多种物质(牛顿流体水、非牛顿流体番茄酱)、柔性工具(海绵)、关节约束物体(水龙头)。成功率 14/20 (70%)。使用 ResNet-34(从头训练)作为视觉编码器的基线策略完全无法执行任务 (0/10)——说明对于复杂的视觉感知任务,CLIP 预训练的 ViT 视觉编码器是必不可少的。

3.4.2 泛化实验(In-the-Wild Generalization)

在 30 个物理位置(家庭、办公室、餐厅、户外)用 15 个不同杯子收集 1400 个演示后,UMI 策略在两段完全未见的环境中测试:

关键发现: 使用与 in-the-wild 策略相同的 CLIP ViT 视觉预训练模型,但仅用实验室窄域数据训练的基线策略,在泛化测试中成功率为 0%。这表明仅微调大型预训练模型是不够的——收集多样化、野外环境中的真实数据对于有效的泛化仍然至关重要。

3.4.3 数据采集效率

在 15 分钟的时间内,UMI 夹爪的演示收集速度是传统 Spacemouse 遥控操作的 3 倍以上(杯子整理任务:UMI 35 次/15min vs 遥控 11 次/15min)。在动态投掷任务中,遥控操作完全无法产生一个成功的演示——传统遥控界面的延迟和人体工学限制使得快速动态动作几乎不可能被演示。


四、算法与代码实现

4.1 3D-Conditioned Diffusion Policy — 完整 PyTorch 实现

本节实现 UMI 的核心组件:一个以 3D 视觉表征(多视图鱼眼图像特征)为条件的扩散策略,用于生成相对轨迹动作序列。

"""
3D-Conditioned Diffusion Policy for Visual Imitation Learning
基于: Cheng Chi et al., "UMI: Universal Manipulation Interface," RSS 2024

核心设计要素 (标注来源):
- PD2: 相对轨迹动作表示 (§2.3.1)
- 扩散策略: DDPM 训练 + DDIM 推理 (Chi et al., RSS 2023)
- 视觉编码器: CLIP ViT 预训练 + 微调 (§3.3)
- 观测窗口: Ho=2 的历史观测 (§3.3)
- 多视图: 主视图 + 左镜面 + 右镜面 (HD3, §2.2.2)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Tuple, Optional, Dict
from einops import rearrange


# ==============================================================================
# 模块 1: 相对轨迹动作表示 (PD2, §2.3.1)
# ==============================================================================

def relative_trajectory_to_actions(
    ee_poses: torch.Tensor,      # [B, T, 7] — (x,y,z, qw,qx,qy,qz) 四元数
    reference_idx: int = 0
) -> torch.Tensor:
    """
    将绝对 SE(3) 位姿序列转换为相对轨迹动作序列.
    第 k 个动作为: a_k = T_{ref}^{-1} @ T_{ref+k}
    
    Args:
        ee_poses: [B, T, 7] — 绝对末端执行器位姿 (xyz + quaternion xyzw)
        reference_idx: 参考时刻的索引 (通常为 0)
    
    Returns:
        相对轨迹: [B, T, 10] — (相对位置3, 6D旋转6, 夹爪宽度1)
    """
    B, T, _ = ee_poses.shape
    
    # 提取参考位姿
    T_ref_pos = ee_poses[:, reference_idx, :3]    # [B, 3]
    T_ref_quat = ee_poses[:, reference_idx, 3:7]  # [B, 4]
    
    # 在 SE(3) 上计算相对变换
    rel_positions = []
    rel_rotations_6d = []
    
    for k in range(T):
        T_k_pos = ee_poses[:, k, :3]
        T_k_quat = ee_poses[:, k, 3:7]
        
        # 相对位移: p_rel = R_ref^T @ (p_k - p_ref)
        R_ref = quaternion_to_matrix(T_ref_quat)  # [B, 3, 3]
        delta_pos = T_k_pos - T_ref_pos            # [B, 3]
        rel_pos = torch.bmm(R_ref.transpose(1, 2), delta_pos.unsqueeze(-1)).squeeze(-1)
        rel_positions.append(rel_pos)
        
        # 相对旋转: R_rel = R_ref^T @ R_k
        R_k = quaternion_to_matrix(T_k_quat)
        R_rel = torch.bmm(R_ref.transpose(1, 2), R_k)
        # 转换为 6D 连续表示 [Zhou et al., 2019] — 使用前两列避免不连续性
        r_6d = R_rel[:, :, :2].reshape(B, 6)
        rel_rotations_6d.append(r_6d)
    
    rel_positions = torch.stack(rel_positions, dim=1)  # [B, T, 3]
    rel_rotations_6d = torch.stack(rel_rotations_6d, dim=1)  # [B, T, 6]
    
    # 拼接: [rel_pos (3) | rel_rot_6d (6)] = 9维
    rel_poses = torch.cat([rel_positions, rel_rotations_6d], dim=-1)  # [B, T, 9]
    
    return rel_poses


def quaternion_to_matrix(quat: torch.Tensor) -> torch.Tensor:
    """
    将四元数 (w,x,y,z) 转换为 3x3 旋转矩阵.
    公式来源: 标准 SO(3) 参数化.
    """
    B = quat.shape[0]
    qw, qx, qy, qz = quat[:, 0], quat[:, 1], quat[:, 2], quat[:, 3]
    
    R = torch.zeros(B, 3, 3, device=quat.device, dtype=quat.dtype)
    R[:, 0, 0] = 1 - 2 * (qy**2 + qz**2)
    R[:, 0, 1] = 2 * (qx * qy - qz * qw)
    R[:, 0, 2] = 2 * (qx * qz + qy * qw)
    R[:, 1, 0] = 2 * (qx * qy + qz * qw)
    R[:, 1, 1] = 1 - 2 * (qx**2 + qz**2)
    R[:, 1, 2] = 2 * (qy * qz - qx * qw)
    R[:, 2, 0] = 2 * (qx * qz - qy * qw)
    R[:, 2, 1] = 2 * (qy * qz + qx * qw)
    R[:, 2, 2] = 1 - 2 * (qx**2 + qy**2)
    return R


# ==============================================================================
# 模块 2: 正弦位置编码 (用于扩散时间步和观测序列位置)
# ==============================================================================

class SinusoidalPosEmb(nn.Module):
    """
    扩散时间步 k 的正弦位置编码.
    参考: DDPM [Ho et al., 2020] 和 Transformer [Vaswani et al., 2017]
    """
    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: [B] — 扩散时间步整数 (0 到 K-1)
        Returns:
            [B, dim] — 正弦位置编码
        """
        device = x.device
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = x.unsqueeze(-1).float() * emb.unsqueeze(0)
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
        return emb


# ==============================================================================
# 模块 3: 多视图视觉编码器 (HD3, §2.2.2)
# ==============================================================================

class MultiViewVisualEncoder(nn.Module):
    """
    多视图视觉编码器: 处理主视图 + 左右镜面反射视图.
    
    设计依据:
    - HD3: 侧镜提供隐性立体视觉, 三视图构成等效三目系统 (§2.2.2)
    - 镜面内容需先做数字反射 (水平翻转) 以避免方向歧义
    - ViT 编码: CLIP 预训练的 ViT-B/16 或 ViT-L/14 (§3.3, 论文 Table A1)
    
    注: 对于双臂任务, 有两个 GoPro 相机, 因此总视图数为 2×3=6.
    这里以单臂(3视图)实现为主, 双臂通过两个实例扩展.
    """
    def __init__(
        self,
        feature_dim: int = 256,          # 输出特征维度
        backbone: str = "vit_base",      # [UNSPECIFIED] — 论文使用具体 ViT 变体, 此处提供占位
        pretrained: bool = True,         # CLIP 预训练权重
        freeze_bn: bool = True,          # 冻结 BatchNorm
    ):
        super().__init__()
        self.feature_dim = feature_dim
        
        # [UNSPECIFIED] 论文未详细说明 CLIP ViT 的具体架构参数
        # 此处使用 ResNet-18 作为轻量占位替代, 实际训练时替换为 ViT-B/16
        import torchvision.models as models
        self.backbone_main = models.resnet18(pretrained=False)
        self.backbone_mirror_left = models.resnet18(pretrained=False)
        self.backbone_mirror_right = models.resnet18(pretrained=False)
        
        # 移除分类头, 保留特征提取器
        resnet_feat_dim = 512
        self.backbone_main.fc = nn.Identity()
        self.backbone_mirror_left.fc = nn.Identity()
        self.backbone_mirror_right.fc = nn.Identity()
        
        # 融合三个视图特征
        self.fusion_proj = nn.Sequential(
            nn.Linear(resnet_feat_dim * 3, feature_dim * 2),
            nn.LayerNorm(feature_dim * 2),
            nn.ReLU(),
            nn.Linear(feature_dim * 2, feature_dim)
        )
        
    def forward(
        self, 
        img_main: torch.Tensor,           # [B, C, H, W] — 原始鱼眼图像主视图
        img_mirror_left: torch.Tensor,    # [B, C, H_m, W_m] — 左镜面裁剪 (已数字反射)
        img_mirror_right: torch.Tensor    # [B, C, H_m, W_m] — 右镜面裁剪 (已数字反射)
    ) -> torch.Tensor:
        """
        Returns:
            visual_feature: [B, feature_dim] — 融合的多视图视觉特征
        """
        # 三个视图分别编码
        feat_main = self.backbone_main(img_main)          # [B, 512]
        feat_left = self.backbone_mirror_left(img_mirror_left)
        feat_right = self.backbone_mirror_right(img_mirror_right)
        
        # 拼接并融合 (HD3: 数字反射确保三个视图中物体方向一致)
        feat_concat = torch.cat([feat_main, feat_left, feat_right], dim=-1)
        visual_feature = self.fusion_proj(feat_concat)   # [B, feature_dim]
        
        return visual_feature


# ==============================================================================
# 模块 4: 条件 U-Net 去噪网络 (DDPM 噪声预测器)
# ==============================================================================

class ConditionalUnet1D(nn.Module):
    """
    一维条件 U-Net 噪声预测网络 ε_θ(A^k, k, z_vis, z_proprio).
    
    架构设计 (来自 Diffusion Policy [Chi et al., 2023]):
    - 基于 1D 时间卷积的 U-Net 结构
    - FiLM (Feature-wise Linear Modulation) 条件注入
    - 视觉特征和本体感知作为全局条件
    
    输入:
    - A^k: [B, T_a, 10] — 噪声化的动作序列 (相对轨迹, 9维位姿 + 1维夹爪宽度)
    - k: [B] — 扩散时间步
    - z_vis: [B, D_vis] — 多视图融合视觉特征
    - z_proprio: [B, T_o, 9] — 观测窗口内的相对末端执行器位姿历史
    
    输出:
    - ε_pred: [B, T_a, 10] — 预测的噪声
    """
    def __init__(
        self,
        action_dim: int = 10,              # 动作维度: rel_pos(3) + rel_rot_6d(6) + grip(1)
        action_horizon: int = 6,           # 动作预测窗口 T_a (§3.3, 论文 Table A1)
        obs_horizon: int = 2,              # 观测窗口 T_o
        visual_feature_dim: int = 256,     # 视觉特征维度
        proprio_dim: int = 9,              # 本体感知维度 (不含夹爪宽度)
        diffusion_step_embed_dim: int = 128,
        down_dims: Tuple[int] = (256, 512, 1024),
        kernel_size: int = 5,
        n_groups: int = 8,
    ):
        super().__init__()
        self.action_dim = action_dim
        self.action_horizon = action_horizon
        self.obs_horizon = obs_horizon
        
        # ========== 编码器 ==========
        # 图像特征投影 (作为 FiLM 条件)
        self.visual_proj = nn.Sequential(
            nn.Linear(visual_feature_dim, visual_feature_dim),
            nn.Mish(),
            nn.Linear(visual_feature_dim, visual_feature_dim)
        )
        
        # 本体感知嵌入
        self.proprio_encoder = nn.Sequential(
            nn.Linear(propri_dim * obs_horizon, proprio_dim * obs_horizon * 2),
            nn.Mish(),
            nn.Linear(propri_dim * obs_horizon * 2, visual_feature_dim)
        )
        
        # 扩散时间步嵌入
        self.time_mlp = nn.Sequential(
            SinusoidalPosEmb(diffusion_step_embed_dim),
            nn.Linear(diffusion_step_embed_dim, diffusion_step_embed_dim * 2),
            nn.Mish(),
            nn.Linear(diffusion_step_embed_dim * 2, diffusion_step_embed_dim)
        )
        
        # 全局条件融合
        cond_dim = visual_feature_dim * 2 + diffusion_step_embed_dim
        self.cond_encoder = nn.Sequential(
            nn.Linear(cond_dim, down_dims[0] * 2),
            nn.Mish(),
            nn.Linear(down_dims[0] * 2, down_dims[0] * 2)
        )
        
        # ========== U-Net 下采样路径 ==========
        # 输入投影
        self.input_proj = nn.Conv1d(action_dim, down_dims[0], kernel_size=3, padding=1)
        
        # 下采样块
        down_blocks = []
        in_dim = down_dims[0]
        for out_dim in down_dims[1:]:
            down_blocks.append(ResnetBlock1D(in_dim, out_dim, kernel_size, n_groups, cond_dim=down_dims[0] * 2))
            in_dim = out_dim
        self.down_blocks = nn.ModuleList(down_blocks)
        
        # ========== U-Net 中间瓶颈 ==========
        mid_dim = down_dims[-1]
        self.mid_block = ResnetBlock1D(mid_dim, mid_dim, kernel_size, n_groups, cond_dim=down_dims[0] * 2)
        
        # ========== U-Net 上采样路径 ==========
        up_blocks = []
        for out_dim in reversed(down_dims[:-1]):
            up_blocks.append(ResnetBlock1D(mid_dim + in_dim, out_dim, kernel_size, n_groups, cond_dim=down_dims[0] * 2))
            # 注: 标准 U-Net 中 up block 的输入是 mid/prev 输出 + 对应下采样层的跳跃连接
        self.up_blocks = nn.ModuleList(up_blocks)
        
        # ========== 输出头 ==========
        self.output_proj = nn.Conv1d(down_dims[0], action_dim, kernel_size=1)
        
    def compute_global_condition(
        self, k: torch.Tensor, visual_feat: torch.Tensor, proprio: torch.Tensor
    ) -> torch.Tensor:
        """计算 U-Net 全局条件 (FiLM 参数)"""
        # 扩散步嵌入
        k_emb = self.time_mlp(k)  # [B, D_time]
        
        # 视觉特征
        v_emb = self.visual_proj(visual_feat)  # [B, D_vis]
        
        # 本体感知
        p_emb = self.proprio_encoder(
            rearrange(proprio, 'b t d -> b (t d)')
        )  # [B, D_vis]
        
        # 融合
        global_feat = torch.cat([k_emb, v_emb, p_emb], dim=-1)
        global_cond = self.cond_encoder(global_feat)
        
        return global_cond
    
    def forward(
        self,
        noisy_action: torch.Tensor,   # [B, T_a, action_dim]
        k: torch.Tensor,              # [B]
        visual_feat: torch.Tensor,    # [B, D_vis]
        proprio: torch.Tensor         # [B, T_o, proprio_dim]
    ) -> torch.Tensor:
        """
        预测添加到噪声化动作中的噪声.
        
        Returns:
            noise_pred: [B, T_a, action_dim] — 预测噪声
        """
        # 计算全局条件
        global_cond = self.compute_global_condition(k, visual_feat, proprio)
        
        # 输入投影: [B, T_a, D] -> [B, D, T_a]
        x = rearrange(noisy_action, 'b t d -> b d t')
        x = self.input_proj(x)  # [B, down_dims[0], T_a]
        
        # 下采样路径 + 保存跳跃连接
        skips = [x]
        for block in self.down_blocks:
            x = block(x, global_cond)
            skips.append(x)
        
        # 中间瓶颈
        x = self.mid_block(x, global_cond)
        
        # 上采样路径 (带跳跃连接)
        for i, block in enumerate(self.up_blocks):
            skip = skips[-(i + 2)]  # 从最深到最浅
            x = torch.cat([x, skip], dim=1)
            x = block(x, global_cond)
        
        # 输出投影
        noise_pred = self.output_proj(x)  # [B, action_dim, T_a]
        noise_pred = rearrange(noise_pred, 'b d t -> b t d')
        
        return noise_pred


class ResnetBlock1D(nn.Module):
    """
    U-Net 中的 1D ResNet 块.
    包含 FiLM 条件调制 [Perez et al., 2018].
    """
    def __init__(self, in_dim: int, out_dim: int, kernel_size: int, n_groups: int, cond_dim: int):
        super().__init__()
        self.norm1 = nn.GroupNorm(n_groups, in_dim)
        self.conv1 = nn.Conv1d(in_dim, out_dim, kernel_size, padding=kernel_size // 2)
        self.norm2 = nn.GroupNorm(n_groups, out_dim)
        self.conv2 = nn.Conv1d(out_dim, out_dim, kernel_size, padding=kernel_size // 2)
        
        # FiLM: 从全局条件生成 scale 和 shift
        self.film_scale = nn.Linear(cond_dim, out_dim)
        self.film_shift = nn.Linear(cond_dim, out_dim)
        
        # 跳跃连接投影 (当维度不匹配时)
        self.residual_conv = nn.Conv1d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()
        
    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: [B, C_in, L]
            cond: [B, cond_dim] — 全局条件
        """
        residual = self.residual_conv(x)
        
        # FiLM 调制参数
        scale = self.film_scale(cond).unsqueeze(-1)  # [B, C_out, 1]
        shift = self.film_shift(cond).unsqueeze(-1)  # [B, C_out, 1]
        
        x = self.norm1(x)
        x = F.mish(x)
        x = self.conv1(x)
        
        # FiLM 注入: 仿射变换
        x = x * (1 + scale) + shift
        
        x = self.norm2(x)
        x = F.mish(x)
        x = self.conv2(x)
        
        return x + residual


# ==============================================================================
# 模块 5: 3D-Conditioned Diffusion Policy (完整策略)
# ==============================================================================

class DiffusionPolicy3D(nn.Module):
    """
    3D-Conditioned Diffusion Policy — UMI 的核心策略网络.
    
    完整流程:
    1. 视觉编码: 多视图鱼眼图像 -> 融合视觉特征 (§2.2, HD3)
    2. 本体感知: 历史相对位姿 -> 速度/运动信息 (§2.3, PD2.2)
    3. 扩散去噪: 条件 U-Net 从噪声中恢复动作序列 (§3.3)
    4. 推理时延匹配: 动作执行时序对齐 (§2.3, PD1)
    
    训练: DDPM 前向扩散 + 条件去噪损失
    推理: DDIM 采样 (16步)
    """
    def __init__(
        self,
        action_dim: int = 10,
        action_horizon: int = 6,
        obs_horizon: int = 2,
        visual_feature_dim: int = 256,
        proprio_dim: int = 9,
        num_diffusion_steps: int = 50,      # 训练扩散步数 K_train
        num_inference_steps: int = 16,       # 推理去噪步数 K_infer (DDIM, §3.3)
        noise_schedule_beta_start: float = 1e-4,
        noise_schedule_beta_end: float = 0.02,
        device: str = "cuda"
    ):
        super().__init__()
        self.action_dim = action_dim
        self.action_horizon = action_horizon
        self.obs_horizon = obs_horizon
        self.num_diffusion_steps = num_diffusion_steps
        self.num_inference_steps = num_inference_steps
        self.device = device
        
        # 子模块
        self.visual_encoder = MultiViewVisualEncoder(feature_dim=visual_feature_dim)
        self.noise_pred_net = ConditionalUnet1D(
            action_dim=action_dim,
            action_horizon=action_horizon,
            obs_horizon=obs_horizon,
            visual_feature_dim=visual_feature_dim,
            proprio_dim=proprio_dim
        )
        
        # ====== 噪声调度 (DDPM linear schedule) ======
        betas = torch.linspace(
            noise_schedule_beta_start, noise_schedule_beta_end, num_diffusion_steps
        )
        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)  # alpha_bar_k
        
        self.register_buffer("betas", betas)
        self.register_buffer("alphas", alphas)
        self.register_buffer("alphas_cumprod", alphas_cumprod)
        self.register_buffer("sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod))
        self.register_buffer("sqrt_one_minus_alphas_cumprod", torch.sqrt(1.0 - alphas_cumprod))
        
    def add_noise(self, clean_action: torch.Tensor, timestep: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        前向扩散: A^k = sqrt(alpha_bar_k) * A^0 + sqrt(1 - alpha_bar_k) * epsilon
        
        Args:
            clean_action: [B, T_a, D] — 干净动作
            timestep: [B] — 扩散时间步索引
            
        Returns:
            noisy_action: [B, T_a, D]
            noise: [B, T_a, D] — 添加的真实噪声
        """
        sqrt_alpha_bar = self.sqrt_alphas_cumprod[timestep]  # [B]
        sqrt_one_minus = self.sqrt_one_minus_alphas_cumprod[timestep]  # [B]
        
        # 广播到动作维度
        sqrt_alpha_bar = sqrt_alpha_bar.view(-1, 1, 1)
        sqrt_one_minus = sqrt_one_minus.view(-1, 1, 1)
        
        noise = torch.randn_like(clean_action)
        noisy_action = sqrt_alpha_bar * clean_action + sqrt_one_minus * noise
        
        return noisy_action, noise
    
    def compute_loss(
        self,
        img_main: torch.Tensor,
        img_mirror_left: torch.Tensor,
        img_mirror_right: torch.Tensor,
        proprio: torch.Tensor,       # [B, T_o, 9]
        action: torch.Tensor         # [B, T_a, 10] — 干净动作 (相对轨迹)
    ) -> torch.Tensor:
        """
        DDPM 训练损失 (§3.3):
        L(θ) = E_{k, A^0, ε} [||ε - ε_θ(A^k, k, z_vis, z_proprio)||^2]
        
        这是公式 (3.3) 的直接实现.
        """
        B = action.shape[0]
        device = action.device
        
        # 1. 视觉编码 (§2.2)
        visual_feat = self.visual_encoder(img_main, img_mirror_left, img_mirror_right)
        
        # 2. 随机采样扩散步
        k = torch.randint(0, self.num_diffusion_steps, (B,), device=device)
        
        # 3. 前向加噪
        noisy_action, noise = self.add_noise(action, k)
        
        # 4. 条件噪声预测
        noise_pred = self.noise_pred_net(noisy_action, k, visual_feat, proprio)
        
        # 5. MSE 损失
        loss = F.mse_loss(noise_pred, noise)
        
        return loss
    
    @torch.no_grad()
    def ddim_sample(
        self,
        img_main: torch.Tensor,
        img_mirror_left: torch.Tensor,
        img_mirror_right: torch.Tensor,
        proprio: torch.Tensor        # [B, T_o, 9]
    ) -> torch.Tensor:
        """
        DDIM 推理采样 [Song et al., 2021].
        
        从纯噪声开始, 通过 K_infer 步去噪生成动作序列.
        UMI 的默认配置: K_infer = 16, K_train = 50 (§3.3, 论文 Table A1 注释).
        
        Returns:
            action: [B, T_a, 10] — 预测的干净动作序列 (相对轨迹)
        """
        B = proprio.shape[0]
        device = proprio.device
        
        # 视觉编码
        visual_feat = self.visual_encoder(img_main, img_mirror_left, img_mirror_right)
        
        # 初始化纯噪声 (A^T ~ N(0, I))
        action = torch.randn(B, self.action_horizon, self.action_dim, device=device)
        
        # DDIM 等间距时间步子序列
        step_size = self.num_diffusion_steps // self.num_inference_steps
        timesteps = list(range(self.num_diffusion_steps - 1, -1, -step_size))
        
        for i in range(len(timesteps) - 1):
            t = torch.full((B,), timesteps[i], device=device, dtype=torch.long)
            t_prev = torch.full((B,), timesteps[i + 1], device=device, dtype=torch.long)
            
            # 预测噪声
            noise_pred = self.noise_pred_net(action, t, visual_feat, proprio)
            
            # 预测干净动作 (DDIM 公式)
            alpha_bar_t = self.alphas_cumprod[t].view(-1, 1, 1)
            alpha_bar_t_prev = self.alphas_cumprod[t_prev].view(-1, 1, 1)
            
            # x_0 = (x_t - sqrt(1 - alpha_bar_t) * eps_theta) / sqrt(alpha_bar_t)
            pred_x0 = (action - torch.sqrt(1 - alpha_bar_t) * noise_pred) / torch.sqrt(alpha_bar_t)
            
            # DDIM 确定性更新: x_{t-1} = sqrt(alpha_bar_{t-1}) * x_0 + sqrt(1 - alpha_bar_{t-1}) * eps_theta
            action = (
                torch.sqrt(alpha_bar_t_prev) * pred_x0 +
                torch.sqrt(1 - alpha_bar_t_prev) * noise_pred
            )
        
        return action
    
    def forward(
        self,
        img_main: torch.Tensor,
        img_mirror_left: torch.Tensor,
        img_mirror_right: torch.Tensor,
        proprio: torch.Tensor
    ) -> torch.Tensor:
        """
        推理前向传播: 从视觉观察 + 本体感知生成动作.
        
        Returns:
            action: [B, T_a, 10] — 相对轨迹动作序列
        """
        return self.ddim_sample(img_main, img_mirror_left, img_mirror_right, proprio)


# ==============================================================================
# 模块 6: 训练与推理脚本骨架
# ==============================================================================

def train_epoch(
    model: DiffusionPolicy3D,
    dataloader,
    optimizer: torch.optim.Optimizer,
    device: str = "cuda"
) -> float:
    """
    单个训练 epoch.
    
    数据加载器应 yield:
    - img_main: [B, C, H, W] — 原始鱼眼图像
    - img_mirror_left: [B, C, H_m, W_m] — 左镜面裁剪 (已数字反射)
    - img_mirror_right: [B, C, H_m, W_m] — 右镜面裁剪 (已数字反射)
    - proprio: [B, T_o, 9] — 历史相对末端执行器位姿
    - action: [B, T_a, 10] — 目标相对轨迹动作
    """
    model.train()
    total_loss = 0.0
    
    for batch in dataloader:
        img_main = batch["img_main"].to(device)
        img_left = batch["img_mirror_left"].to(device)
        img_right = batch["img_mirror_right"].to(device)
        proprio = batch["proprio"].to(device)
        action = batch["action"].to(device)
        
        optimizer.zero_grad()
        loss = model.compute_loss(img_main, img_left, img_right, proprio, action)
        loss.backward()
        
        # 梯度裁剪 — [UNSPECIFIED] 论文未提及具体阈值, 使用常见 default
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        
        optimizer.step()
        total_loss += loss.item()
    
    return total_loss / len(dataloader)


def deploy_policy_loop(
    model: DiffusionPolicy3D,
    camera_stream,          # 实时相机流 (提供三视图)
    robot_controller,       # 机器人控制器接口
    obs_horizon: int = 2,
    latency_camera_ms: float = 35.0,     # 相机延迟 (PD1.1, §A1)
    latency_execution_ms: float = 120.0,  # 执行延迟 (PD1.2, §A4)
    control_freq_hz: float = 10.0         # 控制频率 (Table A1)
):
    """
    策略部署循环 — 集成推理时延匹配 (PD1, §2.3.3).
    
    注: 这是一个骨架代码, 真实部署需要与具体机器人控制栈集成.
    相机延迟、执行延迟的具体值需通过物理测量获得 (详见论文 §A).
    """
    import time
    
    proprio_history = []  # 存储历史相对末端执行器位姿
    
    while True:
        # 1. 获取当前图像 (PD1.1: 观测时延匹配)
        t_capture = time.time()
        img_main, img_left, img_right = camera_stream.get_frames()
        
        # 观测时延补偿: 图像在 t_capture - latency_camera 时刻被实际采集
        t_eff = t_capture - latency_camera_ms / 1000.0
        
        # 2. 获取当前末端执行器位姿并转换为相对表示
        current_pose = robot_controller.get_ee_pose()  # SE(3) 位姿
        
        if len(proprio_history) == 0:
            # 第一个观测: 以自身为参考, 相对值为零/单位变换
            rel_pose = torch.zeros(1, 9)
        else:
            ref_pose = proprio_history[-1][0]  # 上一个观测窗口的最后一个位姿
            rel_pose = compute_relative_pose(current_pose, ref_pose)
        
        proprio_history.append(rel_pose)
        if len(proprio_history) < obs_horizon:
            continue  # 等待积累足够的观测历史
        proprio_history = proprio_history[-obs_horizon:]
        
        # 3. 策略推理: 生成动作序列
        proprio_tensor = torch.stack(proprio_history, dim=0).unsqueeze(0)  # [1, T_o, 9]
        
        with torch.no_grad():
            action_seq = model(
                img_main.unsqueeze(0),
                img_left.unsqueeze(0),
                img_right.unsqueeze(0),
                proprio_tensor.to(model.device)
            )  # [1, T_a, 10]
        action_seq = action_seq[0]  # [T_a, 10]
        
        # 4. 动作时延匹配 (PD1.2): 丢弃过时的动作, 只执行未来的动作
        t_now = time.time()
        dt_per_step = 1.0 / control_freq_hz
        exec_offset_steps = int((latency_execution_ms / 1000.0) / dt_per_step)
        
        # 执行时间戳在 t_now + latency_execution 之后的所有有效动作
        for step_k in range(exec_offset_steps, len(action_seq)):
            desired_action = action_seq[step_k]  # 相对轨迹的第 k 步
            
            # 将相对轨迹转换回绝对目标位姿
            target_timestamp = t_now + (step_k - exec_offset_steps) * dt_per_step
            
            # 发送给机器人控制器
            robot_controller.send_target_pose(
                relative_action=desired_action,
                current_pose=current_pose,
                target_time=target_timestamp
            )
        
        # 等待下一个控制周期
        elapsed = time.time() - t_capture
        sleep_time = max(0, dt_per_step - elapsed)
        time.sleep(sleep_time)


def compute_relative_pose(current_pose, reference_pose):
    """计算 relative SE(3) transform, 返回 [9] 向量."""
    # [UNSPECIFIED] 具体实现依赖机器人控制栈的位姿表示方式
    # 对于 SE(3), 相对位置: R_ref^T @ (p_cur - p_ref)
    # 相对旋转: R_ref^T @ R_cur -> 6D 表示
    pass


# ==============================================================================
# 使用示例
# ==============================================================================

if __name__ == "__main__":
    # 初始化模型
    model = DiffusionPolicy3D(
        action_dim=10,
        action_horizon=6,
        obs_horizon=2,
        visual_feature_dim=256,
        proprio_dim=9,
        num_diffusion_steps=50,
        num_inference_steps=16
    ).cuda()
    
    # 前向传播 shape 测试
    B = 4
    dummy_img_main = torch.randn(B, 3, 224, 224).cuda()
    dummy_img_left = torch.randn(B, 3, 112, 112).cuda()
    dummy_img_right = torch.randn(B, 3, 112, 112).cuda()
    dummy_proprio = torch.randn(B, 2, 9).cuda()
    dummy_action = torch.randn(B, 6, 10).cuda()
    
    # 训练损失
    loss = model.compute_loss(
        dummy_img_main, dummy_img_left, dummy_img_right,
        dummy_proprio, dummy_action
    )
    print(f"Training loss: {loss.item():.4f}")
    
    # 推理生成
    action = model(
        dummy_img_main[:1], dummy_img_left[:1], dummy_img_right[:1],
        dummy_proprio[:1]
    )
    print(f"Generated action shape: {action.shape}")  # Expected: [1, 6, 10]
    print("All checks passed — 3D-Conditioned Diffusion Policy is operational.")

4.2 关键实现说明

多视图视觉融合 (MultiViewVisualEncoder): 论文中的视觉编码器基于 CLIP 预训练的 ViT(ViT-B/16 用于大部分任务,ViT-L/14 用于大规模 in-the-wild 任务)。上述代码中使用 ResNet-18 作为占位实现,原因有二:(1) 完整 ViT 的实现代码量过大,不利于教学展示;(2) 论文明确指出 ResNet-34 在某些任务(如动态投掷、双臂叠衣)中效果可接受,ViT 主要对视觉复杂任务(如洗碗)至关重要。实际部署时应替换为 CLIP ViT 并加载官方预训练权重。

[UNSPECIFIED] 标记汇总:

实现细节 状态 默认选择
ViT 架构具体参数(head数、层数、patch size) UNSPECIFIED 使用 ViT-B/16 的默认参数
梯度裁剪阈值 UNSPECIFIED max_norm=1.0
学习率预热策略 UNSPECIFIED 无预热(论文未提及)
EMA(指数移动平均)权重 UNSPECIFIED 未实现(论文未提及)
数值反射预处理的具体实现 UNSPECIFIED 假设在数据预处理阶段完成
机器人控制接口 UNSPECIFIED 骨架代码仅提供接口签名

训练建议(基于论文 Table A1):


五、核心 Takeaway 与延伸阅读

5.1 五个核心洞察

1. 硬件设计与算法设计是一体两面的同一个问题。 UMI 最深刻的洞察不在于任何一个独立的算法创新,而在于它证明了:通过精心的硬件-算法协同设计(co-design),可以用极简的硬件(一个 GoPro + 3D 打印夹爪,总成本约 370 美元)实现传统方法需要昂贵遥操作设备和专业实验室才能达到——甚至无法达到——的性能。鱼眼镜头的选择直接决定了策略的视觉上下文范围,侧镜的几何布局决定了隐性立体信息的质量,相对轨迹表示的选择使得策略天然具有硬件无关性。在具身智能研究中,忽视硬件接口的设计等价于在错误的抽象层次上优化。

2. 相对化是一切可迁移性的基础。 UMI 反复强调”相对”(relative)的重要性:相对轨迹作为动作表示(PD2.1)、相对位姿作为本体感知(PD2.2)、相对夹爪间位姿作为双臂协调信号(PD2.3)。这不是巧合——消除对全局坐标系的依赖意味着消除对具体部署环境的依赖。在机器人基座位姿变化(Fig. 10a)和跨平台部署(UR5 vs Franka FR2)场景中,相对化设计的价值被直接验证。

3. 时延不是实现细节,而是算法设计的核心约束。 大多数机器人学习论文将时延视为工程问题而非科学问题。UMI 的消融实验给出了一个明确的纠正:在动态投掷任务中,禁用推理时延匹配导致成功率从 87.5% 骤降至 57.5%。在毫秒级的物理交互中,数据采集和策略部署之间的时域不匹配会系统性地破坏策略的行为质量。这个洞察适用于所有需要精确手眼协调的机器人学习系统。

4. 多模态数据中的”更多”是双刃剑。 行为克隆中的多模态动作分布——对于同一观测,存在多种同样有效的动作选择(例如杯柄可以顺时针或逆时针旋转)——通常被视为需要”学习覆盖”的目标分布。UMI 提醒我们,过多的多模态也可能是策略未能充分感知环境状态的表现:当移除鱼眼镜头后,即便物体在视野内,策略也因为无法区分不同情境而变得过度抖动(jittery)。感知能力不足会自动转化为虚假的动作多模态。

5. 野外数据是泛化能力的不可替代来源。 一个重要的反直觉发现是:使用与 in-the-wild 策略完全相同的 CLIP ViT 预训练视觉编码器,但仅用实验室窄域数据训练的基线策略,在野外环境中泛化成功率为 0%。”更多数据”和”更好的预训练模型”都不能替代数据的分布多样性——因为预训练模型虽然在语义空间上有泛化能力,但机器人操控需要的不仅是语义理解,更是对物理交互动态的覆盖。在塑料桌上、木桌上、金属桌上、水面上的杯子操作,各有其独特的接触动力学和视觉外观,这些差异无法仅仅通过更好的视觉编码器来弥补。

5.2 局限性

作者在论文 §VIII 中坦诚地指出了三个主要局限:

  1. 运动学数据过滤是粗粒度的: 当前方法简单地丢弃运动学不可行的演示,而不是学习如何将有效但运动学不可行的动作适配到新平台上。这损失了一部分有价值的数据。
  2. SLAM 依赖纹理: 在纯白墙壁等纹理匮乏的环境中,视觉 SLAM 失效。可能解决方向:使用静态第三人称视角相机 + 夹爪上的额外 fiducial 标记。
  3. 数据采集效率仍不如人手演示: 手持夹爪的重量(780g)和体积限制了演示的速度和灵巧度。未来可探索更轻的材料或直接从人手动作迁移到灵巧手。

5.3 推荐延伸阅读

方向 推荐文献 关联
扩散策略的理论基础 Chi et al., “Diffusion Policy: Visuomotor Policy Learning via Action Diffusion,” RSS 2023 第12章; UMI 的动作生成核心
SLAM 系统 Campos et al., “ORB-SLAM3: An Accurate Open-Source Library for Visual, Visual-Inertial, and Multimap SLAM,” IEEE T-RO, 2021 UMI 的 SLAM 基础
双臂遥操作 Zhao et al., “Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware,” RSS 2023 (ALOHA) 另一种低成本的演示采集方案
从人类视频学习 Shaw et al., “VideoDex: Learning Dexterity from Internet Videos,” CoRL 2023 无手持设备的视觉模仿路线
跨具身泛化 Shafiullah et al., “On Bringing Robots Home,” arXiv 2023 (Dobb-E) 与 UMI 类似的”野外采集+部署”范式
机器人学习放大 Brohan et al., “RT-1: Robotics Transformer for Real-World Control at Scale,” RSS 2023 从大规模多任务数据中学习
视觉-语言-动作模型 Physical Intelligence, “pi-0.5: a Vision-Language-Action Model,” CoRL 2025 第17章; 语言引导的通用操作

六、练习与思考

6.1 (SLAM 理论基础) 给定一个在平坦桌面上以恒定速度 \(v\) 水平移动的 UMI 手持夹爪,GoPro 以 30Hz 记录鱼眼图像和 IMU 数据。推导在直线匀速运动(无旋转)假设下,IMU 预积分项 \(\Delta \mathbf{v}_{ij}\) 和 \(\Delta \mathbf{t}_{ij}\) 的解析表达式。讨论为什么即使在这种简单运动中,仅靠 IMU 积分的位姿漂移在 10 秒内可能超过 10cm。

6.2 (侧镜几何学) 假设左侧镜面位于相机坐标系中的平面 \(x = -d\) 处(法向量指向 \(-x\) 方向),右侧镜面位于 \(x = +d\) 处(法向量指向 \(+x\) 方向)。相机光心位于原点,鱼眼镜头水平 FoV 为 155°。计算:(a) 两个虚拟相机的光心位置;(b) 对于一个位于 \((0, 0.5\text{m}, 0.3\text{m})\)(正前方 0.5m,高出桌面 0.3m)的物体,它能否同时出现在三个视图(主视图、左镜面、右镜面)中?如果能,在各视图中分别对应哪个方向的投影?

6.3 (Cross-Embodiment 运动学适配) 你有一个在 UR5 机器人(六轴工业臂,工作半径 850mm)上训练的 UMI 策略。现在你想将其部署到 Franka FR2(七轴协作臂,工作半径 855mm)上。FR2 的最后一个关节(关节 7,对应末端执行器 pitch)的运动范围被限制在 \([-50°, +50°]\),而 UR5 手腕 3 的对应范围是 \([-360°, +360°]\)(连续旋转)。设计一个运动学数据过滤方案,用于在部署前判断训练数据集中哪些演示序列在 FR2 上可行。你的方案需要处理 SE(3) 位姿轨迹的时序连续性问题。

6.4 (扩散策略的多模态建模) 考虑杯子整理任务中的多模态动作分布:当杯子把手朝向正对机器人时,顺时针和逆时针旋转把手的两种动作序列在对数概率空间中形成两个分离的模式。分析:(a) 为什么 MSE 回归无法同时建模两种模式(会学到两者的均值——一个无效的动作)?(b) 扩散模型如何通过迭代去噪过程自然地解决多模态问题?给出一个形式化的论证:在什么条件下,DDPM 的后验分布 \(p_\theta(\mathbf{A}^0 \mid \mathbf{A}^K, \mathbf{o})\) 可以同时为两个模式分配相近的高概率?

6.5 (时延匹配的设计空间) UMI 的时延匹配方案要求为每个机器人平台物理测量各个组件的延迟。设计一个自动化的延迟标定流程,使得机器人在首次部署到新平台时能够自行测量相机流延迟、关节本体感知延迟和执行延迟,而无需人工干预。讨论你的方案在实际部署中可能遇到哪些挑战(提示:考虑时钟同步、系统负载变化、温度对传感器延迟的影响)。

6.6 (端到端思考) UMI 通过将视觉 SLAM(用于动作提取)和扩散策略(用于动作生成)分离为两个独立的系统,实现了模块化的设计。讨论端到端方法的替代方案:如果你试图训练一个直接从原始鱼眼视频序列 \((\mathbf{I}_{1:T})\) 预测动作序列 \(\mathbf{A}_{1:T}\) 的单一神经网络,它会面临哪些根本性的困难?SLAM 作为显式中间表示提供了哪些归纳偏置(inductive biases)是端到端方法难以从数据中学到的?


本章完成于 2026 年 6 月。UMI 的代码和硬件设计已在 https://umi-gripper.github.io 开源。建议读者在阅读本章后访问项目主页观看实验视频——许多操作行为的质量和泛化能力仅通过文字描述难以充分传达。


← 第12章 行为克隆
↑ 目录
第14章 运动控制案例研究 →