Occupancy Reward Shaping: Improving Credit Assignment for Offline Goal-Conditioned Reinforcement Learning

📄 arXiv: 2604.20627v1 📥 PDF

作者: Aravind Venugopal, Jiayu Chen, Xudong Wu, Chongyi Zheng, Benjamin Eysenbach, Jeff Schneider

分类: cs.LG, cs.RO

发布日期: 2026-04-22

备注: ICLR 2026

🔗 代码/项目: GITHUB | PROJECT_PAGE

💡 一句话要点

提出Occupancy Reward Shaping，改善离线目标条件强化学习中的信用分配问题

🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱二：RL算法与架构 (RL & Architecture)

关键词: 离线强化学习 奖励重塑 信用分配 世界模型 最优传输 目标条件强化学习 占用度量

📋 核心要点

  1. 现有离线强化学习方法在稀疏奖励环境中面临信用分配难题，难以学习到有效的目标导向策略。
  2. 论文提出Occupancy Reward Shaping (ORS) 方法，利用世界模型中的时间信息，通过最优传输构建奖励函数。
  3. 实验表明，ORS在多种长时程任务中显著提升性能，并在真实托卡马克控制任务中验证了其有效性。

📝 摘要（中文）

在离线目标条件强化学习中，动作与其长期结果之间的时间滞后使得信用分配成为一项挑战。生成式世界模型能够捕捉智能体可能访问的未来状态的分布，表明它们已经捕获了时间信息。本文探讨了如何提取世界模型中存储的时间信息来进行信用分配。我们形式化地描述了世界模型中存储的时间信息如何编码世界的底层几何结构。利用最优传输，我们从学习到的占用度量模型中提取这种几何结构，并将其转化为一个能够捕捉目标到达信息的奖励函数。我们提出的Occupancy Reward Shaping (ORS) 方法在很大程度上缓解了稀疏奖励环境中的信用分配问题。ORS在理论上不会改变最优策略，但在13个不同的长时程运动和操作任务中，经验性地将性能提高了2.2倍。此外，我们还在真实的托卡马克控制任务中验证了ORS的有效性。

🔬 方法详解

问题定义：离线目标条件强化学习旨在从预先收集的数据集中学习策略，尤其是在奖励稀疏的环境下。传统的强化学习方法难以将最终目标与早期动作联系起来，导致信用分配困难，学习效率低下。现有的方法通常依赖于模仿学习或奖励重塑，但这些方法可能受到数据质量或奖励函数设计的限制。

核心思路：论文的核心思想是利用世界模型学习到的占用度量（occupancy measure）来指导奖励重塑。占用度量反映了智能体在不同状态下访问的频率，蕴含了环境的几何结构和时间信息。通过提取这种几何结构，可以构建一个奖励函数，引导智能体朝着目标状态前进。

技术框架：ORS方法的整体框架包括以下几个步骤：1）使用离线数据集训练一个生成式世界模型，学习环境的动态特性和占用度量。2）利用最优传输理论，计算当前状态到目标状态之间的距离，该距离反映了从当前状态到达目标状态的难易程度。3）基于计算出的距离，构建一个奖励函数，奖励智能体朝着目标状态移动。4）使用重塑后的奖励函数训练策略。

关键创新：ORS的关键创新在于利用占用度量来指导奖励重塑。与传统的奖励重塑方法不同，ORS不需要手动设计奖励函数，而是通过学习到的世界模型自动提取环境的几何结构，从而更有效地进行信用分配。此外，ORS方法在理论上保证不会改变最优策略，确保了学习的稳定性。

关键设计：在实现ORS时，需要选择合适的生成式世界模型，例如变分自编码器（VAE）或生成对抗网络（GAN）。最优传输的计算可以使用Sinkhorn算法等高效的求解器。奖励函数的具体形式可以设计为距离的负指数函数，以确保奖励值在合理范围内。此外，还需要仔细调整世界模型的训练参数和奖励函数的缩放因子，以获得最佳性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

ORS方法在13个不同的长时程运动和操作任务中，经验性地将性能提高了2.2倍，显著优于现有的离线强化学习方法。此外，ORS还在真实的托卡马克控制任务中取得了成功，证明了其在实际应用中的有效性。代码已开源。

🎯 应用场景

该研究成果可应用于机器人导航、操作、游戏AI等领域，尤其是在奖励稀疏、任务复杂的环境中。例如，可以用于训练机器人完成复杂的装配任务，或训练游戏AI在策略游戏中取得胜利。此外，该方法在托卡马克控制上的成功应用表明其在复杂物理系统控制方面具有潜力，未来可能应用于其他工业控制领域。

📄 摘要（原文）

The temporal lag between actions and their long-term consequences makes credit assignment a challenge when learning goal-directed behaviors from data. Generative world models capture the distribution of future states an agent may visit, indicating that they have captured temporal information. How can that temporal information be extracted to perform credit assignment? In this paper, we formalize how the temporal information stored in world models encodes the underlying geometry of the world. Leveraging optimal transport, we extract this geometry from a learned model of the occupancy measure into a reward function that captures goal-reaching information. Our resulting method, Occupancy Reward Shaping, largely mitigates the problem of credit assignment in sparse reward settings. ORS provably does not alter the optimal policy, yet empirically improves performance by 2.2x across 13 diverse long-horizon locomotion and manipulation tasks. Moreover, we demonstrate the effectiveness of ORS in the real world for controlling nuclear fusion on 3 Tokamak control tasks. Code: https://github.com/aravindvenu7/occupancy_reward_shaping; Website: https://aravindvenu7.github.io/website/ors/