In-Context Imitation Learning via Next-Token Prediction
作者: Letian Fu, Huang Huang, Gaurav Datta, Lawrence Yunliang Chen, William Chung-Ho Panitch, Fangchen Liu, Hui Li, Ken Goldberg
分类: cs.RO, cs.AI
发布日期: 2024-08-28 (更新: 2024-09-27)
💡 一句话要点
提出ICRT，通过上下文学习实现机器人零样本模仿学习。
🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱二：RL算法与架构 (RL & Architecture)
关键词: 机器人模仿学习 上下文学习 Transformer 下一token预测 零样本学习
📋 核心要点
- 现有机器人模仿学习方法通常需要大量训练数据或复杂的奖励函数，限制了其泛化能力和适应性。
- ICRT通过上下文学习，利用Transformer模型直接预测感觉运动轨迹，无需显式训练即可完成新任务。
- 实验表明，ICRT在真实机器人上能够成功模仿新任务，并在多任务环境中显著优于现有方法。
📝 摘要（中文）
本文探索如何增强下一token预测模型，使其能够在真实机器人上执行上下文模仿学习。机器人通过解释输入阶段提供的上下文信息来执行新任务，而无需更新其底层策略参数。我们提出了上下文机器人Transformer（ICRT），这是一种因果Transformer，它对感觉运动轨迹执行自回归预测，而不依赖于任何语言数据或奖励函数。这种公式使得在测试时能够灵活且无需训练地执行新任务，这是通过使用包含图像观察、动作和状态元组的新任务的感觉运动轨迹来提示模型实现的，这些轨迹是通过人工遥操作收集的。使用Franka Emika机器人的实验表明，ICRT可以适应提示指定的新任务，即使在与提示和训练数据不同的环境配置中也是如此。在多任务环境设置中，ICRT在推广到未见任务方面显著优于当前最先进的机器人下一token预测模型。
🔬 方法详解
问题定义：现有机器人模仿学习方法通常需要针对特定任务进行训练，泛化能力较差。此外，依赖奖励函数的方法需要人工设计奖励，过程繁琐且难以适应复杂任务。因此，如何使机器人能够像人类一样，通过少量演示快速学习并执行新任务，是一个重要的挑战。
核心思路：本文的核心思路是利用Transformer模型的上下文学习能力，将机器人模仿学习问题转化为下一token预测问题。通过将人类演示的感觉运动轨迹作为上下文输入，模型可以预测后续的动作，从而实现模仿学习。这种方法无需显式训练，即可使机器人适应新任务。
技术框架：ICRT (In-Context Robot Transformer) 的整体架构是一个因果Transformer模型，它接收一系列感觉运动轨迹作为输入，并预测下一个感觉运动轨迹。输入序列包括图像观察、动作和状态元组。模型通过自回归的方式，逐个预测动作序列，从而控制机器人执行任务。
关键创新：ICRT的关键创新在于将上下文学习应用于机器人模仿学习，并提出了一个无需训练的框架。与传统的模仿学习方法相比，ICRT不需要针对每个新任务进行训练，而是通过上下文信息直接指导机器人的行为。此外，ICRT不依赖于任何语言数据或奖励函数，简化了学习过程。
关键设计：ICRT使用标准的Transformer架构，并针对机器人控制任务进行了优化。输入序列被编码为嵌入向量，然后输入到Transformer编码器中。解码器使用因果掩码，确保模型只能依赖于过去的token进行预测。损失函数采用交叉熵损失，用于衡量预测动作与真实动作之间的差异。具体的参数设置（如Transformer的层数、头数、嵌入维度等）需要根据具体的任务进行调整。
🖼️ 关键图片
📊 实验亮点
实验结果表明，ICRT在真实Franka Emika机器人上能够成功模仿新任务，即使在与训练数据不同的环境中也能表现良好。在多任务环境中，ICRT显著优于现有的下一token预测模型，例如Perceiver-Actor和Multi-Game Decision Transformer，在泛化到未见任务方面表现出更强的能力。具体性能数据可在论文及相关代码仓库中找到。
🎯 应用场景
该研究成果可应用于各种机器人自动化场景，例如：快速部署新任务、处理异常情况、人机协作等。通过提供少量演示，即可使机器人适应新的工作流程，降低了机器人使用的门槛，提高了生产效率。未来，该技术有望应用于家庭服务机器人、工业机器人等领域。
📄 摘要（原文）
We explore how to enhance next-token prediction models to perform in-context imitation learning on a real robot, where the robot executes new tasks by interpreting contextual information provided during the input phase, without updating its underlying policy parameters. We propose In-Context Robot Transformer (ICRT), a causal transformer that performs autoregressive prediction on sensorimotor trajectories without relying on any linguistic data or reward function. This formulation enables flexible and training-free execution of new tasks at test time, achieved by prompting the model with sensorimotor trajectories of the new task composing of image observations, actions and states tuples, collected through human teleoperation. Experiments with a Franka Emika robot demonstrate that the ICRT can adapt to new tasks specified by prompts, even in environment configurations that differ from both the prompt and the training data. In a multitask environment setup, ICRT significantly outperforms current state-of-the-art next-token prediction models in robotics on generalizing to unseen tasks. Code, checkpoints and data are available on https://icrt.dev/