In-Distribution Imagination for Model-Based Offline Reinforcement Learning

arXiv:2609.38673v1 Announce Type: new Abstract: 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.

Related

• Mind Dreamer: Untethering Imagination via Active Latent Intervention on Latent Manifolds
• On Training in Imagination
• Dreaming Smoothly and Sample Efficiently with Gradient Penalized Latent Dynamics
• Dreamer-SAC: Off-Policy Learning in Latent World Models for Sample-Efficient Autonomous Driving

Source: arXiv cs.LG | 2026-10-01