LASER: Latent Space Adjoint Matching for Support-Constrained Entropy-Regularized Offline RL

AI summary

Overview

Research area: Offline reinforcement learning (RL), specifically support-constrained policy optimization with flow-based generative policies and entropy regularization.

Technical level: Advanced. The paper assumes familiarity with Markov decision processes, actor-critic value estimation, flow matching, pushforward/pullback mappings between action and latent spaces, and adjoint sensitivity methods.

Scope: The paper argues that entropy regularization is a necessary ingredient for stable policy extraction in latent-space offline RL and introduces LASER, an algorithm that performs entropy-regularized latent-space policy improvement with flow policies while keeping actions inside the dataset support.

What This Paper Is About

Offline RL must improve on a behavior policy using only a fixed dataset, but pushing the policy toward actions the dataset never saw leads to unreliable value estimates and out-of-distribution (OOD) actions. A recent strategy learns a behavior-cloning decoder with flow matching, then optimizes a policy inside the compact latent space that the decoder maps into the dataset's support. LASER's claim is that staying in support is only half the problem: without entropy regularization, the latent policy collapses onto brittle modes or exploits sharp artifacts in the learned critic, so the authors add an explicit entropy-regularized policy improvement objective computed via adjoint matching.

Key Contributions

1. The authors identify entropy regularization as a central ingredient for stable support-constrained policy extraction when the policy is a flow-matching model in offline RL.
2. They reinterpret recent support-constrained flow policy methods, including DSRL and ReFORM, through the lens of entropy regularization.
3. They introduce LASER (Latent Space Adjoint Matching for Support-Constrained Entropy-Regularized Offline RL), which performs entropy-regularized offline policy improvement with latent-space flow policies, preserves dataset-support constraints, and avoids backpropagation through time.
4. They report experiments across 40 challenging OGBench tasks with varying dataset qualities, showing LASER achieves state-of-the-art performance using one fixed set of method-specific hyperparameters across all tasks, outperforming evaluated baselines including those with task- and dataset-specific tuning.

Main Findings

• Support constraints alone are insufficient: The paper argues that once optimization is restricted to the dataset support, the remaining difficulty is policy extraction — selecting high-value actions without collapsing to brittle modes or exploiting sharp critic artifacts.
• Entropy regularization is essential: Without entropy regularization, the policy becomes nearly deterministic and clusters at a false maximum, causing agents to get stuck (illustrated with agents moving up and hitting a wall in Figure 1).
• Critic artifacts appear without entropy regularization: Figure 2 shows the learned Q-function exhibits strange artifacts and a poorly shaped landscape when entropy regularization is absent.
• Target policy smoothing is not enough: Adding noise to the policy in TD learning (target policy smoothing) fixes the Q-function landscape but leaves the policy nearly deterministic and suboptimal.
• Entropy regularization spreads policy samples: With entropy regularization, policy samples spread out more and produce a better-regularized Q-function, letting agents move around the wall and reach the target in the illustrative task.
• Strong reported benchmark performance: Across 40 OGBench tasks with varying dataset qualities, LASER is reported to achieve state-of-the-art performance, and to outperform evaluated baselines even though those baselines used task- and dataset-specific tuning while LASER used fixed method-specific hyperparameters.
• Specific numeric results: Beyond the count of 40 OGBench tasks and the qualitative failure/success descriptions in Figures 1 and 2, the available paper content does not report per-task scores or other numerical values.

Methodology in Plain English

The approach starts from the observation that regularizing a statistical distance between the learned and behavior policies is a blunt instrument: some distances (such as KL divergence) make the policy too conservative, while others (such as Wasserstein or MMD distance) still permit OOD actions.

Instead, LASER follows the latent-space recipe. A deterministic decoder is trained with flow matching to imitate the behavior policy, mapping samples from a simple latent prior (the paper mentions a uniform distribution on a d-ball as an example) into the dataset's action distribution. If the latent policy's support is contained in the prior's support, then any action it produces through the decoder lies within the behavior policy's support by construction. This converts the hard action-space support constraint into a simple latent-space constraint that can be enforced with techniques like tanh squashing, clipping, or projection. Once inside that latent space, any off-policy optimizer could in principle be used without worrying about OOD actions.

