在线强化学习改进VLA模型
在线强化学习改进VLA模型
2025年1月,来自清华大学、伯克利分校和上海姚期智研究院的研究团队发布了一篇题为《Improving Vision-Language-Action Model with Online Reinforcement Learning》的论文。该研究探讨了如何通过在线强化学习(RL)进一步改进视觉-语言-动作(VLA)模型,以实现更高效的机器人控制。
VLA模型的挑战与机遇
最近的研究已成功地将大型视觉-语言模型(VLM)通过使用专家机器人数据集进行监督微调(SFT)集成到低级机器人控制中,从而产生视觉-语言-动作(VLA)模型。虽然VLA模型功能强大,但如何在与环境交互的过程中改进这些大型模型仍是一个悬而未决的问题。本文探讨如何通过强化学习(RL)进一步改进这些VLA模型,强化学习是一种常用的大型模型微调技术。然而,直接将在线RL应用于大型VLA模型存在重大挑战,包括严重影响大型模型性能的训练不稳定性,以及超出大多数本地机器能力的计算负担。
iRe-VLA框架
为了解决上述挑战,研究团队提出了iRe-VLA框架,该框架在强化学习和监督学习之间进行迭代,有效改进VLA模型,利用RL的探索性优势,同时保持监督学习的稳定性。
模型架构
VLA模型将视觉输入o和自由形式语言指令i转换为低级机器人动作a,表示为O×L→A。该模型包含一个预训练的大型VLM和一个轻量级动作头,如图左侧所示。使用BLIP-2 3B模型作为主干VLM。由于预训练的VLM在语言空间中输出文本token,因此设计一个动作头来产生低级的控制动作。这些动作通常包括末端执行器姿态和夹爪状态的变化。遵循[11, 34]中提出的设计,将VLM的全连接层替换为一个初始化的动作头。
训练流程
训练中,首先对VLA模型进行基于机器人数据集的监督微调(阶段0),然后迭代进行在线RL(阶段1)和监督学习(阶段2)。
阶段0:基于专家数据集的监督学习
首先使用专家机器人数据集D/e={(o/1,l1/,a/1),(o/2,l/2,a/2),…,(o/i,l/i,a/i)}对VLA模型π/θ进行标准的监督微调。损失函数为MSE:
经过监督微调,得到初始的VLA模型π0/θ,φ。π0/θ,φ的性能与专家数据集D/e的规模和质量高度相关。然后开始在线RL来改进π^0/θ,φ。
阶段1:使用冻结VLM的在线RL
引入一个与动作头结构相同的批评头,但输出维度设置为1。为了防止模型崩溃并加速学习过程,在这一阶段冻结VLM参数θ。因此,只有动作头的参数ϕ被优化:
在线RL后,机器人可能会发现新的轨迹x/i来解决新任务。然后将这些成功的轨迹收集到一个在线数据集D/RL=D/RL∪x/i中。
阶段2:基于专家数据和在线收集数据的监督学习
第一阶段,当智体对新任务进行RL时,它有可能会忘记先前学习的任务。因此,在第二阶段,使用新收集的在线数据D/RL和原始专家数据集D/e来监督整个模型,以减轻灾难性遗忘[49]。目标可以写成:
迭代阶段1和阶段2
阶段1中的智体探索新任务的新解决方案,而在阶段2中,它模仿所有可用的成功轨迹。通过在阶段1和阶段2之间交替,大型VLA模型逐渐解决更广泛的任务,同时也能防止在已知任务上发生灾难性遗忘。整个流程在如下算法概述。
实验结果
在两个模拟基准Metaworld和FrankaKitchen以及现实世界的Panda操作任务中进行实验,以验证iRe-VLA框架的有效性。
实验设置如下。使用单个文本条件VLA模型来解决一个域中的所有任务。每个域都涉及分为三组的任务(如图所示):演示数据集观察的专家任务、通过在线RL增强的RL训练任务、以及在先前训练中未见过的保留任务。
实世界实验遵循SERL中描述的设置,如图所示,这是一个用于真实世界RL的有用软件套件。首先在一个包含2000条人类收集的专家数据(涵盖各种任务类别,包括抓取、放置、按下按钮、电缆布线和打开抽屉的集合上训练一个VLA模型。
由于VLA模型的泛化能力,学习的VLA模型在未见过目标上显示出确定的成功率。然后采用在线RL来进一步提高在未见过目标上的成功率。在VLA模型的背景下,实施了一些关键的设计选择,以提高样本效率并确保计算的可负担性。为了提高样本效率,采用SACfD算法。具体来说,当引入一项新任务时,最初利用零样本迁移的VLA模型来收集一个包含20条成功轨迹的演示缓冲区。在训练过程中,从演示缓冲区和在线缓冲区分别采样50%的转移,如[52]中所述。为了控制计算成本,每个图像观察只由VLM处理一次,并将生成的潜在输出存储在缓冲区中。随后,在该潜在空间中实现SACfD算法。