1 Introduction↩︎

Safe reinforcement learning (RL) aims to deploy autonomous agents in unknown environments where exploration must be restricted due to critical physical or operational constraints. In high-stakes applications such as autonomous robots, medical systems, and industrial control, the inherent trial-and-error nature of an RL agent becomes a significant liability. In these domains, the cost of a single catastrophic failure can outweigh the benefits of thousands of successful runs [1], [2]. Consequently, the objective of Safe RL must move beyond simple reward maximization to ensure that safety violations are not merely penalized, but rigorously prevented.

Safe RL problems are formulated as Constrained Markov Decision Processes [3] and safety is typically enforced in an average manner: the expected cumulative cost is required to be controlled below a given threshold [3], [4]. While computationally convenient, these methods suffer from a fundamental “expectation myopia”. Regulating only the average of the cost distribution allows a policy to satisfy safety limits by balancing rare, catastrophic failures against a high volume of safe episodes. This risk-masking behavior is a critical flaw; a policy may appear safe on paper while remaining dangerously prone to extreme tail events that standard safety mechanisms ignore. This “expectation myopia” has motivated a growing interest in risk-sensitive formulations of Safe RL, including chance constraints and coherent risk measures such as Value-at-Risk (VaR) and conditional VaR [5], [6]. More recently, extreme value theory has been applied to model the behavior of heavy-tailed distributions directly [7], [8].

Despite their theoretical appeal, enforcing these constraints in practice is challenging because rare-event estimation is intrinsically unstable in online RL. Policy updates continuously shift the rollout distribution, while trajectory correlations reduce the effective sample size [9]. Moreover, episodic training, termination rules, and cost clipping often induce structural artifacts, such as mass accumulation near zero cost or at violation thresholds, that complicate classical tail modeling assumptions. As a result, direct tail estimation can exhibit high variance and brittleness, leading to oscillations between over-conservatism and catastrophic failure during training.

This paper proposes a novel and fundamentally different perspective for tail-sensitive safe RL: rather than explicitly modeling the tail of the cost distribution, we certify tail risk via distributional comparison. Our key idea is to upper-bound a tail-sensitive risk functional [10] under the current policy by comparing its rollout cost distribution to a designated safe reference distribution [11]. The reference distribution specifies acceptable cost behavior under the desired safety budget and serves as the target of distributional certification. Intuitively, if the current policy’s distribution remains close to a certified safe reference, then its tail risk is provably bounded. This converts safety monitoring into a problem of measuring distributional deviation under non-stationary, sample-based conditions.

To theoretically apply this idea to Safe RL, we leverage Stein’s method, a functional-analytic tool for bounding the difference between expectations under an unknown distribution and a target reference [12], [13]. Specifically, we utilize the Kernelized Stein Discrepancy (KSD) [14] to quantify the divergence between the agent’s current rollout distribution \(q_\pi\) for policy \(\pi\) and a pre-certified safe baseline. KSD uniquely fits our needs for safe RL because it avoids the need for explicit density estimation. While RL policy rollouts provide an abundance of cost samples, the underlying rollout distribution is generally implicit and intractable. KSD bypasses this bottleneck by comparing empirical samples from \(q_\pi\) against the score function (\(\nabla \log p\)) of a chosen reference distribution \(p\).

By combining Stein’s method with a tail-sensitive test function, we are able to derive a computable upper bound on the probability of constraint violations. This bound can serve as a rigorous distributional safety certificate, allowing for real-time monitoring of the policy’s tail risk. Unlike standard expected costs, this certificate provides the agent with a formal guarantee that its behavior remains aligned with the safe reference, preventing hazardous distributional shifts before catastrophic failure occurs.

We integrate this distributional certificate into the policy optimization loop via a mechanism we call SteinGate. Standard Lagrangian methods incorporate safety as a “soft” penalty, often leading to training instabilities [15], [16] or situations in which the reward objective can effectively overpower the constraint when cost penalties are not sufficiently enforced (i.e., constraint trade-offs governed by the Lagrange multiplier) [16]. In contrast, SteinGate eliminates this risk by using the certified upper bound as a bang-bang controller:

1. Reward-Maximization Mode: Optimization proceeds only when the certificate verifies that the tail risk remains within the prescribed budget.

2. Safety-Recovery Mode: If the bound is violated, the algorithm immediately switches to prioritize cost reduction, ignoring reward gradients entirely.

Our approach ensures that the agent never sacrifices safety for high rewards; it avoids the instability of traditional methods that rely on tuning safety penalties by balancing the trade-off between reward and cost. We further prove that, under a trust-region condition on the margined policy update, each improvement step preserves the prescribed tail-risk budget with high probability. Our contributions are summarized as follows:

• Boundary-Aware Hybrid Stein Certificate: We derive a tail-risk upper bound for normalized clipped cost distributions in Safe RL. Our framework introduces a hybrid discrete-continuous reference model and split discrepancy that explicitly accounts for boundary atoms induced by cost clipping, enabling rigorous distributional certification without parametric tail fitting.

• SteinGate Mechanism and Safety Guarantees: We integrate this certificate into a lexicographic switching controller that enforces hard safety priority. We provide concentration results for the empirical certificate and establish a trust-region condition under which policy updates preserve the prescribed risk budget with high probability.

• Empirical Stability and Performance: On continuous-control benchmarks, SteinGate significantly reduces violation rates and improves training stability compared to expectation-based and tail-modeling baselines, while maintaining competitive reward performance. Ablations further show stable behavior across cost thresholds, sample sizes, and reasonable reference choices.

2 Problem Setup↩︎

Safe RL problems are typically formulated as a finite-horizon Constrained Markov Decision Process (CMDP). A CMDP is defined by the tuple \(\mathcal{M}=(\mathcal{S},\mathcal{A},P,r,c,H,\mu_0)\), where \(\mathcal{S}\) and \(\mathcal{A}\) represent the state and action spaces, \(P(\cdot\vert s,a)\) denotes the stochastic transition kernel, \(r(s,a)\) is the reward function, \(c(s,a)\geq 0\) is the cost function, \(H\) is the horizon, and \(\mu_0\) is the initial state distribution. For a stochastic policy \(\pi(\cdot\mid s)\), we consider the induced trajectory \(\tau=(s_0, a_0, \dots, s_{H-1}, a_{H-1}, s_H)\), where \(s_0 \sim \mu_0\), \(a_t \sim \pi(\cdot\vert s_t)\), and \(s_{t+1} \sim P(\cdot\mid s_t, a_t)\). The expected cumulative reward, or reward value function, is defined as: \[\begin{align} V_R^\pi(s) = \mathbb{E}_{\pi} \left[ \sum_{t=0}^{H-1} r(s_t, a_t) \vert s_0=s \right]. \end{align}\] The cumulative episode cost for a specific trajectory \(\tau\) is: \[\begin{align} C(\tau)\;\mathrel{\vcenter{:}}=\;\sum_{t=0}^{H-1} c(s_t,a_t). \end{align}\] Given a predefined cost limit \(d > 0\), the standard CMDP objective is to maximize the expected cumulative reward subject to an expectation-based safety/cost constraint: \[\max_{\pi \in \Pi} \mathbb{E}_{s\sim\mu_0} [V_R^\pi(s)] \quad \text{s.t.} \quad \mathbb{E}_{\tau \sim \pi}[ C(\tau) ] \leq d, \label{eq:cmdp}\tag{1}\] where \(\Pi\) denotes the set of feasible policies.

