🤖 AI Summary
This work addresses the risks of out-of-distribution actions and latent policy collapse in offline reinforcement learning by proposing the LASER algorithm. The method introduces flow matching and entropy regularization within the latent space to stabilize training, while employing an adjoint matching technique to circumvent the computational overhead associated with backpropagation through time. Experimental evaluations demonstrate that LASER achieves state-of-the-art performance across 40 OGBench tasks. Notably, it significantly outperforms baseline methods that require meticulous hyperparameter tuning, even when operating under a fixed set of hyperparameters. These results validate the critical role of entropy regularization in facilitating robust offline learning.
📝 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/.