Command Palette
Search for a command to run...
Transformer Transformer:面向运动条件机器人协同设计的统一模型
Transformer Transformer:面向运动条件机器人协同设计的统一模型
Huy Ha Karen Liu Shuran Song
摘要
机器人操作性能中一个常被忽视的因素是机器人自身的形态。受此问题启发,我们研究了运动条件机器人协同设计,其目标是生成能够跟踪目标末端执行器轨迹(来自人类演示)并同时优化用户定义奖励的完整机器人设计。我们提出了 Transformer Transformer,一种在 RoboTokens(一种对机器人形态、状态和动作的统一分词表示)上训练的扩散 Transformer。同一架构可用于不同的形态空间(例如,轮式双臂机器人、四足机器人、人形机器人)和用例(形态生成、跨形态控制器)。Transformer Transformer 并非针对单一奖励函数过拟合,而是一个动力学模型,其与奖励无关的状态和动作预测可被转换为特定于奖励的价值预测。这些价值预测通过我们称为动力学自引导的过程,用于将形态扩散引导至高价值的机器人设计。跨多个设计空间的实验展示了对未见奖励和轨迹的零样本优化,相比进化基线方法提升了性能和运行速度。最后,我们制造了一个经过优化的 ALOHA 设计,其跟踪误差相比原始设计降低了 70% 以上。
一句话总结
斯坦福大学和哥伦比亚大学的研究者提出了 Transformer Transformer,一种基于统一 RoboTokens 对机器人形态、状态和动作进行 token 化的扩散 transformer,该模型作为一个奖励无关的动力学模型,并利用 Dynamics Self-Guidance 将形态扩散引导至高价值设计,实现了对未见奖励和轨迹的零样本优化,并促成了一个制造出的 ALOHA 设计,其跟踪误差降低了超过 70%。
核心贡献
- RoboTokens 提供了对机器人形态、状态和动作的统一 token 化,使得单个扩散 transformer 能够覆盖多种形态,如轮式双臂、四足和人形机器人。
- Dynamics Self-Guidance 在推理时将 Transformer Transformer 的奖励无关动力学预测转换为奖励特定的价值分数,将形态扩散引导至高价值设计,无需重新训练。
- 在三个设计空间中,该模型零样本优化了未见奖励和轨迹,并且一个制造出的优化 ALOHA 设计相比原始设计,将真实世界的跟踪误差降低了超过 70%。
引言
作者解决了机器人协同设计问题,即自动生成机器人的物理形态以最大化在给定任务上的性能。先前的工作通常依赖于有限的设计表示,过拟合于单一奖励函数,或将生成器、评价器和控制器分离为难以联合优化的独立组件。作者提出了 Transformer Transformer,一个将这些三个角色统一到单一模型中的扩散 transformer。它引入了 RoboTokens,一种灵活的 token 化方案,可以表示任何关节式机器人的形态、状态和动作,并经过训练以建模跨多种形态的通用动力学。在推理时,该模型通过将其奖励无关的动力学预测转换为扩散过程的奖励特定引导,零样本优化用户指定的奖励函数,从而实现快速、GPU 并行化的协同设计。
数据集
作者构建了一个将多种机器人设计与目标运动轨迹配对的数据集,使模型能够评估给定设计跟踪所需技能的优劣。该数据集由三个主要组成部分构建:一组记录的末端执行器运动、三个组合式机器人设计空间,以及控制器生成的状态-动作日志。
数据集构成和来源的关键细节:
- 目标运动数据集:使用基于 ARKit 构建的 UMI 夹爪扩展收集了 76 条运动轨迹。这些轨迹涵盖了单个末端执行器技能,如抛掷、拧紧、打开抽屉和擦洗。划分使用 56 条轨迹进行训练,20 条进行验证。来自 UMI 的双臂洗碗演示也被包含在内,用于移动双臂设计空间。
- 运动学设计空间:基于 Trossen ViperX 300S,该空间变化安装姿态、自由度数量、关节方向和连杆长度。它测试对固定基座机械臂运动学可达性的推理。
- 动力学控制设计空间:一个四足机械臂空间,变化腿部设计(串联 vs. 弹簧加载连杆)、臂部安装方式(固定 vs. 滑动导轨)以及躯干、腿部和臂部的长度。该空间着重考验全身控制性能、稳定性和惯性权衡。
- 任务复杂度设计空间:一个移动双臂空间,具有脊柱变化(固定、滑动、弯曲)和长度变化。它捕捉长周期、协调的双臂任务,如洗碗。
数据生成和处理方式:
- 对于非腿式机器人(固定基座和移动双臂空间),作者使用 Mink,一个用于 MuJoCo 的微分逆运动学库。固定的 80 毫秒前瞻时间补偿了关节位置控制滞后,并包含跟踪、姿态正则化和阻尼项。
- 对于腿式机器人,训练强化学习策略作为全身控制器。并非为每个机器人训练一个策略,而是为每个离散设计变体(例如,腿部类型)训练一个策略,并接收连续变体(例如,腿部长度)作为观测。这产生了 128 个 RL 专家策略,覆盖了 7 个二元设计选择的所有组合。这些策略生成用于数据收集和设计验证的状态和动作数据。
生成的数据集为每个跟踪目标运动的机器人设计提供了状态和动作 RoboTokens,构成了协同设计评估模型的训练基础。
方法
作者提出了 Transformer Transformer,一个旨在建模机器人形态和动力学联合分布的统一框架。总体流程包括通过末端执行器轨迹演示任务,使用奖励引导扩散生成优化的机器人设计,并使用跨形态控制器验证这些设计。
如下图所示:
为实现这一点,作者首先引入了 RoboToken,一种将任何机器人描述转换为可供学习使用的表示的 token 化方案。RoboToken 捕获了时不变的形态属性和时变的动力学。token 化器解析机器人仿真,提取连杆、关节和电机的属性,将它们连接成连续值向量。形态 token 表示物理参数,如基元类型、尺寸和惯性,而动力学 token 捕获特定时间步的状态和动作。为支持多样化的运动学,token 包含指向连接组件的 ID,允许表示任意连接性和被动关节。
token 化过程如下图所示:
这种表示非常高效。如下图所示,RoboToken 序列比传统的基于文本的格式(如 MJCF)紧凑得多,通常需要的 token 数量少 27 到 110 倍。
为在这些 token 上训练模型,作者采用了一种能够处理可变序列长度的高精度连续数据的 Diffusion Transformer (DiT) 架构。对于每个 token,一个特定于 token 类型的线性投影将向量映射到 DiT 潜在空间,并辅以针对 token ID 的学习嵌入。由于该领域因重力而具有 SE(2) 等变性,应用了平面变换增强。每个 token 类型序列被填充到最大长度,并连接成一个单一的多模态序列。模型使用 DDIM 调度器进行训练,并对 episode 时间步进行子采样以形成固定长度的序列。
统一架构及其应用如下图所示:
模型使用不同的掩码建模方案进行训练,以支持两个主要应用。首先,对于 Motion-to-Robot Optimization,模型在目标运动轨迹的条件下同时扩散形态和动力学 token。这种非自回归公式允许模型通过评估依赖于形态和动力学的奖励函数来评估新生成的设计,避免了自回归模型典型的误差累积。形态生成器运行并行扩散过程,并使用指定的奖励函数对候选设计进行排序。其次,对于 Cross-Embodiment Control,模型充当一个形态感知的控制器。在形态、当前状态和目标运动的条件下,它预测专家动作。RoboToken 的灵活性允许模型通过改变上下文中的 token 数量来处理异构的状态和动作空间,从而泛化到未见过的形态。
为在具有挑战性的优化环境中提高采样效率,作者引入了 Dynamics Self-Guidance。这种一阶扩散引导方法利用模型噪声奖励预测的梯度来优化噪声形态预测。在每个扩散步骤,DiT 预测噪声以对 token 去噪。即使完整的扩散过程不可微分,每步的奖励梯度也可以引导样本趋向更高奖励的候选设计。这是遵循分类器引导的 DDIM 来结合的,允许模型为未见过的奖励函数进行优化,而无需训练额外的网络。
模型在不同优化环境中生成多样化设计的能力如下图所示:
实验
评估涵盖了多个机器人设计空间,包括固定基座机械臂、四足机器人、双臂移动机械臂和人形机器人,以评估基于扩散的 token 化和生成框架在统一形态表示、生成物理上合理的设计以及优化未见奖励函数方面的表现。该模型生成的设计保持在训练流形上,尊重跨属性一致性,并支持零样本优化,其并行化评估在速度和奖励上均优于 CMA-ES。它还展示了用于多轨迹优化的扩散组合、测试时计算扩展,以及通过统一控制器自我验证设计的能力,并通过真实世界的 ALOHA 设置证实了跟踪性能的提升和关节速度的降低。