SteinGate: Tail-Sensitive Safe Reinforcement Learning via Stein Discrepancy

2026-07-14

AI Summary

Traditional safe reinforcement learning relies on expected cost constraints, which often fail to mitigate rare but catastrophic tail risks. This work proposes SteinGate, a boundary-aware distributional safety certification mechanism grounded in kernelized Stein discrepancy. By performing a nonparametric goodness-of-fit test, SteinGate determines whether the cost distribution of a candidate policy rollout is consistent with a safe reference distribution, dynamically switching between optimization and recovery behaviors accordingly. The method avoids explicit modeling of tail distributions and explicitly accounts for boundary atoms induced by cost truncation, enabling sensitive detection and response to tail risks. Evaluated on continuous control benchmarks, SteinGate substantially reduces both the frequency and severity of constraint violations during training while maintaining return performance comparable to state-of-the-art baselines.

Abstract

Safe reinforcement learning typically enforces safety by bounding expected cumulative costs, a criterion that often fails to detect rare but catastrophic tail events. To overcome these limitations, this paper introduces SteinGate, a boundary-aware distributional safety certificate that replaces fragile tail fitting with a robust consistency check using Kernelized Stein Discrepancy while accounting for boundary atoms induced by clipped costs. SteinGate evaluates whether observed policy rollout costs remain consistent with a safe reference distribution, providing a non-parametric safety certificate. This certificate is used to dynamically adapt the learning regime: favoring reward-improving policy updates when rollouts remain consistent with the safe reference and switching to recovery behavior when the cost tail deviates. Experiments on continuous-control benchmarks demonstrate that SteinGate significantly reduces both the frequency and severity of constraint violations during training while maintaining competitive returns relative to state-of-the-art baselines.

Problem

• safe reinforcement learning
• tail events
• constraint violations
• distributional safety
• cost tail

Innovation

• Stein Discrepancy
• Safe Reinforcement Learning
• Tail Risk
• Distributional Safety Certificate
• Non-parametric Consistency Check