Geometric Control of Out-of-Distribution Shift in Safe Offline RL

Abstract

Safe offline reinforcement learning (RL) requires optimizing policies within the support of static datasets while satisfying strict safety constraints. Although recent latent generative policies achieve strong empirical performance, they rely heavily on implicit regularization and lack systematic control over distributional shift during policy improvement. In this work, we propose a geometric control framework that leverages the bijective structure of conditional normalizing flows to provide a tractable mechanism to regulate distributional deviation of the policy. By constraining divergence in the latent base space, we derive tractable upper bounds on the induced Wasserstein distance and total variation of the policy distribution, establishing an analyzable connection between latent geometry and downstream behaviors. This insight motivates a decoupled architecture: a flow prior shapes a feasibility-weighted latent manifold using Hamilton--Jacobi reachability signals, while a latent refiner performs geometrically constrained optimization directly in the base space. Across multiple safe RL benchmarks, our method achieves robustly low violation rates with competitive returns, highlighting the benefits of structured geometric regularization.