In many safety-critical applications, the common constraint on expected cost is insufficient. This is because expectation formulation allows a policy to meet safety limits on average by balancing rare, catastrophic failures against a large number of safe runs. To prevent these unacceptable risks, we define the policy risk \(C_h^\pi(\mu_0)\) of a policy \(\pi\) as the expected value of a smooth surrogate function \(h\): \[C_h^\pi (\mu_0) \mathrel{\vcenter{:}}= \mathbb{E}_{\tau\sim\pi} [ h(C(\tau))] = \mathbb{E}_{\pi}\left[ h\left(\sum_{t=0}^{H-1} c(s_t,a_t)\right)\right]. \label{eq:risk}\tag{2}\] We denote \(C_h^\pi (\mu_0)\) as \(C^\pi_h\) and \(\mathbb{E}_{\tau\sim\pi} [ h(C(\tau))]\) as \(\mathbb{E}_{C\sim q_\pi} [ h(C)],\) where \(q_\pi\) is the distribution of cumulative costs of trajectories induced by a policy \(\pi\) to simplify the notation. In the following sections, for a given policy, we also use \(C\) to denote \(C(\tau)\), which is a random variable without ambiguity. The surrogate function \(h\) here serves as a risk-shaping operator that transforms the cumulative costs into a specific safety metric and provides the flexibility to define different safety objectives. Specifically, \(h\) can be implemented as a soft threshold to serve as a smooth proxy for violation probability [17], or as a barrier penalty (e.g., an exponential or power function) to heavily discourage trajectories that approach high-cost regions [18]. It can also be modeled as a survival function to map cost accumulation to a bounded risk of failure [19]. Ultimately, this flexibility allows the risk functional \(C_h^\pi\) to target specific portions of the cost distribution, providing a more precise safety signal than the standard expected cost safety.

Therefore, in this paper, we consider a more general constrained RL objective: \[\label{eq:risk-dp} \max_{\pi \in \Pi} \mathbb{E}_{s\sim\mu_0} [V_R^\pi(s)] \quad \text{s.t.}\quad C_h^\pi(\mu_0)\leq \varepsilon.\tag{3}\]

3 Stein-Based Distributional Certification↩︎

When solving the optimization problem in 3 , directly estimating the risk function \(C_h^\pi\) through Monte Carlo rollouts is often unreliable, particularly in rare-event regimes. When safety violations are infrequent, the empirical estimate of \(\mathbb{E}_{C\sim q_\pi} [ h(C)]\) suffers from high variance. This instability is further aggravated by the non-stationary distributional shifts that occur during online policy updating, making the estimate too brittle for reliable safety guarantees. Below, we provide the necessary background on Stein’s method and describe how to apply it to address this challenge.

Quantifying the discrepancy between a policy’s rollout distribution \(q_\pi\) and a reference \(p\) is inherently difficult, as both distributions are typically implicit and lack a closed-form density. Stein’s method [12] provides a principled way to compare a (generally implicit) rollout distribution \(q_\pi\) to a reference distribution \(p\) by characterizing \(p\) through a linear operator. Let \(p\) be a target distribution on an open set \(\mathcal{C}\subseteq\mathbb{R}^d\) with continuously differentiable density and score function \(s_p(x)\mathrel{\vcenter{:}}=\nabla_c \log p(x)\). A Stein operator \(\mathcal{A}_p\) is any linear operator satisfying the Stein identity \(\mathbb{E}_{C\sim p}[(\mathcal{A}_p f)(C)]=0\) for all functions \(f\) in a suitable class where \(C\) is a random variable. A canonical choice for absolutely continuous targets is the Langevin–Stein operator [13] \[(\mathcal{A}_p f)(c) \;:=\; \langle s_p(c), f(c)\rangle \;+\; \nabla\!\cdot f(c), \label{eq:stein-op}\tag{4}\] where \(f:\mathbb{R}^d\to\mathbb{R}^d\) is vector-valued and \(\nabla\!\cdot f\) denotes the divergence [20]. For any given test function \(h\) in an exceedance function class \(\mathcal{H}\), the Stein equation seeks a solution \(f_h\) such that: \[(\mathcal{A}_p f_h)(C) = h(C)-\mathbb{E}_{C\sim p}[h(C)]. \label{eq:stein-eq}\tag{5}\] Thus, we are able to obtain a Stein-type upper bound in the following lemma, which establishes a core safety certificate. It allows us to bound the unknown risk by measuring how much the current policy drifts from a given baseline without knowing the rollout distribution \(q_\pi\).

We can observe that this certificate is one-sided and monotone: as the policy drifts further from the safe reference (increasing the discrepancy), the upper bound on risk grows accordingly. The proof is in 10.1 in the Appendix.

The Stein certificate in 6 provides a rigorous approach for applying Stein’s method to solve Safe RL problems, specifically for controlling hazardous distributional shifts, a severe issue where an evolving online policy drifts into unsafe cost regimes [4]. In complex environments, standard monitoring tools are often impractical because the rollout distribution \(q_\pi\) induced by the current policy \(\pi\) is implicit. While we can generate cost samples via simulation, we lack an explicit formula for its density. Consequently, traditional metrics such as the Kullback-Leibler (KL) divergence become computationally expensive or unstable, as they typically require density ratios or auxiliary density estimation. Stein’s method bypasses these bottlenecks. It only requires the score function (\(\nabla \log p\)) of the designated safe reference \(p\) and empirical samples from the current policy. By leveraging this property, we establish a differentiable and sample-efficient mechanism for constraining policy behavior. This approach ensures that the agent’s cost distribution remains aligned with safety requirements without ever requiring the calculation of normalizing constants or explicit probability densities.

While the Stein certificate is theoretically sound, the supremum in Lemma 1 is generally intractable. To obtain a computable metric, we restrict \(\mathcal{F}\) to the unit ball of a Reproducing Kernel Hilbert Space (RKHS) with base kernel \(k(\cdot,\cdot)\). This yields the Kernelized Stein Discrepancy (KSD), defined in its squared population form as: \(\mathrm{KSD}^2(q_\pi\|p) = \mathbb{E}_{C,C'\sim q_\pi}\big[u_p(C,C')\big],\) where \(u_p(\cdot,\cdot)\) is the induced Stein kernel determined by \(p\) and \(k\). In practice, given \(m\) rollout samples \(\{c_i\} \sim q_\pi\), we employ the consistent V-statistic estimator [21]: \(\widehat{\mathrm{KSD}}^2_V :=\frac{1}{m^2}\sum_{i=1}^m\sum_{j=1}^m u_p(c_i,c_j)\). By monitoring this discrepancy from a safe reference \(p\), we can upper-bound tail-sensitive risk without requiring parametric modeling or explicit density estimation.

It is useful to distinguish the role of Stein discrepancy from the source of tail sensitivity in our framework. The Stein discrepancy is a general distributional discrepancy: it measures how far the rollout-induced cost distribution deviates from a designated reference distribution through samples from (\(q_\pi\)) and the score function of (\(p\)). Its main computational advantage in online RL is that it avoids both density estimation for the implicit rollout distribution and parametric fitting of rare exceedance events, which are unstable under non-stationary policy updates. Tail sensitivity enters through the certificate in Lemma 1 by choosing a risk-shaping test function (h) that emphasizes high-cost events and by comparing rollouts to a reference distribution that encodes admissible safe behavior. Thus, KSD provides a tractable mechanism for detecting distributional drift, while the safety semantics of the certificate are determined by the choice of (h), the reference distribution, and the boundary-aware treatment of clipped costs described in the next section.

