PaperlayerView exact revision on arXiv

Source-linked AI summary

In-Distribution Imagination for Model-Based Offline Reinforcement Learning

Mintae Kim, Koushil Sreenath

arXiv:2609.38673v1cs.LG

TL;DR

Model-based offline RL can generate unrealistic trajectories when model error accumulates, while many existing methods control rollouts using transition-level uncertainty. The paper proposes IDI, which estimates trajectory support and truncates unsupported rollouts, and reports better rollout-failure prediction and improved performance in limited-data settings.

  • Problem

    In MBORL, accumulated model error can drive imagined trajectories outside the offline data distribution, while existing rollout control primarily relies on transition-level uncertainty.

    Source: §1 · p. 1

  • Method

    IDI estimates support for learned latent trajectory representations and truncates unsupported rollouts, while TRAC regularizes the policy-induced trajectory distribution toward offline trajectories.

    Source: §1 · p. 1

  • Results

    Trajectory support predicts rollout failure better than transition-level uncertainty, and IDI consistently improves offline RL performance in limited-data settings.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results suggest reliable imagination requires trajectory-level rollout control rather than relying only on transition-level uncertainty.

    Source: §1 · pp. 1–2

Abstract

from arXiv · show

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.

1 Introduction

The introduction frames rollout reliability as a trajectory-level property: locally plausible transitions can accumulate into unrealistic imagined trajectories. It proposes IDI to estimate support in latent trajectory space and truncate unsupported rollouts, and reports improved performance in limited-data settings.

1 Introduction

  • Locally plausible transitions can accumulate into unrealistic trajectories, so transition-level uncertainty alone may not determine rollout reliability.The overview contrasts conventional local reliability estimates with IDI's trajectory-level support estimation.Source: §1 · p. 1, Figure 1 · p. 2
  • IDI estimates trajectory support in a learned latent space and truncates imagination when trajectories leave the offline manifold.It is combined with TRAC, which regularizes the policy-induced trajectory distribution toward the offline trajectory distribution.Source: §1 · pp. 1–2
  • The paper identifies rollout reliability as a trajectory-level property and presents latent-space support as the basis for IDI's rollout control.The introduction summarizes these as central contributions.Source: §1 · p. 2
  • Trajectory support predicts rollout failure better than transition-level uncertainty, and IDI improves overall MBORL performance.The reported experiments focus on limited-data settings where model-generated trajectories are useful.Source: §1 · p. 2

2 Preliminaries

The preliminaries formulate offline RL as learning from a static dataset of trajectories and define the discounted-return objective. They describe MBORL as learning dynamics from that dataset, generating imagined trajectories, and using imagined transitions in policy optimization.

2 Preliminaries

  • Offline RL learns from a static dataset of trajectories, with the objective of maximizing discounted return.The setup specifies an infinite-horizon discounted MDP and a policy over actions given states.Source: §2 · p. 2
  • MBORL learns a transition dynamics model from offline data and uses it to generate imagined trajectories for policy optimization.Imagined transitions are combined with offline data for policy optimization.Source: §2 · p. 2
  • The preliminaries describe MBORL dynamics models as an ensemble whose members predict Gaussian next-state distributions and provide uncertainty estimates.The passage specifies state predictions and uncertainty estimates from the ensemble.Source: §2 · p. 2

3 In-Distribution Imagination for Trajectory-Regularized Offline RL

IDI evaluates imagined transitions in the context of recent rollout history, using learned trajectory-window representations to reject transitions that lack offline support. Combined with TRAC, it filters synthetic data while the policy is trained with a behavior prior.

