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

1 Introduction

• Offline RL must improve on the behavior policy using static data while choosing actions with reliable dataset supervision, since unsupported actions can receive spuriously high predicted values and perform poorly.Such actions may be exploited by policy optimization and yield out-of-distribution (OOD) behavior.Source: §1 · p. 1
• Support constraints do not by themselves ensure reliable policy extraction: optimization within support can still collapse to brittle modes or exploit sharp, locally overestimated critics.The paper identifies entropy regularization as a way to discourage premature collapse and smooth the optimization landscape by regularizing the critic.Source: §1 · p. 2
• LASER uses adjoint matching to optimize expressive latent-space flow policies with an explicit entropy-regularized objective while preserving learned latent-space support constraints.This addresses the difficulty of optimizing a flow policy's entropy without backpropagation through time.Source: §1 · p. 2
• The paper identifies entropy regularization as central to stable support-constrained flow-policy extraction and reframes methods including DSRL and ReFORM through that perspective.Source: §1 · p. 2
• LASER achieves the strongest performance profiles among evaluated flow-policy baselines while using fixed method-specific hyperparameters.The experiments cover 40 challenging OGBench tasks with varying dataset qualities, as stated in the paper context.Source: §1 · p. 2

2.1 Offline RL

• Because offline RL uses a fixed dataset, unfamiliar actions can receive unreliable, potentially large Q-values, leading policy optimization toward out-of-distribution actions.The learned policy is therefore constrained to the behavior policy's support.Source: §2.1 · pp. 2–3
• Statistical-distance regularization may still permit out-of-distribution actions or restrict policy improvement, and its weight requires task- and dataset-specific tuning.Source: §2.1 · p. 3

2.2 Mitigating OOD Actions with Latent Space RL

• A latent-space policy is decoded into actions by a behavior-cloning flow map, with latent samples supplying policy stochasticity.The decoder's pushforward of the latent prior approximates the behavior policy.Source: §2.2 · pp. 3–4
• Restricting the latent policy to the prior's simple support makes the action-support constraint easier to enforce and decouples OOD avoidance from policy optimization.The decoder maps the latent policy into action space, where the induced policy remains constrained by the learned behavior model's support.Source: §2.2 · p. 3
• Flow matching learns expressive decoders that can model complex, multimodal behavior distributions.Source: §2.2 · p. 3

3.1 Entropy Regularization is Essential for Latent Space RL

• Explicit entropy regularization lets the policy reach the target in the toy environment after low-entropy optimization gets stuck at a wall.The learned critic had artifacts; target policy smoothing corrected its landscape but yielded a suboptimal policy, whereas entropy regularization improved both policy and Q-function.Source: §3.1 · pp. 4–5
• Flow policies with adjoint matching Pareto-dominate squashed Gaussian latent policies across entropy and Q-value operating points.For every entropy and Q-value achieved by the Gaussian policy, the flow policy has an operating point with both higher entropy and higher Q-value.Source: Figure 3 · p. 5, §3.1 · p. 5

3.2 Efficient Entropy-Regularized Policy Optimization via Adjoint Matching

• Backpropagating through the flow map and entropy requires differentiating through ODE solutions, which is computationally expensive and numerically unstable.LASER instead derives the optimal latent policy as a Q-tilted distribution under finite-support and bounded-critic assumptions; this is a fixed-critic, statewise optimum, not a convergence guarantee for joint actor-critic updates.Source: §3.2 · p. 6
• Adjoint matching trains the latent flow by regressing its velocity field toward an estimated optimal control, avoiding backpropagation through the flow.It requires solving an ODE to compute the adjoint state used for the regression target.Source: §3.2 · pp. 6–7

3.3 Algorithm

• LASER jointly trains a flow-matching behavior decoder, a flow for the uniform latent distribution, a critic, and an entropy-regularized latent policy optimized with adjoint matching.The decoder maps uniform samples from a d-ball to dataset actions, while latent policy optimization maximizes the latent-space Q-function.Source: §3.3 · p. 7

• Offline RL methods address distribution shift by regularizing policies, learning pessimistic value functions, or restricting improvement to supported actions.LASER follows the support-constrained approach and uses an expressive flow policy to enforce support by construction.Source: §4 · p. 8
• Generative policies use diffusion or flow matching to model potentially multimodal actions, while critic-guided methods select, reweight, clone, or directly improve generated actions.Unlike approaches relying on statistical regularization, LASER enforces support through its latent action parameterization.Source: §4 · p. 8
• LASER revisits policy optimization after latent-space methods have constrained actions to the behavior-policy support, emphasizing entropy to prevent collapse and stabilize extraction.The paper distinguishes this standalone role of entropy from implicit regularization arising from behavior-policy constraints.Source: §4 · p. 8

5.1 Setup

