伯克利Digit——基于下一个token预测技术预测机器人动作token：从带RL到不带RL的自回归预测

v_JULY_v

6745人浏览 · 2024-01-28 00:00:55

v_JULY_v · 2024-01-28 00:00:55 发布

前言

这两天，我在朋友圈说道，本质上来讲，斯坦福mobile aloha就是大模型机器人领域里的缩略版ChatGPT时刻，虽然aloha本身并没有像ChatGPT那么强大(毕竟才三人团队)，但会让大模型机器人成为今2024年最大的爆发点，为何这么讲呢？

一方面，ChatGPT对于各大公司最大的意义是，让大家看到了在大模型上投入的希望，而mobile aloha则让大家看到了在大模型机器人上投入的希望，更会加大：大家在RL上的投入与决心
二方面，虽说视频生成也会很火，但还没到能改变许多行业的地步，估计得明后年，至于无人驾驶还得再搞几年

而第二大爆发点，我觉得是agent，这个“大”更多是对社会的影响力，当然，agent是一个框架流程，本身也会嵌入在大模型机器人的设计里头

由于我今年特别看好大模型机器人的发展，且作为一个创业者我希望不断挑战难度更大的事情，而作为一个技术研究者或大模型开发者我希望把各个领域的技术给串起来(比如NLP大模型和CV大模型)，综合为一个大场景而服务，毕竟作为大模型应用而言，场景第一数据第二模型第三

第一部分 Digit：基于类似下个token预测策略 + RL训练机器人

1.1 Digit所使用的动作预测模型：Causal Transformer

目前，尽管一些人形机器人在特定环境下表现良好，但一般存在泛化和适应新环境方面的问题

为了更好的解决这个问题，来自加州大学伯克利分校的研究人员于23年3月通过此篇论文《Real-World Humanoid Locomotion with Reinforcement Learning》提出了运用Causal Transformer的方法，其通过自回归的方式从观察-动作历史信息中预测下一个动作，也就是模型的输出(预测的动作)只依赖于其输入(观察-动作历史信息)中的先前信息(看到这，我惊呆了，它竟然把动作预测干成了类似于：对NLP句子中下一个token的预测)

具体来说，在处理“观察-动作对”时，模型会将每个“观察-动作对”作为一个token，并通过自注意力机制来学习这些token之间的关系
在自注意力计算中，模型会为每个token分配权重，这些权重反映了在预测当前动作时，序列中其他token的重要性
由于Causal Transformer限制了自注意力只能考虑前面的token，它能够捕捉到序列中的因果依赖，即当前动作的决策是基于之前观察-动作的历史信息

最终，Causal Transformer能够使机器人在复杂和动态的环境中，在没有未来信息的情况下做出适应性更强的决策，从而使得机器人可以在室内外各种环境中稳健行走，应对不同地形，甚至还能背个书包，提袋垃圾

顺带说两点你可能意想不到的点

Our controller is trained entirely in simulation and deployed to the real world zero-shot，from page 5
This setting is relevant since our robot is blind and may find itself in such situations during deployment，form page13

1.2 Digit的原理与框架

如下图所示

(A)训练分为两个步骤
$\rightarrow$ 首先，假设环境是完全可观察的，并训练一个「教师状态策略」 $\pi_{s}\left(a_{t} \mid s_{t}\right)$
$\rightarrow$ 其次，采用教师模仿和强化学习相结合的方法来培养「学生观察策略」
即，The policy learning includes two steps: teacher state policy training and student observation policy learning
(B)利用Isaac Gym模拟器，在4个A100 GPU和数千个随机环境中进行大规模并行训练，模拟了机器人的刚体和接触动力学
最后，将该策略迁移到真实机器人上
(C)Causal Transformer通过对下一个动作进行自回归预测来从观察和行动历史中进行训练。他们认为观察-行动历史包含有关世界的有用信息，而强大的Transformer模型可以利用这些信息来调整上下文中的动作

1.2.1 问题建模：可观察的马尔可夫决策过程POMDP

他们将控制问题描述为马尔可夫决策过程MDP，为离散时间决策过程提供了一个数学框架

