SHAPO: Sharpness-Aware Policy Optimization for Safe Exploration
Authors: Kaustubh Mani
Yann Pequignot
Vincent Mai
Liam Paull
Venue: ICLR 2026 Poster (2026)
Keywords: Safe Exploration
Sharpness Aware Minimization
Epistemic Uncertainty
Primary area: reinforcement learning
TL;DR: SHAPO makes policy updates pessimistic under epistemic uncertainty, reweighting gradients toward rare unsafe actions to enable safer exploration and improved safety–performance trade-offs

Abstract: Safe exploration is a prerequisite for deploying reinforcement learning (RL) agents in safety-critical domains. In this paper, we approach safe exploration through the lens of epistemic uncertainty, where the actor’s sensitivity to parameter perturbations serves as a practical proxy for regions of high uncertainty. We propose Sharpness-Aware Policy Optimization (SHAPO), a sharpness-aware policy
update rule that evaluates gradients at perturbed parameters, making policy updates pessimistic with respect to the actor’s epistemic uncertainty. Analytically we show that this adjustment implicitly reweighs policy gradients, amplifying the
influence of rare unsafe actions while tempering contributions from already safe ones, thereby biasing learning toward conservative behavior in under-explored regions. Across several continuous-control tasks, our method consistently improves both safety and task performance over existing baselines, significantly expanding their Pareto frontiers.

Published as a conference paper at ICLR 2026
SHAPO: SHARPNESS-AWARE POLICY OPTIMIZATION FOR SAFE EXPLORATION
Kaustubh Mani
Yann Pequignot Université Laval I Institut Intelligence Données
Vincent Mai LawZero
Université de Montréal Mila Québec Al Institute
Liam Paull
Université de Montréal Mila Québec AI Institute CIFAR AI Chair
ABSTRACT
Safe exploration is a prerequisite for deploying reinforcement learning (RL) agents safety-critical domains. In this paper, we approach safe exploration through the lens of epistemic uncertainty, where the actor's sensitivity param- eter perturbations serves as a practical proxy for regions of high uncertainty. We propose Sharpness-Aware Policy Optimization (SHAPO), sharpness-awer pol- icy update rule that evaluates gradients at perturbed parameters, making policy updates pessimistic with respect the actor's epistemic uncertainty. Analytically we show that this adjustment implicitly reweighs policy gradients, amplifying the influence of rare unsafe actions while tempering contributions from already safe ones, thereby biasing learning toward conservative behavior in under-explored re- gions. Across several continuous-contie tasks, our method consistently improves both safety and task performance over existing baselines, significantly expanding their Pareto frontiers.
INTRODUCTION
Deploying reinforcement learning (RL) agents in safety-critical environments poses unique chal- lenges due to the need for safe exploration: the process by which an agent collects informative ex- perience while avoiding catastrophic failures (Bharadhwaj et al. [2200] Achiam et al.) 2017; Sootla iet al. 2022). A fundamental difficulty arises during early training, when the agent inevitably visits states about which it has limited knowledge, i.e., regions of high epistemic uncertainty. In these gions, the predictions of learned policies and value functions are unreliable, and exploratory actions may cause unsafe irreversible outcomes.
In practice, safe exploration depends on the policies an algorithm visits while learning the task, so the choice of update rule matters: each gradient step should hedge against unreliable estimates in parts of the state-action space visited under the current policy. This connects safe exploration to the agent's s epistemic cnceeraintyyuucereee data are scarce-suggesting that policy updates themselves should be pessimistic with respect to that epistemic uncertainty.
about model parameters that is largest where on-policy
Prior work has largely handled uncertainty on the critic side (ensembles (Lakshminarayanan et al., 2017 Osband al. 20166] Bykovets et al., 2022). ) dropout (Gal & Ghahramani (2016). or distribu- tional methods (Bellemare et al. 2017 Luis et al. 2024) combined with risk measures), which helps quantify variability returns but does not expose a practical, optimization-reade notion of actor's epistemic uncertainty. Maintaining an explicit posterior over policy parameters is generally infeasi- ble for deep on- on-policy RL, making actor-side risk aversion difficult practice. As result, while critic-side approaches common, observe gap in methods for policy updates that carefully account for the actor's parameter uncertainty.
"Corresponding author: kaustubb33995egmm com

---

