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

2026-02-02

AI Summary

This work addresses the challenges of value overestimation and policy degradation in offline reinforcement learning caused by distributional shift, particularly in data-sparse regions. To mitigate these issues, the authors propose the Manifold-Constrained Energy-based Transition Model (MC-ETM), which trains a conditional energy-based model in latent space by integrating manifold projection with diffusion-based negative sampling. The method sharpens the energy landscape by generating near-manifold hard negative samples via Langevin dynamics. Furthermore, an energy-guided truncation mechanism combined with pessimistic Bellman backups is introduced to establish a hybrid pessimistic MDP framework. Experimental results demonstrate that MC-ETM significantly improves multi-step dynamics fidelity and normalized returns on standard offline control benchmarks, outperforming existing approaches—especially in scenarios involving irregular dynamics and sparse data.

Technology Category

• Machine Learning: Learning with Manifolds
• Multiagent Systems: Multiagent Learning
• Search and Optimization: Learning to Search

Application Category

• Economics, Online Markets and Human Computation: The sharing economy
• User Modeling, Personalization and Recommendation: Accountability, Transparency, and Ethics for personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• offline reinforcement learning
• distribution shift
• model error
• value overestimation
• manifold constraint

Innovation

• energy-based model
• manifold learning
• offline reinforcement learning
• distributional robustness
• pessimistic policy optimization