MDP包括
状态空间S、动作空间A
转移函数 $P\left(s_{t+1} \mid s_{t}, a_{t}\right)$ 和标量奖励函数 $R\left(s_{t+1} \mid s_{t}, a_{t}\right)$ ，分别确定在时间步长 $t$ 采取行动 $a_t$ 后从状态 $s_t$ 转移到状态 $s_{t+1}$ 的概率，以及对智能体行动质量的反馈
解决MDP问题的方法是通过强化学习，旨在找到最优策略，以有限或无限的时间范围内最大化预期累积奖励。实际应用中，由于无法准确估计真实环境中潜在状态，在存在噪声观测空间时需要修改MDP框架以反映观测不确定性(In practice, estimating true underlying state of an environment is impossible for real-worldapplications. In the presence of a noisy observation space, the MDP framework needs to bemodified to reflect the uncertainty in the observations)
故通过引入观测空间 $O$ 和观测函数 $Z\left(o_{t} \mid s_{t}\right)$ 来实现，该函数确定将状态 $s_t$ 观察为 $o_t$ 的概率(This can be done by introducing an observation space O and an observation function Z(ot|st), which determines the probability of observing state st as ot)
此时，MDP变成部分可观察的马尔可夫决策过程POMDP，智能体必须根据噪声观测而非真实环境状态做出决策(The MDP now becomes a Partially Observable Markov DecisionProcess (POMDP), where the agent must make decisions based on its noisy observations ratherthan the true state of the environment.)

1.2.2 Causal Transformer基于观察-动作序列确定行动策略

接下来的目标是找到一个策略 $\pi _o$ ，用于在POMDP问题中进行实际部署

他们的策略将由“观察-动作”组成的长度为 $l$ 的历史轨迹作为输入，该历史轨迹表示为 $o_{t}, a_{t-1}, o_{t-1}, a_{t-2}, \ldots, o_{t-l+1}, a_{t-l}$ ，并输出下一个动作 $a_t$

为了实现这一目标，他们采用Transformer(50)来进行序列轨迹建模和动作预测，Transformer最关键的特征是使用自注意力机制，在计算输出时能够权衡每个输入元素的重要性。自注意力机制通过自注意力函数实现，该函数接收一组查询Q、键K和值V，并输出加权和，计算方式如下所示：
$\operatorname{Attention}(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^{T}}{\sqrt{d_{k}}}\right) V$
其中 $d_k$ 代表键的维度
通过自注意力机制，Transformer能够捕获输入元素之间较长距离上的依赖关系。他们将运动轨迹中每个观察-动作对表示为一个token。Transformers可以通过反复这样一个过程：向每个token(Q和K上使用softmax)分配权重并将token(V)映射到特征空间，从而有效地突出相关观察和动作信息，并推断步态和接触状态等重要信息
We represent each observation-action pair in the locomotion trajectory as a token. Transformers are able to extract the structural information of these tokens through a repeated processof assigning weights to each token (softmax on Q and K) in time, and mapping the tokens (V )into features spaces, effectively highlighting relevant observations and actions and thus enablingthe inference of important information such as gait and contact states.
他们利用多层感知机(MLPs)将每个观察-动作对嵌入到特征空间中去(We employ Multi-LayerPerceptrons (MLPs) to embed each observation-action pair into a feature space)
为了捕捉序列中每个token位置信息，在特征中添加正弦位置编码器(sinepositional encoding)
且通过限制自注意力机制只关注前面的token，在利用时间依赖关系推断观察与行动之间产生causal Transformer(We leverage the temporal dependencies among the observations and actions byrestricting the self-attention mechanism to only attend to preceding tokens, resulting in a causaltransformer)

1.2.3 确定目标函数：基于模仿学习和RL联合训练

在强化学习中，智能体必须通过试错不断积累经验并更新策略以优化决策过程。然而，在复杂的高维环境中，获得有用的奖励信号可能需要大量交互和模拟步骤，因此直接使用RL来优化策略在观测空间中是缓慢且资源密集的。为了克服这些限制，他们采用了以下方法

首先，假设环境是完全可观测的(模拟环境中一切可观测)，并通过模拟训练一个教师状态策略 $\pi_{s}\left(a_{t} \mid s_{t}\right)$ ，该训练快速且资源高效。他们调优奖励函数(如步态参数)，直到在模拟中获得最佳状态策略
First, we assume that the environment is fully observable and train a teacher state policy πs(at|st) using simulation. Thistraining is fast and resource-efficient, and we tune the reward functions, such as gait-parameters,until an optimal state policy is obtained in simulation.
接下来，通过Kullback-Leibler(KL)发散将学习到的状态策略提炼为观测策略(Next, we distill the learned state policy to an observation policy through Kullback-Leibler (KL) divergence)
同时考虑状态空间和观测空间之间差异可能导致次优决策问题，采用一种联合优化方法将RL损失与状态策略监督相结合定义目标函数 $L\left(\pi_{o}\right)=L_{R L}\left(\pi_{o}\right)+\lambda D_{K L}\left(\pi_{o} \| \pi_{s}\right)$
Joint optimization with reinforcement learning.The discrepancy between the state spaceand the observation space can result in suboptimal decision-making if relying solely on state-policy supervision, as policies based on these separate spaces may have different reward man-ifolds with respect to the state and observation representations.
To overcome this issue, weutilize a joint optimization approach combining RL loss with state-policy supervision

