Occupancy Reward Shaping: Improving Credit Assignment for Offline Goal-Conditioned Reinforcement Learning
On the map: World models → Latent dynamics models and model-based RL
also in 1 more
Abstract (arXiv)
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/
Key facts. Cited by 1 paper in the map (data through 14 Oct 2026). Placed in Latent dynamics models and model-based RL, one of 11 directions in World models. Builds on 46 of its 67 references that are in the map. First posted on arXiv 22 Apr 2026.
Builds on
46 of 67 references.
References · 46
- 2022Flow Matching for Generative Modelingcited by 5,702
- 2018Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actorcited by 4,885
- 2022Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flowcited by 3,700
- 2016Gaussian Error Linear Units (GELUs)cited by 3,502
- 2020D4RL: Datasets for Deep Data-Driven Reinforcement Learningcited by 1,427
- 2021Offline Reinforcement Learning with Implicit Q-Learningcited by 1,293
- 2017Curiosity-driven Exploration by Self-supervised Predictioncited by 1,194
- 2021A Minimalist Approach to Offline Reinforcement Learningcited by 899
- 2024Flow Matching Guide and Codecited by 301
- 2022Contrastive Learning as Goal-Conditioned Reinforcement Learningcited by 253
- 2024OGBench: Benchmarking Offline Goal-Conditioned RLcited by 238
- 2018Semi-parametric Topological Memory for Navigationcited by 209
34 more
- 2025Flow Q-Learningcited by 184
- 2019Goal-conditioned Imitation Learningcited by 162
- 2022Goal-Conditioned Reinforcement Learning: Problems and Solutionscited by 158
- 2023HIQL: Offline Goal-Conditioned RL with Latent States as Actionscited by 141
- 2021Goal-Conditioned Reinforcement Learning with Imagined Subgoalscited by 134
- 2022Offline Reinforcement Learning with Realizability and Single-policy Concentrabilitycited by 123
- 2023Optimal Goal-Reaching Reinforcement Learning via Quasimetric Learningcited by 121
- 2024Foundation Policies with Hilbert Representationscited by 94
- 2022Rethinking Goal-conditioned Supervised Learning and Its Connection to Offline RLcited by 89
- 2024Is Value Learning Really the Main Bottleneck in Offline RL?cited by 89
- 2019Dynamical Distance Learning for Semi-Supervised and Unsupervised Skill Discoverycited by 68
- 2022Efficient Planning in a Compact Latent Action Spacecited by 57
- 2021Temporal Abstraction in Reinforcement Learning with the Successor Representationcited by 54
- 2023Dual RL: Unification and New Methods for Reinforcement and Imitation Learningcited by 48
- 2020Generative Temporal Difference Learning for Infinite-Horizon Predictioncited by 45
- 2017Multi-step Reinforcement Learning: A Unifying Algorithmcited by 42
- 2021Adversarial Intrinsic Motivation for Reinforcement Learningcited by 40
- 2025Flow-Based Policy for Online Reinforcement Learningcited by 37
- 2023What is Essential for Unseen Goal Generalization of Offline Goal-conditioned RL?cited by 36
- 2023Contrastive Difference Predictive Codingcited by 36
- 2019Learning to Reach Goals via Iterated Supervised Learningcited by 26
- 2023Learning Goal-Conditioned Policies Offline with Self-Supervised Reward Shapingcited by 25
- 2025Reward Models in Deep Reinforcement Learning: A Surveycited by 25
- 2025Temporal Difference Flowscited by 24
- 2025Option-aware Temporally Abstracted Value for Offline Goal-Conditioned Reinforcement Learningcited by 23
- 2024Highly Efficient Self-Adaptive Reward Shaping for Reinforcement Learningcited by 22
- 2023SMORE: Score Models for Offline Goal-Conditioned Reinforcement Learningcited by 21
- 2024Can a MISL Fly? Analysis and Ingredients for Mutual Information Skill Learningcited by 21
- 2023Distance Weighted Supervised Learning for Offline Interaction Datacited by 18
- 2023f-Policy Gradients: A General Framework for Goal Conditioned RL using f-Divergencescited by 14
- 2025Flattening Hierarchies with Policy Bootstrappingcited by 11
- 2024Full Shot Predictions for the DIII-D Tokamak via Deep Recurrent Networkscited by 9
- 2025Automatic Reward Shaping from Confounded Offline Datacited by 6
- 2024Bellman Diffusion Modelscited by 5
Built on by
Citing papers per year: 2026: 1.
1 paper.
Similar, not cited
Close in wording, no citation either way
- 2026Reward-free Pretraining for Reinforcement Learning via Occupancy Coverage Maximization
- 2024Dense Dynamics-Aware Reward Synthesis: Integrating Prior Experience with Demonstrations
- 2026Occupancy Modeling via Drifting: One-Step Generation Enables Direct Policy Extraction
- 2026Reward Prediction with Factorized World States
- 2026Do Better Goal Representations Improve Goal-Conditioned Reinforcement Learning?
- 2026TRACE: Turn-level Reward Assignment via Credit Estimation for Long-Horizon Agents
- 2026Reward as An Agent for Embodied World Models
- 2024ORSO: Accelerating Reward Design via Online Reward Selection and Policy Optimization