Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Manifold-Constrained Energy-Based Transition Models for Offline Reinforcement Learning

Zeyu Fang, Zuyuan Zhang, Mahdi Imani, Tian Lan

TL;DR

The paper tackles offline model-based reinforcement learning under distribution shift, where policy-driven rollouts can venture into underrepresented regions and cause value overestimation. It introduces Manifold-Constrained Energy-Based Transition Models (MC-ETM) that learn a conditional energy function $E_\theta(s,a,s')$ using a manifold projection-diffusion (MPD) approach to generate near-manifold hard negatives via latent perturbations and Langevin dynamics, sharpening the energy landscape around data support. Energy is used as a unified reliability signal, truncating rollouts when $\min_{s'} E_\theta(s,a,s') > \delta$ and applying pessimistic penalties based on $Q$-value dispersion across energy-guided samples within a hybrid pessimistic MDP framework, with a bound that separates in-support evaluation error from truncation risk. Empirically, MC-ETM improves multi-step dynamics fidelity and yields higher normalized returns on D4RL MuJoCo benchmarks, particularly in irregular dynamics and sparse data settings, while ablations validate the necessity of MPD, truncation, and energy-based penalties for stability and performance.

Abstract

Model-based offline reinforcement learning is brittle under distribution shift: policy improvement drives rollouts into state--action regions weakly supported by the dataset, where compounding model error yields severe value overestimation. We propose Manifold-Constrained Energy-based Transition Models (MC-ETM), which train conditional energy-based transition models using a manifold projection--diffusion negative sampler. MC-ETM learns a latent manifold of next states and generates near-manifold hard negatives by perturbing latent codes and running Langevin dynamics in latent space with the learned conditional energy, sharpening the energy landscape around the dataset support and improving sensitivity to subtle out-of-distribution deviations. For policy optimization, the learned energy provides a single reliability signal: rollouts are truncated when the minimum energy over sampled next states exceeds a threshold, and Bellman backups are stabilized via pessimistic penalties based on Q-value-level dispersion across energy-guided samples. We formalize MC-ETM through a hybrid pessimistic MDP formulation and derive a conservative performance bound separating in-support evaluation error from truncation risk. Empirically, MC-ETM improves multi-step dynamics fidelity and yields higher normalized returns on standard offline control benchmarks, particularly under irregular dynamics and sparse data coverage.

Manifold-Constrained Energy-Based Transition Models for Offline Reinforcement Learning

TL;DR

The paper tackles offline model-based reinforcement learning under distribution shift, where policy-driven rollouts can venture into underrepresented regions and cause value overestimation. It introduces Manifold-Constrained Energy-Based Transition Models (MC-ETM) that learn a conditional energy function using a manifold projection-diffusion (MPD) approach to generate near-manifold hard negatives via latent perturbations and Langevin dynamics, sharpening the energy landscape around data support. Energy is used as a unified reliability signal, truncating rollouts when and applying pessimistic penalties based on -value dispersion across energy-guided samples within a hybrid pessimistic MDP framework, with a bound that separates in-support evaluation error from truncation risk. Empirically, MC-ETM improves multi-step dynamics fidelity and yields higher normalized returns on D4RL MuJoCo benchmarks, particularly in irregular dynamics and sparse data settings, while ablations validate the necessity of MPD, truncation, and energy-based penalties for stability and performance.

Abstract

Model-based offline reinforcement learning is brittle under distribution shift: policy improvement drives rollouts into state--action regions weakly supported by the dataset, where compounding model error yields severe value overestimation. We propose Manifold-Constrained Energy-based Transition Models (MC-ETM), which train conditional energy-based transition models using a manifold projection--diffusion negative sampler. MC-ETM learns a latent manifold of next states and generates near-manifold hard negatives by perturbing latent codes and running Langevin dynamics in latent space with the learned conditional energy, sharpening the energy landscape around the dataset support and improving sensitivity to subtle out-of-distribution deviations. For policy optimization, the learned energy provides a single reliability signal: rollouts are truncated when the minimum energy over sampled next states exceeds a threshold, and Bellman backups are stabilized via pessimistic penalties based on Q-value-level dispersion across energy-guided samples. We formalize MC-ETM through a hybrid pessimistic MDP formulation and derive a conservative performance bound separating in-support evaluation error from truncation risk. Empirically, MC-ETM improves multi-step dynamics fidelity and yields higher normalized returns on standard offline control benchmarks, particularly under irregular dynamics and sparse data coverage.
Paper Structure (44 sections, 4 theorems, 39 equations, 4 figures, 6 tables, 2 algorithms)

This paper contains 44 sections, 4 theorems, 39 equations, 4 figures, 6 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Model-free offline RL.
  4. Model-based offline RL.
  5. Energy-based transition modeling.
  6. Our approach.
  7. Preliminaries
  8. General RL
  9. Offline MBRL
  10. ETMs
  11. Methodology
  12. Manifold-Constrained Energy-Based Modeling
  13. Resolving Uncertainties with Energy-based Models
  14. Theoretical Analysis
  15. Experiments
  16. ...and 29 more sections

Key Result

Theorem 4.3

Let $J(\pi)$ denote the expected return of policy $\pi$ in the true MDP. Under Assumption ass:consistency, the performance gap between the optimal policy $\pi^*$ and the policy $\hat{\pi}$ learned by MC-ETM is bounded by:

Figures (4)

  • Figure 1: An illustrative example on fitting a discontinuous transition function.
  • Figure 2: Conceptual visualization of energy landscapes on a 2D slice of the high-dimensional state space
  • Figure 3: Correlation between MC-ETM energy values and prediction MSE. (Left) Scatter plot showing a positive correlation; higher energy implies higher model error. (Right) Histogram of energy values for In-Distribution (ID) vs. Out-of-Distribution (OOD) samples, showing a clear separation that enables effective truncation. The threshold $\delta$ is set to 0 in this case.
  • Figure 4: Visualized rollout trajectories with ground truth, MC-ETM, standard ETM and MLP in different Mujoco tasks.

Theorems & Definitions (8)

  • Definition 4.1
  • Theorem 4.3: Energy-Constrained Performance Bound
  • Lemma 3.1: Pessimism on the manifold
  • proof
  • Lemma 3.2: Truncation loss bound via hitting time
  • proof
  • Theorem 3.3: Energy-Constrained Performance Bound (restated)
  • proof