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

2026-04-22

AI Summary

This work addresses the credit assignment challenge in offline goal-conditioned reinforcement learning, particularly under sparse rewards, where temporal delays between actions and their long-term consequences hinder learning. The authors propose a novel approach based on generative world models that leverages occupancy measures and optimal transport theory to embed temporal information into the geometric structure of the environment. This enables the construction of a goal-reaching shaping reward that preserves the optimality of policies. Evaluated across 13 long-horizon locomotion and manipulation tasks, the method achieves an average performance improvement of 2.2× over baselines. Furthermore, it demonstrates practical efficacy by successfully transferring to three real-world tokamak nuclear fusion control tasks, validating both its effectiveness and real-world applicability.

Technology Category

• Machine Learning: Reinforcement Learning
• Search and Optimization: Learning to Search
• Humans and AI: Human-Aware Planning and Behavior Prediction

Application Category

• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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/

Problem

• credit assignment
• offline reinforcement learning
• goal-conditioned reinforcement learning
• sparse reward
• temporal lag

Innovation

• Occupancy Reward Shaping
• Credit Assignment
• Optimal Transport
• Goal-Conditioned Reinforcement Learning
• World Models