4 SteinGate Framework↩︎

While the Stein certificate provides a principled upper bound on tail risk, its application to RL must account for the compact support of cost distributions. In Safe RL problems, cumulative costs are typically bounded, often exhibiting significant probability mass at the boundaries, such as clusters of zero-cost “perfect” episodes. More importantly, safety specifications are almost always defined relative to a fixed threshold \(d\). This creates a saturation of risk: once the budget is exceeded, the exact magnitude of the cost becomes secondary to the fact of the violation itself.

To align our mathematical framework with this reality, we utilize a threshold-normalized cost representation. By clipping and scaling costs to the interval \([0, 1]\), we ensure that all budget exceedances are treated as uniformly catastrophic events. This mapping ensures that the distributional certificate remains practically relevant by focusing the “sensitivity” of the Stein operator on the region where the budget is consumed, while maintaining the mathematical consistency required for a bounded function class.

4.1 Boundary Atoms↩︎

To encode budget exceedances, we define the normalized clipped cumulative cost \(\tilde{C}(\tau) = \min\{C(\tau)/d,\,1\}\), then \(\tilde{C}(\tau)\in[0,1]\) normalized and the boundary value \(\tilde{C}(\tau)=1\) aggregates all budget violations into a single event. While this representation is ideal for defining a stable reference on \([0,1]\), it creates a structural mismatch with standard Stein machinery. Specifically, the rollout distribution of \(\tilde{C}(\tau)\) is often a mixed distribution. It frequently contains discrete probability mass (atoms) at the boundaries, a “safe” atom at \(\tilde{C}(\tau)=0\) and a “violation” atom at \(\tilde{C}(\tau)=1\) with a continuous component on the interior \((0,1)\) as shown in Figure 1. This structure conflicts with the standard Langevin-Stein operator, which assumes a smooth density on an open domain.

The threshold-normalized setting presents two challenges. 1) Boundary Atoms: Stein identities rely on integration by parts, which requires boundary terms to vanish. However, clipping induces discrete probability mass (atoms) at \(\{0, 1\}\), which standard continuous operators fail to capture with spikes at the boundaries. 2) Endpoint Singularity: Score functions for common interior references (e.g., Beta distributions) often become singular near the boundaries. This can amplify variance and cause numerical instability in KSD estimators. To resolve these issues, we propose a hybrid Stein operator that explicitly accounts for boundary mass while maintaining the stability of the interior reference.

4.2 Hybrid Reference and Discrepancy↩︎

4.2.0.1 Safe reference.

To align with the mixed structure of threshold-normalized and clipped costs, we model admissible behavior using a hybrid reference distribution that explicitly accounts for boundary atoms: \[p_{\text{ref}} = a_0^\star\,\delta_{0} + a_1^\star\,\delta_{1} + \bigl(1-a_0^\star-a_1^\star\bigr)\,p_{\mathrm{int}}, \label{eq:psafe}\tag{7}\] where \(\delta_0,\delta_1\) denote point masses at the boundaries, \(a_0^\star,a_1^\star\in[0,1]\) are the reference boundary masses, and \(p_{\mathrm{int}}\) is a smooth interior reference density on \((0,1)\) (e.g., a Beta distribution). The reference distribution \(p_{\mathrm{ref}}\) should be interpreted as a safety specification rather than as an estimate of the unknown cost distribution induced by an optimal safe policy. The optimal safe rollout distribution depends on the environment dynamics and the learned policy, and is generally not available a priori. The role of \(p_{\mathrm{ref}}\) is instead to encode admissible behavior in the normalized cost space. The boundary masses \(a_0^\star\) and \(a_1^\star\) specify acceptable mass at zero cost and at the violation boundary, while the interior density \(p_{\mathrm{int}}\) specifies the desired shape of costs that remain below the threshold.

For any distribution \(q\) on \([0,1]\), we define the empirical boundary masses and the total interior weight as: \[a_0(q) \mathrel{\vcenter{:}}= \Pr(C=0\vert C\sim q), a_1(q) \mathrel{\vcenter{:}}= \Pr(C=1\vert C\sim q),\] Let \(q_{\mathrm{int}}\) denote the conditional distribution of \(q\) on the interior \((0,1)\). Furthermore, let \(w(q) \mathrel{\vcenter{:}}= 1 - a_0(q) - a_1(q).\) This decomposition allows us to analyze the safety of a policy by separately evaluating its performance at the boundaries and its behavior within the threshold. The reference need not be exactly realizable by every MDP. If the specification is overly restrictive or structurally misaligned with the achievable cost distributions, this appears as a persistent discrepancy between \(q\) and \(p_{\mathrm{ref}}\). From the perspective of the certificate, such mismatch is indistinguishable from unsafe distributional behavior: both indicate that the observed rollout distribution does not satisfy the prescribed safety specification. This is intentional, since the certificate is defined relative to the user-provided reference rather than to an unknown optimal distribution. In practice, the choice of \(p_{\mathrm{ref}}\) controls conservatism: stricter references induce larger discrepancies for the same rollout distribution, while more relaxed references admit a wider range of interior cost.

4.2.0.2 Discrepancy with split penalties.

To address the mixed structure of cost distributions, we define a hybrid discrepancy that separately manages boundary mass shifts and interior distribution mismatches: \[\begin{align} &\mathrm{SD}_{\mathrm{tot}}(q\|p_{\mathrm{ref}}) \mathrel{\vcenter{:}}= \lambda_{\mathrm{disc}}\, \mathrm{SD}_{\mathrm{disc}}(q\|p_{\mathrm{ref}}) \\ &\quad\quad\quad\quad\quad+ \lambda_{\mathrm{int}}\, w(q)\, \mathrm{SD}_{\mathrm{int}}(q_{\mathrm{int}}\|p_{\mathrm{int}}) \\ \le &\; \lambda \;\Big( \mathrm{SD}_{\mathrm{disc}}(q\|p_{\mathrm{ref}}) + w(q)\cdot\mathrm{SD}_{\mathrm{int}}(q_{\mathrm{int}}\|p_{\mathrm{int}}) \Big), \end{align} \label{eq:sd-tot}\tag{8}\] where \(\lambda\mathrel{\vcenter{:}}=\max\{\lambda_{\mathrm{disc}},\lambda_{\mathrm{int}}\} >0 .\) The discrete term \(\mathrm{SD}_{\mathrm{disc}}(q\|p_{\mathrm{ref}})\) can be written as \(\mathrm{SD}_{\mathrm{disc}}(q\|p_{\mathrm{ref}}) :=\bigl(a_0^\star-a_0(q)\bigr)_+ + \bigl(a_1(q)-a_1^\star\bigr)_+,\) where \((x)_+=\max(x,0)\). This term is one-sided: it penalizes excess violation mass (\(a_1(q)>a_1^\star\)) and loss of safe mass (\(a_0(q)<a_0^\star\)), while shifts in the harmless direction, fewer violations or more zero-cost episodes than the reference, incur no penalty. The interior term \(\mathrm{SD}_{\mathrm{int}}\) is a standard Stein discrepancy relative to \(p_{\mathrm{int}}\) and can be computed on the interior conditional \(q_{\mathrm{int}}\) using an appropriate Stein operator for \(p_{\mathrm{int}}\) (e.g., the Langevin–Stein operator on a clipped interior).

Together with the Stein certificate of Lemma 1, the hybrid discrepancy in 8 yields a certificate that is sensitive to safety-critical violation mass at the boundary while still controlling continuous interior mismatch.

