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

TL;DR

This paper introduces MC-ETM, a novel energy-based transition model with manifold constraints that enhances offline reinforcement learning by improving model reliability and performance under distribution shift.

Contribution

MC-ETM employs a manifold projection and diffusion-based negative sampling to sharpen the energy landscape, leading to more accurate and robust policy learning in offline RL.

Findings

Improves multi-step dynamics fidelity.

Achieves higher normalized returns on benchmarks.

Performs well under irregular dynamics and sparse data.

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.