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

arXiv: https://arxiv.org/abs/2604.20627 · HTML: https://mlsanity.com/paper/2604.20627/

## Key facts

- Cited by 1 papers 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.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)
- Offline RL and imitation learning from data / [Offline RL algorithms](https://mlsanity.com/offlinerl/offline-rl-algorithms/)

## Builds on (most cited first, 30 of 46)

- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Curiosity-driven Exploration by Self-supervised Prediction](https://mlsanity.com/paper/1705.05363/) (2017)
- [A Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06860/) (2021)
- [Flow Matching Guide and Code](https://mlsanity.com/paper/2412.06264/) (2024)
- [Contrastive Learning as Goal-Conditioned Reinforcement Learning](https://mlsanity.com/paper/2206.07568/) (2022)
- [OGBench: Benchmarking Offline Goal-Conditioned RL](https://mlsanity.com/paper/2410.20092/) (2024)
- [Semi-parametric Topological Memory for Navigation](https://mlsanity.com/paper/1803.00653/) (2018)
- [Flow Q-Learning](https://mlsanity.com/paper/2502.02538/) (2025)
- [Goal-conditioned Imitation Learning](https://mlsanity.com/paper/1906.05838/) (2019)
- [Goal-Conditioned Reinforcement Learning: Problems and Solutions](https://mlsanity.com/paper/2201.08299/) (2022)
- [HIQL: Offline Goal-Conditioned RL with Latent States as Actions](https://mlsanity.com/paper/2307.11949/) (2023)
- [Goal-Conditioned Reinforcement Learning with Imagined Subgoals](https://mlsanity.com/paper/2107.00541/) (2021)
- [Offline Reinforcement Learning with Realizability and Single-policy Concentrability](https://mlsanity.com/paper/2202.04634/) (2022)
- [Optimal Goal-Reaching Reinforcement Learning via Quasimetric Learning](https://mlsanity.com/paper/2304.01203/) (2023)
- [Foundation Policies with Hilbert Representations](https://mlsanity.com/paper/2402.15567/) (2024)
- [Rethinking Goal-conditioned Supervised Learning and Its Connection to Offline RL](https://mlsanity.com/paper/2202.04478/) (2022)
- [Is Value Learning Really the Main Bottleneck in Offline RL?](https://mlsanity.com/paper/2406.09329/) (2024)
- [Dynamical Distance Learning for Semi-Supervised and Unsupervised Skill Discovery](https://mlsanity.com/paper/1907.08225/) (2019)
- [Efficient Planning in a Compact Latent Action Space](https://mlsanity.com/paper/2208.10291/) (2022)
- [Temporal Abstraction in Reinforcement Learning with the Successor Representation](https://mlsanity.com/paper/2110.05740/) (2021)
- [Dual RL: Unification and New Methods for Reinforcement and Imitation Learning](https://mlsanity.com/paper/2302.08560/) (2023)
- [Generative Temporal Difference Learning for Infinite-Horizon Prediction](https://mlsanity.com/paper/2010.14496/) (2020)
- [Multi-step Reinforcement Learning: A Unifying Algorithm](https://mlsanity.com/paper/1703.01327/) (2017)
- [Adversarial Intrinsic Motivation for Reinforcement Learning](https://mlsanity.com/paper/2105.13345/) (2021)
- [Flow-Based Policy for Online Reinforcement Learning](https://mlsanity.com/paper/2506.12811/) (2025)

## Built on by (most cited first, 1 of 1)

- [SpikeCredit: Temporal Credit Carrier for Reinforcement Learning with Sparse Rewards](https://mlsanity.com/paper/2609.35268/) (2026)

## 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/