Transformer Transformer: 面向运动条件机器人协同设计的统一模型

# 面向运动约束的机器人协同设计统一模型 来源：https://transformer-transformer.github.io/ ## Transformer Transformer Huy Ha (https://huy-ha.github.io/),Karen Liu (https://tml.stanford.edu/),Shuran Song (https://shurans.github.io/) ## 并非所有机器人都生而平等 并非所有机器人都生而平等——但如果你能针对特定任务设计一个呢？我们提出**Transformer Transformer**，这是一个统一模型，恰好做到这一点：给它一个操作演示，它就能生成一个完整的机器人——每个连杆、关节、电机和惯性参数——针对该运动进行优化。我们制造了这样一个设计，用于在ALOHA2 (https://transformer-transformer.github.io/#ref-2)双臂平台上进行抛布操作；相比原始设计，追踪误差降低了73%，最大关节速度降低了30%。 这一成果背后是一个基于**RoboTokens**训练的扩散变换器——**RoboTokens**是对机器人具身、状态和动作的统一标记化。相同的架构涵盖了具身空间（轮式双臂、四足、人形）和用例（具身生成、跨具身控制）。它并非过拟合于某个奖励函数，而是一个动力学模型，其与奖励无关的预测在推理时被转换为与奖励相关的价值预测，然后通过一种我们称为**动力学自引导**的过程来引导具身扩散。在三个设计空间上的实验表明，该方法能够零样本优化未见过的奖励和轨迹，在性能和运行时间上均优于进化基线。 ## 技术总结视频 一段17分钟的讲解——或滚动查看图文并茂的版本。 ## 对于给定的操作任务，什么是最好的机器人？ 机器人的具身决定了它能出色完成哪些任务。我们投入大量精力让机器人更智能——更大的数据集、更好的算法、更强大的策略架构——但糟糕的具身会拖累甚至最优秀的策略。你可以收集世界上所有的投掷数据，但如果你机器人的形状远非最优，它可能会把自己投出去而不是球。 所以，让我们具体地定义这个问题。给定一个目标末端执行器运动和一个奖励函数，我们希望生成一个完整的机器人具身用于该任务——所谓*完整*，是指每个连杆、关节、电机和惯性参数，外加一个驱动整个系统的控制器。我们称之为**运动约束的机器人协同设计**。 *一个随机过程生成的机器人尝试目标运动：它还没完成就倒下了 🙃* ## 演示、生成、验证 我们的框架将机器人协同设计重新定义为一个三步过程：你**演示**所需的末端执行器运动（例如，来自UMI3 (https://transformer-transformer.github.io/#ref-3)的人类演示），我们的模型**生成**一个优化的具身，然后同一个模型通过直接控制来**验证**该设计。同一个网络扮演三个角色——生成器、批评者和跨具身控制器——通过在机器人的统一标记化上联合训练实现。 三步流程：演示、生成、验证。**三步，一个模型。**给定目标末端执行器运动（例如，来自人类演示）和用户定义的奖励，Transformer Transformer*生成*一个完整的机器人——几何、运动学、惯性、电机——然后*控制*该机器人追踪该运动。无需单独的优化器、批评者或控制器管线。论文包含三个部分：统一的**机器人表示**（RoboTokens）、统一的**架构**（Transformer Transformer）和统一的**训练目标**（动力学自引导）。在架构部分之前，你有时间猜出这个双关语 🥸 第一部分 ## RoboTokens：统一的机器人表示 RoboTokens标记化方案。**每个机器人变成一个类型化标记序列。**蓝色标记编码具身（连杆、关节、电机）；橙色标记编码动力学（状态和动作）。相同的分词器处理四足机器人、人形机器人或灵巧手——无需针对每个机器人的适配器——因此单个模型可以在所有类型上进行学习。如果我们想要一个模型协同设计*任何*机器人，网络需要一个跨越所有具身的共享词汇表。RoboTokens就是那个词汇表：针对每个连杆、关节、电机、状态和动作的类型化标记——组织方式使得单个序列既能描述机器人的*具身*（时不变），又能描述其*动力学*（时变）。 RoboTokens具有以下关键设计目标： - **完整。**一个RoboToken序列捕获了关于刚性铰接机器人的一切：5种具身标记类型（连杆、固定关节、旋转/滑动关节、球关节、电机）加上状态和动作标记。 - **灵活。**标记通过ID相互引用——一个关节标记知道它连接的两个连杆标记——因此相同的格式能处理6自由度手臂或35自由度双足。 - **一致。**分词器规范化冗余的空间偏移，并将惯性/变换数据存储在单个框架中，减少了模型需要学习的方差。 - **可扩展。**新的标记类型无需更改架构即可插入——我们为轨迹追踪添加目标末端执行器姿态标记，通过指向相关的末端执行器和时间步ID，无需对机器人其余部分重新标记化。 - **可优化。**与MJCF等自回归文本格式不同，扩散连续值的RoboTokens既提供了*全局可控性*（后面的标记在去噪步骤中影响前面的标记），也提供了*可微性*——这是在推理时基于梯度的奖励优化所需的两个要素。 我们对MuJoCo Menagerie10 (https://transformer-transformer.github.io/#ref-10)中的11个机器人进行了标记化，涵盖了两个数量级的质量（0.65千克的灵巧手到67.5千克的四足机器人），以及6到35个主动关节。每个机器人都变成了28–101个RoboTokens的序列——足够小以便端到端扩散。 RoboTokens相对于MJCF文本标记的紧凑性。**RoboTokens比MJCF文本紧凑27–110倍。**一个完整的机器人描述只需几十个类型化标记（蓝色计数），而语言分词器则需要数千个（灰色计数）——足够小，变换器可以直接扩散它们，而不是自回归地编写XML。 第二部分 ## Transformer Transformer：统一的架构 该模型是一个在RoboTokens上训练的扩散变换器（DiT1 (https://transformer-transformer.github.io/#ref-1)），采用DDIM4 (https://transformer-transformer.github.io/#ref-4)噪声调度。噪声被添加到每个输入标记*除了*我们条件化的那些标记——通过改变*哪些*标记被掩蔽，同一个网络扮演不同角色。掩蔽所有标记得到无条件机器人生成器。掩蔽动作得到跨具身控制器5 (https://transformer-transformer.github.io/#ref-5)、6 (https://transformer-transformer.github.io/#ref-6)。掩蔽具身得到运动约束的机器人设计器。相同的权重，联合训练在动力学上。 （关于双关语：第一个“Transformer”指的是*机器人*，其具身可以改变。第二个指的是自注意力*架构* 🥸） Transformer Transformer架构。**一个架构，两个用例。**（a）运动到机器人优化联合扩散具身和动力学标记。（b）跨具身控制扩散动作，以具身和目标运动为条件。在推理时，模型自身的预测被转换为奖励预测，其梯度指导生成朝向更高奖励的设计。### 具身感知的控制器 由于控制器将完整的具身作为输入，它可以调整控制策略以适应它所驱动的机器人。改变连杆惯性、关节范围或电机增益，同一个策略仍能保持追踪——无需针对每个机器人重新训练。我们通过行为克隆，从RL专家7 (https://transformer-transformer.github.io/#ref-7)、8 (https://transformer-transformer.github.io/#ref-8)、9 (https://transformer-transformer.github.io/#ref-9)（在过程化四足空间上每个设计训练一个）获得了全身控制器。 **一个策略，多种四足机器人。**同一个控制器在具有连续变化（连杆尺寸）和离散变化（自由度、膝盖方向、弹簧加载与串联腿）的四足机器人上追踪相同的轨迹。 ### 跨不同机器人类别的生成器 在没有任何条件的情况下运行模型，它会从高斯噪声中扩散出完整的机器人。在来自Menagerie10 (https://transformer-transformer.github.io/#ref-10)的各种机器人上训练后，单个Transformer Transformer能产生固定基座手臂、四足机器人、人形机器人和灵巧手。由于扩散是随机的，每次运行都会产生不同的机器人。 **机器人，被扩散。**单个模型从相同的噪声调度生成固定基座手臂、四足机器人、人形机器人和灵巧手。 第三部分 ## 动力学自引导：零样本奖励优化 现在进入真正的协同设计问题：我们如何为一个模型在训练时从未见过的奖励函数生成高奖励机器人？ 这里有个技巧。同一个模型同时预测具身*和*其动力学，因此我们可以将这些预测输入用户的奖励函数，并获得一个*预测奖励*。有了预测奖励，我们可以并行运行模型多次，预测每个候选的奖励，并返回最佳的一个。这就是最佳N采样，但在GPU速度下。 我们可以做得更好。由于预测奖励是具身标记的可微函数，我们可以取其梯度并将其注入到扩散过程中——在每一步去噪中将样本推向更高的奖励，遵循分类器引导的DDIM11 (https://transformer-transformer.github.io/#ref-11)。实际上，模型在问自己：“这个机器人应如何改变以增加奖励？”我们称之为**动力学自引导**。 **为什么是“自引导”：**将动力学纳入扩散引导的相关工作训练了一个单独的动力学模型12 (https://transformer-transformer.github.io/#ref-12)或依赖于可微模拟器13 (https://transformer-transformer.github.io/#ref-13)。这里，模型使用*自身的*动力学预测来引导自己。 攀爬景观。**一个梯度，整个回合。**Transformer Transformer并行预测整个回合时间范围内的状态和动作标记；然后奖励梯度通过所有这些流回到具身标记。硬件在长时域动力学上进行优化，而不是一次一个时间步。### 奖励输入，机器人输出 以一个随机生成的四足机器人为例，它在动态投掷过程中摔倒。让我们的模型仅为同一轨迹生成一个机器人，仅使用追踪奖励，它会降低质心并加宽支撑多边形——这些稳定的设计选择是它试图最大化从未见过的追踪奖励时发现的副作用。 现在添加一个扭矩惩罚项，优化景观就发生了变化。模型返回一个更小、更轻的机器人——仍然追踪运动，但平均扭矩降低了30%。添加一个尺寸惩罚，景观再次变化。这正是一些人所说的*零样本奖励优化*：在动力学建模上训练，在推理时为你能写出的任何奖励进行优化。 不同条件生成的不同设计。**相同模型，不同条件，不同机器人。**切换目标轨迹（地板擦拭→动态投掷）或添加尺寸惩罚会重塑优化景观，生成的具身相应变化——无需重新训练。离散选择（腿设计、自由度）和连续选择（连杆长度、安装点）都随之响应。### 组合扩散模型用于整个数据集 操作不是一个运动，而是一组运动的*分布*。我们想要为一个完整任务（例如洗碗）优化机器人，而不是单个轨迹。扩散组合14 (https://transformer-transformer.github.io/#ref-14)、15 (https://transformer-transformer.github.io/#ref-15)让我们恰好做到这一点：以不同轨迹作为条件并行运行模型，在每个扩散步骤平均预测噪声，你就得到了一个同时针对所有轨迹优化的生成器。 我们在公开可用的UMI双臂洗碗数据集3 (https://transformer-transformer.github.io/#ref-3)上进行了测试，留出了26个验证轨迹。Transformer Transformer生成了针对整个留出分布优化的机器人——既导航了离散空间（脊柱设计：固定、滑动、弯曲）也导航了连续空间（手臂尺寸、安装偏移）——然后同一个模型作为跨具身控制器来验证每个生成的设计。 **一个机器人，整个任务。**在26个留出的UMI洗碗轨迹上进行组合后，模型返回一个针对整个分布优化的单一具身——然后同一个网络驱动它完成每个运动。 ## 更好的设计，更快的速度 我们与两个基线进行比较：**随机**（从过程化设计空间中采样一个机器人）和**CMA-ES**16 (https://transformer-transformer.github.io/#ref-16)，一种广泛用于协同设计的黑箱优化器。CMA-ES在MuJoCo17 (https://transformer-transformer.github.io/#ref-17)中展开候选机器人，使用设计空间的控制器（固定基座/双臂使用Mink18 (https://transformer-transformer.github.io/#ref-18)；腿式机器人使用RL专家7 (https://transformer-transformer.github.io/#ref-7)），然后将下一代偏向更高奖励的样本。 我们给所有方法同样的任务：为目标轨迹和奖励生成一个机器人，控制它，报告实现的奖励。 ### 测试时计算扩展了结果 就像大型语言模型随着更长的思考时间变得更好19 (https://transformer-transformer.github.io/#ref-19)，Transformer Transformer在获得更多推理种子进行采样时也能得到更好的设计。性能最终在大约一分钟后达到平台期——这不同于LLMs从更长推理中获得的看似无尽的收益。我们在局限性中讨论这一点。 三个设计空间上的测试时扩展结果。**更高奖励，数量级更快。**在三个设计空间和几个奖励函数上，我们的零阶方法和动力学自引导变体都在允许更多并行样本时产生更高奖励的设计——并且两者都在几秒钟内达到CMA-ES级别的质量，而非数小时。对于多轨迹双臂优化（最右），CMA-ES耗时超过3小时；我们的模型在一分钟内完成。**为什么快这么多？**两个原因相乘。(1) 候选设计的GPU并行化。(2) 非自回归扩散并行评估整个展开，而不是一次一个时间步。典型的已学习动力学模型一次预测一个时间步；我们一次预测整个回合。 ## 现实世界：ALOHA上的抛布操作 **原始ALOHA vs Transformer Transformer优化的ALOHA。**相同任务，相同目标轨迹，两个机器人。优化的设计展开了布料，而原始机器人无法做到。 为了在硬件上对整个系统进行压力测试，我们选择了ALOHA2 (https://transformer-transformer.github.io/#ref-2)上的动态布料展开。抛布20 (https://transformer-transformer.github.io/#ref-20)通过运动学（手臂必须高速越过布料）和未建模的动力学（空气阻力、布料重量、摩擦）来惩罚具身错误。原始的ALOHA设计甚至难以追踪轨迹。 我们要求模型优化ALOHA设计，以获得“追踪速度”奖励，然后制造了生成的设计。与原始设计相比，有两处变化： - **更长的连杆长度。**刚好足以运动学上够到抛布的轨迹，但又足够轻以保持在D