$\rightarrow$ 相当于训练目标包括两项，一项是基于教师监督策略的模仿学习，一项是强化学习两个方面(an imitation learning term based on teacher policy supervision and a reinforcement learning term based)，其中 $\lambda$ 代表状态策略监督的权重因子， $L_{R L}\left(\pi_{o}\right)$ 表示强化学习损失， $D_{K L}\left(\pi_{o} \| \pi_{s}\right)$ 是观察策略 $\pi _o$ 和状态策略 $\pi _s$ 之间的KL散度
$\rightarrow$ 在训练过程中，权重因子 $\lambda$ 逐渐退火至零，在训练视界的中点通常达到零，这使得观测策略能够早期从老师那里获益(即状态策略是观测策略的老师)，并最终超越它(The weighting factor λ is gradually annealed to zero over the course of the training process,typically reaching zero at the mid-point of the training horizon, which enables the observation policy to benefit from the teacher early on and learn to surpass it eventually)

值得注意的是，他们的方法不需要任何预先计算的轨迹或离线数据集，因为状态策略监督和强化学习监督都通过策略学习进行优化
It is importantto note that our approach does not require any pre-computed trajectories or offline datasets, asboth the state-policy supervision and RL-supervision are optimized through on-policy learning.
实验中使用了actor-critic方法来训练RL策略，并且没有共享权重
动作空间由16个驱动关节PD设定值和8个驱动腿关节预测PD增益组成。他们未训练任何控制四个脚趾电机的策略，而是使用固定PD增益将电机设置为默认位置。这种方法在基于模型控制中被广泛采用(55,56)。奖励函数受人类行走生物力学研究启发，并经过试错调整。在奖励设计中没有预先计算步态库

1.3 迁移到现实中后的表现

为了模拟机器人的闭链动力学，引入了“虚拟弹簧”模型。在模拟中随机化机器人的动态属性、控制参数和环境物理属性，以及添加噪声和延迟到观察中

从模拟到现实的转移方面，研究人员在机器人公司Agility Robotics(前身是Dynamic Robotics Laboratory)提供的高保真度模拟器中验证策略，该模拟器准确模拟了Digit机器人的动态和物理属性

经过实验，Digit机器人能够在多种环境中可靠行走，展现出对外部干扰的鲁棒性
以及在不同地形和载荷条件下的适应性(左：脚被绊住，右：下坡)
顺带保持手臂摆动协调有力，不顺拐

第二部分不用RL而直接类似Next Token Prediction训练机器人

上一部分的论文《Real-World Humanoid Locomotion with Reinforcement Learning》，其作者是Ilija Radosavovic, Tete Xiao, Bike Zhang, Trevor Darrell, Jitendra Malik, Koushil Sreenath
24年2.29，上一团队(Ilija Radosavovic, Bike Zhang, Baifeng Shi, Jathushan Rajasegaran, Sarthak Kamat, Trevor Darrell, Koushil Sreenath, Jitendra Malik)通过此篇论文《Humanoid Locomotion as Next Token Prediction》，把机器人动作的预测干成类似下一个token预测的工作进一步推进了下

本质区别在于后者去掉了RL

2.1 根据YouTube视频且基于Transformer做自回归预测

最新这篇论文，我个人觉得比较重要的一个价值在于，它可以把比如YouTube上大量的人类走路的视频都作为训练数据来源之一，很有我之前介绍过的Google genie的感觉(详见此文：Google发布Genie硬杠Sora：通过大量无监督视频训练最终生成可交互虚拟世界)

他们使用由Agility Robotics开发的全尺寸Digit类人机器人