3 In-Distribution Imagination for Trajectory-Regularized Offline RL

  • IDI scores each candidate transition using a post-transition window of recent transitions, then terminates rollouts when the window is unsupported by offline trajectories.The window encoder is trained on offline trajectory windows and remains fixed during rollout generation.Source: §3 · pp. 2–3
  • Contrastive learning makes the window representation capture temporal compatibility, using nearby windows from the same trajectory as positives and unrelated trajectories as negatives.The objective uses cosine similarity and a temperature parameter.Source: §3 · p. 3
  • The support threshold is based on offline windows' leave-one-out scores, with the query embedding excluded to avoid artificially small support scores.The support score uses distances to nearest offline embeddings; larger values indicate incompatibility with observed trajectory segments.Source: §3 · p. 3
  • IDI initializes each imagined rollout with preceding offline transitions, progressively replaces that history with imagined transitions, and rejects an unsupported candidate before it enters the synthetic dataset.Accepted transitions are added to the synthetic dataset and imagination continues; rejected candidates terminate the rollout.Source: §3 · p. 3
  • TRAC supplies a behavior-policy action prior, while IDI-filtered transitions augment the policy-training data; IDI retains approximately 55–70% of generated transitions.The behavior prior remains trained solely on the original offline dataset, and IDI truncation is not treated as an environment terminal condition.Source: §3 · pp. 3–4

4 Experiments

In limited-data MuJoCo experiments, IDI outperforms fixed-horizon and uncertainty-based rollout control. Trajectory support tracks accumulated rollout error more closely than ensemble disagreement, while accepted rollout horizons vary across environments.

Experimental setup

  • The experiments use 100-episode subsamples of 1000-episode MuJoCo datasets and compare full-data TRAC, limited-data TRAC, fixed horizons, uncertainty truncation, and IDI.All methods use the TRAC backbone and learned dynamics model.Source: §4 · p. 4

Limited-data performance

  • HalfCheetah normalized return rises from approximately 67 to 88 with IDI, approaching approximately 93 for full-data TRAC.IDI also achieves the strongest performance across the limited-data benchmark comparisons.Source: §4 · pp. 4–5
  • Average accepted rollout horizons are 6.2 steps on HalfCheetah, 4.8 on Hopper, and 3.7 on Walker2d, indicating that no single rollout horizon is optimal across environments.IDI adapts rollout length to rollout reliability, unlike fixed-horizon rollouts.Source: §4 · p. 5

Rollout failure detection

  • Rollout error increases from approximately 0.03 to 0.14 over ten HalfCheetah rollout steps, while trajectory support rises nearly monotonically from approximately 0.66 to 0.76 and ensemble disagreement changes modestly.The authors interpret this pattern as evidence that accumulated rollout drift is not fully reflected by local uncertainty estimates.Source: §4 · p. 5

5 Conclusion

The conclusion presents IDI as history-aware rollout control that filters unsupported imagined transitions, and pairs it with TRAC's behavior-supported policy regularization. Limited-data MuJoCo experiments show improved performance and a stronger relationship between support and accumulated rollout error than between error and instantaneous ensemble disagreement.

5 Conclusion

  • IDI evaluates candidate transitions using learned latent representations of recent trajectory segments and terminates imagination before unsupported transitions enter policy training.The conclusion frames reliable imagination as sequential data selection based on compatibility with the trajectory prefix.Source: §5 · p. 5
  • IDI and TRAC address complementary distribution-shift sources: IDI controls the reliability of model-generated data, while TRAC regularizes policy improvement toward behavior-supported actions.Source: §5 · p. 5
  • In limited-data MuJoCo settings, IDI consistently improves performance over fixed-horizon and uncertainty-based rollout control.Source: §5 · p. 5
  • The learned support score relates more strongly to accumulated rollout error than instantaneous ensemble disagreement, suggesting recent rollout history contains information local uncertainty misses.Source: §5 · p. 5

A Trajectory-Regularized Actor-Critic (TRAC)

TRAC formulates offline RL as finite-horizon KL control, regularizing policy-induced trajectory distributions toward behavior trajectories. Replacing entropy regularization’s uniform action reference with a data-dependent behavior prior lets policy improvement favor behavior-supported actions.