4.3 SteinGate: A Bang-bang Controller↩︎

We are now ready to introduce our empirical hybrid certificate to govern policy optimization. For a given reference \(p_{\mathrm{ref}}\) and a parametrized policy \(\pi_\theta\), we define the empirical Stein certificate as: \[U(\theta) = \mathbb{E}_{C\sim p_{\mathrm{ref}}}[h(C)] + \lambda\, \widehat{\mathrm{SD}}_{\mathrm{tot}}(q_{\pi_\theta} \| p_{\mathrm{ref}}), \label{eq:hybrid-certif}\tag{9}\] where the total discrepancy is estimated from the most recent rollout samples.

Instead of treating safety as a soft penalty, we implement this certificate as a bang-bang controller compared against a safety threshold \(\varepsilon\). The agent’s behavior is determined by a strict switching rule: \[\text{Mode}(\theta) = \begin{cases} \text{reward-maximization}, & \text{if } U(\theta) \le \varepsilon, \\ \text{safety-recovery}, & \text{if } U(\theta) > \varepsilon. \end{cases}\]

The scalar \(\lambda\) in \(U(\theta)\) does not play the same role as a Lagrange multiplier in constrained policy optimization. In Lagrangian methods, the multiplier weights the cost term inside the optimization objective, so reward and safety gradients are combined in a single update direction. In SteinGate, by contrast, \(\lambda\) only scales the distributional certificate used by the gate. Once the mode is selected, the policy update is either reward-driven or recovery-driven; the two gradients are not linearly combined. Thus, \(\lambda\) controls the conservativeness of the safety test, and hence the point at which switching occurs, rather than mediating a continuous trade-off between reward maximization and constraint satisfaction.

The same gating interpretation also determines how reference misspecification is handled algorithmically. If the certificate remains above the threshold across successive iterations, this may indicate either genuinely unsafe rollout behavior or a reference specification that is too restrictive for the environment. The controller treats both cases conservatively, since both correspond to failure to satisfy the prescribed distributional safety certificate. Persistent selection of the recovery mode can therefore serve as a practical diagnostic for reference misspecification, suggesting that boundary masses or the interior reference require relaxation.

This framework, summarized in Algorithm 2, is what we refer to as “SteinGate.” This design provides two critical technical advantages over standard Lagrangian penalty approaches. First, Lagrangian approaches optimize a linear combination of objectives. In this setting, the relative weighting between reward and cost is governed by a learned multiplier, and poorly tuned or unstable updates can cause the agent to prioritize reward performance while violating safety constraints, since the cost term is only enforced in expectation. SteinGate establishes strict lexicographic priority: when the rollout distribution drifts into the unsafe tail region, the reward gradient is ignored entirely, forcing the agent to prioritize recovery. Second, it circumvents the well-known instability of dual-variable tuning, replacing oscillatory dual ascent updates with a robust, boundary-aware distributional check. This design effectively stabilizes the learning process, as evidenced by the empirical evaluations in Section 5.2.

Although the controller uses a discontinuous switching rule, SteinGate does not switch on instantaneous costs. The gate is evaluated using a batch-level distributional certificate computed from recent rollouts, so isolated high-cost trajectories are averaged through the empirical discrepancy rather than immediately changing the update mode. Moreover, the trust-region update used by the backbone optimizer limits the per-iteration change in the policy, which in turn limits abrupt changes in the induced cost distribution. The two modes are also asymmetric: reward updates are allowed only while the certificate remains below the prescribed budget, whereas recovery updates explicitly act to reduce the cost violation once the budget is exceeded. Together, these features restrict switching signal evolution to the distributional scale of policy updates rather than individual trajectory noise.

4.4 Safety Guarantee for SteinGate↩︎

Before presenting our main results, we first make the following assumption on the empirical certificate error.

We are now ready to present the main results. The following theorem provides a sufficient condition for ensuring a safe policy after each policy improvement step. The detailed proof can be found in Appendix 10.4.

Theorem 4.3 gives a sufficient condition under which a policy update preserves the prescribed risk budget. The lower bound on \(\lambda_k\) specifies how conservative the certificate must be relative to the size of the policy change, as measured by the trust-region distance. Since the certificate is one-sided, increasing \(\lambda_k\) tightens the gate by assigning greater weight to distributional deviation from the safe reference. This differs from a Lagrangian trade-off parameter: \(\lambda_k\) does not determine the relative magnitude of reward and cost gradients in a shared objective, but only the conservativeness of the safety test used before selecting the update mode. Under the stated trust-region condition, any value above the bound is sufficient to keep the post-update risk within the budget with high probability.

5 Experiments and Results↩︎

This section empirically evaluates SteinGate on a suite of continuous-control safety tasks. Our experiments are designed to examine whether SteinGate can alleviate the trade-off between task performance and safety violations that arises in constrained RL. We further study the robustness of SteinGate under varying cost thresholds for the safety constraint and analyze its sensitivity to the choice of the interior reference distribution used in the safety updates.

5.1 Experimental Setup↩︎

Benchmarks. We employ the Safety Gymnasium and Safety MuJoCo suites [22], [23]. Safety Gymnasium tasks (PointGoal, PointCircle, CarGoal, CarCircle) simulate autonomous navigation where agents must reach targets while avoiding hazards or constraining movement patterns. Safety MuJoCo tasks (Ant, HalfCheetah, Humanoid, Swimmer) focus on locomotion control with velocity limits, testing the agent’s ability to maintain stability under state-dependent constraints.

Baselines. We compare SteinGate against multiple safe RL methods. 1) CPO [4] represents the standard for expectation-based constraints, enforcing safety on the average cost. Notably, we utilize CPO as the base optimization engine for SteinGate. 2) EVO [7] utilizes extreme value theory to model the cost distribution tail, explicitly targeting rare, high-impact violations. 3) Saute [24] and 4) Simmer [25] employ state augmentation to strictly enforce constraints, with Simmer additionally inducing a curriculum on the safety budget.

Evaluation Protocol. All methods are trained with identical interaction budgets (10M steps) and environmental configurations. We report performance over six random seeds, visualizing the mean and standard deviation of episodic return and cost. We introduce two quantitative metrics to assess convergence: 1) Tail Feasibility measures the fraction of episodes satisfying the cost constraint during the final 20% of training, and 2) Feasible Return reports the average episodic return conditioned on constraint satisfaction, isolating the performance of valid policies.

Implementation Details. To guarantee a fair comparison, all algorithms utilize identical policy and value network architectures and share core hyperparameters; this ensures that observed performance differences stem solely from the safety mechanisms. The practical estimation procedure of Stein certificate \(U(\theta)\) is detailed in Appendix 11. We employ the OmniSafe implementations for standard baselines and the official public codebase for EVO. Detailed hyperparameter configurations are provided in the Appendix 13.3.

5.2 Results and Discussion↩︎

