Research map / Reinforcement learning / Offline RL
Offline Transition Modeling via Contrastive Energy Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Energy-based Transition Models (ETM), which learn environment dynamics as a scalar energy function trained with InfoNCE and Langevin negatives, and applies them to off-policy evaluation and model-based offline RL (EMPO).
Problem
Forward (Gaussian) transition models have an inductive bias toward smooth regressors, so they struggle to fit discontinuous or high-curvature transition dynamics and generalize poorly out of distribution, hurting model-based offline evaluation and control.
Approach
Model the next-state distribution implicitly through a scalar energy function trained contrastively with InfoNCE (shown to be a robust variant of maximum likelihood), with negatives drawn by Langevin MCMC and a gradient penalty; sampling is done by iterative gradient-based search.
Key results
- On the 32-task DOPE benchmark, ETM attains the lowest mean absolute error and best rank correlation and regret among OPE methods, outperforming FTM and FQE/DR/IS/DICE/VPM.
- EMPO (ETM with SAC and ensemble-std penalty) achieves average normalized score 82.0 over 12 D4RL tasks, vs. 80.0 for MOBILE (previous best).
- EMPO solves 7 of 12 D4RL tasks (score >= 95), while the best prior method solves at most 5.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use energy-based transition models instead of forward transition models | Transition modeling on OOD D4RL hopper and walker2d transitions · D4RL hopper and walker2d | mean absolute error | Helped |
| Use energy-based transition models instead of forward transition models | Off-policy evaluation · DOPE benchmark, 32 tasks | mean absolute error | Helped |
| Use energy-based transition models instead of forward transition models | Off-policy evaluation · DOPE benchmark, 32 tasks | rank correlation | Helped |
| Use energy-based transition models instead of forward transition models | Offline policy optimization with SAC · D4RL Gym-Mujoco tasks | normalized score: 80.0 -> 82.0 | Helped |
| Use energy-based transition models instead of forward transition models | Offline policy optimization with SAC · D4RL Gym-Mujoco tasks | solved tasks: 5/12 -> 7/12 | Helped |
| Use energy-based transition models instead of forward transition models | In-distribution transition prediction · D4RL tasks | absolute error: 16/20 | Hurt |
| Increase the number of negative samples from 16 to 8 or 24 (ablation) | ETM · DOPE evaluation on four tasks | absolute error, rank correlation, and regret | No clear effect |
| Use ETM instead of FTM for inference | Transition model inference · Not benchmark-specific | inference time: approximately five times | Hurt |
Surprising findings
- Forward transition models have smaller in-distribution errors on 16 of 20 MuJoCo tasks, even though ETMs generalize better out of distribution.
Limitations
- Iterative gradient-based sampling at inference makes ETMs about five times slower than forward models.
- EMPO's overall learning time is about 40% higher than MOPO.
- Forward models are more accurate in-distribution on many MuJoCo tasks.
Experimental setup
Datasets: D4RL hopper D4RL walker2d D4RL halfcheetah D4RL ant D4RL Adroit DOPE benchmark
Benchmarks: DOPE D4RL Gym-Mujoco
Metrics: mean absolute error rank correlation regret@1 normalized average return solved tasks (score >= 95)
Baselines: CQL TD3+BC EDAC MOPO COMBO Trajectory Transformer RAMBO MOBILE FQE DR IS DICE
Builds on: MOPO Soft Actor-Critic InfoNCE DOPE benchmark Implicit behavioral cloning (energy-based implicit models)
Closest papers by meaning
- Analytic Energy-Guided Policy Optimization for Offline Reinforcement Learning (NeurIPS 2025 · similarity 0.90)
- Reward-Consistent Dynamics Models are Strongly Generalizable for Offline Reinforcement Learning (ICLR 2024 · similarity 0.89)
- Temporal Distance-aware Transition Augmentation for Offline Model-based Reinforcement Learning (ICML 2025 · similarity 0.88)
- ContraDiff: Planning Towards High Return States via Contrastive Learning (ICLR 2025 · similarity 0.88)
- Energy-Weighted Flow Matching for Offline Reinforcement Learning (ICLR 2025 · similarity 0.88)
- Critic-Guided Decision Transformer for Offline Reinforcement Learning (AAAI 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.