# In-Distribution Imagination for Model-Based Offline Reinforcement Learning

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Offline RL algorithms, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 18 of its 19 references that are in the map.
- First posted on arXiv 29 Sep 2026.

## Placed in

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

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

- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [MOReL : Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2005.05951/) (2020)
- [COMBO: Conservative Offline Model-Based Policy Optimization](https://mlsanity.com/paper/2102.08363/) (2021)
- [1000 Layer Networks for Self-Supervised RL: Scaling Depth Can Enable New Goal-Reaching Capabilities](https://mlsanity.com/paper/2503.14858/) (2025)
- [InfoNCE: Identifying the Gap Between Theory and Practice](https://mlsanity.com/paper/2407.00143/) (2024)
- [Double Check Your State Before Trusting It: Confidence-Aware Bidirectional Offline Model-Based Imagination](https://mlsanity.com/paper/2206.07989/) (2022)
- [Is Model Ensemble Necessary? Model-based RL via a Single Model with Lipschitz Regularized Value Function](https://mlsanity.com/paper/2302.01244/) (2023)
- [SUMO: Search-Based Uncertainty Estimation for Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2408.12970/) (2024)
- [Learning-based Trajectory Tracking for Bird-inspired Flapping-Wing Robots](https://mlsanity.com/paper/2411.15130/) (2024)
- [Estimation of Aerodynamics Forces in Dynamic Morphing Wing Flight](https://mlsanity.com/paper/2508.02984/) (2025)
- [RoVerFly: Robust and Versatile Implicit Hybrid Control of Quadrotor-Payload Systems](https://mlsanity.com/paper/2509.11149/) (2025)
- [Finite Memory Belief Approximation for Optimal Control in Partially Observable Markov Decision Processes](https://mlsanity.com/paper/2601.03132/) (2026)
- [WOMBET: World Model-Based Experience Transfer for Robust and Sample-efficient Reinforcement Learning](https://mlsanity.com/paper/2604.08958/) (2026)
- [Robust Adversarial Policy Optimization Under Dynamics Uncertainty](https://mlsanity.com/paper/2604.10974/) (2026)
- [Trajectory-Regularized Stochastic Optimal Control via KL Divergence](https://mlsanity.com/paper/2607.22201/) (2026)

## Abstract (arXiv)

Model-based offline reinforcement learning (MBORL) improves sample efficiency through model-generated trajectories. However, accumulative model error can drive imagined trajectories outside the offline data distribution, leading to unrealistic synthetic data and unstable policy optimization. Many existing methods primarily control rollouts using transition-level uncertainty. We propose in-distribution imagination (IDI), a rollout control framework that estimates trajectory support in a learned representation space and adaptively truncates rollouts that leave the offline trajectory manifold. Combined with trajectory-regularized RL, an extension of entropy-regularized RL, IDI consistently improves performance in limited-data settings. Experiments show that trajectory support predicts rollout failure substantially better than transition-level uncertainty, highlighting the importance of trajectory-level rollout control in MBORL.