SteinGate: Tail-Sensitive Safe Reinforcement Learning via Stein Discrepancy

Safe reinforcement learning (RL) often relies on bounding the expected cumulative cost. However, this approach suffers from "expectation myopia," where a policy may satisfy average safety constraints while still being prone to rare but catastrophic tail events. This paper addresses the challenge of enforcing safety in high-stakes environments by developing a mechanism that can rigorously detect and mitigate these dangerous tail risks during training.

The authors introduce SteinGate, a framework that replaces fragile tail-fitting methods with a distributional consistency check. Instead of modeling the tail of the cost distribution directly, SteinGate uses Kernelized Stein Discrepancy (KSD) to measure the divergence between the current policy's rollout distribution and a pre-certified safe reference distribution. To handle the structural artifacts of RL—such as cost clipping and mass accumulation at zero or violation thresholds—the authors propose a hybrid Stein operator. This operator treats boundary atoms (discrete probability mass at the limits of the cost range) separately from the continuous interior density, allowing for a more accurate and stable safety certificate.

SteinGate integrates this safety certificate into a lexicographic switching controller. The agent operates in a "Reward-Maximization Mode" when the certificate confirms that tail risk remains within the prescribed budget. If the certificate detects a deviation, the agent switches to a "Safety-Recovery Mode," prioritizing cost reduction and ignoring reward gradients. Experimental results on continuous-control benchmarks demonstrate that this approach significantly reduces both the frequency and severity of constraint violations compared to traditional expectation-based and tail-modeling baselines, all while maintaining competitive reward performance.