In-Distribution Imagination for Model-Based Offline Reinforcement Learning

Mintae Kim; Koushil Sreenath

cs.LGSep 29, 2026

In-Distribution Imagination for Model-Based Offline Reinforcement Learning

Authors: Mintae Kim, Koushil Sreenath

Organizations: Hybrid Robotics, BAIR, UC Berkeley

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 \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.

Figures & tables

Figure 1 : Conceptual overview of in-distribution imagination (IDI). (a) Conventional MBORL relies on transition-level reliability, but locally plausible transitions may still accumulate into unrealistic trajectories due to compounding model error. (b) IDI instead measures trajectory support in a learned latent trajectory space and truncates imagination when trajectories leave the offline manifold. This converts rollout control from local uncertainty estimation to trajectory-level support estimation.

Figure 2 : Normalized return learning curves on MuJoCo locomotion benchmarks under the limited-data setting. TRAC is trained using only

100

episodes, while the dashed line denotes TRAC trained on the full dataset (

1000

episodes). Fixed-horizon and uncertainty-based truncation improve over the limited-data baseline, but their effectiveness varies across environments. IDI consistently achieves the strongest performance through history-aware rollout control.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Method	HalfCheetah	Hopper	Walker2d
TRAC (1000 traj.)	$\mathbf{93.2\pm 1.1}$	$\mathbf{91.5\pm 2.3}$	$\mathbf{88.7\pm 1.9}$
TRAC (100 traj.)	$67.4\pm 3.6$	$61.2\pm 5.4$	$58.7\pm 4.2$
Fixed $H=1$	$74.5\pm 2.8$	$67.9\pm 4.7$	$65.1\pm 3.6$
Fixed $H=3$	$81.7\pm 2.5$	$74.3\pm 3.8$	$71.8\pm 3.1$
Fixed $H=5$	$84.6\pm 2.1$	$79.8\pm 3.1$	$75.2\pm 2.7$
Fixed $H=10$	$78.9\pm 8.7$	$72.1\pm 10.3$	$69.4\pm 9.1$

Appendix

Table 2 : Final normalized return (mean

\pm

std) under the limited-data setting.

Method	HalfCheetah	Hopper	Walker2d
Fixed $H=5$	0.67	0.61	0.55
Uncertainty truncation	0.74	0.68	0.64
IDI	0.80	0.80	0.77

Appendix

Table 3 : Fraction of lost performance recovered relative to full-data TRAC.

Explore similar work

On Training in Imagination
May 7, 2026Nadav Timor, Ravid Shwartz-Ziv, Micah Goldblum +2Model-Based Reinforcement Learning Rollout
Weizmann Institute of Science · New York University · Columbia University
IADD-TR: Intervention-Aware Dynamics Decoupling with Targeted Regularization for Model-Based Reinforcement Learning
Aug 11, 2026Zefeng Liang, Jie Qiao, Ruichu Cai +2Model-Based Reinforcement Learning Regularization
School of Computer Science, Guangdong University of Technology, Guangzhou, China · College of Science, Shantou University, Shantou, China
Counterfactual Transport Flows for Offline Conservative Trajectory Refinement
Jun 8, 2026Lena Krieger, Xuan Zhao, Zhuo Cao +3Model-Based Reinforcement Learning Long Refinement Trajectories
1IAS-8, Forschungszentrum J ulich, Germany · 2LMU Munich, Munich Center for Machine Learning (MCML), Germany · Department of Computer Science, Aarhus University, Denmark.

May 7, 2026, cs.LGJ/K move · Enter open · S save

On Training in Imagination

Nadav Timor, Ravid Shwartz-Ziv, Micah Goldblum, Yann LeCun+1

Weizmann Institute of Science · New York University · Columbia University

Model-Based Reinforcement Learning Rollout Imagination

Open paper AskPDF arXiv

State-of-the-art model-based reinforcement learning methods train policies on imagined rollouts. These rollouts are trajectories generated by a learned dynamics model and are scored by a learned reward model, but without querying the true environment during policy updates. We study this training paradigm by quantifying how errors in learned dynamics and reward models affect returns and policy optimization. First, we extend the analysis of Asadi et al. (2018) to MDPs with learned reward models, and derive the optimal sample allocation--the ratio of dynamics samples to reward samples that minimizes a bound on return error under power-law scaling assumptions. We identify lower Lipschitz constants of the learned dynamics, reward, and policy as a representation desideratum that tightens this bound, and we connect this perspective to the temporal-straightening objective of Wang et al. (2026). Second, we examine how policy optimization with REINFORCE tolerates noisy rewards, which are often cheaper to obtain. We show that zero-mean reward noise leaves the gradient estimator unbiased and adds at most a variance term that decreases with the number of rollouts. This introduces a practical tradeoff: given a fixed budget, should one buy more rollouts with cheaper but noisier rewards, or fewer rollouts with more expensive but less noisy rewards? We reduce this choice to a one-dimensional optimization problem and characterize the optimum.

Abstract

Explore similar work

On Training in Imagination

IADD-TR: Intervention-Aware Dynamics Decoupling with Targeted Regularization for Model-Based Reinforcement Learning

Counterfactual Transport Flows for Offline Conservative Trajectory Refinement