SteinGate vs. Baselines. Figure 3 illustrates the training dynamics across the Safety Gymnasium suite. SteinGate exhibits a robust ability to converge to the feasible set, driving episodic costs below the threshold while sustaining continuous reward improvement. In contrast, CPO prioritizes reward maximization at the expense of safety, resulting in frequent and sustained violations (particularly in CarGoal1). Conversely, EVO demonstrates strict safety adherence but suffers from significant conservatism, often plateauing at suboptimal return levels. This confirms that relying solely on tail modeling (EVO) or expectation constraints (CPO) leads to suboptimal trade-offs, whereas SteinGate effectively identifies policies that satisfy the constraint while achieving strong task performance. Table 1 quantifies this advantage over the final training phase. On the PointGoal1 and CarGoal1 tasks, SteinGate matches the near-perfect tail feasibility of EVO (\(>0.98\)) while delivering substantially higher feasible returns (e.g., nearly \(3\times\) higher on PointGoal1 and \(4\times\) on CarGoal1). This indicates that SteinGate resolves the over-conservatism inherent in EVO’s extreme-value approach. On PointCircle1, SteinGate outperforms CPO by achieving equivalent returns with strictly higher feasibility (\(0.99\) vs. \(0.77\)), eliminating the safety violations common in primal-dual methods. Finally, compared to state-augmentation baselines (Saute and Simmer), SteinGate consistently achieves superior feasible returns, demonstrating that direct policy regulation is more efficient than curriculum-based or augmented-state approaches. We observe similar trends on the velocity benchmarks and against additional baselines; results are reported in Appendix 12.

2pt

Sensitivity to Cost Thresholds. Figure 4 illustrates SteinGate’s behavior under varying cost limits \((\kappa \in \{10, 25, 40\})\). As the constraint is tightened, SteinGate adaptively regulates its policy to maintain episodic cost near the prescribed threshold while modulating the achievable return. For moderate thresholds (\(\kappa\) = 25, 40), the learned policies consistently stabilize strictly below the cost limits. Under the highly restrictive setting (\(\kappa\) = 10), the episodic cost fluctuates tightly around the threshold, indicating operation at the boundary of the feasible set with only marginal violations. Concurrently, episodic return increases monotonically as the cost limit is relaxed, demonstrating that SteinGate effectively capitalizes on the expanded safety budget to maximize task performance. Collectively, these results demonstrate that SteinGate scales smoothly with stringency of cost constraint, maintaining a stable balance between reward maximization and safety.

Sensitivity to Reference Distribution. Figure 5 examines the impact of the interior reference distribution on PointGoal1, comparing the default \(\mathrm{Beta}(2,5)\) against a conservative \(\mathrm{Beta}(1,10)\) and a logit-normal alternative. Across all three specifications, the costs remain relatively consistent, suggesting that the mechanism for constraint enforcement is not highly sensitive to the specific choice of reference family on this task. The primary difference lies in the efficiency of reward acquisition: the conservative \(\mathrm{Beta}(1,10)\) induces a more cautious utilization of the safety budget, resulting in lower task performance. Conversely, both the default \(\mathrm{Beta}(2,5)\) and the logit-normal references converge to higher final returns. These results indicate that while the reference distribution modulates the degree of conservatism, SteinGate can maintain effective constraint satisfaction across different reasonable parametric choices. The observed robustness across different reference families suggests that SteinGate depends primarily on the specification of an admissible safety profile rather than on a particular parametric form. More conservative references induce correspondingly conservative policies, while comparable safety behavior across multiple reference families indicates that the framework is not tied to a specific distributional assumption.

Sensitivity to Sample Size. Finally, we investigate the robustness of SteinGate to the number of cost samples available for estimating the certificate, in Figure 6. Unlike tail-modeling approaches (e.g., EVO) that often require off-policy data augmentation to accurately fit distributions, SteinGate operates strictly on-policy. We evaluated SteinGate on CarGoal1 with effective batch sizes of \(N \in \{8, 16, 32\}\) samples per epoch. Despite the scarcity of data in the low-sample regime, SteinGate consistently enforced safety across all settings. We observed a natural trade-off: larger batches (\(N\)=32) resulted in slightly more conservative policies, while smaller batches (\(N\)=8) exhibited higher variance but remained safe. These results demonstrate that the SteinGate mechanism is highly sample-efficient and effective even within the limited interaction budgets of standard on-policy training loops.

This behavior is consistent with the role of KSD as a sample-based distributional certificate rather than a parametric tail estimator: the method does not require fitting a rare-event model from a large buffer of exceedances, but instead monitors whether the observed rollout costs remain consistent with the prescribed reference.

6 Related Work↩︎

Expectation-Based Safety. Safe RL is commonly formulated as a CMDP problem, where safety is enforced by constraining the expected cumulative cost. Representative methods include trust-region methods like constrained policy optimization (CPO) [4], FISOR [26], SEVPO [27], projection-based approaches such as PCPO [28] and CUP [29], first-order constrained optimization in policy space (FOCOPS) [23], multi-timescale Lagrangian methods such as RCPO [30] and PID-Lagrangian [15], and reductions to unconstrained RL[31], [32]. While these approaches provide control over expected cost, they do not directly regulate the distribution of outcomes; policies with identical means may exhibit severe tail violations.

Risk-Sensitive and Distributional Safety. To address this limitation, several works impose constraints on distributional functionals of cost. Common choices include Value-at-Risk (VaR) and Conditional Value-at-Risk (CVaR) [33]. In RL, CVaR constraints have been implemented via state augmentation and Lagrangian optimization [6], [34], while actor–critic variants such as WCSAC [35] and QCPO [36] estimate quantiles or parametric cost distributions. Extreme value theory has also been applied to model high-cost exceedances using generalized Pareto distributions, as in extreme value policy optimization (EVO) [7] and related approaches [8]. These methods explicitly target tail behavior but rely on estimating rare-event statistics. Alternatively, Sauté RL [24] and Simmer [25] enforce safety almost surely via budget state augmentation.

Stein Discrepancies in RL. Kernelized Stein discrepancy (KSD) [14] offers a nonparametric measure of distributional deviation, previously applied in RL for detecting transition shifts [37] or inducing pessimism in offline settings [38]. In contrast, we employ KSD to certify tail-sensitive safety in online RL via a boundary-aware formulation that monitors the cost distribution during optimization.

7 Conclusion↩︎

This paper introduced and studied SteinGate, a distributional safety framework for tail-sensitive constrained reinforcement learning. By leveraging kernelized Stein discrepancy, SteinGate certifies safety through deviation from a designated hybrid reference distribution, avoiding direct estimation of rare-event statistics. The boundary-aware formulation accounts for probability mass induced by cost clipping and episodic termination, enabling theoretically-sound safety monitoring under non-stationary policy updates. Empirical results across multiple benchmarks demonstrate that SteinGate reduces violation frequency and severity during training while maintaining strong task performance. These findings suggest that distributional certification provides a principled and practical approach to mitigating extreme violations in safe reinforcement learning.

8 Acknowledgments↩︎

The authors gratefully acknowledge the in part support by USDA-NIFA funded AgAID Institute award 2021-67021-35344 and NSF ECCS-2534263 grant. The views expressed are those of the authors and do not reflect the official policy or position of the USDA-NIFA and the NSF.

9 Background of Stein’s Method and Kernelized Stein Discrepancy↩︎

Stein’s method yields identities and discrepancy measures between distributions.

