# Mitigating Distribution Shift in Model-based Offline RL via Shifts-aware Reward Learning

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Offline RL algorithms, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 19 of its 33 references that are in the map.
- First posted on arXiv 23 Aug 2024.

## Placed in

- Offline RL and imitation learning from data / [Offline RL algorithms](https://mlsanity.com/offlinerl/offline-rl-algorithms/)

## Builds on (most cited first, 19 of 19)

- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [A Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06860/) (2021)
- [Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction](https://mlsanity.com/paper/1906.00949/) (2019)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [Behavior Regularized Offline Reinforcement Learning](https://mlsanity.com/paper/1911.11361/) (2019)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemble](https://mlsanity.com/paper/2110.01548/) (2021)
- [Model-Based Value Estimation for Efficient Model-Free Reinforcement Learning](https://mlsanity.com/paper/1803.00101/) (2018)
- [Algorithmic Framework for Model-based Deep Reinforcement Learning with Theoretical Guarantees](https://mlsanity.com/paper/1807.03858/) (2018)
- [Off-Dynamics Reinforcement Learning: Training for Transfer with Domain Classifiers](https://mlsanity.com/paper/2006.13916/) (2020)
- [State Alignment-based Imitation Learning](https://mlsanity.com/paper/1911.10947/) (2019)
- [Revisiting Design Choices in Offline Model-Based Reinforcement Learning](https://mlsanity.com/paper/2110.04135/) (2021)
- [State Regularized Policy Optimization on Data with Dynamics Shift](https://mlsanity.com/paper/2306.03552/) (2023)
- [Conservative Bayesian Model-Based Value Expansion for Offline Policy Optimization](https://mlsanity.com/paper/2210.03802/) (2022)

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

- [Generative OOD-regularized Model-based Policy Optimization](https://mlsanity.com/paper/2605.24405/) (2026)
- [A Unified Causal-Origin Taxonomy of Distributional Shifts in Reinforcement Learning](https://mlsanity.com/paper/2606.16933/) (2026)

## Abstract (arXiv)

Model-based offline reinforcement learning trains policies using pre-collected datasets and learned environment models, eliminating the need for direct real-world environment interaction. However, this paradigm is inherently challenged by distribution shift~(DS). Existing methods address this issue by leveraging off-policy mechanisms and estimating model uncertainty, but they often result in inconsistent objectives and lack a unified theoretical foundation. This paper offers a comprehensive analysis that disentangles the problem into two fundamental components: model bias and policy shift. Our theoretical and empirical investigations reveal how these factors distort value estimation and restrict policy optimization. To tackle these challenges, we derive a novel shifts-aware reward through a unified probabilistic inference framework, which modifies the vanilla reward to refine value learning and facilitate policy training. Building on this, we develop a practical implementation that leverages classifier-based techniques to approximate the adjusted reward for effective policy optimization. Empirical results across multiple benchmarks demonstrate that the proposed approach mitigates distribution shift and achieves superior or comparable performance, validating our theoretical insights.