LASER's distinguishing step is replacing implicit or architecture-dependent regularization with an explicit entropy-regularized policy improvement objective. The technical obstacle is that computing and optimizing the entropy of a flow policy in latent space is nontrivial. The authors use adjoint matching, which enables entropy-regularized optimization of expressive flow policies in latent space while preserving the support constraints induced by the learned latent space, and which avoids backpropagation through time.

Why This Matters

Impact on research: The paper reframes the offline RL discussion away from "which divergence to regularize" and toward "how to extract high-value actions safely from the dataset support." By explicitly arguing that entropy regularization is necessary in latent-space RL and by reinterpreting existing methods such as DSRL and ReFORM through that lens, it gives the community a unifying diagnostic for why some support-constrained methods are brittle. It also connects flow-matching generative modeling with adjoint-based entropy computation, a combination that may be reusable beyond offline RL.

Real-world applications (bullets):

• Robotics and autonomous driving, where policies are trained from logged demonstrations and exploration during deployment is costly or unsafe.
• Industrial control and process optimization, where historical sensor/actuator logs are abundant but online experimentation is disruptive.
• Healthcare treatment policy learning from retrospective patient records, where out-of-distribution actions are unacceptable.
• Recommendation and sequential decision systems trained on historical interaction logs, where deploying untested actions risks poor user outcomes.

Industry relevance: Offline RL is attractive precisely because it avoids costly or risky online interaction. LASER's reported property of working with a single fixed set of method-specific hyperparameters across 40 tasks is directly relevant to practitioners, since per-task hyperparameter tuning is often the practical blocker for deploying offline RL in production pipelines.

Future Directions

• Scaling and domain transfer: The reported evaluation is confined to OGBench. Whether the fixed-hyperparameter result holds on other benchmarks, on real robot or industrial datasets, or on higher-dimensional action spaces is not established in the provided content.
• Entropy estimation accuracy: Since the method relies on adjoint matching to optimize the entropy of a flow policy in latent space, the accuracy and cost of that estimation in large latent dimensions remain open questions.
• Relation to other regularizers: The paper reinterprets DSRL and ReFORM as implicit entropy-regularized methods; a natural follow-up is a systematic comparison of explicit versus implicit entropy regularization under matched compute and tuning budgets.
• Critic regularization mechanism: The paper claims entropy regularization smooths the policy optimization landscape by regularizing the critic. Isolating how much of the gain comes from this critic-smoothing effect versus from policy exploration would clarify the underlying mechanism.
• Interaction with support-constraint enforcement: The paper mentions tanh squashing, clipping, and projection as ways to enforce the latent support constraint; how each choice interacts with the entropy objective is not resolved in the available content.

Target Audience

This paper is most useful to offline RL researchers and graduate students working on policy extraction, support constraints, or flow/diffusion-based policies. It also suits practitioners building decision-making systems from fixed logged datasets who need reliable behavior without online exploration. Readers need a solid background in reinforcement learning and generative modeling to follow the mathematical development; beginners would find the paper difficult without prior exposure to flow matching and entropy-regularized RL.

Authors’ abstract

While offline reinforcement learning (RL) enables policy optimization from static datasets without costly online interaction, it remains bottlenecked by the risk of executing out-of-distribution (OOD) actions. Recent approaches mitigate this by learning a behavior-cloning policy through flow matching and then performing RL within its constrained latent space. However, naively optimizing the latent policy can easily cause the policy to collapse into a brittle mode or exploit sharp artifacts of the learned critic. In this work, we find that entropy regularization is essential in latent-space RL for addressing these challenges. We introduce LASER, a novel offline RL algorithm that applies latent-space adjoint matching to achieve entropy-regularized latent-space RL with expressive flow policies while avoiding backpropagation through time. Through comprehensive experiments on 40 challenging OGBench tasks with varying dataset qualities, we show that LASER achieves state-of-the-art performance. Notably, LASER uses fixed method-specific hyperparameters across all tasks and outperforms the evaluated baselines, including those with task- and dataset-specific tuning, which highlights the robust applicability of LASER. Project website: https://mit-realm.github.io/laser/.