Octo

发表于 更新于 分类于 阅读次数： 本文字数： 2.8k 阅读时长 ≈ 10 分钟

Octo: An Open-Source Generalist Robot Policy[1]

作者是来自UCB、Stanford、CMU和DeepMind的Sudeep Dasari等人，论文引用[1]:Team, Octo Model et al. “Octo: An Open-Source Generalist Robot Policy.” ArXiv abs/2405.12213 (2024): n. pag.

Time

  • 2024.May

Key Words

  • input tokenizer
  • transformer backbone
  • readout head: diffusion denoising
  • 能够微调到new robot setups

总结

  1. 在多种robot datasets上预训练的large policies有潜力转变robotic learning: 不是从零开始训练新的Policies，这样的通用的robot policies只需要一点的in-domain data就能fine-tune，泛化地很好。然而，为了在多个robotic learning的场景的应用，这样的policies需要处理多个sensors和action spaces，适应常用的robotic platforms，有效地微调到新的domains，在本文中，作者旨在为开发开源的、广泛应用的和通用的robotic manipulation奠定基础。作者提出了Octo，是一个大的基于transformer的policy，从Open X-Embodiment dataset上的800k trajectories上训练的，这是目前最大的robot manipulation dataset，可以通过language command或者goal images进行instruct，能够用新的sensory inputs和action spaces，在标准的消费级GPUs上有效的微调。

  2. 常用的robotic learning的方式是使用针对特定机器人和当前任务收集的数据集来训练策略。从零开始训练需要大量的针对每个task的数据，最后的policies通常泛化性很小，原则上，搜集其它robots和tasks的经验提供了一个可能的solution，将models置于多种robotic control的问题中，可能提高泛化性和在下游任务上的表现，然而，即使是一个通用的models在NLP和CV中很常见，通过类似的通用的robot model，控制很多的robots用于多个tasks被证明是有挑战的，训练一个统一的control policy in robotics展示出了独特的挑战，要求处理不同的robot embodiments、sensor setups、action spaces、task specifications等。 在这个方向，一些工作提出了robotic foundation models，直接将robot observations映射到actions，在新的domain和robots提供zero-shot或者few-shot的泛化性，作者将这些models称之为generalist robot policies(GRPs)，强调了它们在多个tasks, environments上的low-level visuomotor control。例如，GNM model泛化到不同的robotic navigation场景，RobotCat model对于goal-conditioned tasks处理不同的robot embodiment，RT-X model在5个robot embodiment上执行language-conditioned manipulation，尽管这些models是走向true 通用目的的robot model上展示出了重要的steps，它们在很多重要的方面有局限：它们将下游用户限制在预定义和有限制的input observations，例如single camera stream，缺乏对新的domain的有效的finetuning，重要的是，这些models中最大的并未对公众开放。

    作者设计了一个system，用于预训练通用的robot policies，在下游robotic applications上多种接口更加适用，model的核心是一个transformer架构，将任意的Input tokens映射到output tokens，可以在多种robots和tasks的数据集上训练，没有额外的training，这个policy能够接受不同的camera的configurations，能够控制不同的robots，通过language command或者goal images进行guide，所有功能仅需调整输入模型的token即可实现，更重要的是，这些model有了新的sensory inputs、action spaces或者增加合适的适配器的形态后，能够适应，用很小的target domain dataset和合理的compute budge进行微调。

    作者的主要贡献是：Octo，是一个基于transformer的policy，在最大的robot manipulation dataset上预训练的，Octo是第一个GRP，能够有效地微调到新的observations和action spaces，是第一个generalist robot manipulation policy，全开源，包括training pipeline, model checkpoints和data，最后，组成Octo的individual components，一个transformer backbone支持language和goal image，一个diffusion head，来model expressive action 分布，将这些components结合为一个powerful generalist robot policy是unique和novel的

    作者展示了, Octo能够用作finetuning到unseen steups的有效地初始化。

  3. 许多研究采用从机器人收集的大量轨迹数据集来训练策略，从早期通过自主数据采集扩展策略训练的研究，到近期探索现代基于Transformer的策略与大型demonstration dataset结合的工作。这些工作主要聚焦于single embodiment，然而，Octo通过在多种机器人形态上构建的数据集进行策略训练，既有效扩大了训练数据集的规模，又支持针对不同机器人配置的微调适配。最近，一些工作聚焦于提高robot policies的泛化能力，很多工作利用多种non-robot data或者预训练的vision-language foundation models来提高policy泛化到新的scenes和tasks的能力，和Octo最相关的是在多个robot embodiments上训练robot policies，GNM model 泛化到不同的robot navigation setups，RoboCat和RT-X控制多个single-arm manipulation robots，然而这些models有impressive policy learning results，一个关键问题是缺乏flexibility，通常要求用户必须沿用预训练阶段使用的传感器输入和action spaces，且无法适应新的观测和动作空间。另外，最大的model仍然是对公众不开放的，Octo在多个方面和这些工作不同：它在一个更大,更diverse的robot data上训练的，通过有效的微调至新的机器人配置，支持多种下游应用；完全开源。

  4. Octo的设计是受最近的robot imitation learning和scalable transformer training的启发，包括用denoising diffusion objectives for action decoding，prediction of action chunks，sequences of future actions，model layouts。作者的工作是第一个在学习cross-embodied generalist policies的环境中利用这些方法的。 对于训练一个通用的robot policies的key ingredient是robot training data，相比于vision和language data，能够从web上获取，得到大规模的robot data是有挑战的，通常涉及硬件上的投资。

  5. Octo至此多种常用的robot，sensor configurations，actions，同时提供能够在大量数据上训练的generic和scalable recipe，它也支持自然语言instructions, goal images, observation histories, 多模态, 通过diffusion decoding的chunked action prediction。另外，作者涉及了Octo，使得能够微调至新的robot setups，包括不同action spaces的robots，不同cameras和proprioceptive information的结合，这个设计使得Octo是一个灵活、能够广泛应用的generalist robot policy，能够用于多个下游robotic 应用和research projects.

    • 架构： Octo是要给基于Transformer的policy \(\pi\)，包含三个关键的key parts: input tokenizers，将自然语言instructions \(l\), goals \(g\) 和 observations sequences \(o_1,..., o_H\) 变为tokens \([T_l, T_g, T_o]\)，一个transformer backbone处理tokens和产生embeddings, \(e_l, e_g, e_o = T(T_l, T_g, T_o)\), 和readout heads \(R(e)\)，产生desired outputs actions \(a\)。task and observation tokenizers: 将task definitions(language instructions l and goal image g) 和observation o，用modality-specific tokenizers转换成常用的tokenized format。

      • language inputs: are tokenized，然后通过一个预训练的transformer，产生一系列的language embedding tokens，用\(t5-base\) model
      • Image observations and goals: 通过一个浅的convolution stack，然后split成一系列的flattened patches 通过给task和observation tokens增加一个可学习的position embeddings p，组成transformer的输入sequence，然后按顺序进行arrange。
      • Transformer backbone and readout heads：一旦inputs转换成统一的token sequence，它们通过一个transformer，这类似于之前的工作，在observations和actions的sequences上训练基于transformer的policies，Octo transformer的attention pattern是block-wise masked： observation tokens仅attend 相同或者更早的time steps \(T_{o,0:t}\) ，还有task tokens \(T_T\), 不存在的observations的tokens完全被masked out，这个modular design使得在微调的时候，作者能够增加或者rmeove observations or tasks，除了这些input token blocks，作者也插入了一个learned readout tokens \(T_{R,t}\), 一个readout token \(T_{R,t}\) attend to observation和task tokens，但是不被任何observation or task token attend，因此，它们只能被动地read和处理内在的embeddings，readout tokens和BERT中的[CLS] token类似，作为一个observation sequence的compact vector embedding，一个轻量的action head，执行diffusion process，用于readout tokens的embeddings，这个action head预测一些连续的actions。

      作者的设计使得在下游任务微调的时候，能够灵活地增加task和observation inputs or action output heads，当增加新的tasks、observations, 或者loss functions的时候，完全保留预训练的transformer的weights，仅增加新的positional embeddings，一个新的lightweight encoder或者根据规格变更所需的新增头部参数，这和之前的架构不一样，增加或者去掉一个image input 或者改变task specification的时候，需要重新初始化或者重新训练预训练model的large components

      这个灵活性使得Octo是一个真正的通用的model，因为不能覆盖所有的roboto sensor和action configurations during pretraining，在微调的时候适应Octo的inputs和outputs使其是一个多功能的tool for robotics community，之前采用了标准的transformer backbone，或者融合将visual encoders和MLP output heads融合的设计会锁定模型所期望的输入类型与顺序，相比之下，切换observation或者Task对于Octo，不需要重新初始化。

    • Training objective: 用了一个conditional diffusion decoding head来预测连续的，多模态的action分布，**重要的是，每个action predicition只执行一次transformer backbone的forward pass，随后所有multi-step的denoising process完全在小的diffusion head中完成，作者发现，这个policy parameterization在zero-shot和finetuing evaluations上超过了用MSE action heads或者discretized action distribution训练的policies，为了生成一个action，作者采样一个高斯noise vector，然后用学习到的denoising network(依赖于之前denoising step的输出，当前step的索引k，和Transformer action readout output embedding e)。

    \[x^{k-1} = \alpha \left( x^k - \gamma \theta_\theta(x^k, e, k) + \mathcal{N}(0, \sigma^2 \mathbf{I}) \right)\] 超参数 \(\alpha, \gamma, \sigma\) 对应noise schedule，用标准的cosine schedule，训练diffusion head的时候，用标准的DDPM Objective，给dataset actions增加高斯noise，然后训练denoising network来重建原始的action。

    在微调的时候，用相同的diffusion training objective然后更新full model，这一方法显著优于冻结部分预训练参数的策略。所有微调实验均遵循统一方案：基于包含约100条轨迹的小规模目标领域数据集，通过余弦退火学习率衰减（含线性预热阶段）进行50,000步训练

    • Training Details：训练了两个变体：Octo-Small，有一个transformer backbone，和ViT-S的size一样，一个Octo-Base，和ViT-B的size一样。训练的时候，对每个training example，随机zero out language instruction 或者goal image，使得Octo能够基于语言指令或目标图像进行条件化控制。

Architecture \(Fig.1^{[1]}\) Left: Octo将task descriptions和input observations，用一个预训练的Language model和一个轻量的CNN进行tokenize，Top: Transformer backbone处理task和observation tokens的sequences，产生readout tokens，然后通过output heads得到actions。 Bottom: block-wise attention 结构使得能够增加或者去掉inputs and outputs during finetuing：例如：可以增加新的observations或者action spaces，不需要修改任何预训练的参数。