Published as a conference paper at ICLR 2026
We address this gap by introducing novel approach to policy update guided by gradients evalu- ated at an adjusted parameter: Sharpness Aware Policy Optimization (SHAPO). Our method builds upon sharpness-awasd aware optimization (Foret et al. 2020; Kim et al. 2022b), which incorporates the sharpness of the loss landscape to enhance generalization. We reinterpret this adjustment through the perspective of epistemic uncertainty, demonstrating that it induces systematically pessimistic bias with respect to the actor's performance. In a simplified analytical setting, we show that SHAPO modifies the treatment low-probability actions: those deemed unsafe are assigned greater weight in the policy update, whereas those identified as safe are downweighted. Empirically, SHAPO ex- pands the aafety-efficiency Pareto frontier across range of on-policy Safe RL baselines in the Safety-Gym (Ray et al. 2019) and MuJoCo (Todorov et al. (012) environments, while simultane- ously reducing cumulative failures and mitigating heavy-tailed episodic cost distributions. Ablation studies further indicate that perturbations defined in the Fisher metric are superior to their Fuclidean counterparts, and that applying sharpness-awars optimization to the actor is more consequential for safe exploration than its application to the critic.
2 BACKGROUND
2.1 CONSTRAINED MARKOV DECISION PROCESSES
Constrained Markov Decision Process (CMDP) (Altman 2021 is tuple
M A, T, c, 7, po),
with state space S, action space A, transition kernel T, reward function r: SxA R, nonnegative
cost function : S x A + t.,, discount factor 7 E [0, 1), and initial state distribution μoo- A stationary policy induces discounted returns
J.(T) = E Σ
('r(ss, un) Ja(x) = E.
7o(ss, mt)
1= 0
0
and the objective is
π = argmax (,(T), I1: J((TT S 3}.
(I)
xElle
for constraint threshold 3 2 0.
common approach consists introducing a Lagrangian objective,
JA(TO) (mo) X( (JJ((To) - 3),
with multiplier& adapted to enforce the constraint.
We write (s, a) Ex (R) 9o 9, = a] for the action-value function associated with n, where [x(T) = Ee= 7 (r(s), 4z) - )c(s),a))) Similarly, we write VT² Ea-a( [+) (QX a)] for the value function. Finally, for a policy n, we write d= (s) for the discounted state distribution given by d= (8) >) Ei Pr(s sT
2.2 POLICY OPTIMIZATION
A standard way to learn policies that maximize expected return consists of iteratively improving a policy until convergence optimum. We provide here some notations and results from Kakade & Langford (2002): Schulman et al. (2015). Given current policy we seek to find a policy that improves on mo. The improvement of a policy и over No can expressed as:
(x(T) - (mo) Es-a" Ea-s(als) AQ (s,
where (s, a) = 2 a) (s)(s) quantifies the advantage of taking action a compared to the average action under policy mo.
This improvement is an expectation over the state distribution of the unknown policy π, so, in practice, replaced by the the following quantity that can be estimated using only experience collected with the current Np:
AS L (T) = Es-arr Eawro(als) To(a|s)
TT (als).
(3)
2

---

Published as a conference paper at ICLR 2026
However in order for positive L (T) effectively represent an improvement, need to ensure that π does not differ too much from xo As a consequence, to generate an updated policy म, the natural gradient is found as the solution the following optimization problem: problieemm...roommm.mm..
maximize L (T)
π subject to DKt. (Tu, n) 5
where DRL (To, TT) max, (L((T(-|ss is the maximum of Kullback-Leible divergence over states and d > 0 specifies the trust region constraint.
We recall the following result adapted from (Schulman et al. 2015 Theorem 1) that bounds the improvement (or deterioration)
policy æ over the current policy To in terms of L (T):
Theorem 1, For a = Di. (Tu x), we have
47 L (x) Ba S JA(T) (A(T)) S L20 (x) Ba with B, = aa (1 7)2
max (A) a)l.
s,a
Intuitively, this result indicates that the further is from (in the sense of DKL the more LA (T) must deviate from L (TO.]] to ensure true improvement (or deterioration).
We now consider parametrized policy wo, let mo Th and overload our previous nota- tion by using instead of me when there is no risk of confusion. Since the maximum in t0 KK (To, πe) is intractable, the expectation DRL (To, πT) E [[( is used
in practice. The second-order approximation of this function of around 6g is given by DRi.
x Fo, where Fe, denotes the Fisher Information Matrix defined by F&, Eewdro (a ()) (Vo log (a s))e==,)) log ( (Kullback 1997 p. 26). Now, ing linear approximation LA (0) g² of the loss function around 0g leads the following optimization problem:
UF., (g. 6) = arg max (g, c). 2 SSA
(5)
TRPO (Schulman et al. 2015 uses Vo (0) for the parameter update. We observe that relying on the gradient of the loss L may result in an overconfident update, as it does not take into account the epistemic uncertainty associated with mg-
2.3 SHARPNESS- AWARE OPTIMIZATION
Sharpness-Awe Aware Minimization (SAM) (Foret et al. 2020) is a general framework that aims to bias solutions towards flat regions of the loss utility landscape. In general terms, for utility function [(T) that we seek to maximize for π E IT, SAM framework proposes the following modified optimization problem:
max rell #EN(T)
min C(#),
where ((T) (T) is some choice of neighborhood of т. Given a parametrization with 0 E θ Eu- clidean neighborhoods in parameter space represent a simple choice that was studied by Foret al. (2000) and has also applied to RL (Lee & Yoon (025). Other neighborhoods based on functional similarity have also been proposed in the supervised learning setting. Fisher-SAM (Kim et al.) 2022b) proposes to define the neighborhood of a predictive distribution
x[|[x) in terms of expected Kullback-Leibh divergence (T) = { Ez ( < 8} for some 8 In practice, maximization this worst-case value
neighborhood is achieved by suc- cessive ascent parameter 0 following the gradient of the loss at L(I + c) where is such that the next section, we adapt the sharpness- aware optimization approach for policy optimization. Furthermore, we show that this method is equivalent to optimizing risk-aware objective in light of the epistemic uncertainty surrounding the current policy.
Tg4e R arg min (xx L(*) (see Alg. 2 appendix).
Note that T., (To) ESsdrT Banwal (e) A8, a) since (R(s,a) Q2(s, a) Barwug (+) 121 a)]
3

---

Published as a conference paper at ICLR 2026
3 SAFE EXPLORATION
On-policy deep RL algorithms for CMDPs begin with a random policy and iteratively update to
TT; T2: Ilss The objective to maximize (T(m) subject to ( S B. However, during
this process the agent incurs a cumulative training cost proportional to Σi Jc(TT)) The primary objective of safe exploration is therefore minimize this cumulative cost while still learning a policy that yields high expected reward.
A complementary objective is to avoid policies n; that, even if low-cost in expectation, occasion- ally produce trajectories with very high cost. Writing the discounted trajectory cost as Je(T) = 2t= 2'((st, az), the goal is control the tail of the distribution of Je(T), thereby reducing the frequency of rare but catastrophic failures.
We observe that these two objectives extend beyond merely finding a solution to the CMDP. They emphasize not only the importance of maximizing expected rewards while adhering to cost con- straints but also the critical need for the policies explored during training to avoid rare high-cost outcomes. In practice, selecting different values of À in the equation 2 results in a trade-off be- tween expected return and total cost. The focus of this work on safe exploration, which involves advancing this Pareto front.
Policy updates are crucial to the effectiveness of these methods. Given the inherent uncertainties associated with these updates, adopting pessimistic approach-especialle in the context of the agent's epistemic uncertanty-emmerge as a logical strategy for achieving safe exploration, which implement in this work. Specifically, we believe that safe exploration framework must effec- tively address rare and potentially unsafe events during the policy update process. By prioritizing this consideration, we aim to enhance both the reliability and safety of the learning process and its outcomes.
4 SHARPNESS AWARE POLICY OPTIMIZATION (SHAPO)
In the context of policy optimization, SAM framework (Foret et al. 2020 suggests to look for policy æ that maximizes
min (x( REN(TT
instead of just (A(x). This shift in objectives can be intuitively justified at high level follows. We may prefer policy over even if (A(T)) < Jx(22) if Ti achieves a higher value for the worst-case objective This preference may arise from the understanding that the expected return of policies within the neighborhood ((T:) can provide more accurate estimate of the actual perfor- mance of TT. This particularly relevant because our learning problem may not fully encapsulate real-world complexities, such environmental dynamics the real-world implementation of policy. By focusing on the worst-case expected return of nearby policies, we aim to better align our estimates with the true performance outcomes in realistic settings.
4.1 THE POLICY UPDATE
As outlined in the background section, goal is to lear policy that maximizes (x) by gradually updating the current policy Te,- This is achieved by solving the optimization problem under trust region constraint Sup and using the gradient g = V (0) e= We present here our proposed approach, which differs primarily in that it substitutes the direction g with the gradient of LA computed 0g (Down where Down is carefully chosen adjustment.
Given the current policy Te. we now outline our proposed method for updating it within the frame- work of SAM. We have at our disposal the utility function L (0) which estimates the relative improvement (or deterioration) of nearby policies Ty. Our first step therefore consists of solving the minimization problem
L (#). (EO(
min
We note that choosing an isotropic Euclidean neighborhood in parameter space for N (TT) presents certain challenges. A small distance in parameter space (II 0oll2 small) does not necessarily
4

---

Published as a conference paper at ICLR 2026
Algorithm SHAPO(B, L, F, dDown
Require: parameter f,, objective function L, Fisher matrix F, constraint dDwn
1: Compute the gradient at 8: 9 (o)(o)lo
2: Compute the perturbation: éDown = Ür(-9, 8Down)
3: Compute the gradient at perturbed parameter: 9 f VoL 9 F-wa
4: Return SHAPO gradient ğ
guarantee that my is similar to To in terms of DRL Consequently, when L (0) < 0, does not ensure that Jx(TT) < (TT). In contrast, while selecting neighborhood of distributionally similar policies better reflects the underlying geometry of the statistical manifold of the parameters (Kim) ct al. 2022b). it also contributes to ensure the validity of utility function L in virtue of Theorem
Consequently, to identify policy से that is distributionally close to To but performs worse than To. propose to formulate the inner minimization in manner analogous to the optimization problem namely as
minimize I. (#t
subject to DRL. (TT. F) < 6Dww
for some trust region constraint ōDwn 2 O. We therefore estimate the solution to the inner minimiza- tion with the policy मे = temws where the perturbation CDown UF,, (-V VoL (0))e=6,> dDwwn) is the solution to the optimization problem where Fa, is the Fisher information matrix of our policy parametrization evaluated Op.
Next we proceed to update To based on the direction provided by the gradient of I.A,, at our found perturbed parameter 0g Down- This step involves solving the maximization problem [5] where the gradient V VoL
L. (0) is now evaluated at 0p + EDown This results in a parameter update given by
'Up = UF., L 2 (0) òUp) where dup 2 0 is the trust region constraint. The pseudo-
code for computing the SHAPO gradients is given Algorithm [
4.2 REINTERPRETING FISHER SAM AS PESSIMISM IN FACE OF EPISTEMIC UNCERTAINTY
In effect, our proposed approach differs from classical local optimization techniques in that the update of the current parameter 0g is guided by the gradient evaluated at Og + <Dwnn instead of lg- We now provide different perspective the perturbation <Down, arising from pessimistic viewpoint in the context of epistemic uncertainty.
We propose to model the uncertainty about the parameter Oo. which defines the current policy wo Te, using a multivariate normal distribution (0)) = (N,,, Σ) with mean 0g and covari- ance matrix
F where F-1 denotes the inverse of the Fisher Information Matrix evaluated at 6g. The choice to model epistemic uncertainty this way is motivated by the asymptotic properties established by the Bernstein-von Mises Theorem (van der Vaart 1998) where n repre- sents the number of independent samples used to estimate the parameters of the model. However, it should be noted that the case of RI. this not simple to compute due to the correlation of samples, gather
Fo,
as we will elaborate in Section 4.4 This approach captures the intuitive notion that more data, our uncertainty about the parameter decreases, and the distribution of our estimates be- comes more concentrated around the true value. This choice not only aligns with theoretical founda- tions but also offers a practical starting point for exploring the implications of epistemic uncertainty in the context of policy optimization.
We observe that this uncertainty about the current parameter Op translates into an uncertainty about estimate
the expected return JA (To) of our current policy, which
(A( JA (To.) + LA (B)
when To is close to To in distribution. Writing = VOLA, (0))e=o, and assuming a first- order approximation L (0 + E) eTg around 6. the deviation of expected return is given by Y LA (0),0 (00, 1F-1) that follows normal distribution A/ o2) with variance
o² 1L-
5

---

Published as a conference paper at ICLR 2026
Log (9)
8
Opven
Figure 1:
SHAPO directs the policy update using the gradient I at an adjusted parameter (Right). Two perspectives on this adjustment are illustrated: (Left) The adjustment aims to minimize the expected return while staying within the trust region defined by duewn. (Middle) Parameter uncertainty translates uncertainty the expected return, resulting in maximum likelihood adjustment that remains below the o-quantile ya.
While L (0n) = represents the mean this distribution, parameters for which L (0) falls at, for instance, the 5th percentile of the distribution Y provide expected return. This highlights that, although the mean may indicate a certain level of performance, there are credible scenarios in which performance could be significantly lower. Therefore, adopting a pessimistic standpoint, it is justified to prioritize improving a more conservative expected return when updating Pg We demonstrate that evaluating the gradient at 6g + Down to inform the update of Oo can be interpreted as optimizing an estimate of the expected return that is pessimistic in face of the parameter uncertainty.
nuanced perspective the current
For a E (0, }), let denote the a-quantile of the random variable Y, and 2a denote the a-quantile of the normal distribution (). Note that under our first-order approximation of LA, the con- dition LA (9g + e) S Ya corresponds to e" S Ya- The following proposition links this epistemic uncertainty perspective with the parameter perturbation described in the previous section. It is illus- trated in Fig.
Proposition Lei Q(0) (Oo, ! F For every a E (0, 5),, the solution to to the optimization problem
maximize log ((u +
subject to e g S Va
coincides with the adjustment make in SHAPO for the trust constraint pDown
- namely
=2 Ea = Ur((g 2m
Intuitively, this result states that the adjusted parameter 0g + CDwn considered in SHAPO can be viewed as the most likely parameter under the parameter distribution Q(0) that falls in the lower tail the a-quantile. Our approach SHAPO, which updates the current policy To by using the gradient calculated at this adjusted parameter, can therefore be understood as promoting an enhancement of To in a worst-case or pessimistic scenario, considering the existing epistemic uncertainty sur- rounding mo. By applying SHAPO with ōDown > 0 to update the actor, we are therefore effectively enforcing pessimism in face of epistemic uncertainty about our current policy.
Furthermore, by establishing the relationship 5pown 2 this proposition allows to interpret the constraint 6Down as measure of confidence in the expected return of the current policy in the presence of epistemic uncertainty. More details and proof in Appendix B
2
4.3 ANALYSIS OF SHAPO ON SIMPLE GAUSSIAN POLICY
Our proposed approach to update the current policy To is guided by the gradient ğ of LA (0) evalu- ated at 0n Down. while classical approaches use the gradient g evaluated at 0g. Here we study how these two gradients g and ğ differ in simplified setting.
We assume that the policy some state is given by a 1D Gaussian policy 1,; Na(a; p, 1) parametrized by the mean μl and that the current policy is (a) = (a; 0, specified by = O. We propose to study the discrepancy between the two gradients g and ğ with respect to l in terms of an observed action and an advantage value A under a constraint opowa. Our utility function L and its gradient can now be expressed as:
A n T(a VaLluga,Ar A)= To(a) L(p; a, A) To(a)
A
v (a;(4;

---

Published as a conference paper at ICLR 2026
Current policy distribution
With advantage A
With advantage A-10.0
0.40
sanndrd gradient SHAFO gradient
66
standard gradient
0.25
SHAPG gradient
20-
0.2
50
0.20
-20
0.20
-20
0.55
-40
action caserved
action observed
Figure 2: 1D Gaussian policy: The gradient utilized in SHAPO exhibits larger amplitude for rare and unsafe actions, while it comparatively smaller rare safe actions, compared the classical gradient.
The standard gradient g(a, A) V a, 1)-o and the gradient ő(a, A) V Inia a, A) at j = 140 + DDwn are as given in the following Proposition (See Appendix C for details).
Proposition 3. Under the 1 I Gaussian model, we have
a -2
with A Sa p Aa <0
A(a A)= =A and J(a, A) = exexp
a a+p P otherwise,
where p Down-
Figure 2 show plots of the two gradients functions of the action observed a for two values of advantage (A = 10 and A 10) with = 0.1 (i.e. 6Dow = 0.005).
We that SHAPO gradients differ from standard gradients for actions |a| > 1. When the ad- vantage is negative, indicating an unsafe action, we see that SHAPO's gradient larger compared to the standard gradient. contrast, for positive advantage SHAPO's gradient is smaller than the stan- dard gradient. The analysis of this simplified setting shows that, under certain conditions, SHAPPOs policy update promotes different treatment of rare actions depending on the sign of the corre- sponding advantage. When unsafe, rare actions are taken very seriously and an important update of the policy is promoted. In comparison. SHAPO opts for smaller updates when facing safe actions, even when they have large advantage. Overall, this behavior aligns well with our safe exploration objectives and helps explain why our approach favors policies that infrequently incur high costs (cf. Figure 5)
4.4 SHAPO FOR ON-POLICY RL. ALGORITHMS
As presented in this section, Sharpness Aware Policy Optimization (SHAPO) modifies the TRPO algorithm with a Lagrange objective and pseudo-code is provided in Algorithm 3 We abstract this modification in Algorithm [ this code returns a parameter direction for updating the current policy Ta based on the objective function L, the Fisher matrix F of the policy parametrization and constraint bound Down- Formulated this way, method can be applied to most online policy RL algorithms and we provide pseudo-code in Appendix for for CRPO and CPO.
While Proposition 2 the trust-region constraint 6Dwn to the percentile z, of the normal distribu- tion (where a is the level of pessimism in face of uncertainty), the effective sample size n cannot be easily estimated in continuous-control reinforcement learning where data are temporally correlated and only subset of collected samples provide independent information. Furthermore, due to the complex coupling between the actor and critic, together with the challenges posed by exploration, it is difficult to determine the appropriate level of pessimism (i.e. the value of a)) to adopt at each stage of the training.
In this context, adopting constant value of 9Down during training and treating this value tunable hyperparameter allows us to directly control the desired safety- efficien tradeoff without relying on uncertain estimates of n. This makes the constant-Opewn strategy both simple and effective, offering a stable mechanism for regulating pessimism throughout training. We further discuss this approach and other options to schedule the level of pessimism in Subsection 6.1 and in Appendix E
Finally, in our method, we also incorporate Sharpness Aware Minimization with Euclidean neighbor- hoods for the critic, proposed by Foret al. (2020) for supervised learning task. By considering
7

---

Published as a conference paper at ICLR 2026
PointGoal1
PointButton1
Ant-p4
Walker26-v4
12
5400
25.3
10
4008
Retum
Return
2009
2003
1530
1039
1003
10
Cost Rate
1e-3
12
Cost Rate
1a
3.
653
133
750
200
Tallures)
2508
500
1500
4003
faluurss)
SC0
CRPO
PiDLag
CPO
SauteRL
Vanilla
SHAPO
Figure 3: Performance of SHAPO on four different tasks across four different baselines over 5 training seeds. We report Cost Rate (Total Cost / Total Env Steps) for Safety Gym environments after training for 10M vironment steps and Total Cost MuloCo environments after 5M environment steps. SHAPO (shown with a star) significantly improves the performance of the baseline algorithms (shown with a circle) on both safety and efficiency thus improving the safe exploration capabilities of the baseline algorithms. Width and height of the cllipses represent the standard error over 5 seeds.
PointGoal
PointButton1
Vanilla
Vanilla
Ant
Walker2d
Vanila
Recum
Return
Epteade Cost
te
31
to
1.3
18
Steps
Mean
80th Percentile
95th Percentile
Avs-10
Avs-100
Figure 5: SHAPO results in policies with episodic cost distributions with smaller tails (80th percentile lambda (Ainie) for PIDLag algorithm with and with- and 95th percentile) crucially for safe exploration dur- out SHAPO. We can see that SHAPO improves upon ing the learning phase. the Vanilla PIDLag for all configurations.
Figure 4: Results for different values initial
the sharpness of the loss landscape during the training of the critic, we aim to focus on flatter regions that yield more stable value function approximations.
EXPERIMENTAL SETUP
Environments. We evaluate SHAPO on Safety Gym Ray et al., 2019) (PointGoalI-vO, PointButton/-vO) and MuloCo (Todorov et al., 2012) (Ant-v4, Walker2d-v4), In Poin an agent must reach a random goal while avoiding unsafe regions that accumulate cost; in PointButton v0, it must press buttons in the correct order, with penalties for wrong presses and additional hazards. In both tasks, the objective is to succeed while keeping cost under B = 10. In MuJoCo, the agent must run fast without falling, where each fall is a constraint violation.
Baselines. Although Sec. 4 more specifically describes how our approach modifies TRPO (Schul- man al. 2015) with Lagrangian objective, SHAPO virtually applies any on-policy method in straight manner. We therefore also study the effect of SHAPO on CPO (Achiam et al. 20171. PIDLag (Stooke et al., 2020) (Lagrangian TRPO with PID control of A). CRPO (Xu et al, 2021 (alternating primal updates) and SauteRL. (Sootla et al., 2022) (state augmentation with con- straint budget with TRPO (Schulman al. 2015) as base algorithm). Pseudo-code for the SHAPO versions of the baseline algorithms is provided in Appendix A
Metrics. In MuJoCo we report episodic return (efficiency) and total failures/cost (safety). In Safety Gym track episodic return and cost during training, plus final average return and cost-rate (total cost / steps) as defined in (Ray et al. 2019). All results are generated by averaging over seeds, the standard error in (Cost Rate 7 Total Cost) and y (Return) are represented by the width and height of the ellipses, respectively.
Implementation Details. When adding SHAPO, baseline's hyperparameters frozen and only 6 and Peritic are tuned. Thus the safety-efficieee tradeoff is determined by the baseline, while SHAPO provides orthogonal improvement. More details in the Appendix D
8

---

Published as a conference paper at ICLR 2026
Antov4
Walker
PointGoall
PointButton1
1300
40:11
2.3
Retum
2500
1200
2000
03-
Otal Cost
2500
3000
Total Cost
3000
4230
4000
-03
RDDaag
Cost Rate
Ix-2
Rate
Le-2
Verilla
SAM Citcc
SHAPO Actor
SHAPO
Figure 6: We compare results for SHAPO when using perturbation in the euclidean space E-SAM or the Fisher Metric space (SHAPO). We can see that SHAPO is consistently better than E-SAM because of perturba- tion that ensure worsening of the inner objective.
Figure 7: SHAPO leads to significant improvement in cost-rate, even without SAM critic. While SAM on critic leads to improvement return, only adopting SAM for the critic does not promote safe exploration in PointGoall and PointButton
RESULTS
Fig. summarizes performance across four baseline algorithms and their SHAPO counterparts on four tasks. In PointGoall and PointButtonI, we plot cost-rate (x-axis) against return (y-axis), captur- ing the safety-efficiency tradcoff during training. For Ant-v4 and Walker2d-v4, we report cumulative failures (total cost). Across all benchmarks, SHAPO consistently improves either safety, task per- formance, or both. For instance, SauteRL. (Sootla et al. (022) achieves the lowest cost-rate on PointGoall and SHAPO further reduces this cost without hurting returns. Since SauteRL relies on augmenting states with the remaining budget, is inapplicable in MuJoCo (where 3 = 0 means any incurred cost violates the constraint), so we omit those results. Detailed results have been presented
in the Appendix Fig. 14 16
The safety- efficeeeffccieney tradeoff in Lagrangian methods is shaped by hyperparameters such as the initial multiplier dinit- Fig. 4 shows results for PIDLag across different Aénit values. In every setting, SHAPO expands the Pareto frontier of return versus cost, indicating that it consistently delivers better safety-efficiency tradeoffs than the vanilla algorithm.
To assess safety beyond mean statistics, Fig. 5 shows episodic cost distributions. While both vanilla baselines and SHAPO reduce average cost, only SHAPO consistently suppresses heavy tails (95th percentile in Ant-v4, 80th in Walker2d-v4.. This tail suppression significantly enhances the agent's safety and aligns with our safe exploration objective. Consistent with our analysis in Sec. 4.2) SHAPO's pessimism when faced with parameter uncertainty reduces catastrophic events, yielding fewer cumulative failures (i.e. Total Cost) (Fig. 3). Detailed plots are included in the Appendix.
6.1 ABLATIONS & ANALYSIS
Fisher versus Euclidean perturbations: Fig. 3 compares perturbations in Euclidean (E-SAM) and Fisher (SHAPO) spaces. SHAPO consistently outperforms E-SAM across tasks and baselines, vali- dating that Fisher perturbations, being aligned with the local geometry of on-policy data, yield more principled and effective risk-sensitive updates.
Actor versus critic perturbations: Fig. 7 isolates the effect of applying SHAPO to different com- ponents. When SAM is applied only to the critic (SAM Critic), perturbations do not enforce risk- aversion and can even increase cost, since squared-error perturbations treat under- and overesti- mation symmetrically. In contrast, applying SHAPO only to the actor (SHAPO Actor) markedly improves safety by steering exploration conservatively. Finally, Using SHAPO (SHAPO Actor SAM on critic) yields the strongest gains in both safety and efficiency.
2 we consider for each update step t the number ne of state-action pairs collected so far by our agent as a proxy for the effective sample size n from Subsection 4.2 This allows us to interpret the bound hown at step / in terms of n, and the percentile Za: of the distribution (0, 1) at the level of pessimism d₁. Our simple strategy using a fixed ō therefore can be interpreted as increasing the level of pessimism as the agent learns. Indeed, we then have 2., = - / 22n and thus a₁ = Φ( - 22n00uwn). where & is the cumulative distribution
Pessimism schedule: In light of Proposition
Wal is more challenging, with the th percentile remaining near instead.
report the 80th

---

Published as a conference paper at ICLR 2026
function of WN 1). Another simple approach instead consist of choosing fixed value of = pessimism a E (0, 2) for the whole learning process, and then compute the bound Bown = 2n, at step t, which effectively decreases as we progress through learning. Fig. 8 compares these two scheduling pessimism. Figure X shows that both of these schedules allow SHAPO to consistently outperform the vanilla algorithm, highlighting the robustness of the proposed adjustment with respect the choice of Down which controls the level of pessimism. More details and analysis are provided in Appendix
different approaches
7 RELATED WORK
PointGoall
PointButtoon1
Uncertainty RL: RL agents must con- tend with aleatoric (outcome) and epis- temic (model/parameter) uncertainty (Os-) band et al. 2016, Ghavamzzdeh et al.) 2015). Bayesian model-based methods propagate posterior uncertainty for explo- ration and risk assessment (Ghavamzadeh al. 2015); model-free approaches use ensembles/randdmizze value functions for epistemic estimates (Osband et al.] 2016 2018 Mai et al.) (022) and distributional critics for risk-sensitive control (Bellemare et al. 2017] Dabney et al., (018)) Pes- simistic objectives further bias learning to- ward conservative estimates (Kumar et al.) 2020
2.5
Return
Return
1.0
1005 Lo00 3.02 1:300 106 1.993
L065
107
1.08
1.00
1.10
3.11
Cost Rate
Cost Rate
Vanilla
SHAPO Ifixed Baseel
SHAPO (fixed
Figure 8: Pessimism schedule. We compare the selection
a constant ōDwa with the selection of a value a E (0, 0.5)
using the schedule & =2 where n,, is the number of state-action pairs seen by the agent. Both approaches out- perform the vamilla approach.
Sharpness-Aware Optimization: Sharpness-awarr objectives minimize the worst-case (or stochas- tic) loss in local noisy parameter neighborhood (Rahn et al. 2023), yielding solutions that are insensitive to weight perturbations and thus more robust to parameter uncertainty (Foret et al. 2020 Andriushchenko Flammarion 2022, Möllenhoff & Khan 2022). Geometry-aware variants (won al. [222]] Kim et al. (2022a) tailor the neighborhood to better approximate uncertainty directions, while noisy-neighborhood extensions explicitly inject randomness into the perturbation to probe loss under parameter noise and improve robustness (Back al. 2024). (Lee & Yoon 2025) shows that applying SAM to PPO significantly improves robustness to aleatoric variations of the environment test time. In this paper, focus instead the potential of sharpness aware optimization for safe exploration.
Safe Exploration: Safe exploration deals with epistemic uncertainty associated with data scarce regions of the state-space(Achhan et al., 2017] Liu et al.] [002]] SSotla et al. [202]] Stooke et al. 2020 Mani et al., 20255] Gu et al. (2024). Strategies include Bayesian model-based control with
uncertainty-aware planning (Berkenkamp) 2019), Lagrangian world models (Huang et al., 2023)
safety critics that filter risky actions or impose constraints (Fisac al, 2019: Srinivasan al. 2020 Kang et al.] (2022; [Bharadhwaj et al. (020)) and state augmentation with accumulated cost as a risk proxy (Sootla et al., 2022] Jiang et al. 2023). In this paper, we are interested specifically in epistemic uncertainty of the actor and in designing update rule to be pessimistic with respect to that uncertainty.
CONCLUSION
Safe exploration is ultimately about shaping updates to respect uncertainty, not only enforcing con- straints on the final policy. By instantiating SAM the actor's Fisher/KL geometry, we treat sharp- ness- ness-the sensitivity of the policy surrogate to tiny parameter changes a practical proxy for epistemic uncertainty. When small nudges can flip outcomes, the local neighborhood is under- supported by data; SHAPO biases learning toward locally stable policies whose nearby variants behave similarly, tempering brittle, risk-seeking moves and emphasizing improvements that reliably reduce cost. This creates closed loop: safer updates yield safer trajectories, which produce cleaner data in risky regions, reducing epistemic uncertainty and naturally relaxing conservatism as learning progresses.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
All code, configuration files, and exact hyperparameters required to reproduce our results publicly available 813E
ftttss://anoymmone www.onmoous.ooen.......... sorpenssccenne//ee/ Experiments were run in standard, publicly ac- cessible environments (OpenAT Gym, Safety Gym), with version details documented in the repository.
Safe-Policy.
ACKNOWLEDGEMENTS
This work is supported by the DEEL Project CRDPJ 537462-18 funded by the Natural Sciences and Engineering Research Council Canada (NSERC) and the Consortium for Research and Innovation in Aerospace in Québec (CRIAQ), together with its industrial partners Thales Canada inc, Bell Textron Canada Limited, CAE inc and Bombardier inc
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303. 08774, 2023.
Joshua Achiam, David Held, Aviv Tamar, and Pieter Abbeel. Constrained policy optimization. In International Conference Machine Learning, pp. 22-31. PMLR, 2017.
Eitan Altman. Constrained Markov decision processes. Routledge, 2021.
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-awee mini- mization. In International conference on machine learning, pp. 639-668. PMLR, 2022.
Christina Back, Zico Kolter, and Aditi Raghunathan. Why is SAM robust to label noise? arXiv preprint arXiv:2405.0361 006666666,, 2024.
Mare Bellemare, Will Dabney, and Rémi Munos. distributional perspective on reinforcement learning. In International Conference Machine Learning, pp. 449 458. PMLR, 2017.
Felix Berkenkamp. Safe exploration in reinforcement learning: Theory and applications in robotics. PhD thesis, ETH Zurich, 2019.
Homanga Bharadhwaj. Aviral Kumar, Nicholas Rhinehart, Sergey Levine, Florian Shkurti, and Ani- mesh Garg. Conservative safety crities for exploration. arXiv preprint arXiv:2010.144: 2020.
Eugene Bykovets, Yannick Metz, Mennatallah El-Assady, Daniel A. Keim, and Joachim M. Buh- mann. How to enable uncertainty estimation in proximal policy optimization. arXiv preprint 1wXiv:221000336:4 2022.
Will Dabney, Georg Ostrovski. David Silver, and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. In International conference on machine learning, pp. 1096- 1105. PMLR. 2018.
Jaime Fisac, Neil Lugovoy, Vicenç Kubies-Royo Shromona Ghosh, and Claire Tomlin. Bridg- ing hamilton-Jacee Jacobi safety analysis and reinforcement learning. In 2019 International Conference on Robotics and Automation (ICRA), pp. 8550-8556. IEEE, 2019.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. arXiv preprint arXiv:2010.01412, 2020.
Yarin Gal and Zoubin Ghahramani. Dropout bayesian approximation: Representing model uncertainty in deep learning. In Proceedings of the 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 1050-1059. PMLR, 2016.
https://deel. quebec
11

---

Published as a conference paper at ICLR 2026
Mohammad Ghavamzadeh, Shie Mannor, Joelle Pineau, Aviv Tamar, al. Bayesian reinforcement learning: A survey. Foundations and Trends@ in Machine Learning, 8(5-6):359 483, 2015.
Shangding Gu, Long Yang, Yali Du, Guang Chen, Florian Walter, Jun Wang, and Alois Knoll. A review safe reinforcement learning: Methods, theories and applications. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.
Weidong Huang, Jiaming Ji, Borong Zhang, Chunhe Xia, and Yaodong Yang. Safe dreamerv3: Safe reinforcement learning with world models. arXiv preprint arXiv:2307.07776. 2023.
Jiaming Ji, Jiayi Zhou, Borong Zhang, Juntao Dai, Xuehai Pan, Ruiyang Sun, Weidong Huang, Yiran Geng, Mickel Liu, and Yaodong Yang. OmniSafe: An infrastructure for accelerating safe reinforcement learning research. Journal of Machine Learning Research, 25(285):1-6, 2024.
Hao Jiang, Tien Mai, Pradeep Varakantham, and Minh Huy Hoang. Solving richly con-
strained reinforcement learning through state augmentation and reward penalties. arXiv preprint arXiv:2301.11592, 2023.
Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning. In Proceedings of the nineteenth international conference on machine learning, pp. 267-274, 2002.
Sham M Kakade. A natural policy gradient. Advances in neural information processing systems, 14, 2001.
Katie Kang, Paula Gradu, Jason J Choi. Michael Janner, Claire Tomlin, and Sergey Levine. Lya- punov density models: Constraining distribution shift in bearning-based control. In International Conference on Machine Learning, pp. 10733. PMLR, 2022.
Minyoung Kim, Da Li, Shell X Hu, and Timothy Hospedales. Fisher sam: Information geometry and sharpness aware minimisation. In International Conference on Machine Learning, pp. 11148 11161. PMLR, 2022a.
Minyoung Kim, Da Li, Shell Hu, and Timothy Hospedales. Fisher sam: Information geometry and sharpness aware minimisation. In International Conference on Machine Learning, pp. 11148- 11161. PMLR, 2022b.
Solomon Kullback. Information theory and statistics. Courier Corporation, 1997.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances neural information processing systems, 33:1 1979919991 2020.
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness-
aware minimization for scale-invariant learning of deep neural networks. In International confer- ence on machine learming, pp. 5905-5914 14. PMLR, 2021.
Balaji Lakshminarayanann Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances in Neural Information Processing Systems, volume 30, 2017.
Hyun Kyu Lee and Sung Whan Yoon. Flat reward policy parameter space implies robust rein- forcement learning. In The Thirteenth International Conference Learning Representations, 2025.
Zuxin Liu, Zhepeng Cen, Vladislav Isenbaev, Wei Liu, Steven Wu, Bo Li, and Ding Zhao. Con- strained variational policy optimization for safe reinforcement learning. International Confer- ence on Machine Learning, pp. 13644-13668. PMLR, 2022.
Carlos E Luis, Alessandro G Bottero, Julia Vinogradska, Felix Berkenkamp, and Jan Peters. Value- distributional model- based reinforcement learning. Journal of Machine Learning Research, 25 1298):1-42, 2024.
Vincent Mai, Kaustubh Mani. and Liam Paull. Sample efficient deep reinforcement learning via uncertainty estimation. arXiv preprint arXiv:2201.0166m 2022.
12

---

Published as a conference paper at ICLR 2026
Kaustubh Mani, Vincent Mai, Charlie Gauthier, Annie Chen, Samer Nashed, and Liam Paull. Safety representations for safer policy learning. arXiv preprint arXiv:2502.20341, 2025.
Thomas Möllenhoff and Mohammad Emtiyaz Khan. Sam as an optimal relaxation of bayes. arXiv preprint arXiv:2210.01620, 2022.
lan Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep exploration via bootstrapped DQN. Advances in Neural Information Processing Systems, 29, 2016.
Ian Osband, John Aslanides, and Albin Cassirer. Randomized prior functions for deep reinforcement learning. Advances in Neural Information Processing Systems, 31, 2018.
Nate Rahn, Pierluca D'Oro, Harley Wiltzer, Pierre-Luc Bacon, and Marc Bellemare. Policy opti- mization in noisy neighborhood: On return landscapes in continuous control. In A. Oh, T. Nau- mann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 30618-30640. Curran Associates, Inc., 2023.
Alex Ray, Joshua Achiam, and Dario Amodei. Benchmarking Safe Exploration in Deep Reinforce- ment Learning. 2019.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp- 1889-1897, Lille, France, 07-09 Jul 2015. PMLR.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.005
2017.
Aivar Sootla, Alexander Cowen-Rivers, Taher Jafferjee, Ziyan Wang. David Mguni, Jun Wang, and Haitham Ammar. Sauté RL: Almost surely safe reinforcement learning using state augmen- tation. In International Conference on Machine Learning, pp. 20423-20443. PMLR, 2022.
Krishnan Srinivasan, Benjamin Eysenbach, Sehoon Ha, Jie Tan, and Chelsea Finn. Learning to be safe: Deep RL with safety critic. arXiv preprint arXiv:2010. 14603, 2020.
Adam Stooke, Joshua Achiam, and Pieter Abbeel. Responsive safety reinforcement learning by PID Lagrangian methods. In International Conference on Machine Learning, pp. 9133-9143. PMLR, 2020.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference intelligent robots and systems, pp. 5026-5033. IEEE, 2012.
Mark Towers, Ariel Kwiatkowski, Jordan Terry, John U Balis, Gianluca De Cola, Tristan Deleu, Manuel Goulão, Andreas Kallinteris, Markus Krimmel, Arjun KG, et al. Gymnasium: standard interface for reinforcement learning environments. arXiv preprint arXiv:2407.17032. 2024.
AW van der Vaart. Asymptotic statistics. cambridge series in statistical and probabilistic mathemat- 3. cambridge univ. press, cambridge. 1998.
Tengyu Xu, Yingbin Liang, and Guanghui Lan. CRPO: new approach for safe reinforcement learning with convergence guarantee. In International Conference on Machine Learning, pp- 11480-1 1491. PMLR, 2021.
13

---

Published as a conference paper at ICLR 2026
Algorithm 2 Sharpness-aware optimization
Initialize /
2: while not converged do
3: Compute such that TØ+e arg min (
4: Calculate the gradient at perturbed parameter: g V +
5: Update parameters: 0 <- 0 + ag {a is the learning rate]
6: end while
Algorithm SHAPO (TRPOLag): Sharpness Aware Policy Optimization (TRPOLag)
1: Initialize parameters 0
2: Set trust region constraint dup and SHAPO constraint 9
3: while not converged do
4: Compute the Fisher matrix F at 0
5:
Compute the gradient: g <- VoL mo (0)10
6:
Compute the perturbation: <Down UF( 15((-9. 9, ŌDown)
7:
Compute the gradient at perturbed parameter: f- Vo VoL (0)
8:
Compute the parameter update: étip = (~ Öup)
9:
Update parameters: 0 + CUp
10:
end while
Return policy m
Algorithm 4 SHAPO (PPO): Sharpness-Aware Policy Optimization (PPO)
1: Initialize policy parameters
2: Set clipping parameter and SHAPO constraint 9
3: Set learning rate
4: while not converged do
5: Compute the Fisher matrix at 0
6: Compute SHAPO gradient for the clipped objective: ğ k C SHAPO LPO. FF, ddown)
7: Update parameters with the SHAPO gradient: 0 + 0 + 17..
8: end while
A PSEUDO-CODE
Pseudo-code for Sharpness Aware Optimization presented in Algorithm 2 and for our Sharpness Aware Policy Optimization for Lagrangian TRPO (SHAPO TRPOLag) in Algorithm 3
Recall that (r(g, 5) stands for the (estimated) solution to the optimization problem
Ur(g. 8) = arg max(g, c). F:58
(10)
where 5 >> is vector and F is symmetric positive definite matrix. In TRPO at current do the gradient of L (0) evaluated at 6g and is the Fisher information matrix Fs, of the policy parametrization Te. 0 E θ, evaluated at 6g-
We provide pscudo-code for SHAPO in Algorithm I Formulated this way, method can be ap- plied to most on-policy RL algorithms. It returns a parameter direction for updating the current policy Na based the objective function, the Fisher matrix of the policy parametrization and con- straint bound p This function effectively abstracts lines present how our method applies other algorithms.
Algorithm and allows us to
While TRPOLag combines cost and reward into Lagrangian objective LA, at the current policy mo (ef. Equation other algorithms like CPO CRPO consider distinct objectives. Given the
14

---

Published as a conference paper at ICLR 2026
Algorithm 5 SHAPPO-RPPO
Initialize parameters 0
2: Set trust region constraint ÓUp and SHAPO constraint O
3: while not converged do
4: Compute the Fisher matrix F at 0
5: if Expected Cost constraint threshold (3) then
6: Compute SHAPO gradient on Cost Surrogate 4 SHAPO(0, 1 F, ō
7: else
8:
Compute SHAPO gradient Reward Surrogate + SHAPO(0, Lay F, dDww))
9:
end if'
10:
Compute the parameter update: cUp Ür(0, dup)
11:
Update parameters: 0 +- + CUp
end while
13:
Return policy T
Algorithm 6 SHAPO-CPO
Initialize parameters 0
2: Set trust region constraint dup and SHAPO constraint 9
3: while not converged do
4: Compute Fisher matrix F' at 0
5:
Compute SHAPO gradient on Cost Surrogate Gr <- -SHAPO(O,, L &Down)
6:
Compute SHAPO gradient on Reward Surrogate Ge +- SHAPO(0,, L F,, ōDown)
7:
Resolve gradient update ÉUp <- CPO Update(ūr. Dc. dup)
8: Update parameters: 0 UUp
9:
end while
10: Return policy π
current policy To we define the two following objectives similarly to Equation 3
4 a) To (a (a(a)))
L (0) = Eg E. (a
(11)
A a) (a(a
L (0) Eg (awsn(als)
ITe (als),
(12)
where d denotes the discounted state distribution under vo and A, s) (resp.A9) quantifies the advantage in reward (resp. cost) of taking action in state compared to the average action under mo in state 8. More precisely, letting (s, E [T) = 8, ag = a] for reward have (s)
action-value function associated with To, where R,(T) Ei=o ( at),
(s, () Fo-raa (e) (s, a)) and similarly for cost with Re(T) = Eie 7 7"((st, ar).
While CRPO's update relies on a single gradient direction for either improving reward decreasing cost (see Algorithm 5 each step, CPO relies on both pieces of information through an update function that we denote by CPO_Update (see Algorithm (6) both cases, our method SHAPO provides gradient directions for both reward and cost following the same steps, thereby incorporating pessimism in face of epistemic uncertainty of the actor into these algorithms. Since SauteRL (Sootla any underlying policy gradient algorithm. We' ve used TRPO (Schulman et al. 2015) as the base algorithm for SauteRL. and the SHAPO version be derived from Algorithm In Algorithm 4 propose a simple approach to apply SHAPO on the popular PPO algorithm that uses the following clipped objective (for some > 0):
al. 2022) is state-augmentatton technique,
T. (PP(() = Eeruyo Eawnolale) min (r(0,a, a, (8, clip(r(0, 8), 1,, 1 A
where s))
xo(a|s)
Ra(a(s)*
15

---

Published as a conference paper at ICLR 2026
B REINTERPRETING FISHER SAM AS PESSIMISM IN FACE OF EPISTEMIC UNCERTAINTY
For the sake of completeness, prove first the following well known result (cf. Kakade (2001):
Schulman et al. ):15)):
Lemma 4. The solution to the optimization problem
ÚF(g,0)
arg (a, E)
(13)
3e Fe56
given by Ur(g,ö) symmetric positive definite.
22 1- g where F-1 is the inverse of which is assumed to be
Proof.
We introduce Lagrange multiplier for the constraint 3/²T Fe S å and define the following
Lagrangian:
A( (g. NC,eT TF Fc- 6).
We then find the stationary point L(e, A). First, we have:
0 v.L(E, = AFe
which leads to = F Next we have:
- 5,
Vx = Fe
that leads to the constraint _ F. = 5. Substituting F-1g into the constraint, and using the fact that F-1 is also symmetric, find:
92
g
212 =
7
X
26
Consequently, as desired:
28 F
=
F
Recall that Q()) is assumed to follow a multivariate normal distribution N 1F-1) where F-1 is the inverse of the Fisher matrix at Og. Recall that Y = g" (0 0o) with 0 ~ Q with a-quantiles denoted by ya- Finally 2a- denotes the q-qunntile for (0, 1). We now provide the proof Propo- sition 2 which is illustrated in Figure
Proposition 5. Let Q(+) = ((o.. )) For every a E (0, 2) the solution ca to the optimization problem
maximize log Q(0 + +
(14)
subject to g" c S yo
coincides with the aljustment we make in SHAPO, namely we have eDow = Up( Spown) when
-2
8pown
2n
Proof. Note that the precision matrix of Q(#) is equal to nF and
log Q(0o + E) - n Fe constant.
We introduce a Lagrange multiplier for the constraint gf S yo and define the following La- grangian:
n T Fe --0 -a).
L(e, A) 2
16

---

Siees recommended to decrease åDown according to ōDown -2 reflectnng the increasing trust in the estimates as more data becomes available. This strategy not only enhances the robustness of the Published as a conference paper at ICLR 2026
Next, we look for the stationary points of L. First,
v.L(e, -nFe Ag
which implies c c= -- Next we see that
= VxL(e,
7 g = ya.
Therefore,
A Va g9¹
1g Fr1 and
X
1yya
g² F-1g
Consequently, we obtain
Ya
F 9.
g
9
Now using the fact that Ya = aza where o with Q(0), get
9"
F
F-15
g is the standard deviation of = g'c
v
Za
F Frg
ng" F F-19
By Lemma we have
25
Up(- 9.6)
V g" -
F
9.
g
which is equal to e when
V28, namely when 5 = as desired.
22 2n
We observe that the relation dDwnn 2n offers an attractive perspective on the trust constraint for the inner minimization in SHAPO. This equation expresses the bound in terms the a-quantile of the standard normal distribution N(0, 1) and n. which represents the sample size used for estimating dg- For given sample size n, selecting confidence level-such the 5th percentile-yvildé, specific value for down. This approach implies that one can initially set Down = by determining an appropriate confidence level E (0, 2). As the actor gains more experience, it would appear
minimization process but also aligns the trust constraint with the actor's growing experience.
C ANALYSIS OF SHAPO ON A SIMPLE GAUSSIAN POLICY
Our proposed approach to update the current policy wo is guided by the gradient of T. (0) eval- uated at 0 + <Down. while classical approaches use the gradient g evaluated at 0g. Here we study how these two gradients g and differ in simplified setting, We observe that the contribution of single state-action pair (s, a) towards LA, is given by:
A 0) To(a)s) T (a|s),
L (0 a)
whose gradient with respect 0 is given by
A a)
V VoL (0)
(1s.))oorll (ano(als)
Ta(a)s)
Each observed state-action pair therefore contributes towards the gradient, to increase or decrease the likelihood of action a depending on the sign of A a) (in way that is inversely proportional to the likelihood that action under To))
Here we study how the two gradients (standard) and ğ (SHAPO) differ in a simplified setting. We assume that the policy at some state is given by a ID Gaussian policy (a; μ4, 1) = Na 1) parametrized by the mean μ, with probability density function given by:
1
(a 11)
<(a; ()) =
exp
/2T
2
17

---

Published as a conference paper at ICLR 2026
The gradient of (Ta;; p)) with respect to μ is expressed as:
d <(a μ4) =(a ()T((;;))
0p
Assuming that the current policy is Tn(a) = x (a: 0, specified by je = 0, we study the discrepancy between the two gradients 9 and with respect to p in terms of observed action a and an advantage value A. Our utility function L and its gradient can now be expressed as:
A L (1; a, A) TT(a)
A
((;
V Vuli (u:a, A)
TT(a) V (a
The gradient g(a, A) at fr becomes:
(a (o) = Aa g(a A) V V (ula))wo To(a)
Therefore, this situation, the standard gradient is simply the product the advantage with the action With SHAPO the gradient is instead computed at 1i <Down- where <Down parameter perturbation crafted to minimize L Note that with our choice of parametrization
KL(T( Mo, ( 5/00 + <Dww. 1)) = fhan which is half the square of the Euclidean distance between mean parameters. Therefore, the approach based on the Kullback-Leibler neighborhood (KL(T(: ; Ao. 1) |T( 140 + CDown, 1) s ōDww)) coincides with the one using the Euclidean metric on parameters (CDown $ p) and can be expressed in terms of = 22 The perturbation is therefore given by:
2
Down (a, A) p g(a, A)
(g(a, A)
This adjustment contributes to increase the likelihood of action when the advantage A is negative (unsafe action) and to decrease it when A positive. The adjusted parameter jaaa, A) μo + (Down (a, A) is given by:
i(aa A) - psign(Aa)
So, for example, this adjusted mean is negative when A < and a < 0, thereby leading to a
likelihood <(a; ji, larger than (a; 0, 1).
Now we compute the gradient 9 at jt:
A To(a)
J(a, A) V T.2 (pla))u=0
(a (a)
As desired, this leads to the following expression for ğ(a, A)
A(a - pxp
a² aa-pa=
if Aa < 0
g(a, A)
d² d'=pa+))²
A(aa + expp
if Aa 2
D IMPLEMENTATION DETAILS & ANALYSIS
D.1 HYPERPARAMETER DETAILS
For the base algorithms we use the optimized hyperparameters available in open-source implemen- tations like Omnisafe (ji et al. 2024 SHAPO hyperparameters ōrown and Peritic are optimized by performing grid search between range [0.001, 0.00001] for opown and {0.01, 0.05, 0.005) for Perêtic while keeping the hyperparameters of the base algorithm were kept fixed. Full list of hyperpa- rameters for each baseline and the corresponding environment have been provided here (https:) //anonynous. 4open. se www.eecceee/eeeeeeeeeee/eeee Pww.Poollyyyy--o Opt.ty---iiiieeooo im za -on- 813E). Most com- mon choice of SHAPO hyperparameters that gave the best performance across tasks and baselines
https://gi
gww.ingtggnnnnnnnnmneeeee safe
18

---

Published as a conference paper at ICLR 2026
Sensitivity to down
Sensitivity to Pcnitic
12
10
10
Return
Return
8
6
4
0.0108
0.0110
0.0112
0.0114
0.0108 0.0109 0.0110 0.0111 0.0112 0.01133
Cost Rate
Cost Rate
Vanilla
SHAPO 8yswn Sp Se-06
Vanilla
SHAPO Purk- -.01
SHAPO Sma #8
SHAPO 6asm 15-05
1e-05
SHAPO Pow0=0.001
SHAPO Purti - 0.05
58
Basan
SHAPO A.ro:00.005
SHAPO A.1e-0011
SHAPO Sa 1e:07
SHAPO 6 0.0001
SHAPO Samas 5e-07
SHAPO 6psam 0.0005
SHAPO Sw 1e-06
SHAPO
Gasen 0.000
Figure 9: Hyperparameter Sensitivity: We evaluate the sensitivity of SHAPO to Down (left) and Peritic (right). Very small ōDown effectively disables SHAPO and yields performance similar the vanilla baseline, while overly large values destabilize training. Moderate values provide the most consistent improvements. For Perities performance is stable across broad range, as long as the value is not excessively large.
was Down = 0.0001 and Peritic = 0.01. There a very clear monotonous hyperparameter sensi- tivity SHAPO i.e. if the ōDown value is small, the policy optimization is unchanged from the base algorithm thus the performance also resembles the base algorithm's 's performance, while the Down is too high, the date-disttibuttonss of the original policy d that generated the data and that of perturbed policy d* are no longer bounded and thus its no longer possible to provide a tight bound on the deterioration of the policy required for the inner SHAPO objective, which leads to unstability and inferior performance.
D.2 HYPERPARAMETER SENSITIVITY
Metric
Vanilla
SHAPO
Total Runtime (mins) Update Time (s)
85 min
100 min
To study the robustness of SHAPO the choice of hyperparameters, we conduct an analysis (Fig. 9 varying the two key hyperparameters, Sdown and Parities while keeping all other settings fixed. This setup allows to isolate the effect of each hyperpa- rameter on safety and efficiency of the resulting pol- icy. The results show that very small ddown values effectively disable the perturbation step in SHAPO, causing the algorithm to behave like the vanilla base- line. Conversely, excessively large &down values lead to unstable updates and degraded performance. We observe clear band of intermediate values that sistently improves both learning speed and final performance. In contrast, Peritic exhibits much less sensitivity: performance remains stable across wide range of values, with degradation only occur- ring when the regularization becomes extremely large. These trends confirm that SHAPO is robust reasonable choices of Peritic and requires only mild tuning ddown achieve reliable gains.
4
5.5
Table Table shows the total runtime of Vanilla (PIDLag) and SHAPO (PIDLag) PointGoall task as well as time taken for individual updates. As can see that the compute overhead is negligible in compari- sion the overall time it takes to train the RL agent and most of the runtime is taken up by collecting data through policy rollouts.
D.3 RUNTIME ANALYSIS
Here we analyze the computational overhead introduced by SHAPO during training across different tasks. In experiments, overall training time is dominated by environment interaction rather than
19

---

Published as a conference paper at ICLR 2026
Fixed a
Fixed down
10a
alpha-0.01
Daown
alphe=0.05
Sasee -0.0005
10
alpha=0.1
S 1e-05
alpha-0.2
- 55-05
10-1
alpha=0.3
6oown 0000
Doown
alpha-0.499
alpha=0.4
0.3
10.6-
0.2
10-9.
10-10
10-12
0.0
0.0
0.2
0.4
0 0.8
0.6
1.0
10A
10a
10
107
1e7
Comparison of pessimism schedules under SHAPO. Left: Fixed a. Holding the pessimism level a constant causes the corresponding bound 8mann / to shrink rapidly as nu grows, reducing the robustmess margin over time. Right: Fixed ŌDown. Keeping 9Dow constant makes the implied percentile ae 4 (22,, decrease with rus meaning the agent becomes This desirable for safe exploration: early on, when data limited, pessimism remains low to avoid overly conservative behavior, while increased pessimism later helps guard against epistemic uncertainty once sufficient experience has been gathered.
Figure 10:
risk-averse training progresses.
by policy or critic updates. Consequently, the additional TRPO step required for SHAPO's parameter adjustment contributes only marginally to the total wall-clock time, which includes data collection, policy updates, and critic updates. Table [ reports the runtime for training RL agent for 10M timesteps on the PointGoall task averaged across 5 different training runs. As shown, SHAPO incurs only small increase in total runtime, since most of the compute is spent collecting policy rollouts. The per-update computation time for the actor and critics increases by roughly 30 percent, but this overhead is minor compared to the dominant cost of data collection. All the experiments were performed on CPU with 10 cores and 30 GB of RAM. To increase the efficiency of data-collection, we use vectorized environments provided with gymnasium (Towers et al. 2024) and safety-gymnasium Ray et al. 2019), with 5 environments running in parallel.
E PESSIMISM SCHEDULE
In order to compare different ways of controlling the level of pessimism used SHAPO, we use at each update step the number of collected state-action samples n, proxy for the ideal sample size n introduced in Subsection 4.2 It is important to note that the true value of n-the number of informative. approximately independent samples governing posterior contraction-is unknown in practice. The proxy n. may therefore increase much faster than the effective statistical sample size, which has direct implications for how pessimism evolves during learning.
Increasing affects the pessimism level differently depending on whether we fix &Down or fix the target percentile a. When dDwa is kept constant, Proposition implies that the induced percentile satisfies
= - 22,
0x (
so that ag decreases n4, grows. Thus, the agent becomes increasingly pessimistic, focusing on more extreme tail events as training progresses. This behavior is desirable for safe exploration: early in learning, when data scarce and uncertainty is large, low pessimism avoids overly conservative behavior that would suppress exploration; more experience is gathered, increasing pessimism helps guard against refined estimates of unsafe regions.
By contrast, if we fix a pessimism level a E (0, 1 then the corresponding 6
8 Down
2x
2n
20

---

Published as a conference paper at ICLR 2026
Venilla
Fixed
Fixed Oyee
PointGoal2
PointButton1
Fixed O
Fixed Ø
32-
Dotum
Retur
Returns
$.000
Goite
33110
38501 0.300 GOILE GOL23
10
Cost Rate
Cost Rate 3e-2
SHANp Cocce
Shann bo
RegAO and
Aus=75
Figure 12: Both fixed drown and fixed schedules
Axym10
A1y0100
lead significant improvement performance over the vanilla bascline PointGoal/ task. Although the choice of a is not straightforward as both very low (a = 0.01) and high alphas (a = 0.4) give good per- formance.
Figure 11: We compare two pessimissm schedules SHAPO by using a fixed 8nows throughout the train- for fixed
ing calculating 6Down using 6Dows a, where n assumed he all the state-action pairs experienced by the agent. Both schedules expand the Pareto frontier of the Vanilla (PIDLag) baseline.
decays rate 1/m,. Because nt may grow much faster than the true effective sample size, this schedule can drive oown to zero prematurely, effectively nullifying any pessimism the later stages of learning. This mirrors the Bernstein-von Mises posterior contraction, but may not model appro- priately the epistemic uncertainty state of an actor trained in a reinforcement learning setting, where epistemic uncertainty persists due partial observability, non-stationarity, or limited coverage of the state space. Fig. 10 (illuustrttes this phenomenon: the fixed-a schedule quickly collapses the down as ny grows, whereas fixing dDow maintains nontrivial degree of pessimism throughout training.
Fig. I and 12 compare these two approaches to scheduling pessimism. Both strategies- down or fixing a enable SHAPO to outperform the vanilla baseline (PIDLag in this case). Fig. shows that both pessimism schedules expand the Pareto frontier of the vanilla algorithm enhancing safety-efficiency tradeoff' across different Asnit values. Fig. 12 compares the performance of SHAPO compared to the vanilla baseline (in grey) for different choices of ODown or We can see that there are many choices of fixed a's and 6owa 's that lead to improvement over the vanilla PIDLag, but the choice of that leads to best performance is not that obvious as it impacts the exploration of the algorithm.
Since the effective sample size cannot be reliably estimated in continuous-contre reinforcement learning-whhii wwhre data are temporally correlated and only a fraction of collected samples contribute independent information -fiiin-iinggg 6Dww provides a practical and robust alternative. Treating d as a tunable hyperparameter allows practitioners to directly control the desired safety-efficiency tradeoff without relying on uncertain estimates of n. This makes the fixed-Opown strategy both simple and effective, offering stable mechanism for regulating pessimism throughout training.
F SHAPO FOR UNCONSTRAINED RL
SHAPO is general mechanism that be incorporated into any on-policy reinforcement learning algorithm, and its utility extends beyond explicitly constrained RL settings. To illustrate this ver- satility, we evaluate SHAPO in three standard MuloCo (Todorov et al., 2012) continuous-control environments: Ant-v4, Walker2d-v4, and HalfCheetah-vv. The first two environments (Ant and Walker2d) include termination conditions in which the episode ends immediately if the agent falls. Although they do not provide an explicit cost signal or negative penalty for falling, this termination logic implicitly induces lower tail of poor outcomes: each fall truncates the episode, prevents fur- ther reward collection, and therefore constitutes an undesirable event in the return distribution. In contrast, HalfCheetah has no such termination or failure mode-the episode continues regardless of how the agent moves -yieldig a much smoother return distribution without pronounced lower tail. We compare PPO (Schulman et al. 2017) to its CHAPO-aagmennt variant (SHAPO-PPO), implemented described in Algorithm 4 From this perspective, SHAPO's role becomes clear: it modulates the actor's s sensitivity to tail events in the loss landscape, whether these arise from explicit constraints from inherent structural properties of the environment.
21

---

Published as a conference paper at ICLR 2026
Ant-v4
Walker-v4
HalfCheetah-v4
5000
SHAPO-PPO
5000
PPO
SHAPO-PPO
6000
SHAPO-P9O
PPO
4000
5000
4000
4000
Return
3000
2000
2000
2000
1000
1000
1000
0
0
1000
2
Steps
Steps
Steps
1.0
1.0
PPO
Rate Termination
SHAPO-PPO
0.9
SHAPO-PPO
0.04
SHAPO-PPO
0.8
0.02
0.7-
0.6
0.00
-0.04
2
Steps
Steps
Steps
5000
FPO
SHAPO-PPO
0.04
SHAPO-PPO
Total Terminations
6000
4000
5000
0.02
3000
4000
0.00
3000
2000
2000
1000
PPO
1000
SHAPO-PPO
01
0
Steps
Steps
Steps
Figure 13: SHAPO on Standard RL: Columns correspond to Ant-v4, Walker2d-v4, and HalfCheetah-
v4. Top row: episodic return; middle row: rate of termination; bottom row: total terminations. SHAPO yields substantial gains on Ant-v4 and Walker2d-v4, where falling triggers early episode termination and creates a pronounced lower tail in the return distribution. In HaljCheetah--44 which lacks termination failure modes, the return distribution has meaningful lower tail, and SHAPO consequently provides limited improvement.
Fig. 13 illustrates how the impact of SHAPO depends the presence absence of tail events in the environment. In Ant and Walker2d, SHAPO consistently improves the sample efficiency of the PPO baseline by redueing premature terminations; this allows agents to remain in the episode longer, resulting in fewer resets and more reward collected per rollout. In HalfCheetah, however, where no catastrophic failure or termination mode exists, the return distribution lacks a meaningful lower tail, and accordingly, SHAPO offers limited additional benefit. These results underscore that while SHAPO provides tangible advantages in settings with significant tail risk, its effect is naturally less pronounced in environments where undesirable events do not produce clear tail behavior in the return landscape.
G LLM USAGE
ChatGPT (Achiam al. 2023) was used to polish writing and also tool correct grammar and spelling mistakes.
22

---

Published as a conference paper at ICLR 2026
CPO
CRPO
PID
SauteRL
1.5
12
SHAPO
SHAPO
10-
SHAPO
SHAPO
Vanilla
8-
2-
-0.5
-1.0
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.35
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Steps
ls7
Steps
1n7
Steps
ls7
Steps
l
80-
Verill SHAPP
90
Venilla SHAPO
Varille SHAPO
Venille SHAPO
70
Episodic Cost
60
70
50
40
30-
30
20
10
10f
10'
10F
10'
10a
10a
10a
Steps
Steps
Steps
Steps
Figure 14: PointGoall
CPO
CRPO
PID
SauteRL
2.5
0.25
SHAPO
SHAPO
SHAPO
SHAPO
0.00
3-
1.5
-.225
1
0.50
0.75
0.0
-1.25
-1.0
1.55
0.00
0.25
0.50
0.75
1:00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
107
167
Steps
167
107
Vanilla
Vanilla SHAPO
Varilla SHAPO
SHAPO
SHAPO
70
Episodic Cost
60
80
60
50
50
40-
30
20-
10
10
10'
10'
10'
10"
Stegs
Steps
Steps
Steps
Figure 15: PointButtonI
23

---

Published as a conference paper at ICLR 2026
CPO
CRPO
PIDLag
2000
Vanilla SHAPO
4000
Vanilla SHAPO
5000
2000
4000
3000
Return
-2000
2000
2000
-4000
4000
1000
-6000
Vanilla SHAPO
Steps
1e6
Steps
1e6
Steps
1e6
109-
Vanilla
10a
Vanilla SHAPO
Vanilla SHAPO
SHAPO
Episodic Cost
10-1-
10-1-
10-1
10
10-2-
10"4
10
Steps
1e6
Steps
1e6
Steps
1e6
800
Vanilla SHAPO
800
700
600
600
600
500
Total
400
300
200
200
200
Vanilla SHAPO
Vanilla SHAPO
100
0
2
2
5
0
1
2
Steps
1e6
Steps
1e6
Steps
1e6
Figure 16: Ant-v4
24

---

Published as a conference paper at ICLR 2026
CPO
CRPO
PIDLag
Vanilla
2000
Vanilla
4000
1500
1750
3500
1250
1500
3000
2500
Return
1000
2000
1500
500
500
1000
250
250
500
Vanilla SHAPO
0
0-
Steps
1e6
Steps
1e6
Steps
1e6
10a
Vanilla SHAPO
109
Vanilla SHAPO
10a
SHAPO
6x-1
6K10-1
Episodic Cost
1x
10-
10
1x-1
10
10-1
2x
2 2100
-0
Steps
1e6
Steps
1e6
Steps
1e6
4500
Vanilla
Vanilla SHAPO
Vanilla SHAPO
4000
SHAPO
3500
4000
3000
3000
Total Cost
3000
2500
2000
2000
2000
1500
1000
1000
500
500
Steps
1e6
Steps
1e6
Steps
Figure 17: Walker2d-v4
25


Reviews:
[{"review_id": "3EIEpRPd7I", "reviewer": "Reviewer_2kQM", "summary": "This paper introduces Sharpness-Aware Policy Optimization (SHAPO) to address the challenge of safe exploration in safety-critical RL. The core idea is to use the actor's sharpness as a practical proxy for epistemic uncertainty. The method implements a pessimistic policy update. In each step, it first solves an inner-loop optimization to find the \"worst-case\" perturbed parameters ($\\theta_0 + \\epsilon_{Down}$) that minimize the policy objective within a trust region which is defined by the Fisher Information Matrix (KL divergence). It then computes the policy gradient ($\\tilde{g}$) at this worst-case point and uses this pessimistic gradient for the final policy update .\nAnalytically, the authors show this process implicitly reweighs gradients, amplifying the effect of rare, unsafe actions and tempering the effect of rare, safe actions.", "questions": "- Could you clarify whether the adaptive scaling rule was implemented or tested in any form?\nIf not, how do you reconcile this discrepancy between the theoretical motivation and the actual implementation?\n- Could you prove that the core safety intuition from Figure 2  holds for the Fisher-based update, not just the simplified Euclidean case  shown in the appendix?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper tackles the critical challenge of safe exploration in a practical and useful manner, which is a prerequisite for real-world RL deployment.\n\n- The analytical insights, including Proposition 3 on gradient adjustments for rare actions and the reinterpretation of perturbations as pessimistic quantiles (Proposition 2), are clear and well-supported.\n\n- Evaluations across multiple baselines and environments demonstrate consistent gains, with Pareto improvements and reduced catastrophic failures", "weaknesses": "-  Major Theory-Implementation Mismatch\n\nThe paper's theoretical justification for being \"uncertainty-aware\" (presented in L249-256)  hinges on the idea that the inner trust region $\\delta_{Down}$ should adapt based on the amount of data $n$. However, the actual implementation described in Appendix D uses a fixed, grid-searched $\\delta_{Down}$.\n\n- Gap in Intuitive Justification\n\nThe paper argues its strength comes from using the Fisher metric . However, the core intuition for why it is safe (Section 3.3)  is derived using a simplified Euclidean metric (Appendix C, L778: \"For simplicity, we consider here a Euclidean perturbation...\").", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IjWcbgwbwt", "reviewer": "Reviewer_oXHA", "summary": "The paper studies the application of sharpness awareness minimization (SAM) in policy optimization. SAM aims at optimizing the model such that the worst-case loss in its neighboring parameters is low, encouraging the optimizers to find low-loss solutions that are also flat. In particular, the authors extend Fisher-SAM, a prior SAM method that leverages Fisher information matrix to estimate the local geometry to efficiently estimate the local loss-maximizing parameter, to trust-region policy optimization for safe exploration. The paper provides an interpretation of Fisher-SAM as a pessimistic estimation of the loss subject to the uncertainty quantified by the Fisher matrix. This in turn allows the authors to interpret the worse-case expected return in the neighboring parameters informed by the Fisher matrix as an uncertainty lower-bound on the actual expected return. The paper evaluates the proposed method, SHAPO, on safety gym and shows that the proposed method outperforms prior methods in terms of the return/cost trade-off.", "questions": "- In Figure 3, 6 and 7, how are the shapes of the circles around the mean/average performance determined?\n- In algorithm 2, how is $U_\\theta$ computed?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is easy to read and the proposed method is well-motivated and technically sound. \n- The empirical results show statistically significant improvements of the proposed method over prior safe RL methods. Ablations are thorough and show that major components of the algorithms all contribute to the effectiveness of the proposed method.\n- Even though SAM/Fisher-SAM is not new, the application of it in the context of safe RL is new and seems effective.", "weaknesses": "- The proposed method requires solving the natural gradient direction that maybe very expensive. \n   - I tried to look for the implementation details in the paper but could not find it. It would be good if the authors could include more implementation details in the paper for better reproducibility. \n   - I checked the linked anonymous codebase briefly and it seems that the implementation involves running an iterative conjugate gradient procedure (https://anonymous.4open.science/r/Safe-Policy-Optimization-813E/safepo/single_agent/shapo.py). This procedure can incur a non-trivial amount of computation overhead. It would be good if the authors could include more details on the run-time of the algorithm compared to prior methods in the paper. \n\n- Hyperparameter sensitivity analysis is missing from the empirical evaluations. In Appendix, the authors mentioned that \"Most common choice of SHAPO hyperparameters that gave the best performance across tasks and baselines was $\\delta_{\\mathrm{Down}} = 0.0001$ and $\\rho_{\\mathrm{critic}} = 0.01$\", but it is unclear how different hyperparameters influence the performance of the algorithm. Including some analysis on the sensitivity of the performance with respect to these hyperparameters (especially $\\delta_{\\mathrm{Down}}$ could help gain a better understanding of the robustness and effectiveness of the algorithm.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Eou93ERA1y", "reviewer": "Reviewer_tvhm", "summary": "The paper explores the use of sharpness aware optimization for safe RL. Sharpness aware optimization introduces a max-min problem instead of maximization to flatten the optimization profile. In my understanding this max-min problem would aim to prefer a flatter maximum to a sharp peak maximum in the parameter space. This approach aims to improve robustness of the optimization solution to epistemic uncertainty. The authors conduct a series of experiments that show an improvement of performance of different algorithms with and without SHAPO.", "questions": "* The approach makes total sense and it seems relevant not only for safe RL, but also for RL. What happens if SHAPO is applied without a safety constraint? Did the authors perform this ablation study? \n* Can the authors provide similar Figures to Fig 6 and 7, but in terms of episode return vs episode cost? I think their decision to focus on cost rate is correct, but having a full picture (maybe in appendix) would be good.\n* Can the authors elaborate what’s the base method for Saute RL? As the authors are aware, Saute RL can be applied to any RL algorithm.   \n* Safety gymnasium has many more environments that could be useful for your future research and evaluations. Safety Gymnasium: A Unified Safe Reinforcement Learning Benchmark https://arxiv.org/abs/2310.12567\n* What’s the additional computational burden that the approach adds to the algorithm?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.The idea is novel to my best knowledge and quite interesting. \n2. The paper contains a deep discussion on my sharpness aware optimization is a good fit for ML problems in general. In section 3.3, the authors dive deeply into why SHAPO can be a good fit for RL.\n3. The experiments and ablation study are great", "weaknesses": "I couldn't come up with many weaknesses, but here are a couple \n* The policy update section 3.1 is building the solution bottom-up, but maybe a top-down approach would be slightly easier to read. \n* It would be good to explain how the policy update is actually implemented. The authors present a quadratic optimization problem in  TRPO style update and state that SHAPO update can be applied to any on-policy algorithm. Can the authors elaborate how SHAPO can be applied to a PPO update?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ofn7WTHROV", "reviewer": "Reviewer_huer", "summary": "This paper introduces Sharpness-Aware Policy Optimization, a novel approach to enhance safe exploration in Reinforcement Learning by addressing the actor's epistemic uncertainty. The core idea is to use the \"sharpness\" (sensitivity of the policy to parameter perturbations) as a practical proxy for uncertainty, particularly in regions with scarce data.\nSHAPO leverages principles from Sharpness-Aware Minimization (SAM). The proposed update rule operates in a pessimistic manner: it first finds the worst-case policy parameter θ within a local neighborhood that minimizes the performance objective (maximizes risk), and then computes the policy gradient at this \"pessimistic\" point. This results in a modification of the policy gradient that analytically promotes conservative behavior by amplifying the penalty for rare, unsafe actions while attenuating the reward for rare, safe actions. The method is implemented as a plugin on top of existing safe RL algorithms (TRPO) and demonstrates improved performance on the safety-efficiency Pareto frontier across various Safety Gym tasks.", "questions": "1. Given the theoretical prescription that $\\delta _{Down}$ should decay with the number of samples n, why was a fixed, tuned hyperparameter used in the implementation? Could the authors provide an ablation study comparing the fixed hyperparameter approach with the theoretically suggested annealing schedule?\n2. Please clarify and correct the discrepancy in the definition of the posterior/likelihood distribution Q(θ) between Section 3.2 and Appendix B. A corrected, rigorous proof for Proposition 2 (or 5) is required.\n3. To solidify the contribution of using \"sharpness\" as an uncertainty proxy, could the authors compare SHAPO's performance against policy gradient methods that explicitly model actor uncertainty or risk, such as an Ensemble Actor approach or a CVaR-based policy optimization method?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.  Applying Sharpness-Aware Minimization (SAM), typically used for improving generalization in supervised learning, to Safe RL is a  conceptually elegant contribution. This provides a new lens for addressing epistemic uncertainty in the actor network.\n2. The analysis (Section 3.3) and empirical results support the claim that the SHAPO gradient effectively introduces a pessimistic bias. The method consistently achieves a lower frequency of catastrophic events and generally improves the safety performance of multiple SOTA safe RL baselines when used as a plugin.\n3. The update rule is straightforward and can be readily integrated with different Policy Optimization algorithms without fundamental changes to their core optimization objective, showcasing high modularity.", "weaknesses": "1. The paper suffers from a crucial disconnect between the theory and implementation. Proposition 2 defines the perturbation magnitude \\delta_{Down}​\t as a function of the sample size n and confidence level α. However, in the practical implementation and hyperparameter search (Appendix D), \\delta_{Down}​ is treated as a fixed hyperparameter, which contradicts the theoretical guidance for an annealing schedule based on n.\n2. The modeling of the posterior distribution Q(θ) is solely \"motivated\" by the BvM theorem. This represents a strong, unproven assumption, as the strict regularity conditions required for BvM are unlikely to hold in high-dimensional, non-convex deep RL settings, thus introducing a significant theoretical gap.\n3. There is an explicit inconsistency in the definition of the covariance matrix for the distribution Q(θ) between the main text and the appendix proof of Proposition 5, which uses a distribution corresponding to a precision matrix proportional to $\\sqrt n$. This error undermines the rigor of the derived relationship for $\\delta _{Down}$.\n​4. The paper focuses on comparing SHAPO's performance against standard constrained optimization methods (CPO, CRPO, etc.). To fully validate the claim of solving epistemic uncertainty in the actor, the baselines should ideally include methods that explicitly handle uncertainty or risk sensitivity, such as: Policy Gradient methods utilizing Dropout or Ensembles on the actor. Risk-Sensitive Policy Optimization methods (e.g., those based on CVaR or other risk measures).\n5.  While $\\delta _{Down}$ is the critical new hyperparameter of the method, the paper provides only a qualitative discussion of its sensitivity (Appendix D) and lacks a comprehensive, quantitative ablation study.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)