A Trajectory-Regularized Actor-Critic (TRAC)

  • TRAC maximizes undiscounted finite-horizon return while regularizing the policy trajectory distribution toward the behavior trajectory distribution.The trajectory KL reduces to action-distribution regularization because the policy and behavior distributions share initial-state and transition dynamics.Source: A Trajectory-Regularized Actor-Critic… · p. 6
  • The regularization strength λ is positive, and policy action distributions are assumed absolutely continuous with respect to the behavior policy on reachable states.This assumption ensures the trajectory KL is finite.Source: A Trajectory-Regularized Actor-Critic… · p. 6
  • The Gibbs variational principle yields a behavior-prior Boltzmann policy that reweights behavior-supported actions by value and excludes actions outside behavior support.The practical actor update projects onto this policy using reverse KL.Source: A Trajectory-Regularized Actor-Critic… · pp. 6–7
  • The behavior prior is learned from the original offline dataset and remains fixed while the actor and critics train on offline plus accepted imagined data.IDI controls imagined-data reliability; TRAC separately regularizes optimization toward behavior-supported actions.Source: A Trajectory-Regularized Actor-Critic… · p. 7
  • The critic uses an undiscounted finite-horizon target with terminal rewards, and double target critics mitigate value overestimation.The target value uses the learned behavior prior; target critics are updated with Polyak averaging.Source: A Trajectory-Regularized Actor-Critic… · p. 7
  • With a uniform action prior, TRAC recovers finite-horizon maximum-entropy RL up to a policy-independent additive constant.TRAC replaces entropy regularization’s uniform reference with a data-dependent behavior prior.Source: A Trajectory-Regularized Actor-Critic… · p. 7

B Implementation Details

IDI encodes recent trajectory windows, scores their support against offline windows, and accepts candidate transitions only while they remain supported. The implementation uses shared offline data for its models and representation, then trains the policy on accepted imagined data alongside the offline dataset.

B Implementation Details

  • The model, behavior prior, and encoder use the same 100-trajectory subsample, while actor and critics train on offline plus accepted imagined data.The behavior prior is learned only from the original dataset; actor and critic parameters are shared across finite-horizon stages.Source: B Implementation Details · pp. 7–8
  • The encoder represents windows of eight transitions as nine states and eight actions, using a two-layer MLP with 64-dimensional latent output.It is trained on offline windows with InfoNCE and frozen for support scoring and rollouts.Source: B Implementation Details · p. 8
  • IDI uses 10 nearest neighbors and a support threshold at the 95th percentile of offline leave-one-out scores.The query embedding is excluded from the reference set when scoring an offline window.Source: B Implementation Details · p. 8
  • Rollouts begin with preceding offline transitions as history, which accepted imagined transitions progressively replace.The initial imagined window contains seven offline transitions and one candidate transition.Source: B Implementation Details · p. 8
  • A candidate transition is added to the imagined dataset only when its post-transition window’s support score is at or below the threshold; otherwise it is discarded and rollout stops.The candidate is included in the scored window before acceptance.Source: B Implementation Details · p. 8
  • Rollouts are capped at 10 steps, and IDI accepts about 55–70% of generated transitions across environments.Average accepted horizons are 6.2, 4.8, and 3.7 steps on HalfCheetah, Hopper, and Walker2d, respectively; truncation is not an environment terminal condition.Source: B Implementation Details · p. 8
  • The implementation averages results over five seeds and uses Adam with a 3 × 10−4 learning rate and batch size 256.The TRAC regularization coefficient is learned automatically, and bootstrapping stops only at the true finite-horizon terminal state.Source: B Implementation Details · p. 8
  • The accompanying tables report final limited-data normalized returns and the fraction of lost performance recovered relative to full-data TRAC.Source: Table 2 · p. 9, Table 3 · p. 9

C Additional Performance Results

IDI performs best among the rollout-control methods across the tested environments, while fixed-horizon performance is sensitive to rollout length and uncertainty-based truncation is less effective. Recovery results support rollout quality, rather than quantity, as the primary determinant of synthetic-data benefit.

C Additional Performance Results

  • IDI achieves the strongest performance among rollout-control methods across environments, while fixed horizons are sensitive and H = 10 often destabilizes training through compounding model error.Uncertainty-based truncation reduces some rollout failures but is less effective than trajectory-level support estimation.Source: C Additional Performance Results · p. 9
  • Recovery ratios quantify the fraction of lost performance recovered relative to full-data TRAC; results suggest synthetic-data benefit depends primarily on rollout quality rather than quantity.IDI recovers a larger fraction by preventing unreliable long-horizon imagination.Source: C Additional Performance Results · p. 9
Loading 2609.38673v1…