机器人强化学习:如何用一个数字代替整个动作矩阵?
有一个问题困扰了做机器人强化学习的研究者好几年:你训练好一个会生成动作的流匹配模型,想用强化学习再优化它一下,让它表现更好。但这个模型生成一个动作要走很多步,就像你写一篇文章要打几十次草稿才能定稿。强化学习的评价信号只在最后一稿给出,可你想告诉前面每一次草稿修改方向,这信号要怎么传回去?这篇论文要讲...
做机器人强化学习的研究者被一个问题卡了好几年:动作生成用流匹配模型训练好了,想再用强化学习优化一把,让它表现更好。麻烦在于,这个模型生成一个动作要走很多步,好比写一篇文章得打几十次草稿才定稿。强化学习的评价信号只在最后一稿给出,前面每一次草稿该怎么改,信号怎么传回去?
这篇论文讲的,就是研究者把这件麻烦事变简单的经过。
流匹配模型是什么
先说说流匹配策略。
流匹配是一种生成模型技术,从一堆随机噪声出发,经过很多小步骤,逐渐把噪声"整理"成有意义的输出,比如一个机器人动作。
可以把它想成雕塑家凿石头:粗糙的石料(随机噪声)经过几十次精细凿刻(流步骤),最后成了一座雕像(具体动作)。每一凿都在前一步基础上微调方向,石料不是一下子变成雕像的,是一步步靠近的。
机器人策略如今越来越多用这种方式生成动作,原因是它能表达特别丰富多样的动作分布,老式方法只能生成单一僵硬的输出。Physical Intelligence的π0系列、英伟达的GR00T N1,这些大型机器人基础模型骨子里都是流匹配。
问题来了:这些模型通常先用人类示范数据做模仿学习,训练出一个基础版本,研究者再想用强化学习给它"开小灶",让它在基础能力之上更强,超过示范数据本身的水平。这就要用到一个叫Q函数的东西。
Q函数又被称作价值函数,输入一个状态和一个动作,它给出分数,表示在这个状态下做这个动作大约能换来多少长期回报。强化学习的做法,就是持续调整策略,使策略产出的动作获得更高的Q函数评分。
但要把Q函数的打分信息传给流匹配模型内部的每一步雕刻动作,没那么直接。Q函数只能看到雕刻完成后的最终成品,雕刻过程中间任何一凿是什么状态,它看不到,也不理解。
当时领域里怎么处理这个难题
现在业界有几种办法。
最简单粗暴的一种是干脆不碰这个流匹配模型本身,冻住它,在外面加一个小型修正模块,让新模块去学怎么对最终生成的动作做些微调。这就像不改原来的文章,而是雇一个编辑在你写完之后帮你润色几句。好处是简单安全,坏处是编辑的能力永远受限于你原文的底子,文章天花板本来就在那,编辑顶多帮你擦亮一点,没法真正突破上限。EXPO、DSRL这类方法走的就是这条路。
另一条路更硬核,叫adjoint matching。
Adjoint matching(伴随匹配)在理论上更为严谨。它借助一种名为"伴随常微分方程"的数学工具,把Q函数在最终动作处的梯度信息逐层往回传递,从而指明雕刻过程中每一凿该朝哪个方向下刀。
听上去不错:直接改动模型本身,不像前面那些外挂方案那样被原模型的能力上限卡住。可代价很要命——梯度信息要往回传,每一步都得算一次"向量雅可比积",而这一次运算要完整穿过策略网络。
雅可比矩阵描述一个函数在某一点附近各个输出维度对各个输入维度的敏感程度,可以理解成一张"牵一发动全身"的关联图。
假如你的流匹配模型生成一个动作需要走十步、二十步,那就得做十次、二十次这种高开销的运算;策略网络越大,单次运算也越贵。打个比方,你写文章要改二十遍,而编辑每给一次反馈之前,都得把你写作时脑子里牵扯到的所有知识关联重新捋一遍,想想这有多慢。
论文作者正是卡在这里:adjoint matching能直接改进生成模型本身,这份理论上的好处他们不想丢;可每一步都要做昂贵运算,这个代价又实在吃不消。
他们发现了什么
研究者们做了一件很朴素的事:打开那个雅可比矩阵看看里面长什么样。
他们把预训练好的流匹配模型在一批状态上跑了一遍,算出雅可比矩阵的平均值,再画成热力图。结果出人意料:矩阵上几乎只有对角线一带着色,其余地方基本一片空白。
所谓对角矩阵结构,就是矩阵中仅有从左上到右下那条对角线上有数值,其他位置的数值接近零。这表示每个输出维度几乎只受与之对应的那个输入维度影响,各维度之间牵扯不大。
论文里给出的图很直观,在antmaze、humanoidmaze、puzzle、cube四个完全不同的任务领域上,这个对角线结构都稳定出现,不是偶然现象。研究者后面在附录里把所有十个任务领域都测了一遍,结果一致。
这个发现说明了一件事:那个代价高昂的"向量雅可比积"运算,大半力气都花在算那些近似为零的非对角线项上。真正管用的信息,几乎全挤在对角线这条窄缝里。
反过来看,如果没发现这条规律,研究者只能继续背着随流步数和模型规模线性增长的计算包袱,adjoint matching就始终走不出小模型实验,没法上到π0这类动辄几亿参数的大型机器人基础模型。规律的价值就在这里:它把"必须全算"的问题变成了"可以近似"的问题。
从矩阵到一个数字:标量伴随
雅可比矩阵既然近乎对角,研究者索性做了个大胆简化:整个矩阵用一个数字乘单位矩阵来近似。
单位矩阵*:对角线全是1、其余全是0的矩阵,相当于什么都不做,原样输出。
这个近似听着有些粗暴,数学上却占了大便宜。原本靠数值积分一步步往回传的伴随常微分方程,在这个假设下居然有了析解——可以直接写出公式,不必再逐步数值求解。研究者证明,该假设下解出的结果恰好等于流动时间τ乘以Q函数在最终动作上的梯度。
τ(流动时间)*:表示生成过程进行到哪一步了,从0(刚开始,全是噪声)到1(结束,生成完整动作)。
研究者接着发现,公式里有个指数因子,它只缩放整体大小、不改变方向,于是索性一并省掉,剩下的式子极简单:某一步的调整方向,就是负的τ乘以Q函数的梯度。他们把这套东西命名为标量伴随。
简化之后直接的好处是:流匹配模型内部再也不必穿行做任何向量雅可比积运算。只需在最终生成的动作上算一次Q函数的梯度,再乘上不同的τ值,分发给每一步即可。
打个比方,原来编辑要在你写作的每一步都重新理解整篇文章的所有逻辑关联才能给反馈;现在编辑发现,其实只要在你写完后说一句"整体方向需要往左调一点",你自己按比例把这句话分摊到每次修改里就够了,越靠后的修改分摊的比重越大,因为越靠后的草稿离定稿越近,越要精细调整。不这么简化,编辑每次反馈前都得重梳一遍全文逻辑,你改二十次稿,编辑就得重梳二十次,这工作量显然不成比例。
论文里有一张图专门验证这个简化靠不靠谱。他们拿标量伴随和精确伴随(老老实实做向量雅可比积算出的真实结果)比方向和大小。结果显示,两者在大小的变化趋势上几乎一致,流动时间越接近0,两者的调整幅度都越小,趋势完全对得上。方向上,标量伴随始终与精确伴随保持正相关,没有跑偏。不过研究者也坦承,这个夹角在操作类任务(比如cube、puzzle这种需要精细抓取操作的任务)上比在导航类任务(比如antmaze、humanoidmaze这种走迷宫式的任务)上要大一些,这个细节后面会牵出另一个重要发现。
单独用标量伴随还不够:价值函数的陷阱
研究者本以为找到这个简化公式就万事大吉,结果一测试才发现,光用标量伴随,在操作类任务上效果并不好。
背后的原因挺微妙。精确伴随要穿过雅可比矩阵,Q函数的梯度信息先被这个矩阵"过滤"一道才传到每一步,相当于信息经过了一层缓冲。标量伴随则把Q函数的梯度原封不动地分发出去,没有缓冲。
也就是说,一旦Q函数本身在策略生成的那个具体动作上估得不准,这个误差会被标量伴随毫无保留地放大传出去,没有任何过滤机制来缓冲。
Q函数估计不准的问题,在强化学习里早就不是新鲜事了。研究者们此前提出过各种办法来约束Q函数,防止它在没见过的动作上乱打高分。常见做法大致分两类。
in-sample 这一类方法的代表是 IQL:Q 函数只拿数据集中真实出现过的动作来训练,策略自己生成的动作,它根本不去算 Q 值。行为正则化则不同,ReBRAC 走的就是这条路,它限制策略生成的动作与数据集动作之间的偏离程度,避免策略跑偏太远。
论文里实测发现,这两类方法用在标量伴随上效果都不理想,几乎和完全不加任何约束的效果差不多。
道理并不复杂。标量伴随给出的更新信号,直接取决于 Q 函数在“策略当下生成的那个动作”上的梯度。in-sample 方法压根不评估这个动作上的 Q 值,真正起作用的那个点它覆盖不到。行为正则化管的是动作落在什么位置,而不是 Q 函数给这个动作打出的分数;它约束住了“别跑太远”,却没约束住“你在这个位置上的打分靠不靠谱”。
打个比方:老师批改学生的作文草稿,只按教材范文的标准打分(in-sample),学生自己写的那篇草稿该得几分,从来没看过;或者老师只规定“这次改动字数别太多”(行为正则化),却不看“改完之后这篇文章到底写得好不好”。两种管法,都没落到真正该管的地方——学生草稿本身的质量评分。
于是研究者提出了一个针对性的解法,叫价值惩罚。
价值惩罚*:训练 Q 函数时额外加一项惩罚,专门约束 Q 函数在策略刚生成的那个动作上的打分,要求它不能比在数据集真实动作上的打分高出太多。
这个设计正好对准了标量伴随真正依赖的位置——策略生成动作处的 Q 函数取值。成本也几乎为零:策略生成的那个动作在训练过程中本来就采样好了,直接拿来用即可,不必再额外采样。
研究者还顺带验证了另一个相关猜想:如果惩罚的不是 Q 函数在这个动作上的数值大小,而是它的梯度大小,会怎样?结果同样有效,cube-quadruple 任务上尤其明显。论文还从数学上证明,价值惩罚做一阶展开后,恰好与梯度大小惩罚等价关联,两者本质上是同一个道理的不同表现形式。
为了搞清楚这个解释到底成不成立,研究者又补了一组对照实验。做法是把价值惩罚挪到精确伴随上用,也就是TRQAM那套依赖完整雅可比矩阵的老办法。结果价值惩罚确实也让精确伴随有所改善,可提升的幅度跟它给标量伴随带来的增益比,差得远。这正好对上了前面的推断:精确伴随本身要穿过雅可比矩阵,已经多了一层缓冲,价值惩罚能添的好处自然就有限了。
最终的方法叫SQAM,把标量伴随和价值惩罚拼在一起,就是这篇论文给出的方案。
SQAM(Q-learning with Scalar Adjoint Matching)*是一种让流匹配策略做强化学习微调的算法:核心思路是用一个标量数字顶替整个雅可比矩阵,再配上一个专门盯住策略所生成动作的价值惩罚,算力省下来一大截,性能却没有让步。
算法骨架其实是从此前的TRQAM继承来的,那套框架靠信任域来约束微调幅度,不让它跑得太偏。研究者把这套信任域机制原样留着,只把其中最吃资源的那一步,也就是逐步求解伴随常微分方程,换成了标量伴随的闭式公式。
实验做得不小。论文在OGBench基准上铺开大规模对比,横跨十个任务领域、五十个具体任务,导航类有蚂蚁迷宫、人形迷宫,操作类有方块抓取、拼图、场景摆放。
结果大体上集中印证了这套方法的价值。五十个任务整体看,SQAM平均成功率达到75%,此前最强的基线TRQAM是68%。两个数字差距不算惊人,关键在于增益落在哪些任务上。
| 任务领域 | 最强基线成功率 | SQAM成功率 | 提升幅度 ||---|---|---|---|| antmaze-giant | 41% | **62%** | 21个百分点 || humanoidmaze-large | 36% | **54%** | 18个百分点 || cube-triple | 50% | **72%** | 22个百分点 || cube-quadruple | 19% | **54%** | 35个百分点 |
这四个领域恰好是整个基准里最难啃的四块,数据集最大,最多到一亿条转移样本,任务也最复杂,动作空间维度最高。换句话说,SQAM并不是在哪个任务上都均匀地好一点,它是在最硬的骨头上啃出了明显差距。这个信号挺有意思:任务越难,这套方法的优势反而越明显,而不是简单任务上数字漂亮、难任务上就原形毕露。
额定数字背后还有一处细节:在cube-double这个任务上,价值惩罚反倒让离线阶段的成绩掉了,于是研究者在这个任务上直接把价值惩罚关掉。可见价值惩罚不是放之四海而皆准的万能钥匙,得按具体任务去调,论文也把这个局限诚实地摆了出来。
光在仿真器里好用还不够。研究者把SQAM搬到了真实的双臂机器人上,是一台配两只20自由度灵巧手的人形机器人,底座为RB-Y1,手是Wuji Hand 2,脑子用的是名叫RLDX-1的视觉语言动作预训练模型。
VLA模型(视觉语言动作模型)*是一种大型预训练模型,接收图像和文字指令,输出机器人动作,眼下属机器人领域最热门的技术路线之一。
真实任务一共设计了三个:把塑料袋翻过来、把吸管插进杯子、把水果放进锅里再盖上盖子。每个任务先靠人类示范做监督微调,然后用SQAM跑15000步离线强化学习,再接着跑10000步在线强化学习,也就是让机器人一边跟真实世界互动一边继续学。
三个任务上,SQAM全都压过了纯监督微调。以吸管任务为例,监督微调版本30次里只成功7次,SQAM到离线阶段提到13次,在线继续训练后又升到15次,差不多翻了一倍。拿来对比的另一个方法EXPO-FT走的是冻住基础模型、只学一个外挂修正模块的路子,在这个任务上完全没学会,30次里成功0次。
这组对比很能说明问题。EXPO-FT得从零学一个全新的修正策略,这个策略没有继承基础模型已经掌握的能力,好比让一个新手编辑从零开始摸你的文章风格;SQAM则直接在原有模型上微调,等于让原作者自己改自己的稿子,对全局的理解天然就在。
研究者还细看了一件事:成功的那些回合花了多长时间完成任务。翻塑料袋和放水果盖盖子这两个任务上,SQAM的完成时间都比监督微调版本更短,也就是说成功率上去了,动作本身也更果断利索。
写在后面
读到雅可比矩阵呈对角结构这个发现时,我第一反应是这事巧得让人起疑,会不会是流匹配这类生成模型本身的训练目标就倾向于让网络学出这种结构。论文没有给出理论解释,作者也把这个疑问老老实实留在了未来工作里:他们知道现象在,但不知道为什么在。比起硬凑一个牵强的说法,这种坦白反而更让人愿意信任这篇工作。
还有一处让我来回琢磨:价值惩罚与梯度范数惩罚之间,存在一个一阶展开上的等价关系。一个盯数值,一个盯梯度,看上去是两条完全不同的正则化路子,数学上却能证明二者是同一件事的两种近似写法。这让我想到,不少时候我们以为自己做出了“两套不同的方案”,说到底不过是同一个洞察换了两层皮。
论文收尾那句“问题可能不在策略更新上,而在Q函数本身”,也值得停下来单独品一品。整篇工作起初惦记的是怎样免掉伴随矩阵的计算,到最后却拐回一个更根本的疑问:喂给策略的反馈信号,本身靠不靠谱。从“怎么算得更快”一路走到“算出来的东西对不对”,这个思路上的转身,是我觉得全文最有价值的地方。
反过来想,如果一份粗糙的近似就足以撑起整个策略更新,那真正卡住上限的那堵墙又在哪里?
Q&A
Q1:SQAM是什么?
A:SQAM的全称是Q-Learning with Scalar Adjoint Matching,属于一类让流匹配机器人策略完成强化学习微调的算法。它注意到预训练流模型的速度场雅可比矩阵主要落在对角线上,于是拿一个标量数字去近似整个矩阵,每个流步骤原本都要做的向量雅可比积运算就此省掉;再加上一项专门约束策略所生成动作处Q函数取值的价值惩罚,两者共同把性能往上推。
Q2:SQAM相比之前的adjoint matching方法有什么优势?
A:以往的adjoint matching(例如QAM、TRQAM)在流匹配模型的每个生成步骤上都得做一次向量雅可比积运算,算力开销随流步数与策略规模线性上涨。SQAM改用闭式的标量伴随来顶替这一过程,完全不必穿过策略网络做向量雅可比积,计算量因此大幅下降;同时它在OGBench最难的四个任务领域上,成功率反倒比TRQAM高出18到35个百分点。
Q3:SQAM在真实机器人上效果如何?
A:研究者将SQAM部署到装有双臂和灵巧手的真实机器人上,对一个视觉语言动作预训练模型做微调,去完成三个操作任务。结果是三个任务上SQAM都明显压过纯监督微调:以插吸管任务为例,成功率由7/30升到15/30,而作为对照的EXPO-FT方法在这个任务上压根没学会。