# Learning on One Mode: Addressing Multi-modality in Offline Reinforcement Learning

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Imitation learning methods, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 28 of its 38 references that are in the map.
- First posted on arXiv 4 Dec 2024.

## Placed in

- Offline RL and imitation learning from data / [Imitation learning methods](https://mlsanity.com/offlinerl/imitation-learning-methods/)

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

- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (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)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [AWAC: Accelerating Online Reinforcement Learning with Offline Datasets](https://mlsanity.com/paper/2006.09359/) (2020)
- [Behavior Regularized Offline Reinforcement Learning](https://mlsanity.com/paper/1911.11361/) (2019)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning](https://mlsanity.com/paper/2208.06193/) (2022)
- [Learning Latent Plans from Play](https://mlsanity.com/paper/1903.01973/) (2019)
- [Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemble](https://mlsanity.com/paper/2110.01548/) (2021)
- [Critic Regularized Regression](https://mlsanity.com/paper/2006.15134/) (2020)
- [Keep Doing What Worked: Behavioral Modelling Priors for Offline Reinforcement Learning](https://mlsanity.com/paper/2002.08396/) (2020)
- [Offline RL Without Off-Policy Evaluation](https://mlsanity.com/paper/2106.08909/) (2021)
- [PLAS: Latent Action Space for Offline Reinforcement Learning](https://mlsanity.com/paper/2011.07213/) (2020)
- [Pessimistic Bootstrapping for Uncertainty-Driven Offline Reinforcement Learning](https://mlsanity.com/paper/2202.11566/) (2022)
- [Rethinking Goal-conditioned Supervised Learning and Its Connection to Offline RL](https://mlsanity.com/paper/2202.04478/) (2022)
- [How Far I'll Go: Offline Goal-Conditioned Reinforcement Learning via f-Advantage Regression](https://mlsanity.com/paper/2206.03023/) (2022)
- [What is Essential for Unseen Goal Generalization of Offline Goal-conditioned RL?](https://mlsanity.com/paper/2305.18882/) (2023)
- [Let Offline RL Flow: Training Conservative Agents in the Latent Space of Normalizing Flows](https://mlsanity.com/paper/2211.11096/) (2022)
- [GOPlan: Goal-conditioned Offline Reinforcement Learning by Planning with Learned Models](https://mlsanity.com/paper/2310.20025/) (2023)
- [Supported Trust Region Optimization for Offline Reinforcement Learning](https://mlsanity.com/paper/2311.08935/) (2023)
- [Balancing policy constraint and ensemble size in uncertainty-based offline reinforcement learning](https://mlsanity.com/paper/2303.14716/) (2023)
- [A Behavior Regularized Implicit Policy for Offline Reinforcement Learning](https://mlsanity.com/paper/2202.09673/) (2022)
- [TD3 with Reverse KL Regularizer for Offline Reinforcement Learning from Mixed Datasets](https://mlsanity.com/paper/2212.02125/) (2022)
- [Goal-conditioned Offline Reinforcement Learning through State Space Partitioning](https://mlsanity.com/paper/2303.09367/) (2023)

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

- [Retrosynthesis Planning via Worst-path Policy Optimisation in Tree-structured MDPs](https://mlsanity.com/paper/2509.10504/) (2025)
- [Beyond Penalization: Diffusion-based Out-of-Distribution Detection and Selective Regularization in Offline Reinforcement Learning](https://mlsanity.com/paper/2605.08202/) (2026)
- [Unifying Value Alignment and Assignment in Cross-Domain Offline Reinforcement Learning with Heterogeneous Datasets](https://mlsanity.com/paper/2605.24862/) (2026)

## Abstract (arXiv)

Offline reinforcement learning (RL) seeks to learn optimal policies from static datasets without interacting with the environment. A common challenge is handling multi-modal action distributions, where multiple behaviours are represented in the data. Existing methods often assume unimodal behaviour policies, leading to suboptimal performance when this assumption is violated. We propose weighted imitation Learning on One Mode (LOM), a novel approach that focuses on learning from a single, promising mode of the behaviour policy. By using a Gaussian mixture model to identify modes and selecting the best mode based on expected returns, LOM avoids the pitfalls of averaging over conflicting actions. Theoretically, we show that LOM improves performance while maintaining simplicity in policy learning. Empirically, LOM outperforms existing methods on standard D4RL benchmarks and demonstrates its effectiveness in complex, multi-modal scenarios.