Let \(k(x,x')\) be a positive definite kernel. The spectral decomposition of \(k(x,x')\), as implied by Mercer’s theorem, is defined as \[\label{eq:mercer} k(x,x')=\sum_{j}\lambda_j e_j(x)e_j(x'),\tag{12}\] where \(\{e_j\}\) and \(\{\lambda_j\}\) are the orthonormal eigenfunctions and positive eigenvalues of \(k(x,x')\), respectively, satisfying \[\int e_i(x)e_j(x)\,dx=\mathbb{I}[i=j],\qquad \forall\, i,j.\]

For a positive definite kernel \(k(x,x')\), its related RKHS \(\mathcal{H}\) comprises linear combinations of its eigenfunctions, i.e., \[f(x)=\sum_{j} f_j e_j(x) \quad \text{with} \quad \sum_{j}\frac{f_j^2}{\lambda_j}<\infty,\] endowed with an inner product \[\langle f,g\rangle_{\mathcal{H}}=\sum_{j}\frac{f_j g_j}{\lambda_j} \quad \text{between} \quad f(x)=\sum_{j} f_j e_j(x),\;\;g(x)=\sum_{j} g_j e_j(x).\] Thus this Hilbert space is equipped with a norm \(\|f\|_{\mathcal{H}}\) where \[\|f\|_{\mathcal{H}}^2=\langle f,f\rangle_{\mathcal{H}}=\sum_{j}\frac{f_j^2}{\lambda_j}.\]

One can verify that \(k(x,\cdot)\in\mathcal{H}\) and satisfies the important reproducing property, \[f(x)=\langle f,k(\cdot,x)\rangle_{\mathcal{H}}, \qquad k(x,x')=\langle k(\cdot,x),k(\cdot,x')\rangle_{\mathcal{H}}.\] Every positive definite kernel \(k\) defines a unique RKHS for which \(k\) is a reproducing kernel.

It is easy to check that the RBF kernel \[k(x,x')=\exp\!\left(-\frac{1}{2h^{2}}\|x-x'\|_2^{2}\right)\] is in the Stein class for smooth densities supported on \(\mathcal{X}=\mathbb{R}^{d}\).

10 Theoretical Proofs↩︎

10.1 Proof of Lemma 1↩︎

First, we show that \(B(h,p)\) has an upper bound for specific test-function class \(\mathcal{H}\).

Then we start the proof of 1.

10.2 Stein Inequality↩︎

10.3 Bound for Empirical Error of Stein Certificate↩︎

Given \(n\) i.i.d.rollouts (or approximately independent mini-batch samples) producing \(\{c_i\}_{i=1}^n\sim q\), define empirical atom masses \[\hat{a}_0 := \frac{1}{n}\sum_{i=1}^n \mathbf{1}\{c_i\le\tau_0\},\qquad \hat{a}_1 := \frac{1}{n}\sum_{i=1}^n \mathbf{1}\{c_i\ge 1-\tau_1\},\qquad \hat{w}:=1-\hat{a}_0-\hat{a}_1. \label{eq:pi-hat}\tag{19}\] Let \(\{c^{\mathrm{int}}_i\}_{i=1}^m\subset\{c^{}_i\}_{i=1}^n\) be the interior samples with \(m=n_{\mathrm{int}}\). Let \(\widehat{SD}_{\mathrm{int}}\) be the KSD estimator on interior samples with V-statistic of the Stein kernel after clipping away from the boundary. Define the empirical discrepancy \[\widehat{SD}_{\textrm{disc}} := (a_0^\star-\hat{a}_0)_+ + (\hat{a}_1-a_1^\star)_+, \qquad \widehat{SD}_{\textrm{tot}} := \widehat{SD}_{\textrm{disc}} + \hat{w}\cdot \widehat{SD}_{\mathrm{int}}, \label{eq:sd-hat}\tag{20}\] where \[\widehat{SD}_{\mathrm{int}} := \widehat{\mathrm{KSD}}_V^2 := \frac{1}{m^2}\sum_{i=1}^m\sum_{j=1}^m u_{p_{\mathrm{int}}}(C_i^{\mathrm{int}},C_j^{\mathrm{int}}).\] and the empirical certificate \[U := \mathbb{E}_{p_{\text{ref}}}[h(C)] + \lambda\,\widehat{SD}_{\textrm{tot}}. \label{eq:U-hat}\tag{21}\] Here we precompute \(\mathbb{E}_{p_{\text{ref}}}[h(C)]\) by Monte-Carlo Sampling is a fixed constant independent of \(q\).

10.4 Proof of Theorem 2↩︎

Consider an episodic CMDP \(\mathcal{M}\) with horizon \(H\), state space \(\mathcal{S}\), action space \(\mathcal{A}\), transition kernel \(P(\cdot\mid s,a)\), and per-step cost \(c:\mathcal{S}\times\mathcal{A}\to[0,1]\). For a trajectory \(\tau=(s_0,a_0,\dots,s_{H-1},a_{H-1})\) generated by policy \(\pi\), define the (unnormalized) cumulative cost \[C(\tau)\;:=\;\sum_{t=0}^{H-1} c(s_t,a_t).\] Let \(h:\mathbb{R}\to\mathbb{R}\) be a measurable test function (e.g., sigmoid). Our target functional is \[C_h^{\pi}(\mu_0)\;:=\;\mathbb{E}_{\tau\sim\pi}\big[h(C(\tau))\big].\] This is not additive in \((s_t,a_t)\) in general, hence performance difference lemma does not apply directly.

10.4.0.1 Augmented CMDP.

We define an augmented CMDP \(\bar{\mathcal{M}}\) whose state includes the running cumulative cost. Let the augmented state space be \[\bar{\mathcal{S}} := \mathcal{S}\times \mathbb{R}, \qquad \bar{s}_t := (s_t, z_t), \qquad \bar{\mu}_0:=({\mu}_0, z_0)\] where the auxiliary coordinate \(z_t\) evolves deterministically as \[z_0:=0,\qquad z_{t+1} := z_t + c(s_t,a_t),\quad t=0,\dots,H-1.\] The augmented transition kernel is therefore \[\bar{P}\big((s',z')\mid (s,z),a\big) = P(s'\mid s,a)\cdot \mathbf{1}\{z'=z+c(s,a)\}.\]

Define a terminal cost on the augmented state at time \(H\): \[\bar{c}_H(\bar{s}_H) := h(z_H),\] and define per-step costs for \(t=0,\dots,H-1\) as identically zero: \[\bar{c}_t(\bar{s}_t,a_t):=0,\qquad t=0,\dots,H-1.\] Let \(\bar{V}_c^\pi(\bar\mu_0)\) denote the expected total cost under current policy \(\pi\) in \(\bar{\mathcal{M}}\): \[\bar{V}_c^\pi(\bar\mu_0) := \mathbb{E}_{\tau\sim \pi}\Big[\sum_{t=0}^{H-1}\bar{c}_t(\bar{s}_t,a_t) + \bar{c}_H(\bar{s}_H)\Big].\] By construction, \[\label{eq:Jh95equals95augJ} \bar{V}_c^\pi(\bar\mu_0)=\mathbb{E}_{\tau\sim\pi}[h(z_H)]=\mathbb{E}_{\tau\sim\pi}[h(C(\tau))]=C_h^\pi(\mu_0).\tag{28}\]

10.4.0.2 Applying Performance Difference Lemma on the augmented CMDP.

Now we can apply standard trust-region performance-difference bounds on \(\bar{\mathcal{M}}\), because \(\bar{J}(\pi)\) is an expected return in the augmented CMDP.

Let \(\pi,\pi'\) be two policies on \(\mathcal{S}\) (equivalently on \(\bar{\mathcal{S}}\) by ignoring \(z\) in the action selection). Define the value and \(Q\) functions for the augmented MDP with respect to policy \(\pi\): \[C_h^\pi(\bar{s}_t) := \mathbb{E}_{\pi}\Big[\bar{c}_H(\bar{s}_H)\,\big|\,\bar{s}_t\Big],\qquad Q_h^\pi(\bar{s}_t,a) := \mathbb{E}_\pi\Big[\bar{c}_H(\bar{s}_H)\,\big|\,\bar{s}_t,a_t=a\Big],\] and the corresponding advantage \[A_h^\pi(\bar{s}_t,a):=Q_h^\pi(\bar{s}_t,a)-C_h^\pi(\bar{s}_t).\] Let \(\bar{d}^\pi_t\) denote the state distribution of \(\bar{s}_t\) under \(\pi\) in the augmented CMDP and \(\bar{d}^\pi:=\frac{1}{H}\sum_{t=0}^{H-1}\bar{d}^\pi_t\).

By augmentation, we can apply trust-region performance-difference bounds on \(\bar {\mathcal{M}}\) . We therefore consider the conservative surrogate program whose constraint can be rewritten as: \[\begin{align} \label{eq:stein-tr} &\underbrace{ C_h^{\pi_k}(\bar \mu_0) + \sum_{t=0}^{H-1}\mathbb{E}_{\bar s\sim \bar d^{\pi_k},a\sim\pi}\big[A^{\pi_k}_h(\bar s,a)\big] }_{\text{CPO}} + \underbrace{m_k}_{\text{Stein margin}} \le \varepsilon, \;\; \bar D_{KL}(\pi\|\pi_k)\le \delta, \end{align}\tag{29}\] where the Stein margin is \[m_k\triangleq \lambda_k\cdot \underline{SD}_{\mathrm{tot}}^+(q_{\pi_k}\|p_{\mathrm{ref}})\;\;\ge 0.\]

Fix the current policy \(\pi_k\) and let \(\pi_{k+1}\) satisfy a trust-region constraint. Applying 30 –31 with \((\pi,\pi')=(\pi_k,\pi_{k+1})\) yields \[\label{eq:one95step95aug95bound} C_h^{\pi_{k+1}}(\bar \mu_0) \;\le\; C_h^{\pi_{k}}(\bar \mu_0) + \sum_{t=0}^{H-1}\mathbb{E}_{\bar{s}\sim \bar{d}^{\pi_k}_t,\;a\sim\pi_{k+1}(\cdot\mid s)} \big[A_h^{\pi_k}(\bar{s},a)\big] + 2H\,\epsilon_h^{\pi_{k+1}}\cdot \sum_{t=0}^{H-1}\mathbb{E}_{\bar{s}\sim \bar{d}^\pi_k}\Big[D_{\mathrm{TV}}\big(\pi_{k+1}(\cdot\mid s),\pi_k(\cdot\mid s)\big)\Big].\tag{32}\]

Define the usual surrogate objective on the augmented MDP: \[L_h(\pi_{k+1};\pi_k) := C_h^{\pi_{k}}(\bar \mu_0) + \sum_{t=0}^{H-1}\mathbb{E}_{\bar{s}\sim \bar{d}^\pi_t,\;a\sim\pi_{k+1}(\cdot\mid s)} \big[A_h^{\pi_k}(\bar{s},a)\big].\] Then 32 becomes \[\label{eq:Jh95le95surrogate95plus95TR} C_h^{\pi_{k+1}}(\bar \mu_0) \;\le\; L_h(\pi_{k+1};\pi_k) + 2H\,\epsilon_h^{\pi_{k+1}}\cdot \sum_{t=0}^{H-1}\mathbb{E}_{\bar{s}\sim \bar{d}^\pi_k}\Big[D_{\mathrm{TV}}\big(\pi_{k+1}(\cdot\mid s),\pi_k(\cdot\mid s)\big)\Big].\tag{33}\]

By 29 , our update enforces a marginalized surrogate safety constraint \[L_h(\pi_{k+1};\pi_k)\le \varepsilon - m_{k}.\] Plugging this into 33 yields the one-step bound \[\label{eq:stein95gate95one95step} C_h^{\pi_{k+1}}(\bar \mu_0) \;\le\; d - m_{k} + 2H\,\epsilon_h^{\pi_{k+1}}\cdot \sum_{t=0}^{H-1}\mathbb{E}_{\bar{s}\sim \bar{d}^\pi_k}\Big[D_{\mathrm{TV}}\big(\pi_{k+1}(\cdot\mid s),\pi_k(\cdot\mid s)\big)\Big].\tag{34}\]

Therefore, if the Stein margin satisfies \[\lambda_k\ge \frac{2H\,\epsilon_h^{\pi_{k+1}}\cdot \sum_{t=0}^{H-1}\mathbb{E}_{\bar{s}\sim \bar{d}^\pi_k}\Big[D_{\mathrm{TV}}\big(\pi_{k+1}(\cdot\mid s),\pi_k(\cdot\mid s)\big)\Big]}{\underline{SD}_{\rm tot}^+(q_{\pi_k}\|p_{\mathrm{ref}})}\] then the following relation holds with probability at least \(1-2\delta\), \[C_h^{\pi_{k+1}}(\bar \mu_0)\;\le\; \varepsilon.\]

11 Practical SteinGate↩︎

In subsection 4.2, we provided a decomposition of the tail-risk functional into three components: (i) the reference risk under the hybrid safe distribution, (ii) a boundary mass discrepancy term, and (iii) an interior Stein discrepancy term weighted by the interior mass. This structure naturally suggests a plug-in estimator in which each quantity appearing on  9 is approximated from finite rollout samples.

First, the reference risk \(\mathbb{E}_{p_{\mathrm{ref}}}[h(C)]\) depends only on the chosen safe reference and can therefore be computed offline or by Monte Carlo sampling from the interior reference density \(p_{\mathrm{int}}\). This term establishes a baseline level of admissible tail risk.

Second, the discrete discrepancy \(SD_{\mathrm{disc}}(q \,\|\, p_{\mathrm{ref}})\) depends only on the boundary masses \(a_0(q)\) and \(a_1(q)\). In the bounded-cost setting, these quantities correspond to the probability that trajectories accumulate negligible cost or saturate at the maximal violation level. Estimating these probabilities reduces to counting the fraction of rollout samples falling within small neighborhoods of the endpoints. Importantly, this step preserves the asymmetry of the theoretical bound: only loss of safe mass and excess violation mass contribute to the certificate, while harmless shifts in the opposite direction are ignored.

Third, the interior discrepancy \(SD_{\mathrm{int}}(q_{\mathrm{int}} \,\|\, p_{\mathrm{int}})\) is approximated using a kernelized Stein discrepancy relative to the interior reference density. This requires only a Stein operator for \(p_{\mathrm{int}}\) and a function class over the interior. Restricting the discrepancy computation to the interior samples implements the conditional distribution \(q_{\mathrm{int}}\), while the empirical interior mass \(w(q)\) provides the corresponding weight.

Combining these three empirical components yields a certificate of equation 9 , which is a direct numerical instantiation of the theoretical upper bound in Lemma 1. The only approximation introduced is the replacement of population expectations and discrepancies by their sample analogues. As the number of rollout samples increases and the boundary neighborhoods shrink, these estimates converge to their population counterparts and the empirical certificate approaches the theoretical bound, up to the vanishing boundary approximation term \(\varepsilon_{\mathrm{bdry}}\).

This construction ensures that the practical estimator remains structurally aligned with the theory: boundary mass shifts and interior distributional mismatch are controlled by separate terms with distinct gains, and both contribute additively to the final certificate. No additional modeling assumptions are introduced beyond those required by the hybrid Stein bound itself.

Finally, we incorporate a lightweight empirical guard that replaces the Stein bound with the empirical risk estimate in regimes where the boundary discrepancy vanishes and the empirical risk is already below the reference level. This step does not alter the form of the certificate in the unsafe regime; it merely reduces variance when the distribution is clearly safe and the bound is nearly tight. The resulting procedure preserves the conservativeness of the Stein certificate while improving numerical stability in benign regimes. The implementation can be found at https://github.com/yassineCh/SteinGate.

12 Additional Results↩︎

12.1 Additional Benchmarks↩︎

Figure 8 extends our evaluation to the Safety MuJoCo suite, which tasks agents with locomotion control under velocity-based constraints. Consistent with the navigation results, SteinGate demonstrates superior constraint adherence, rapidly converging to feasible policies while maintaining competitive forward velocities. In contrast, CPO suffers from significant instability, characterized by large cost oscillations and frequent safety violations, particularly in the complex dynamics of Ant and HalfCheetah. While EVO succeeds in maintaining safety, it again exhibits excessive conservatism, resulting in slower convergence and lower asymptotic returns, most notably on Humanoid. Although state-augmentation methods (Saute, Simmer) improve upon CPO’s safety profile, they frequently operate precariously close to the boundary with intermittent violations. These results confirm SteinGate’s robustness, demonstrating that its safety mechanism generalizes effectively from spatial navigation to velocity-constrained locomotion without sacrificing task performance.

12.2 Additional Baselines↩︎

Figure 9 expands our comparative analysis to include Lagrangian methods (PPO-Lag, TRPO-Lag) [39], [40] and PPO-based implementations of state augmentation techniques (Saute-PPO, Simmer-PPO) [24], [25]. While the main text utilizes TRPO-based baselines to align with the backbones of CPO, EVO, and SteinGate, these additional comparisons isolate the impact of the underlying optimizer from the safety mechanism itself. The Lagrangian baselines prioritize reward maximization but fail to strictly enforce constraints, resulting in sustained operation above the cost limit. Although switching the backbone of Saute and Simmer to PPO reduces the magnitude of these violations compared to Lagrangian approaches, they still exhibit frequent excursions beyond the safety boundary. In contrast, SteinGate demonstrates superior constraint adherence without requiring the hyperparameter sensitivity often associated with Lagrangian multipliers, maintaining cost trajectories that tightly track the prescribed limit while delivering competitive task performance.

Comparison with CPO Variants. Figure 10 compares SteinGate with standard CPO and two of its variants, RCPO [30] and PCPO [28], across four Safety Gymnasium tasks. Although all methods share the same trust-region constrained optimization backbone, their behaviors differ substantially. Standard CPO and RCPO achieve relatively high returns but exhibit sustained cost levels close to or above the constraint, indicating frequent violations during training. PCPO also fails to consistently satisfy the constraint and shows noticeable excursions above the cost limit across environments. In contrast, SteinGate maintains systematically lower episodic cost while achieving competitive return. These results indicate that the improvement provided by SteinGate is not simply due to the choice of backbone optimizer, but arises from the proposed Stein-based gating mechanism, which offers more reliable control of constraint satisfaction than alternative variants of CPO.

13 Experimental Details↩︎

13.1 Sample Size and Training Dynamics↩︎

Our experimental framework is grounded in the OmniSafe library [41]. We implement SteinGate by extending the OmniSafe version of CPO. Similarly, we utilize standard OmniSafe implementations for baselines such as Saute and Simmer. For EVO, we adopt the official standalone implementation; however, we note that it is also constructed upon the omnisafe CPO. To ensure a strictly fair comparison, we maintained the default architectural hyperparameters and interaction budgets used for the baselines. In on-policy algorithms, the number of cost samples available for a safety update is determined by the steps_per_epoch (interaction budget per update) and the environment’s episode_length. To accelerate training on our hardware, we used vector_env_nums=8 (eight parallel environments). Given the task parameters, this resulted in an effective sample size of \(N = 16\) cost episodes per epoch for Goal tasks (episode length 1000) and \(N = 40\) for Circle tasks (episode length 500).

A core advantage of SteinGate is its ability to operate without the extensive replay buffers or importance resampling required by EVT-based methods such as EVO. To verify that the safety guarantees of SteinGate hold independently of the specific OmniSafe configuration, we conducted an ablation on CarGoal1 by modulating the steps_per_epoch parameters to obtain batch sizes of \(N \in \{8, 16, 32\}\).

It is important to note that varying \(N\) in this framework implicitly changes the number of policy updates per total environment steps (e.g., larger batches imply fewer total updates for a fixed interaction budget). Despite this confounding factor, the results in paragraph [pg:sample95size95ablation] show SteinGate effectiveness. This ablation confirms that SteinGate is not brittle with respect to sample size and can be integrated into standard on-policy implementations (e.g., CPO) without requiring the collection of off-policy data to stabilize the safety estimator.

13.2 Environment Descriptions↩︎

We evaluate SteinGate on a collection of continuous-control environments drawn from the Safety-Gymnasium benchmark suite, which is built on the MuJoCo simulator [22], [42]. These environments are designed to study reinforcement learning under explicit safety constraints by associating penalties with hazardous states or behaviors.

Point and Car Tasks. The Point and Car agents are evaluated on both Goal and Circle variants. In the Goal tasks, the agent must navigate toward a sequence of target locations while avoiding hazardous regions in the environment. After each successful reach, a new goal position is sampled, requiring continual navigation under safety constraints. In the Circle tasks, the agent is encouraged to follow a predefined circular trajectory. Deviations outside the designated safe region incur costs, creating a trade-off between tracking accuracy and constraint satisfaction. The Car variants introduce additional control complexity through nonholonomic dynamics, making safe navigation more challenging than for the Point agent.

Velocity-Constrained Locomotion Tasks. We additionally consider velocity-limited locomotion tasks for articulated agents, including Ant, HalfCheetah, Swimmer, and Humanoid. In these environments, the agent is rewarded for forward progress while being penalized for exceeding a specified velocity threshold. The safety constraint therefore requires the agent to coordinate its movement to achieve steady locomotion without violating speed limits. These tasks capture a different class of safety considerations, where constraints apply to the dynamics of motion rather than spatial regions.

13.3 Hyperparameters↩︎

All experiments were conducted using the default hyperparameter configurations provided by the OmniSafe library for the corresponding baseline algorithms. We did not perform task-specific tuning or additional hyperparameter searches. This choice ensures a fair and reproducible comparison and isolates the effect of the proposed SteinGate mechanism from confounding optimization or architectural choices. Unless otherwise specified, network architectures, learning rates, batch sizes, and interaction budgets follow the standard OmniSafe settings.

For SteinGate, the practical implementation introduces a number of parameters related to the construction of the Stein certificate. These parameters are shared across all tasks and are reported in Table 2.

For EVO, the official implementation does not fully specify several hyperparameters related to buffer usage and percentile thresholds. In particular, the following parameters were not documented in the released code: use_cost_buffer, use_reward_buffer, cost_buffer_size, reward_buffer_size, initial_distribution_percentile, initial_distribution_reward_percentile, cost_buffer_percentile, reward_buffer_percentile, and oversample_factor. We initialized a subset of these parameters based on the descriptions provided in the EVO paper and performed limited sweeps over the remaining unspecified values. From these runs, we selected the configuration that achieved the best trade-off between reward performance and constraint satisfaction. The final values used in our experiments are summarized in Table 3.

References↩︎

---

1. Equal contribution↩︎