Occupancy Reward Shaping: Improving Credit Assignment for Offline Goal-Conditioned Reinforcement Learning
Aravind Venugopal, Jiayu Chen, Xudong Wu, Chongyi Zheng, Benjamin Eysenbach + 1 more
TLDR
Occupancy Reward Shaping uses world models and optimal transport to improve credit assignment in offline goal-conditioned RL, boosting performance significantly.
Key contributions
- Formalizes how temporal information in world models encodes the underlying geometry of the world.
- Extracts this geometry into a reward function using optimal transport for goal-reaching information.
- Introduces Occupancy Reward Shaping (ORS) to mitigate credit assignment in sparse reward settings.
- Improves performance by 2.2x across 13 tasks and works for real-world nuclear fusion control.
Why it matters
This paper addresses the challenging credit assignment problem in goal-conditioned RL by leveraging temporal information from world models. Occupancy Reward Shaping significantly boosts performance in sparse reward settings, demonstrating its effectiveness across diverse tasks and in real-world nuclear fusion control.
Original Abstract
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/
📬 Weekly AI Paper Digest
Get the top 10 AI/ML arXiv papers from the week — summarized, scored, and delivered to your inbox every Monday.