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/

翻译：行为与其长期后果之间的时间延迟，使得从数据中学习目标导向行为时的信用分配成为一项挑战。生成式世界模型捕捉了智能体可能访问的未来状态分布，表明它们已掌握时间信息。如何提取这些时间信息以进行信用分配？在本文中，我们形式化了世界模型中存储的时间信息如何编码世界的潜在几何结构。利用最优传输理论，我们从学习到的占据测度模型中提取出这种几何结构，并将其转化为一个捕捉目标到达信息的奖励函数。由此产生的方法——占据奖励塑形（ORS），在很大程度上缓解了稀疏奖励设置下的信用分配问题。ORS 理论上不改变最优策略，但在实践中，它在 13 种不同的长时域运动与操作任务中，将性能提升了 2.2 倍。此外，我们还在现实世界中展示了 ORS 在 3 个托卡马克控制任务上用于控制核聚变的有效性。代码：https://github.com/aravindvenu7/occupancy_reward_shaping；网站：https://aravindvenu7.github.io/website/ors/