首先在仿真中收集传感器运动轨迹的数据集。这些包括通过强化学习训练的神经网络策略的完整轨迹(即上一部分的这个工作：Real-world humanoid locomotion with reinforcement learning），以及来自三个不同来源的不完整轨迹：

基于模型预测控制的Agility Robotics控制器
人体动作捕捉
人类的YouTube视频

使用计算机视觉技术重建人类视频，并通过逆向运动学重新定位动作捕捉和YouTube轨迹。然后，训练一个Transformer模型来自回归地预测轨迹。在测试时，自回归地执行动作，并忽略感知预测

最终，研究发现，仅使用离线数据训练的自回归策略在测试环境中与使用强化学习的最先进方法相当(We find that our autoregressive policies trained from offline data alone are comparable to the state-of-the-art approaches that use reinforcement learning (33 ) in tested settings)

2.1.1 训练方法

首先，每个传感器运动轨迹都是一系列感官观察和动作的序列： $\mathcal{T}=\left(o_{1}, a_{1}, o_{2}, a_{2}, \ldots, o_{T}, a_{T}\right)$ ，然后将轨迹分词为K个token，以获得 $t=\left(t_{1}, t_{2}, t_{3}, \ldots, t_{K}\right)$
接下来，他们的目标是训练一个神经网络自回归地建模 $p(t)$ ： $p(t)=\prod_{k=1}^{K} p\left(t_{k} \mid t_{k-1}, \ldots, t_{1}\right)$
然后，通过最小化轨迹数据集上的负对数似然来训练他们的模型 $L=\sum_{t \in \mathcal{D}}-\log p(t)$
假设一个具有恒定方差的高斯分布，并训练一个神经网络来最小化均方误差，即预测token与实际token之间的误差：
$L=\frac{1}{K} \sum_{k=1}^{K}\left(\widehat{t}_{k}-t_{k}\right)^{2}$

2.1.2 模型架构之Transformer

他们的模型是一个基本的Transformer

给定完整或不完整数据的轨迹，首先将轨迹分词为token
为每个模态学习单独的线性投影层，但在时间上共享
为了编码时间信息，使用位置嵌入。假设 $o_{i} \in \mathcal{R}^{m}$ 和 $a_{i} \in \mathcal{R}^{n}$ ，则：
$t_{i}=\operatorname{concat}\left(o_{i}, a_{i}\right)$
$h_{i}^{0}=W t_{i}$
其中 $W \in \mathcal{R}^{d \times(m+n)}$ 是一个线性投影层，用于将连接的观测和动作模态投影到一个d维嵌入向量中，上标 0表示第0层的嵌入，即输入层
当动作不可用时，使用一个掩码token $[\mathrm{M}] \in \mathcal{R}^{n}$ 来替代 $a_i$ ，并且 $[M]$ 被初始化为一个随机向量，并与整个模型一起进行端到端的学习，该模型将嵌入向量序列 $H_{0}=\left\{h_{1}^{0}, h_{2}^{0}, \ldots, h_{t}^{0}\right\}$ 作为输入
Transformer架构包含 L个层，每个层由一个多头自注意模块和一个MLP模块组成。假设第 $l$ 层的输出是 $H_l$ ，则 $H_{l+1}$ 的输出为
$\begin{aligned} \tilde{H}_{l} & =\text { LayerNorm }\left(H_{l}\right) \\ \tilde{H}_{l} & =\tilde{H}_{l}+M \operatorname{HSA}\left(\tilde{H}_{l}\right) \\ H_{l+1} & =\tilde{H}_{l}+M \operatorname{LP}\left(\tilde{H}_{l}\right) \end{aligned}$
在这里，多头自注意力具有因果掩码(the multi-head self-attention has causal masking)，其中token只关注自身和过去的token
一旦token通过所有层进行处理，通过学习线性投影层 $\widehat{W} \in \mathcal{R}^{(m+n) \times d}$ 将嵌入投影到预测的状态和动作(we project the embedding to predicted states and actions, by learning a linear projection layer cW ∈ R(m+n)×d)
$\begin{aligned} \widehat{t}_{i+1} & =\widehat{W} h_{i}^{L} \\ \widehat{o}_{i+1} & =\left(\widehat{t}_{i+1}\right)_{0: m} \\ \widehat{a}_{i+1} & =\left(\widehat{t}_{i+1}\right)_{m:(m+n)} \end{aligned}$
然后使用目标函数 $L=\frac{1}{K} \sum_{k=1}^{K}\left(\widehat{t}_{k}-t_{k}\right)^{2}$ 来训练transformer。在token被屏蔽的情况下，不应用任何损失。他们使用两种类型的数据来训练他们的transformer，如下图所示

这使他们能够使用各种数据源，从而在数据方面实现扩展

在推理时，他们的transformer模型始终可以访问观察-动作对。在这种设置下，为每个观察-动作对token自
回归地应用transformer模型

通过对过去的观察和动作进行条件预测，预测下一个动作(或观察-动作对)并执行该动作
然后从机器人那里获取观察结果并丢弃预测的观察结果，使用观察到的观察结果和预测的动作作为下一组token，并将它们与过去的对连接起来以预测下一个观察-动作对

2.2 数据集

如之前所述，数据集包括来自不同来源的轨迹：

神经网络策略
作为训练轨迹的第一个来源，使用了一个经过大规模强化学习训练的神经网络策略(33)
具体来说，这个策略是在Isaac Gym(Isaac gym: High performance gpu-based physics simulation for robot learning)中的数千个随机环境中使用数十亿个样本进行训练的。在Agility Robotics的模拟器中运行这个策略，并收集了10秒钟的10k个轨迹，这些轨迹都是在平坦的地面上进行的，没有进行域随机化。每个轨迹都以从截断正态分布中采样的速度命令为条件，具体如下：向前的线性速度为[0.0, 1.0]m/s，向侧面的线性速度为[−0.5, 0.5]m/s，旋转角速度为[−0.5, 0.5]rad/s。由于可以访问数据生成策略，能够记录完整的观测结果以及模型预测的准确动作
基于模型的控制器
作为轨迹的第二个来源，使用了由Agility Robotics开发的基于模型的控制器。它是部署在Digit人形机器人上并在Agility Robotics的模拟器中可用的控制器。收集了两组10k个行走轨迹，每个轨迹持续10秒，在平坦的地面上进行。在这两种情况下，对速度指令进行采样如下：
线性前进速度 [−1.0， 1.0]米/秒
线性侧向速度 [−1.0， 1.0]米/秒，以及转向角速度 [−1.0， 1.0]弧
度/秒

他们使用默认的基于模型的配置为一个集合，并对另一个集合的腿长、步行间隙和地板的弹性进行随机化。
由于该控制器输出关节力矩，与他们的关节位置动作空间不一致。只记录观测结果，不记录动作。这些数据作为具有合理观测结果的轨迹的来源，但没有动作。
人体动作捕捉
作为轨迹的下一个来源，他们使用来自KIT数据集(The KIT motion-language dataset)通过AMASS存储库(AMASS: Archive of motion capture as surface shapes)分发的人体动作捕捉(MoCap)记录。这些数据是在实验室环境中使用光学基于标记的跟踪记录的。该数据集包含约4k个轨迹。
他们使用约1k个站立、行走和奔跑的轨迹的子集。除了不包含地面真实动作外，MoCap轨迹还面临着额外的挑战：不同的形态。也就是说，MoCap轨迹捕捉到了人体关键点在3D空间中的位置。为了将这些轨迹用于训练机器人，他们解决了一个逆运动学问题，以找到相应的机器人姿势

他们制定了一个逆运动学优化问题：
其中， $q$ 是广义坐标下的机器人状态， $N$ 和 $d_t$ 是优化的时间范围和采样时间，优化变量包括 $q$ 、 $\dot{\mathbf{q}}$
对于约束条件，(12b)是姿势 $q$ 的欧拉积分，(12c)约束了 $q$ 和 $\dot{\mathbf{q}}$ 在它们的可接受集合 Q和 V内的范围
在成本函数中， φtraj跟踪人类轨迹中的关键点位置，而 φreg表示正则化成本，例如关节速度最小化和平滑性
以及来自YouTube的人体视频
人们进行各种活动的互联网视频是学习人类步态的潜在数据来源。
然而，原始像素没有关于人类状态和行为的信息。为了恢复这一点，首先运行一个计算机视觉跟踪算法PHALP(Tracking people by predicting 3d appearance, location and pose)来提取3D中的人类轨迹。这提供了人体关节SMPL(Smpl: A skinned multi-person linear model)参数的估计和世界坐标系中人体关节的嘈杂估计。他们使用人体关节位置将运动重新定位到人形机器人上，使用上述逆运动学优化
一旦将运动重新定位从互联网视频到人形机器人轨迹，他们使用低优化成本对轨迹进行过滤。请注意，这些数据的规模带来了噪声的代价