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

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 consistently low violation rates with competitive returns, highlighting the benefits of structured geometric regularization.

This paper studies how to train safe decision-making agents from fixed offline data, without allowing risky online trial-and-error. A key challenge is that learned policies may choose actions that are outside the data distribution, which can lead to unsafe behavior. We propose GSCO, a method that learns a structured latent representation of safe behaviors using normalizing flows and then improves actions within this learned space while controlling distributional shift. The method separates reward improvement, safety correction, and geometric regularization, making policy updates more stable and less likely to leave the supported data region. Experiments on safe offline reinforcement learning benchmarks show that GSCO achieves low safety violations while maintaining competitive returns.