Pi0

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

\(\Pi_0\): A Vision-Language-Action Flow Model for General Robot Control[1]

作者是来自Physical Intelligence的Kevin Black 等人。论文引用[1]:Black, Kevin et al. “π0: A Vision-Language-Action Flow Model for General Robot Control.” ArXiv abs/2410.24164 (2024): n. pag.

Time

  • 2024.Nov

Key Words

  • flow matching
  • 一句话：作者的robot policy用了一个预训练的VLM backbone，还有多个embodiment dataset，通过增加一个单独的action expert来做roboto control，通过flow matching产生连续的actions，使得能够实现精确和丝滑的控制，model可以基于prompt，执行tasks，或者在高质量的data上进行微调，实现复杂的多阶段的任务，例如折叠多个衣服。

总结

  1. Robot Learning有很大的前景，解锁灵活、通用和熟练的机器人系统的全部潜力，还能解决AI中的一些最深的问题。然而，将robot learning带到一个有效的real-world systems的要求的通用能力的水平面临一些数据、泛化性和鲁棒性方面的挑战。在本文中，作者讨论了通用的机器人策略(generalist robot policies)如何解决这些挑战，以及作者如何涉及有效的generalist robot policies for complex和highly dexterous tasks。作者提出了一个新的flow matching 架构，建立在一个预训练的VLM基础之上，继承了internet-scale语义信息，作者然后讨论这个model是如何在一个来自多个robot platforms、包括单臂、双臂和移动manipulators上的数据集上进行训练的。作者将从模型通过直接prompting执行任务的能力、遵循人类语言指令的能力，以及遵循高层次VLM策略指令、和通过微调得到新的skills的能力这几个方面来评估我们的模型。

  2. 一个人应当能够换尿布、策划入侵、宰猪、掌舵船只、设计建筑、撰写十四行诗、管理账目、砌墙、接骨、安慰临终者、服从命令、发布命令、协作共事、独立行动、解算方程、分析新问题、扬粪施肥、编程计算机、烹调美食、高效战斗、英勇赴死。专业化是昆虫的宿命，罗伯特.海因莱因，《时间足够你爱》。 人工智能系统形态多样，规模各异，从高级的专业化的systems，解决人类无法难以理解的问题，例如预测蛋白质的comformation，到能够基于texual prompts产生逼真的高分辨率的images or videos，然而，人类智能超越机器智能的核心维度在于其通用性，能够解决多种物理环境中的多样的任务，智能地回应环境的限制、语言指令和没有想到的干扰。或许，人工智能在实现这种通用性方面最显著的进展，体现在LLMs和VLMs中：在大量的web上的images和text上预训练的systems，然后用精心准备的datasets进行微调，旨在引导出期望的行为和响应模式。尽管研究表明此类模型展现出广泛的指令遵循和问题解决能力，但它们并不像人类那样真实置身于物理世界，其对物理交互的理解完全基于抽象描述。如果现有方法希望在开发具备人类物理情境适应能力的AI系统方面取得实质性进展，就必须使用来自具身机器人代理的物理环境数据进行训练。能够用来执行多种机器人行为的灵活、通用的models展示出了巨大的实际效果。能够被赋予执行多种机器人行为任务的灵活通用模型具有深远的实际影响，但它们也可能为当前机器人学习所面临的最严峻挑战提供解决方案：例如数据的可用性，泛化性和鲁棒性。在NLP和CV中，在多个多任务数据上预训练的通用目的的foundation models，往往超过narrowly tailored的和专门的解决方法。例如：如果目标是识别image中的biirds，在多个不同的image-language 关联上预训练，然后微调或者prompt for the bird recognition task，然后仅在bird 数据上进行train。类似地，这样也可能找到specialized robot systems，首先在highly diverse robot data进行预训练，然后微调或者prompt for the deired task，这能够解决数据稀缺性挑战，因为很多数据是适用于generalist model，包括其它任务的、其它robots甚至是non-robots sources的数据。它可能解决robustness和泛化性的挑战，因为多样化数据对观测与操作的覆盖范围更广泛，提供了多样化的场景、修正机制和恢复行为——这些特征在狭窄的专业化数据中可能并不存在。因此，在robot learning上采用大规模预训练方法展示出了解决field's 的挑战，使得实际的learning-enabled robots成为一个现实，同时，促进了对AI中最深问题的理解。

    然而，开发这样的通用的robot policies，例如robot foundation models，涉及很多主要的挑战：首先，任何研究必须是large scale的，因为大规模的预训练的benefits通常不会在small scales上展示出来；其次，它要求合适的模型架构，能够有效地利用多样的数据，同时能够表示复杂和微妙的、和复杂物理场景交互必须的行为；第三，它要求合适的训练方法，这或许是最重要的部分，很多最近的NLP和CV中的进展依赖于精心设计的策略for curating pre-training和post-training data。

    **在本文中，作者提出了一个prototype model和learning的框架，称之为 \(\pi_0\)，为了引入多种数据源，利用预训练的VLM来导入Internet-scale experience，作者的model以VLM为基础，继承了通用的knowledge，semantic reasoning和LLMs和VLMs的问题解决的能力，作者进一步训练model，引入robot actions，将其变成VLA model，为了利用多种robot data sources变得可行，作者采用了cross-embodiment training，来自多个robot types的data结合进同一个model，这些不同的robot types有不同的configuration spaces和action representations，包括single和dual-arm systems，还有移动的manipulators，另外，为了使得执行highly dexterous 和复杂的物理tasks成为可能，作者用了一个action chunking with flow matching的架构，来表示复杂连续的动作分布，这使得作者的model最多能在dexterous tasks上以50Hz控制robots，为了将flow matching和VLMs结合，作者用了一个新的action expert，用flow-based outputs来增强标准的VLM。

    和语言模型一样，作者的model只是方法的一部分，为了灵活地执行复杂的任务，作者需要正确地训练策略，作者的策略模仿了pre-training/post-training separation，这在大规模的LLMs和VLM上很常见，model首先在非常大和多样的corpus上进行预训练，然后在narrow和carefully curated data上进行微调，来引出behaviors的理想的pattern，在作者的case中，dexterity， efficiency和robustness，直觉上，仅在高质量的data上进行训练并不能教model如何从错误中恢复，因为错误在data上很少见，仅在低质量的pretraining data上进行训练，不能教model有效地act，将两者结合能够实现期望的行为：模型会尽可能地模仿高质量数据的行为模式，同时仍具备一系列恢复和纠正机制，可在出现错误时灵活调用。

    作者的工作的contributions包括：一个新的通用的robot policy架构，基于预训练的VLM和flow matching，一个针对此类机器人基础模型的预训练与后训练策略开展的empirical investigation，作者对模型进行了三种方式的评估：直接使用语言指令进行开箱即用的评估、在下游任务上进行微调的评估，以及结合高级语义策略的评估——该策略会输出中间语言指令，以执行复杂且具有时间延续性的任务。同时，作者的model和system利用多个最近的工作的ideas。

  3. 作者的工作基于最近的在大规模robot learning的方法，包括multimodal language models，VLA models(用预训练的VLMs，然后对robot control 进行fine-tuned)，这样的models利用自回归的discretization来表示actions，类似于text tokens，相比之下，作者采用了一个新的设计，将VLMs微调，通过flow matching产生actions，这使得能够处理高频的action chunks和highly dexterous tasks，也展示了这些任务对以往的自回归视觉语言动作模型（VLA）构成了重大挑战。这像最近在diffusion models for action generation上的工作，相比于这些工作，作者的model利用一个预训练的VLM backbone，作者的contribution也是integrative，聚焦于一个framework for robot foundation models，包括不仅模型架构本身，还有pretraining recipe，pretraining/post-training阶段。

    在robot control之外，很多models将预训练的language models和diffusion进行结合，包括specifically hybridize diffusion和自回归的large language models，这样的models通常和image generation有关，但是作者的action generation model建立在之前提出的concepts上，作者通过在单个序列元素上应用扩散风格的（flow matching）损失来训练模型，以此取代仅使用解码器的 Transformer 中常见的交叉熵损失。作者为与扩散过程对应的token使用了一套独立的权重。在将这些概念引入视觉语言动作（VLA）模型的基础上，作者提出了首个基于flow matching的 VLA 模型，能够生成用于灵巧控制的高频action chunks for dexterous control。

    作者的工作也建立在之前的large-scale robot learning上的工作，早期的工作利用自监督或者autonomous data collection，提供了一个易处理的data source，用于简单的任务，但是没有更多more dexterous 的行为，最近，搜集了高质量的数据集用于robot control，使得更广的泛化性，但是，这只是针对更简单的tasks，包含object relocation等。more dexterous tasks是在一个更小的scale被研究，通常只有10s 或者100s的training轨迹，作者目标之一是研究复杂和dexterous 行为，作者利用一个更大的dataset，大概10,000 hours的demos，和开源的OXE dataset互补，剧作者所知，这是目前为止最大哦的robot learning expertiments，在这个scale上，作者展示了一个更复杂的pretraining/post-training recipe是highly effective，类似于用于LLMs的recipes，一个预训练的阶段使得model具有广博的知识基础，在post-training中，用更高质量的curated data进行refine，实现desired 行为。

    任务的复杂性超过了之前的工作，最近的工作阐释了大量更复杂和dexterous 的行为，例如cooking shrimp，作者展示了，作者的framework能够学习long tasks，有时长度上能达到数十分钟，这些任务的行为既包含物理上的灵巧性，也具有组合上的复杂性。例如，laundry folding 任务要求robot来操作多种衣服，在任何configuration下开始，然后依次进行fold。

  4. 在训练framework中，作者首先组合了一个pretraining mixture，包含一个在7个不同的robot configuration下的，用于68个tasks的dexterous manipulation datasets和整个的OXE dataset的加权的combination。预训练阶段用了不同的language labels，包含task names和segment annotations。预训练的目的是训练一个base model，展示出broad capabilities和泛化性，但不一定在任何单一任务上追求高性能。这个base model遵循language commands，对一些任务有基础的执行能力。对于复杂的任务，需要post-training：用高质量的数据将model适应到下游任务。作者研究了小到中等数据量的有效后训练，以及大数据集的高质量后训练，用于复杂任务，如洗衣折叠和移动操作。作者的model是基于PaliGemma VLM，然后进一步用data mixture进行训练，为了将base PaliGemma VLM转换成 \(\pi_0\)，用flow matching，增加了action outputs，来产生连续的action distributions。

  5. \(\pi_0\) 包含一个language model transformer backbone，遵循标准的late fusion VLM 方法，image encoders将robot image observations嵌入到和language tokens相同的向量空间，用robotics-specific inputs和outputs来增强这个backbone，也就是proprioceptive state和robot actions， \(\pi_0\) 用了conditional flow matching，来建模actions的连续分布。Flow matching提供了high precision和多模态的能力，使其能够适应高频的Dexterous tasks，作者的架构受Transfusion的启发，用多个objectives训练single transformer，tokens对应于通过一个flow matching loss监督的连续outputs。基于Transfusion，作者发现，为robotics-specific tokens设置一个单独的weights可以提高性能，这个设计类似于MoE，有两个mixture elements，第一个element用于image和text inputs，第二个用于robotic-specific inputs和outputs，将第二个weights set称之为 action expert。 作者想建模data distribution \(p(A_t|O_t)\)，\(O_t\)是observation，\(A_t\)是对应的future actions的action chunk，这个observation包含多个RGB images，一个language command，robot proprioceptive state，\(O_t = [I^1_t, ..., I^n_t, l_t, q_t]\)，\(I^i_t\) 是第i个image，\(l_t\) 是language token的序列，\(q_t\) 是joint angles的vector，images \(I^i_t\) 和state \(q_t\) 通过对应的encoders和进行编码，然后通过一个linear projection layer投影到和language tokens相同的向量空间。对于action chunk \(A_t\) 中的每个action \(a_t`\) ，有一个对应的action token，通过action expert，在训练期间，用一个conditional flow matching loss来监督这些action。最近在high-resolution image和video合成中的工作表明：当和一个simple linear-gaussian结合的时候，flow matching实现了一个很强的性能。在实际中，网络通过sampling random noise进行训练，计算noisy actions，然后训练网络的输出来匹配denoising vector field，action expert用一个full bidirectional attention mask，因此所有的action tokens能够attend to each other，在训练期间，从一个beta 分布中，采样flow matching timestep \(\tau\)。

    在推理的时候，通过集成learned vector field，产生actions，用前向的Euler集成规则。

    原则上，作者的model可以从零初始化，或者从任意的VLM backbone进行微调，实际中作者用了PaliGemma作为base model，PaliGemma是一个开源的VLM，在size和performance之间提供了convenient trade-off。作者加了300M 参数for action expert

    • Non-VLM baseline model：除了主要的VLA model，作者也训练了一个类似的baseline model，不用VLM 初始化，这个model称之为 \(\pi_0\)-small，没有用VLM初始化。 和LLM的训练类似，有预训练和后训练，预训练的目标是使model得到广泛的applicable和通用的能力，后训练让modele能够有效地执行下游任务。

  6. 很多复杂的任务要求semantic reasoning和高阶的策略，例如table bussing，能从将高阶任务分解成多个中级的subtasks的high-level策略中获益，因为作者的model被训练用来处理language inputs，可以用一个high-level VLM来做语义推理，类似于LLM/VLM planning方法，采用此类高层次策略，协助模型以high-level strategy指导多个实验任务的执行。 作者的dexterous manipulation datasets包含7个不同的robot configurations和68个tasks，分别是 UR5e, Bimanual UR5e, Franka, Bimanual Trossen， Bimanual ARX &bimanual AgileX，Mobile Trossen & mobile ARX, Mobile Fibocom。

framework \(Fig.2^{[1]}\) 先是pre-training mixture，包含作者自己的dexterous manipulation datasets和开源数据，作者用这个mixture来训练flow matching VLA model，包含一个更大的VLM backbone和一个更小的action expert，用于处理robot states和actions，VLM backbone weights是从PaliGemma初始化的，提供了从large-scale internet 预训练的representations，\(\pi_0\) model 能够用不同的action spaces来控制多个robot embodiments，来完成多样的任务。