跳到主要内容
返回时间线
arXiv来源发表:

零样本强化学习不再被稠密任务编码拖累:加一行稀疏惩罚,DMC平均回报提高53分

核心概要

在TD-JEPA的任务编码上加一项稀疏惩罚,使推断奖励在目标附近更集中,零样本强化学习在DMC平均回报提高53分、OGBench平均成功率提高10.4个百分点。

Source-provided article image: Sparse-TD-JEPA: Improving Zero-Shot Reinforcement Learning With Sparse Latent Task Representations
Figure 1 · arXiv

深度剖析

零样本强化学习预训练一个任务编码器,测试时只用少量带奖励标注的转移做闭式线性回归就能恢复策略;当编码器输出稠密时,推断出的奖励在状态空间上摊得很开,目标状态与背景之间的对比变弱,策略可跟随的信号随之变弱。 此前基于后继特征与任务编码器的零样本方法把注意力放在如何构造任务表示上,例如用转移图的特征函数、意图条件价值或时间距离来组织隐空间,而编码器输出稠密导致奖励信号弥散这一瓶颈受到的关注较少。 在DeepMind Control与OGBench共65个任务上,稀疏版本在全部13个域级比较中都优于稠密基线TD-JEPA;DMC平均回报提高53分,OGBench平均成功率提高10.4个百分点,其中antmaze导航与拼接域提高13.4至18.0个百分点。

做法是在训练时对投影后的任务嵌入施加一项稀疏惩罚,让每个嵌入把质量集中到更少的强激活维度上,从而把奖励预测结构压进一个更小的有效子空间。 该惩罚只作用于训练,不改变网络结构,也不改动测试时的闭式推断步骤,因此实现上只是在TD-JEPA的更新步里增加一行代码。 评估覆盖DMC的walker、cheetah、quadruped、pointmass四个域共20个任务,以及OGBench九个域共45个任务,每个任务做10次独立rollout取平均。

分析显示稀疏编码器给出的推断奖励景观更尖锐、更贴近目标,这为零样本性能提升提供了一个机制上的解释。 稠密编码器在目标房间内推断奖励仍然整体偏高,进入房间后难以指出目标的确切位置;稀疏编码器则在目标处形成明显峰值,向房间边缘逐渐下降。 在20个独立重采样的观测批次上,用5折交叉验证的相关系数衡量推断奖励与任务进度信号的一致性,稀疏编码器在九个域中的六个显著更高,且这一差距在标注集规模跨越多个数量级时大致保持平坦。

启示与展望

该结果适用于基于后继特征与任务编码器的零样本强化学习设定:在无奖励离线数据上预训练,测试时用少量带标注转移做闭式线性回归恢复策略。评估覆盖DMC的四个域共20个任务与OGBench的九个域共45个任务,涵盖运动、导航、轨迹拼接、拼图与机械臂操作。稀疏惩罚目前只在TD-JEPA上实例化,作者提出将其应用到其他基于后继特征的零样本方法,并把得到的低有效维度任务嵌入用于理解单个隐维度编码了什么。

稀疏惩罚只约束每个样本自身的嵌入,并不强制不同状态共享同一支撑集,是否形成共享的低维结构是一个经验问题。在距离-奖励相关性分析中,Large Navigate域未达到统计显著;在标注集规模扫描中,puzzle-3x3域对两种模型都没有可靠信号。稀疏编码器是否更好地表示距离目标之外的其他任务维度,作者留作未来工作。