Driving · 2026-10-01
In-Distribution Imagination for Model-Based Offline Reinforcement Learning
Mintae Kim, Koushil Sreenath
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
The authors propose in-distribution imagination (IDI) for model-based offline RL. IDI measures trajectory support in a learned latent representation space using KNN over offline trajectory windows, and adaptively truncates imagined rollouts when they leave the offline trajectory manifold. Combined with trajectory-regularized actor-critic (TRAC), IDI consistently improves limited-data performance and shows trajectory support predicts rollout failure better than transition-level uncertainty.
Why it matters
MBORL can generate unrealistic synthetic trajectories due to compounding model error, causing unstable policy optimization. The paper argues that rollout reliability is a trajectory-level property rather than a transition-level one. By controlling imagination with trajectory support (instead of local transition uncertainty), IDI aims to accept only behavior-supported, offline-manifold-compatible synthetic data and improve offline RL training stability.
Method
- IDI learns a history-aware trajectory encoder from offline trajectory windows and computes a KNN-based support score for a post-transition window in latent space.
- Rollouts are truncated online: a candidate transition is accepted only if its post-transition support score is below a threshold determined from offline leave-one-out support scores; otherwise the rollout terminates.
- IDI is combined with TRAC, which regularizes the learned policy-induced finite-horizon trajectory distribution toward the offline behavior trajectory distribution.
Abstract (from 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 \emph{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 papers
- From Suboptimal to Optimal: One-Step Latent Flow Policies for Offline Reinforcement Learning
- Trajectory-Aware Policy Optimization: Closing the State Distribution Gap in Long-Horizon LLM RL
- Beyond Policy Support: Interaction Constrained Offline Reinforcement Learning for Autonomous Driving
- Higher-Order Action Supervision Makes A Strong Policy Class
- VIGOR: Zero-Shot Visual Generalization via Latent-Space Consistency in Model-Based Reinforcement Learning
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First public release: 2026-09-29 · Indexed by Pub-AI: 2026-10-01 UTC
In-Distribution Imagination for Model-Based Offline Reinforcement Learning
Preprint posted: 2026-09-29 · Original record ↗
- arxiv
- 2609.38673
Explicit source versions (1)
- v1 · 2026-09-29
- Preprint available
Source date: 2026-09-29 · First detected: 2026-10-25 UTC · Historical baseline
Source evidence (1)
- Official evidence ↗ · atom:published (official metadata)
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 \emph{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.