25年5月来自清华大学和南洋理工的论文“VLA-RL: Towards Masterful and General Robotic Manipulation with Scalable Reinforcement Learning”。
近期,高容量视觉-语言-动作 (VLA) 模型通过模拟人类演示,在一系列机器人操作任务中展现出令人印象深刻的性能。然而,在数据分布不均的场景下,利用访问状态有限的离线数据会导致执行失败。直观地说,一种基于探索的方法,在测试时改进在线收集的数据,可以解决这一限制。VLA-RL,是一个算法和系统框架,利用在线强化学习 (RL) 来改进下游任务中预训练的自回归 VLA。在一个统一的视角下,首先引入一种用于自回归 VLA 训练的轨迹级强化学习公式,该公式将通用的机器人操作轨迹建模为多模态多轮对话。为了应对稀疏奖励的挑战,将一个预训练的视觉-语言模型微调为机器人过程奖励模型(RPRM),该模型使用自动提取的任务段上伪奖励标签进行训练。为了扩大规模,确定多项可提升稳定性和效率的实现成果,包括课程选择策略、GPU 平衡的矢量化环境、批量解码和 Critic 预热。VLA-RL 使 OpenVLA-7B 在 LIBERO 中 40 个具有挑战性的机器人操作任务上的表现超越最强的微调基线 4.5%,甚至匹敌 π0-FAST 等先进商业模型的性能。值得注意的是,VLA-RL 受益于测试-时间优化的提升,这表明机器人技术领域在推理规模化定律的早期萌芽。
如图所示:VLA- RL和之前 VLA 比较
前言
通用机器人操作长期以来一直是机器人领域的核心追求。智能体需要基于视觉观察和指定当前任务的人类指令,以交互方式确定机器人的下一个动作(末端执行器的姿态),以执行各种任务。近年来,高容量、预训练的视觉和语言模型已在各种语言条件下的操作任务中展现出普遍性。其中,OpenVLA-7B [39] 是一个领先的开源 VLA 模型,因此它成为了方法的基础模型。其核心是一个自回归 LLM 模型 Llama-2-7B [69],它配备一个双流视觉编码器,该编码器由预训练的 SigLIP [86] 和 DinoV2 [55] 模型组成。在每个时间步 t,它将第三人称摄像机拍摄的图像 o_t 和人类指令 vin_t 作为输入,并输出一个动作 token 序列 vout_t,其中每个动作 token 代表机器人动作空间中一个维度的离散区间。最终的机器人动作使用后处理函数 f 从该序列中提取出来,得到 a_t = f(vout_t)。然而,优化自回归 VLA 在算法和系统方面都面临挑战,包括强化学习的通用操作、稀疏奖励问题以及大规模评估和优化等。
总体流程
VLA-RL 的总体流程如图所示。其开发一个算法和系统框架,用于利用强化学习 (RL) 训练自回归 VLA。该系统包含三个模型:需要在常用的 Actor-Critic 框架中训练的策略和价值模型,以及一个冻结的机器人过程奖励模型(PRM),该模型用于稠密环境给出的稀疏奖励。在算法层面,将自回归 VLA-RL 训练设计为多模态和多轮对话。GPU 平衡的矢量化环境、批量解码、课程选择策略和 Critic 预热等系统性技术进一步提升系统的训练效率和稳定性。最终,训练后的 VLA 模型能够通过优化预期奖励来生成可行的动作,从而成功执行各种操作任务。
通用机器人操作的多轮对话
为了扩展强化学习以优化用于通用操作的自回归 VLA,首先将马尔可夫决策过程设计为多轮对话。令 V 表示离散的、有限的词汇 token 集。空间 Vm 和 Vn 表示可能的输入和输出文本序列,其中 m 和 n 分别指定输入和输出的最大序列长度。将状态空间定义为笛卡尔积 S = O x Vm,其中 O 是图像空间。动作空间由 VLA 生成的所有可能输出话语集合 Vn 给出。因此,由 θ 参数化的 VLA 策略可以形式化为映射 π_θ : O x Vm -> Vn。
在每个时间步,该策略分配一个概率 π_θ(vout_t | o_t, vin_t),以在给定输入图像 o_t 和提示 vin_t 的情况下发出输出序列 vout_t。环境的转换由函数 T : S x A -> S 控制,该函数描述每个动作之后状态如何演变。环境奖励函数 R : S -> R 将每次行动后的结果质量量化为 r_t。在 T 个时间步的轨迹中,目标是最大化折扣奖励总和,Rγ = sum (t) = 0T γt r_t,其中 γ 是折扣因子。采用近端策略优化 (PPO) [62] 进行稳定的策略优化。
展开阶段。首先将更新后的 LoRA [26] 权重与原始检查点合并,并将其广播到推理引擎。然后,智体根据其当前策略 π_θ_old 与环境交互,生成一系列状态、动作和奖励(即轨迹)。动作序列的对数概率可以分解为自回归模型中 token 级对数概率的总和。
学习阶段。PPO 目标函数利用带裁剪的重要性采样来确保稳定更新。对于每个状态,优势 A_t 通过广义优势估计 (GAE) [61] 计算。整个过程总结在算法 1 中。
机器人过程奖励模型(RPRM)
奖励建模是将强化学习应用于一般操作的关键,它需要:(1) 在反馈本身稀疏的环境中提供密集的奖励;(2) 避免奖励攻击,即智体以非预期的方式利用奖励函数。为此提出机器人过程奖励模型,这是一种奖励密集化方法,与 VLA 的 token 生成过程相一致。
奖励建模作为下一个 token 预测。传统的机器人强化学习常常受到稀疏奖励的影响,通常仅在任务完成时提供二进制信号。本文将奖励建模重新表述为下一个 token 预测问题,利用预训练视觉-语言模型的自回归特性。给定状态和动作的轨迹,机器人过程奖励模型 (RPRM) 可以预测成功动作序列的可能性。训练目标是最大化有希望的动作 token 的对数似然,并由指示任务完成进度的伪-奖励信号加权。
自主伪奖励标签生成。为了在无需大量人工标记的情况下有效训练机器人过程奖励模型,其开发一套自主标签生成流程,该流程可根据成功轨迹创建高质量的伪奖励标签:(1) 里程碑分割:从专家演示和之前的模型运行中收集包含各种成功轨迹的数据集。根据夹持器张开度的显著变化将轨迹细分为子任务,因为这些变化通常标志着某个功能步骤的完成。(2) 进度标记:在每个分段子任务中,识别机器人末端执行器速度趋近于零的关键帧。这些点通常对应于稳定状态或细粒度运动的完成。将正伪-奖励分配给通向这些关键帧的 VLA 动作序列。
最终奖励是黄金稀疏奖励与机器人过程奖励模型(RPRM)预测奖励的直接总和。实证分析表明,这种方法在显著加速学习的同时,与实际任务成功率保持着密切的相关性。
VLA-RL 系统
由于 RL 的性能高度依赖于实现细节,如下介绍一些在本项目中采用的技巧,以提高学习效率和稳定性。
课程选择策略。实现一个自适应课程,根据智体的当前能力选择任务。每个任务由一条指令和一个初始状态组成,跟踪成功率 s_j 并计算采样概率:
P (task_j) = exp ((0.5 - s_j )/τ)
其中 τ 控制探索。该方程将成功率约为 50% 的任务优先作为智体能力的前沿,同时保持对已掌握任务和具有挑战性任务的接触,从而提高样本效率和泛化能力。
Critic 预热。当从头开始训练价值模型(Critic)时,它最初会产生不准确的价值估计,这可能会在训练早期误导模型。为了解决这个问题,实现一个 Critic 预热阶段,在这个阶段中,用模拟预训练策略收集初始轨迹,并在开始联合策略-价值优化之前,专门训练价值网络进行多次迭代。
GPU 平衡的矢量化环境。为并行部署实现多个矢量化环境,其中每个训练 GPU 都包含一个环境子集。现代渲染器通常依靠 GPU 进行加速,但随着矢量化环境数量的增加,GPU 内存消耗会显著增加。为了解决这个问题,为每个 GPU 工作器分配一组各自的环境,以便与之交互和学习。同时,用“all_reduce”操作收集所有工作器的环境状态,以供推理引擎使用。
基础设施。PPO 基础设施使用 bfloat16 将模型拟合到内存中。假设总共有 G 个 GPU,分配一个专用的 1 GPU 用于使用 vLLM [41] 加速进行推理,另分配 G - 1 个 GPU 用于 Ray [53] 的学习,就像在 OpenRLHF [30] 和 open-instruct [42] 中所做的那样。在代码库中,在 vLLM 插件中实现 OpenVLA [39],以避免使用原始的 Huggingface Transformers 生成函数,因为该函数在处理大批量时会导致错误的结果。分布式训练过程由 PyTorch 完全分片数据并行 (FSDP) [90] 管理,以支持大规模训练。
实验设置。如何将 VLA-RL 应用于模拟机器人设置和任务,特别是利用 LIBERO 基准 [45]?LIBERO 基准 [45] 包含四个具有挑战性的任务套件:LIBERO-Spatial、LIBERO-Object、LIBERO-Goal 和 LIBERO-Long(或 LIBERO-10),分别关注各种空间关系、物体类别、目标任务和扩展的序列挑战。在实验中,专注于从基础模型 OpenVLA-7B [39] 开始进行强化学习,该模型针对每个任务套件进行了监督微调 (SFT) 训练。在测试阶段,将在每个套件的 500 个回合中对所有对应模型进行评估。如图展示该基准的几个成功样本。
基线。对于强化学习训练,用 [39] 中发布的检查点作为基础 SFT 模型。除了 SFT 基线之外,还报告从零开始训练的扩散策略 [12]、基于扩散的微调 VLA Octo [68] 以及使用直接偏好优化 (DPO) [59] 训练的 GRAPE [89] 的性能,以便更好地参考。在指标方面,报告平均成功率 (SR) 和平均排名 [39]。