• The main suite covers four environments, five tasks per environment, and CLEAN and NOISY datasets, while LASER, ReFORM, and IFQL use fixed method-specific hyperparameters across tasks.The 40 settings span locomotion and manipulation; other baselines' hyperparameters vary by environment and dataset.Source: §5.1 · pp. 8–9

5.2 Results

• LASER outperforms all evaluated baselines in aggregate across 40 settings, reaching at least 95% of best mean success in 34/40 and mean success 100% in 13/40 (global p = 1.35 × 10^-3).The performance profiles aggregate settings across CLEAN and NOISY datasets.Source: Figure 5 · p. 9, §5.2 · p. 9
• Inverse temperatures 1/α ∈[5, 20] yield comparable results; 1/α = 3 drops performance on both datasets, while 1/α = 30 degrades CLEAN performance but remains competitive on NOISY data.The authors report stable results across environments at 1/α = 10.Source: §5.2 · p. 9, Figure 6 · p. 9
• Adjoint matching yields more stable training and stronger performance than BPTT, while training 3.6× faster (123 it/s vs 34 it/s).Adjoint matching avoids BPTT through the flow policy's optimization process.Source: §5.2 · pp. 9–10, Figure 6 · p. 9
• Changing the decoder prior changes behavior regularization while retaining the same latent-space entropy regularizer, and yields comparable final success rates in the tested task.The runs retain a uniform entropy reference on the same latent ball.Source: Figure 6 · p. 9, §5.2 · p. 10

• Support constraints prevent OOD actions, while entropy regularization stabilizes policy extraction within the learned dataset support.The authors say the combination helps prevent mode collapse, smooths the optimization landscape, and avoids local minima.Source: §6 · p. 10
• LASER achieves state-of-the-art performance on the OGBench suite using fixed method-specific hyperparameters.Source: §6 · p. 10
• Sampling from the flow model remains computationally intensive during training for deep neural ODEs, despite adjoint matching avoiding BPTT.The authors suggest distillation as one possible way to reduce this cost.Source: §6 · p. 10
• Support constraints do not guarantee safe behavior because datasets may contain unsafe actions and learned decoders can make errors; physical deployment under distribution shift may cause harm.The authors say safety validation and monitoring remain necessary before physical deployment.Source: §6 · p. 10

• Offline RL methods address distribution shift with policy-distance or value-function regularization, or by penalizing actions outside behavior-policy support, but sample-based proxies and Gaussian approximations may miss multimodal behavior.Source: A Extended Related Work · p. 15
• Flow and diffusion policies offer expressive alternatives for multimodal actions, but iterative sampling makes critic-based optimization more challenging than optimization of standard Gaussian actors.Source: A Extended Related Work · p. 15
• Generative-policy critic-evaluation methods select or reweight sampled actions or clone high-value data, whereas critic-gradient methods can require costly, unstable backpropagation through sampling or sacrifice expressiveness through distillation.Source: A Extended Related Work · p. 15
• LASER uses adjoint matching for critic-gradient optimization but relies on latent support constraints rather than KL regularization toward behavior, which can be too restrictive when optimal actions have low behavior-policy probability.Source: A Extended Related Work · p. 15
• Because latent support control handles support constraints, LASER makes entropy regularization explicit in policy extraction; the method uses an expressive flow policy and reports a better performance-entropy tradeoff than Gaussian policies.Source: A Extended Related Work · p. 16

• The proof assumes supp(qL) has finite positive Lebesgue measure and that πL has a density satisfying the latent-space support constraints.Source: B Proof of Lemma 3.2 · p. 16
• The proof takes the uniform distribution over supp(qL) as its reference and notes that a term in the objective is constant with respect to πL.Source: B Proof of Lemma 3.2 · p. 16
• Writing the objective as KL divergence to πL⋆ plus a constant shows that πL⋆ is the unique optimizer, since the divergence is nonnegative and vanishes only at equality.Source: B Proof of Lemma 3.2 · p. 17

C.1 The Static Flow Issue and Implicit Entropy Regularization in ReFORM

• When a proposed Euler step leaves the ball and velocity aligns with the outward normal, the reflection cancels the velocity, making the corrected step stationary.Source: §C.1 · pp. 17–18
• In Figure 7, ReFORM’s latent mass concentrates along two radii, resulting in a suboptimal action distribution.Source: Figure 7 · p. 18, §C.1 · p. 18
• Premature halting prevents collapse to a single point, unintentionally acting as an implicit entropy regularizer that partly explains ReFORM’s empirical success.Source: §C.1 · p. 18

C.2 The Ground-truth Base Velocity Field from the Base Space to the Latent Space

• Because the Gaussian base and uniform-ball target are spherically symmetric, their optimal transport is radial, and matching radial quantiles gives an analytical mapping.Source: §C.2 · p. 18
• Evaluating the velocity field requires inverting the trajectory’s radial map; because this generally entails numerical inversion, the implementation learns the field by flow matching to avoid repeated solves.Source: §C.2 · p. 19

D Practical Implementation of LASER

• The decoder uses a learned flow-matching velocity field and a uniform latent prior over a bounded ball; sampled latents are Euler-integrated to produce actions.Source: D Practical Implementation of LASER · p. 19
• The critic uses standard TD learning with an ensemble, and some environments benefit from a pessimistic target-value backup.Source: D Practical Implementation of LASER · p. 19
• Policy optimization discretizes the forward SDE and backward adjoint ODE with N = 10; fixed regression targets let gradients pass only through the latent-policy vector field.Terminal latent critic gradients are clipped before adjoint initialization, while the final integration step uses the reference flow and projects onto the latent ball.Source: D Practical Implementation of LASER · pp. 19–20
• Projection can place probability mass on the latent-ball boundary, creating a singular component, so the projected implementation is a practical approximation to the density-based formulation.Source: D Practical Implementation of LASER · pp. 20–21

E Details of the Reach-avoid Example in Figure 1

• In the reach-avoid example, the no-entropy baseline replaces adjoint matching with gradient ascent, while the no-support-constraint baseline uses QAM with a large inverse temperature.Source: E Details of the Reach-avoid Example… · p. 21

F.1 Environments

• The benchmark covers 40 settings across antmaze, cube-single, cube-double, and scene, with five tasks and CLEAN and NOISY datasets per environment.Source: §F.1 · p. 21

F.2 Baselines

• Baselines include Wasserstein-regularized FQL, KL-regularized QAM, IQL-based IFQL, Gaussian-latent DSRL, and support-constrained ReFORM, with distinct policy-extraction approaches.Source: §F.2 · pp. 21–22

F.3 Hyperparameters

• LASER keeps its method-specific hyperparameters fixed across tasks, while common settings can vary by environment and dataset and some baseline-specific settings are tuned per environment.The latent-ball radius is l = √d.Source: §F.3 · p. 23

F.4 Computational Resources

• 12% more training time than ReFORM and 33% more than QAM; LASER’s throughput is comparable to QAM-E (123 vs 125 it/s).Lean-adjoint solving is needed only during offline training; deployment integrates the latent flow and decoder without adjoints or critic gradients.Source: §F.4 · p. 23

F.5 Additional Results

• Both implicit entropy regularization in ReFORM and explicit regularization in LASER significantly improve on NoER, with LASER performing better than ReFORM.The ablation covers 20 cube-single and cube-double tasks across CLEAN and NOISY datasets.Source: §F.5 · p. 24, Table 5 · p. 25
• On cube-double task 2, inverse temperature 1/α = 10 outperforms NoER and TPS on both datasets; 1/α = 3 reduces performance on both.At 1/α = 30, performance is poor on CLEAN but competitive on NOISY; broader action coverage from noisy demonstrations is offered as one possible explanation.Source: §F.5 · p. 24, Table 6 · p. 25
• On cube-double task 2 with CLEAN data, both tested QAM+Mix probabilities, η = 0.03 and η = 0.1, reduce QAM’s mean success rate.QAM+Mix broadens the action-space reference toward a uniform distribution; the authors report that this naïve approach does not perform well.Source: §F.5 · p. 25, Table 7 · p. 26
• LASER achieves high performance on five CLEAN puzzle-4x4 tasks; ReFORM’s latent boundary correction is active at rates of 97.4%–99.0% on cube-double tasks.Puzzle results use the CLEAN (100M) dataset; ReFORM boundary-hit rates cover five cube-double tasks on both CLEAN and NOISY datasets.Source: §F.5 · p. 26, Table 8 · p. 26, Table 9 · p. 26
• LASER is not sensitive to σ0, while enabling terminal-gradient clipping or reducing the parameter-gradient cap from 10 to 1 improves training stability with comparable performance.The clipping comparison is on cube-double task 2 with CLEAN data; the reported variants use a parameter-gradient cap of 10 or 1.Source: Figure 8 · p. 26, §F.5 · p. 27
• LASER yields the best overall performance across 40 main evaluation settings, achieving near-perfect success rates on multiple tasks.Final checkpoints are evaluated over 50 episodes per seed, with results reported as means and standard deviations across 3 seeds.Source: §F.5 · p. 27, Table 11 · p. 28
• A paired comparison across 40 settings gives a global p-value of 1.35 × 10^-3 for improvement over every baseline, under the paired-test assumptions.The analysis treats setting-level differences as independent despite shared environments and datasets, so the authors interpret it as supporting evidence for aggregate suite performance.Source: §F.5 · p. 27, Table 10 · p. 27