Flat Reward in Policy Parameter Space Implies Robust Reinforcement Learning
Authors: Hyun Kyu Lee
Sung Whan Yoon
Venue: ICLR 2025 Oral (2025)
Keywords: Reinforcement learning
Flat Minima
Robust Reinforcement learning
Primary area: reinforcement learning

Abstract: Investigating flat minima on loss surfaces in parameter space is well-documented in the supervised learning context, highlighting its advantages for model generalization. However, limited attention has been paid to the reinforcement learning (RL) context, where the impact of flatter reward landscapes in policy parameter space remains largely unexplored. Beyond merely extrapolating from supervised learning, which suggests a link between flat reward landscapes and enhanced generalization, we aim to formally connect the flatness of the reward surface to the robustness of RL models. In policy models where a deep neural network determines actions, flatter reward landscapes in response to parameter perturbations lead to consistent rewards even when actions are perturbed. Moreover, robustness to action perturbations further enhances robustness against other variations, such as changes in state transition probabilities and reward functions. We extensively simulate various RL environments, confirming the consistent benefits of flatter reward landscapes in enhancing the robustness of RL under diverse conditions, including action selection, transition dynamics, and reward functions. The code for these experiments is available at https://github.com/HK-05/flatreward-RRL.

Published as a conference paper at ICLR 2025
FLAT REWARD IN POLICY PARAMETER SPACE IMPLIES ROBUST REINFORCEMENT LEARNING
Hyun Kyu ,ee', Sung Whan Yoon!.2.
GGaddaae School of Artificial Intelligence and *Department of Electrical Engineering Ulsan National Institute of Science and Technology [dnwidlwl, shyoon8 @nn funist.as. at..ac.......................... ac. .r
ABSTRACT
Investigating flat minima loss surfaces in parameter space is well-documented in the supervised learning context, highlighting its advantages for model gener- alization. However, limited attention has been paid the reinforcement learning (RL) context, where the impact of flatter reward landscapes in policy parameter space remains largely unexplored. Beyond merely extrapolating from supervised learning, which suggests a link between flat reward landscapes and enhanced gen- eralization, we aim to formally connect the flatness of the reward surface to the robustness of RL models. In policy models where a deep neural network deter- mines actions, flatter reward landscapes in response to parameter perturbations lead to consistent rewards even when actions perturbed. Moreover, robustness to action perturbations further enhances robustness against other variations, such changes in state transition probabilities and reward functions. We extensively simulate various RL environments, confirming the consistent benefits of flatter reward landscapes in enhancing the robustness of RL under diverse conditions, including action selection, transition dynamics, and reward functions. The code for these experiments available at https///ithub..om/HHH (ot-KK---//aaaaaad--R-....
INTRODUCTION
Reinforcement Learning (RL) has emerged as a powerful strategy for optimizing sequential decision-making tasks, with applications ranging from roboties to game playing and beyond (Mnih et al., 2015; Silver et al., 2016). Despite its great success, one of the critical challenges in RL remains its susceptibility to variations in the environment, which can significantly degrade perfor- mance (Peng et al., 2018). This challenge has driven interest in Robust Reinforcement Learning (RRL), subfield focused on enhaneing the generalization and resilience of RL algorithms against such environmental variabilities (Morimoto & Doya, 2005b; Rajeswaran et al., 2017: Pattanaik et al., 2018). Traditional RRL approaches, such as robust policy optimization and action robust reinforce- ment learning, have made strides in this direction. However, these methods often involve complex implementations and limited algorithmic diversity, especially when considering the broader land- scape of RL (Pinto et al., 2017; Tessler et al., 2019).
In contrast, recent advancements in supervised learning have highlighted the importance of the land- scape of loss values in the model parameter space. Specifically, flatter minima- with smooth vary- ing loss values around them them-are are preferred as they exhibit strong generalization capabilities and are less sensitive to perturbations in the input distribution (Keskar et al., 2017: Hochreiter & Schmidhu- ber, 1997; Dinh et al., 2017; Foret et al., 2021). This raises a natural question: Does a flat reward landscape imply robustness in RL against environmental changes?
While a few works have begun to explore this question, finding that optimization for flatter reward maxima sometimes shows meaningful gains against varying environments (Lee et al., 2024), and that RL methods with flatter reward surfaces tend to exhibit robust performance (Sullivan et al., 2022), these prior studies often rely on a naive transfer of benefits observed in supervised settings to RL optimization. Moreover, they primarily provide empirical evidence without establishing
Corresponding Author

---

Published as a conference paper at ICLR 2025
rigorous link between flat rewards and robustness. To the best of our knowledge, the potential of flat reward landscapes in RL remains largely unexplored.
To illustrate the practical significance of this concept, consider the following preliminary experi- ment. We ran a 2D navigation task where an agent must move from a fixed start position to a goal position, with very narrow path between slight action perturbations can severely deteriorate performance. We compare a standard Proxi- mal Policy Optimization (PPO) algorithm with our variant of PPO enhanced with Sharpness-Aware Minimization (SAM), denoted as "SAM+PPO," which pursues flatter reward maxima.
obstacles (see Fig. 1). In this action-critical setting,
As shown in Fig. 1, the paths generated by mul- tiple trials are depicted in orange. The standard PPO agent tends to follow the shortest path, bring- ing it dangerously close to the obstacles. This be- havior makes the agent susceptible failure when actions are slightly perturbed, minor deviations can result in collisions. In contrast, the SAM+PPO agent maintains safer margin from the obstacles, avoiding the narrow path and demonstrating more robust decision-makee making in the presence of action perturbations. This preliminary experiment high- lights how pursuing flatter reward maxima can lead to enhanced robustness in RL agents.
Figure Trajectories of agents in the 2D nav- igation task: (a) PPO tends to take the short- est path, risking collision with obstacles; (b) SAM+PPO maintains safer margin, demon- strating robustness to action perturbations.
Clorakk
Treeday
Trajetory
Styn
Gad
Gval
(a) PPO
(b) SAM+PPO
In this paper, we delve deeper into the concept of flat reward landscapes and their implications for robust reinforcement learning. Our main claim is that flatter reward maxima rigorously lead robustness against action perturbations (Proposition 1). We also provide insights into how this robustness extends beyond action robustness to encompass transition probability and reward robustness. To validate our claims, we conduct extensive exper- iments in various MuloCo environments (Todorov et al., 2012), including Hopper, Walker2d, and HalfCheetah, by varying actions, transition probabilities, and rewards.
Our work is outlined as follows: Section 2 reviews related prior works, and Section 3 presents the preliminaries with focus on Robust Markov Decision Processes and SAM. In Section 4, we intro- duce our definition of flat reward maxima and establish the rigorous link between flat rewards and action robustness. Section presents extensive simulations and analyses that validate the effect of flatness the reward surface in various RL benchmarks. We further compare our flat reward policy with existing robust RL. methods, emphasizing gains in performance and computational efficiency (see also Appendix E). Finally, Section 6 concludes the paper.
2
RELATED WORK
2.1 ROBUST REINFORCEMENT LEARNING
Robust Reinforcement learning (RRL) focuses on developing agents that maintain high performance despite uncertainties in the environment. Early approaches in this domain were predominantly model-based, leveraging explicit models of the environment to account for uncertainties. Mori- moto & Doya (2005a) utilized Ho control theory formulate RL as a zero-sum game between the agent and nature, providing robustness through worst-case scenario optimization. Similarly. Iyen- gar (2005) and Nilim& & El Ghaoui (2005) introduced Robust Markov Decision Processes (RMDPss) that model uncertainties in transition probabilities and reward functions, enabling the derivation of robust policies with theoretical guarantees. Several studies have proposed model-based robust RI. algorithms in the tabular domain, with strong emphasis improving sample efficiency (Yang et al., 2022: Zhou et al., 2021; Panaganti & Kalathil, 2022).
However, model based methods often face scalability issues due to the computational complexity involved in high-dimensions spaces (Xu & Mannor, 2010; Wiesemann et al., 2013; Mannor et al., 2016: Goyal& & -aand-Ceeeent, 2023). To overcome these limitations, the focus gradually has been shifted towards model-free approaches, which do not require explicit environmental models, making it more scalable and applicable complex tasks.
2

---

Published as a conference paper at ICLR 2025
Recent research has explored model-free robust RL in settings with unknown uncertainty sets, where the agent pursues robustness without precise knowledge of the environmental uncertainties. Tessler et al. (2019) proposed Action Robust Reinforcement Learning, which models uncertainties in the action space and learns policies robust to adversarial action perturbations. Mankowitz et al. (2019) introduced Robust Reinforcement Learning via Regularized Policy Optimization, applying regular- ization techniques to improve robustness without explicit adversarial training. Derman et al. (2021) developed the TWICE algorithm, model-free method that achieves robustness by accounting for both variance and bias in policy evaluation. Roy et al. (2017) introduced method for robust Q- learning that accounts for uncertainties in transition dynamics by optimizing the worst-case perfor- mance over an ambiguity set, which is estimated online without prior knowledge (Wang & Zou, 2021; Badrinath Kalathil, 2021).
Advancements in posecy- methods have further propelled robust RI. (Kumar et al., 2024; Li et al., 2022; Grand- CranndClément& Kroer, 2021). Pinto et (2017) introduced Robust Adversarial Re- inforcement Learning (RARL), where an adversary is trained alongside the agent to perturb the en- vironment, enhancing the agent's robustness. By directly parameterizing and optimizing the policy, these methods offer advantages in handling continuous and high-dimensional action spaces. Wang & Zou (2022) developed a model-free robust RL algorithm that utilizes a worst-case policy gradient approach to handle unknown uncertainties in the environment. Li et al. (2022) proposed First-Order Constrained Optimization in Policy Space (FOCOPS), which ensures robustness by enforcing safety constraints during policy updates.
Our work aligns with this progression, focusing on model-free, policy-based RL to enhance robust- ness through the optimization of flat reward surfaces. To the best of our knowledge, the flat reward surface of policy-based RL has not yet been deeply explored in the view of robustness.
2.2 FLAT MINIMA AND FLATNESS IN REINFORCEMENT LEARNING
The concept of flat minima in the loss landscape has been extensively studied in supervised learn- ing, where flatter regions are associated with better generalization and robustness to perturba- tions (Hochreiter & Schmidhuber, 1997; Keskar et al., 2017; Dinh et al., 2017; Dziugaite & Roy, 2017; Petzka al., 2021; Andriushchenko et al., 2023). Several methods have been proposed to find flat minima, including Entropy-SGD (Chaudhari et al., 2019), Stochastic Weight Averaging (SWA) (Izmailov al., 2018), and SWA with Densely (SWAD) (Cha al., 2021).
Among the prior works, Foret et al. (2021) introduced Sharpness-Aweer Minimization (SAM), which directly embeds the loss flatness into the objective function, enabling the search for min- ima with flatter surrounding regions. SAM has shown significant improvements in generalization performance across various supervised learning scenarios, including domain generalization (Cha et al., 2021) and federated learning (Qu et al., 2022; Caldarola et al., 2022; Sun et al., 2023).
Lee et al. (2024) applied SAM to policy gradient methods, observing empirical improvements in policy robustness. However, the prior work shows the following limitations: It primarily focuses on exhibiting partial empirical evidence without a formal bridge between the flatness in the reward landscape and the robustness in RL. Also, the effects of reward flatness are not carefully scrutinized in the multiple key perspectives of RL, i.c., action, transition probability, and reward.
Our work is also based the SAM-based policy gradient methods, but we do not focus further pursuing algorithmic advances. Our work distinguishes itself by providing a further understanding of how the flatter reward maxima contribute the robustness of RL agents, which remains unknown yet. Also, we tested how flatness improves the robustness of RL against the perturbations of actions, transition probabilities, and rewards.
2.3 REWARD SURFACE VISUALIZATION
Ilyas et al. (2020) is among the first to visualize reward surfaces to characterize problems with policy gradient estimates. They plotted policy gradient estimates versus uniform random directions, demonstrating through striking visual examples that low-sample estimates of the policy gradient often guide the policy in directions no better than random ones. This work emphasizes the need to better understand the optimization landscape of rewards.
3

---

Published as a conference paper at ICLR 2025
With the visualization methods by Li et al. (2018), Bekci & Gümüsş (2020) utilized loss landscape visualizations to study Soft Actor-Critic (SAC) agents (Haarnoja et al., 2018) trained on inventory optimization tasks. They investigated the impact of policy stochasticity and action smoothing on the curvature of the loss landscapes in several MuloCo environments, providing empirical evidence of how different factors influence the optimization landscape in RI..
Ota et al. (2024) utilized loss landscape visualization methods to compare shallow and deep net- works in RL. It demonstrates that deeper models perform poorly due to increased complexity and curvature in loss landscapes. With the visual intuitions, they developed methods to effectively train deeper networks for RL tasks. Notably, their work plotted the loss function of SAC agents, which includes an entropy regularization term (Haarnoja et al., 2018). Recently, it is observed that an RL method with a flatter reward surface tends show robust performance (Sullivan al., 2022).
In our work, we also visualize the surface of rewards to confirm the flatness of reward functions. In addition, we compute the flatness metric provide quantitative results.
PRELIMINARIES
This section outlines the key preliminaries that underpin our work, including the basics of Markov Decision Process (MDP), Policy Gradient Methods, Robust MDP, and Action Robust MDP. In ad- dition, briefly introduce Sharpness-AWe Aware Minimization (SAM), which Hidely-documenter optimization method to pursue flatter minima in loss surface for the supervised learning case.
3.1 MARKOV DECISION PROCESSES (MDP) AND POLICY GRADIENT METHODS
As well-documented many works, MDP and policy gradient methods are formalized as follows. Our description mainly refers (Sutton & Barto, 2018; Sutton et al., 1999; Williams, 1992).
Markoy Decision Process (MDP) is defined by a tplle (S, A, P. 7), where S represents the set of states, A represents the set of actions, P: S x A x S + [0, 1] is the state transition probability function, r S x A + R is the reward function that assigns a real value to each state-action pair, and y [0, is the discount factor that models the decreasing importance of future rewards.
The objective in an MDP is to find a policy и : S > A that maximizes the expected return. The expected return from state under policy is given by:
(*(()= Σ 7rr(S,, ox) Sp s,
where st and a₂ denote the state and action time t, respectively, and the expectation is the distribution of possible trajectories generated by following x.
erucial class of algorithms in deep reinforcement learning used to optimize policies directly. These methods employ gradients of the expected return with respect the policy parameters perform optimization (Sutton al., 1999). In this case, the policy is commonly a parameterized model, i.e., (a) with learnable weights f. According to the policy gradient theorem, the gradient of the expected return with respect to the policy parameters 0 is:
Policy Gradient Methods
V VoV*(s)= =E, Σ t= V' log Ta
Ta(a (arls:)
where R; = 2 ax) pression facilitates the application of gradient ascent methods iteratively improve the policy by adjusting parameters in the direction that maximizes the expected return (Williams, 1992).
denotes the cumulative reward from time t onwards. This ex-
3.2 ROBUST MDPs AND ACTION ROBUST MDP
Robust MDPs extend the classical Markov Decision Process framework to handle uncertainties in transition probabilities and rewards (Nilim & El Ghaoui, 2005). In a robust MDP, instead of single deterministic transition probability, the model allows for a range of possible outcomes defined by an uncertainty set P. This set encapsulates all plausible variations in the transition dynamics under real- world conditions, and is defined as: P = HP..a {IPs.a.eSS..E Ps.a E P..a. Ys E S, Va E A}. where each Psaa represents compact subset of As, the simplex of probability distributions over the state
4

---

Published as a conference paper at ICLR 2025
space S, which is defined as: As {p Risl
p(s') 0 for all s' E ES Eves p(s') =1}. The
optimization problem in a robust MDP is framed as follows:
00
maxmin Bp.. pEP
(1)
focusing on Max-Min strategy ensure the policy is resilient against the worst-case scenario of transition probabilities within P (lyengar, 2005).
Action Robust MDP specifically targets robustness in the context of actions taken within the MDP framework (Tessler et al., 2019). It considers the worst-case scenario where actions taken deviate from those chosen by the policy, potentially due errors or external disturbances. Thus, the goal is to devise policy that yields the highest minimum reward possible under any allowable perturbation of actions. Formally, the objective is follows:
max min Bp.
where de denotes allowable deviations in action space, upper bounded by 3 2 and y is the discount rate. This setup not only enhances the resilience of the strategy under practical conditions where perfect action execution unfeasible but also shows the formulation with that used Robust MDPs by maintaining the Max-Min structure across transition probability p and reward function
3.3 SHARPNESS-AWARR MINIMIZATION (SAM)
Sharpness-Aware Minimization (SAM) aims to search loss minima that are not only high- performing but also demonstrating the minimal sensitivity to perturbations in the parameter space (Foret et al., 2021). The objective function SAM is given by:
min max <(0+)), E
(3)
where L(0) is the loss function for given batch of samples, represents the model parameters, and is the parameter perturbations within a sphere by the norm p. The Max operation inside the objective function seeks the worst-case performance within the perturbation boundary, and the outer Min operation attempts to find the model parameters that are most robust against such perturbations.
The gradient of the Min-Max problem can approximated by four steps of computations. A de- tailed optimization procedure is provided in Appendix C.: i) Compute the gradient at 0, which maximally increases the loss values, ii) Perturb model parameter to become 0 + e, where is pointing the previously computed gradient direction. iii) Compute the gradient at the perturbed parameter 0 + e. iv) Update 0 by using the gradient computed at 0 + c.
LINKING FLAT REWARD TO ACTION ROBUSTNESS
For reward function r(s. a) with given state-action pair (s,a), a), here define the E-flat reward maxima by rehearsing the definition loss flatness of the supervised learning case (Shi et al., 2021).
Definition 1 (E-flat reward maxima) For a reward function r(s,a) and a policy model To (a(a))) parameterized by 0, a maximum 0* E-flat reward maxima when the following constraints hold:
For all E R. s.l. ||l|| S E, Be-p.orr-- Σ
t (r(ssr,a)) =
There exists R s.t. ||]|| E, By-p.osto-...all) E I'r(se,
<r-
(4)
where T" := Expparrm (a]s) Ee rrr(se, 4:)l and & is positive real number.
Based the objective of the action robustness, we further define Aacctoo robust policy as follows:
Definition (Aacction robust policy) For a reward function r(s, a). a policy model ne- (als) pa- rameterized by 8* is A-action robust when the following constraints hold:
For all 6, IR s.t. ||&el| < A,
E Eg-p.arTo--
1= 2'r(s at + 61) -
There exists 8N RIAI S.L. | > A. Egpppr
En rr at + 8:)
<r", (5)
where x" E Eee 7 "r'(ss, ax)) and A is positive real number.
5

---

Published as a conference paper at ICLR 2025
a+0
(a) E-flat reward maxima
(b) A-action robust policy
Figure 2: Illustration of Definitions I and 2. (a) E-flat reward maxima defined in Definition 1, highlighting the E-ttat reward r" region around 0" (b) A-action robust policy as defined in Defini- tion 2, demonstrating how the policy maintains r" within perturbations of size A.
Built upon the definitions, we here address how the flat reward maxima link to the action robust policy. To formalize the rigorous link, we provide mathematical statement as follows.
Proposition 1 (Flat reward links action robustness) If 0* is E-flat reward maximum, then the policy ng- is -action robust, where:
A ( +
(6)
and J(0") == Volle(s) 0=0- is the Jacobian matrix of the mean action pa(s) with respect to 8, evaluated at 0".
The proof of the proposition is in Appendix A.
Remark 1.1 (A link Max-Min problem of action robustness) For -acactoon robust policy de- rived by E-flat reward maxima 0", the policy directly satisfies the objective of action robust MDP:
0r = arg max
min SSA
Banp.a
7'r(st, as 6.)
(7)
which implies that flatter reward yields the robustness against action perturbations.
Remark 1.2 (An informal link to other factors) We here provide an intuition that the reward sur- face flatness also links to the robustness against the perturbations of other factors, including reward function and transition probability. i) For reward function perturbations, i.e., r(s,a) + f(s, a), directly corresponds the perturbations of loss function in the supervised learning case. Thus, when a reward function has merely slight perturbations, a maximum located at flatter surface tends retain higher rewards. ii) When the MDP's transition probability has perturbations, i.e., (())ss,a) - P(s'(s,, (s,''s,,0) a), even the same action-state pair results in the different transition other states. When the state transition changes, it gives different reward values, thus it relates to the undesired perturbation of rewards even with the same state-action pair. Consequently, it implies that reward flatness also promotes the robustness against the transition probability changes, which eventually links the reward changes.
EXPERIMENTAL RESULTS
MuJoCo tasks: We conduct extensive experiments across three continuous control tasks from the MuloCo environment (Todorov et al., 2012): HalfCheetah-v3, Hopper-v3, and Walker2d-v3.
Algorithms be considered: The evaluations are mainly for comparing the standard PPO algo- rithm and our SAM-enhanced PPO, which is called SAM+PPO. SAM+PPO's objective is:
min max E Ep.ror. Se
(8)
1=
where it solves the Min-Max problem with the perturbation e that maximally worsens the re- ward, i.e., maximizing the minus-signed reward. Our SAM+PPO refers the prior works (Lee

---

Published as a conference paper at ICLR 2025
al., 2024). Also, add recent baseline of robust reinforcement learning: Robust Natu- ral Actor-Critic (RNAC) (Zhou et al., 2024) and Robust Adversarial Reinforcement Learning (RARL) (Pinto et al., 2017).
We evaluate PPO, RNAC, RARL, and SAM+PPO under various pertur- bations, including action noise, transition probability changes, and reward function variations. Also, we visualize the reward surfaces to analyze the flatness achieved by PPO and SAM+PPO and com- pare our method with existing robust reinforcement learning algorithms. Additional evaluations of other RL algorithms and SAM-enhanced version are presented in Appendix D.2
Experiments to be done:
5.1 EXPERIMENTAL SETUP
For PPO, SAM+PPO, RNAC, and RARL, we use multi-layer perceptron (MLP) architecture with three layers for their policy networks. Also, for each experimental case, PPO and SAM+PPO are trained in the nominal (unperturbed) environment using identical hyperparametters for fair com- parison. For RNAC and RARL, we tune the hyperparameters to achieve the performance reported in (Zhou et al., 2024), and (Pinto al., 2017) respectively, specified in the original RNAC and RARL implementation for the environments. Each experiment was conducted over five independent trials, each initialized with different random seed to ensure statistical significance. Furthermore, for each evaluation, we performed 100 evaluation runs and averaged the results to enhance the stability and accuracy of our findings. We added the detailed hyperparameters and settings in Appendix B.
5.2 ACTION ROBUSTNESS EVALUATION
We first evaluate the action robustness, which seems to be the most highly related factor to flat rewards. To simulate the perturbed actions, we add zero-mean Gaussian noise with varying standard deviations da ranging from 0.0 to 0.5 to the agent's s actions. The noisy action agsisy is computed as: Onaisy a + (0, 2 where the original action output by the agent's policy. If the addition of noise results in actions outside the valid range, clip the actions to the allowable bounds of the environment.
Action Robuddess Evalustion
Action Robustnes Evmiuttion
Action Robustnes Evaleriion
6000
PBO
4020
PPO
5000
PPO
RNAC
Average Krward
RNAC
Averey Becwnl
RNAC
KARL
4000
8000
Avenor
2000
3000
2000
0.3
0.3
0.2
Stundar Deviation
Suundard Deviation
Stadlad Devittion
(a) HalfCheetah-v3
(b) Hopper-v3
Walker2d-v3
Figure 3: Action robustness evaluation across different environments. The average return is plotted against the action noise standard deviation da-
Figure presents the results of the action robustness evaluation across three distinct environments: HalfCheetah-v3, Hopper-v3, and Walker2d-v3. In all three environments, the performance of all agents gradually decreases as the action noise increases, which is expected due the increased uncertainty in action execution. Importantly, SAM+PPO consistently outperforms PPO and RNAC across all noise levels, and the rate of performance degradation much slower compared to PPO and RNAC. RARL demonstrates competitive robustness performance, maintaining relatively higher returns compared PPO and RNAC under moderate noise levels. These results indicate that the flat reward achieved by SAM+PPO makes the policy less sensitive to action perturbations, thus coinciding with our claim.
5.3 TRANSITION PROBABILITY ROBUSTNESS EVALUATION
We perturb the transition dynamics by modifying the physical parameters of the environment. By tailoring the given three environments, we vary the torso and friction coefficients.
7

---

Published as a conference paper at ICLR 2025
Variation in Torso Mass: We scale the torso mass of the agent by factors ranging from 0.5 to 1.5 times its nominal value in increments of 0.1. This tests the agent's ' ability to adapt to changes in its own dynamics, such carrying additional weight or structural modifications.
Mass Relnstmess Evaluction
Mass Rollestnees Evallaatinn
Maes R Rohustness Evaluction
6000
Reward
Average Reward
3000
9000
PPU
PPO
PPU KNAC
2000
RARL
RNAC:
2000
RARL
RNAC
2000
RARL
0.6
tb
12
14
04
13
1.2
1.4
6.6
0.8
1.
12
14
Moss Factor
Mass Factor
Mase Factor
(a) HalfChectah-v3
(b) Hopper-v3
(c) Walker2d-v3
Figure 4: Robustness evaluation under torso mass variations. The average return plotted against the mass scaling factor. The nominal torso mass factor is 1.0.
As shown in Figure 4, the performance of all agents decreases as the torso mass deviates from the nominal. However, SAM+PPO consistently achieves higher returns than others. RNAC generally shows better performance than PPO but worse than SAM+PPO. RARL demonstrates competitive performance, maintaining relatively higher returns compared to PPO and RNAC under moderate mass scaling factors. These results suggest that the policies learned by SAM+PPO are more robust to changes in the agent's dynamics, likely due to the flatter reward surface that reduces sensitivity to such perturbations.
Variation in Friction Coefficient: We vary the friction coefficients between the agent's feet and the ground from 0.4 to 1.6. The nominal friction coefficient is 1.0. This simulates different ground conditions, such slippery rough surfaces, which affect the agent's ability move and balance.
Priction Robustness Evalcation
Frictioe: Robestness Evaination
Priction Robustness Evaluction
6000
7000
Reward
PPO
Average Reward
PPU
6200
3000
KARL
6000
RARL
PPO
2000
2000
RARL
3000
0.30
W75
1:00
1.50
0.30
0.75
1.00
135
120
0.30
0.75
130
1:25
130
Friction Factor
Friction Factor
Friction Factor
(a) HalfCheetah-v3
(b) Hopper-v3
(c) Walker2d-v3
Figure 5: Robustness evaluation under friction coefficient variations. The average return plotted
against the friction scaling factor.
Figure presents the friction robustness evaluation. SAM+PPO shows a clear advantage over PPO, maintaining higher performance as the friction coefficient varies. SAM+PPO generally outperforms RNAC. RNAC shows strong performance with larger friction factors for Walker2d-v3, but it severely degrades when the friction decreases (less friction makes the ground surface slippery). RARL main- tains relatively stable performance across range of friction coefficients, particularly excelling in environments with extreme cases of friction where other methods experience significant degrada- tion. These results indicate that the policies learned by SAM+PPO are more adaptable to different surface conditions, which is critical for real-world applications where ground properties can change unpredictably.
Mass and Friction Joint Variations: Herein, we focus on evaluating the agents' performance across the joint variations mass and friction, which shows a comprehensive view robustness.
Figure 6 presents the reward heatmaps for three environments. The widely highlighted regions of SAM+PPO demonstrate that is robust compared to others across a broader spectrum of environ- mental variations. While SAM+PPO shows slightly lower performance in scenarios with extremely high friction, significantly outperforms others in environments with lower friction, such as slip- pery surfaces. RARL exhibits competitive performance in certain regions, particularly maintaining
8

---

Published as a conference paper at ICLR 2025
better performance than PPO and RNAC under specific combined perturbations of mass and fric- tion. Surprisingly, RNAC shows marginal gains over PPO. underscores the superior robustness of SAM+PPO in adapting compounded changes in mass and friction over others.
1004
500
Paciin Paate
12
500
Faciam Paste
1
906
000
150
1000
200
Mias Factu
Macs Factut
Mase Factie
(a) HalfCheetah-v3
(b) Hopper-v3
(c) Walker2d-v3
Figure 6: Reward heatmaps for combined mass and friction variations across different environments. Yellow to dark blue indicates better to worse performance, respectively.
Table I summarizes the performance of PPO, SAM+PPO, RNAC and RARL under various perturba- tion scenarios across different environments. These results collectively demonstrate that SAM+PPO enhances the agent's ability to withstand simultaneous environmental perturbations in mass and friction, providing more robust and reliable performance compared to others.
Table Performance comparison under various perturbations.
Perturbation
Metric
HalfCheetah-v3
Hopper-v3
Walker2d-r3
Nominal
Perturbed
Nominal
Perturbed
Nominal
Perturbed
PPO
4758
1469( 3289))
3217
1467( -1770)
4883
607(
Action Noise
RNAC
5181
-
3115
(321(-2124)
4147
652 (1195)
RARL
4996
-412()1584)
2819
1645 (1174)
4020
764| (62(6))5))
SAM+PPO
6323
1949( 1574)
3766
2312( -145()))551)
5129
2033( (3006)
PPO
4837
3865(-972)
3215
1556( 1659)
4957
782( (4175)
Mass Scale Factor 1.2
RNAC
5185
-811(-611)
3303
(139(-114))
4373
741(-3632
RARL
6561
6016( 5455)
3136
3078( (58)
4303
4343( 50)
SAM+PPO
6562
6210 (352)
3-199
3508(+99)
5205
5284(+79)
PPO
4723
4774(+-1)
3075
1580(-1495)
4996
4756( 240)
Friction Coefficient 0.88
5422 24)
3653
679 (994))
4438
3844( 534)
RARL
6560
6516 44)
3136
4:26
2828( 1598)
SAM+PPO
6562
6539 23)
3500
2728(
(728(-777 772)
5073
5134( +611
(-/+) values mean the performance degradation from 'Nominal' 'Perturbed.',
Boldface is used highest performance 'Nominal', smallest performance degradation in 'Perturbed'
5.4 REWARD FUNCTION ROBUSTNESS EVALUATION
For testing reward function robustness, notable difference this part is that reward perturbations added during training. As done in the previous testing, when the training is done with nominal rewards and then perturbed during testing, the policy thus yields the expectation of the perturbed rewards, leading trivial message. Therefore, how different agents are trained against the reward perturbations. Specifically, we introduced zero- mean Gaussian noise with standard deviation of a, 0.1 to the rewards. After training, evaluated the policies in the nominal environment.
adopt noise rewards during training confirm
Table 2: Performance comparison of agents trained with and without reward noise (o, 0.1)
HallCheetah--3
Hopper-v3
Walker2d-v3
Algorithm
Nominal
Noisy
Nominal
Noisy
Nominal
Noisy
PPO
4820
3150
2945(
205)
4780
2204(-25700
RNAC
5423
4088(
1335)
3211
3035(
176)
4184
3172( 10122
RARL
5620
4617( 1003)
3124
2903
131)
4388
3085( 1303)
SAM+PPO
6530
5990 (5440)
3505
3377
128)
5120
4226( 894)
values means the performance degradation from 'Nominal' 'Noisy.'
As shown in Table 2. the degradation of SAM+PPO is significantly less than that of PPO across all cases. SAM+PPO is less sensitive to variations of rewards, leading to better stability. critical when the reward may not always be accurate consistent in real-world settings.

---

Published as a conference paper at ICLR 2025
5.5 REWARD SURFACE VISUALIZATIONS AND FLATNESS MEASUREMENTS
We visualize the reward surface by using the method in (Sullivan et al., 2022). Also, we first compute the flatness metrics, including the maximum eigenvalue of the Hessian, i.e., Amax and the Low-Pase Filter (LPF) flatness measure (Keskar et al., 2017; Bisla et al., 2022).
HaPCheettt-a3 [Hubco] Mewn Esicodic Reward
Hopperv3 [uuocc Meen Episooic Reward
Wa |Muoco |Neen Episodic Reward
6000
300
4200
1890
1.0
1.
2.5
0.5
3.5
-04
0.0
-05
2.2
-3
5.
(a) PPO HalfCheetah-v3
(b) PPO Hopper-v3
(c) PPO Walker-v3
HatCheetah-w3] |M Mean Esitodic Reward
Hopperiv: Mujoco| Meen Episodic Reward
walker2iv3 Mujoco |eean Episodic Reward
0030
6300
3500
4239
500
230
3000
1100
-10
-10
(d) SAM+PPO HallCheetah-v3
(e) SAM+PPO Hopper-v3
(1) SAM+PPO Walker-v3
Figure 7: Reward surface visualizations of agents. The and y axes represent perturbations along random directions in the parameter space, and the z-axis represents the average return.
Figure 7 shows the reward surfaces for PPO and SAM+PPO in the HalfCheetah-v3 environment. Similar patterns are observed in Hopper-v3 and Walker2d-v3. The SAM+PPO reward surface is noticeably flatter, meaning that the agent's performance is less sensitive to parameter changes. This supports that SAM yields flatter reward maxima, and it is strongly correlated to more robust policies.
As shown in Table SAM+PPO generally achieves lower flatness metries compared PPO, indi- cating that SAM+PPO converges to flatter minima in the policy parameter space.
Table Flatness metries for PPO and SAM+PPO (L: indicates that lower is better).
Metries
Amea J. (Keskar et al., 2017)
LPF J (Bisla al., 2022)
Environment
HalfCheetah--3
Hopper-v3
Walker2d-v3
HalfCheetah--33
Hopper-v3 Walker2d-v3
PPO SAM+PPO
15192.95 275.93
131.07 80.86
7239.59
0.0385
0.00034
0.0269
271.91
0.00097
0,00018
0.00028
6 CONCLUSION
In this paper, we aim to unravel the meaning of flat reward maxima from the perspective of the robustness of reinforcement learning. Theoretically, we provide a formal link between flatness and robustness. Empirically, we broadly confirm our claim, showing that SAM+PPO outperforms stan- dard PPO and existing robust RL algorithms like RARL and RNAC, validating the efficacy of flat rewards. This work emphasizes the importance of flatness in RL's policy space. We believe that it can broaden new avenues for developing robust RL algorithms by combining with prior robust RL methodologies pursue strongly robust RI. in real-world applications.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by the Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No. RS-2020-11201336, Ar- tificial Intelligence Graduate School Program (UNIST)), (No. novative Human Resource Development for Local Intellectualization program), and the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. RS- 2024-00459023..
In-
REFERENCES
Maksym Andriushchenko, Francesco Croce, Maximilian Müller, Matthias Hein, and Nicolas Flam- marion. A modern look at the relationship between sharpness and generalization. arXiv preprint arXiv:2302.07011, 2023.
Kishan Panaganti Badrinath and Dileep Kalathil. Robust reinforcement learning using least squares policy iteration with provable performance guarantees. In International Conference on Machine Learning. pp. 511-520. PMLR, 2021.
Recep Yusuf Bekci and Mehmet Gümüůş Visualizing the loss landscape of actor critic methods with applications in inventory optimization. arXiv preprint arXiv:2009.0239) 2020.
Devansh Bisla, Jing Wang, and Anna Choromanska. Low-pass filtering sgd for recovering flat optima in the deep learning optimization landscape. In International Conference on Artificial Intelligence and Statistics, pp. 8299-8339, PMLR, 2022.
Debora Caldarola. Barbara Caputo, and Marco Ciccone. Improving generalization in federated learning by seeking flat minima. European Conference on Computer Vision, pp. 654-672. Springer, 2022.
Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park. Swad: Domain generalization by seeking flat minima. Advances in Neural Infor- mation Processing Systems, 14:22405-22411 2021.
Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-sgd: Biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019(12):
124018, 2019.
Esther Derman, Matthieu Geist, and Shie Mannor. Twice regularized mdps and the equivalence between robustness and regularization. In Advances in Neural Information Processing Systems, volume 34, pp.
2021.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference Machine Learning, pp. 1019-1028. PMLR, 2017.
Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. arXiv preprint arXiv:1703.11008, 2017.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-awaie minimiza- tion for efficiently improving generalization. In International Conference on Learning Represen- tations. 2021.
Vineet Goyal and Julien Mathematics of Operations Research, 48(1):203-22) 2266 2023.
Robust markov decision processes: Beyond rectangularity.
Julien Grand-Clément and Christian Kroer. Scalable first-order methods for robust mdps. In Pro- ceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 1p..1200 12086-12094, 2021.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- on machine learning, pp. 1861-1870. PMLR, 2018.

---

Published as a conference paper at ICLR 2025
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1):1 19977.
Andrew Hyys, Logan Engstrom, Shibani Santurkar, Dimitris Tsipras, Firdaus Janoos, Larry Rudolph, and Aleksander Madry. closer look at deep policy gradients. In International Conference on Learning Representations (ICLR), 2020. URL https: https///ppeeeeeeeeeeeeee forum??d=ryxxEkBBtR
Garud N Iyengar. Robust dynamic programming. Mathematics of Operations Research, 30(2):
257-280, 2005.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wil- son. Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv: 10X3::8803.05477, 2018.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batcl training for deep learning: Generalization gap and sharp minima. In Proceedings of the International Conference on Learning Representations (ICLR), 2017. URL httpps://ogt http:///pppeeeiieeeee net www.:nneeootmmd/ttttieeeoom forum?id=H1 ooRl1
Navdeep Kumar, Esther Derman, Matthicu Geist, Kfir Levy, and Shie Mannor. Policy gradient for rectangular robust markov decision processes. Advances in Neural Information Processing Systems, 36, 2024.
Hojoon Lee, Hanseul Cho, Hyunseung Kim, Dachoon Gwak, Joonkee Kim, Jaegul Choo, Se- SouYoung Yun, and Chulhee Yun. Plastic: Improving input and label plasticity for sample efficient rein- forcement learning. Advances Neural Information Processing Systems, 36, 2024.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss land- scape of neural nets. Advances in neural information processing systems, 31, 2018.
Yan Li, Guanghui Lan, and Tuo Zhao. First-order policy optimization for robust markov decision process. arXiv preprint 220.:2299911579 2022.
Daniel J. Mankowitz, Abdul Fatir Abbas, Cheng Soon Ong, Shic Mannor, and Marcus Hutter. Ro- bust reinforcement learning for continuous control with model misspecification, In Proceedings of the 36th International Conference Machine Learning (ICML), pp.
PMLR, 2019.
Shie Mannor, Ofir Mebel, and Huan Xu. Robust mdps with k-rectangular uncertainty. Mathematics of Operations Research, 4144):444-1150, 2016.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 8(7540):529-533 2015.
Jun Morimoto and Kenji Doya. Robust reinforcement learning. Neural Computation, 17(2):335- 359, 2005a.
Jun Morimoto and Kenji Doya. Robust reinforcement learning. Neural computation, 17(2):335-359) 2005b.
Ali Nilim and Laurent El Ghaoui. Robust control of markov decision processes with uncertain transition matrices. Operations Research, 53(5):780-798, 2005.
Kei Ota, Devesh K Jha, and Asako Kanezaki. A framework for training larger networks for deep reinforcement learning. Machine Learning, June 2024. doi: 022077/50999--04--45557777777 URL https: /oii. /tg/10.... 1010000//1100 024 06547-6.
Kishan Panaganti and Dileep Kalathil. Sample complexity of robust reinforcement learning with a generative model. In International Conference on Artificial Intelligence and Statistics, pp. 9582- 9602. PMLR, 2022.
Anay Pattanaik, Zhenyi Tang, Shuijing Liu, Gautham Bommannan, and Girish Chowdhary. Ro- bust deep reinforcement learning with adversarial attacks. In 17th International Conference Autonomous Agents and Multiagent Systems, AAMAS 2018, pp. 2040-2042, International Foun- dation for Autonomous Agents and Multiagent Systems (IFAAMAS), 2018.
12

---

Published as a conference paper at ICLR 2025
Xue Bin Peng, Marcin Andrychowicz, Wojciech Zaremba, and Pieter Abbeel. Sim-to-real transfer of robotic control with dynamics randomization. In 2018 IEEE international conference on robotics and automation (ICRA), pp. 3803-3810. IEEE, 2018.
Henning Petzka, Michael Kamp, Linara Adilova, Cristian Sminchisescu, and Mario Boley. Relative flatness and generalization. Advances in neural information processing systems, 34:18420-18432, 2021.
Lerrel Pinto, James Davidson, Rahul Sukthankar, and Abhinav Gupta. Robust adversarial reinforce- ment learning. In International conference on machine learning, pp. 2817-2826. PMLR, 2017.
Zhe Qu, Xingyu Li, Rui Duan, Yao Liu, Bo Tang, and Zhuo Lu. Generalized federated learning via sharpness aware minimization. In International conference on machine learning, pp- 18250- 18280. PMLR, 2022.
Aravind Rajeswaran, Sarvjeet Ghotra, Balaraman Ravindran, and Sergey Levine. Epopt: Learn- ing robust neural network policies using model ensembles. In Proceedings of the International Conference on Learning Representations (ICLR), 2017. URL https: http:::///ppeeeeee....mo
Aurko Roy, Huan Xu, and Sebastian Pokutta. Reinforcement learning under model mismatch. Ad- vances neural information processing systems, 30, 2017.
Guangyuan Shi, Jiaxin Chen, Wenlong Zhang, Li-Ming Zhan, and Xiao-Ming Wu. Overcoming catastrophic forgetting incremental few-shot learning by finding flat minima. Advances in neural information processing systems, 34:6747-6761. 2021.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittiisser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
Ryan Sullivan, Jordan K Terry, Benjamin Black, and John P Dickerson. Cliff diving: Exploring reward surfaces in reinforcement learning environments. In International Conference on Machine Learning, pp. 20744-20776. PMLR, 2022.
Yan Sun, Li Shen, Shixiang Chen, Liang Ding, and Dacheng Tao. Dynamic regularized sharpness aware minimization in federated learning: Approaching global consistency and smooth landscape. In International Conference on Machine Learning, pp. 32991-33013 PMLR, 2023.
Richard S Sutton and Andrew G Barto. Reinforcement Learning: An Introduction. MIT Press, 2nd edition, 2018. URL http: http::/ttccomppp//////// /tt.///ccoppeeeeeeddee....eeeese net sest//bbo/thheebbbottet book- 2nd. html.
Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
Chen Tessler, Yonathan Efroni, and Shie Mannor. Action robust reinforcement learning and appli- cations in continuous control. In International Conference Machine Learning, pp. 6215-6224. PMLR, 2019.
Emanuel Todorov. Tom Erez, and Yuval Tassa. Mujoco: A physies engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pp. 5026-5033. IEEE, 2012.
Yue Wang and Shaofeng Zou. Online robust reinforcement learning with model uncertainty. Ad- vances in Neural Information Processing Systems, 14:7193-7206,, 2021.
Yue Wang and Shaofeng Zou. Policy gradient method for robust reinforcement learning. In Inter- national conference machine learning, pp. 13484-23526, PMLR, 2022.
Wolfram Wiesemann, Daniel Kuhn, and Berç Rustem. Robust markov decision processes. Mathe- matics of Operations Research, 38(1):153-188, 2013.
13

---

Published as a conference paper at ICLR 2025
Ronald J Williams. Simple statistical srudient-following algorithms for connectionist reinforcement learning. Machine learning, 8:229- 256, 1992.
Huan Xu and Shie Mannor. Distributionally robust markov decision processes. Advances in Neural Information Processing Systems, 23, 2010.
Wenhao Yang, Liangyu Zhang, and Zhihua Zhang, Toward theoretical understandings of robust markov decision processes: Sample complexity and asymptotics. The Annals of Statistics, 50(6): 3223-3248, 2022.
Ruida Zhou, Tao Liu, Min Cheng. Dileep Kalathil, PR Kumar, and Chao Tian. Natural actor-critic for robust reinforcement learning with function approximation. Advances in neural information processing systems, 36, 2024.
Zhengqing Zhou, Zhengyuan Zhou, Qinxun Bai, Linhai Qiu, Jose Blanchet, and Peter Glynn. Finite- sample regret bound for distributionally robust offline tabular reinforcement learning. In Interna- tional Conference on Artificial Intelligence and Statistics, pp. 3331-3339. PMLR, 2021.
A PROOF OF PROPOSITION
We aim to show that an E-flat reward maximum 0* (as per Definition 1) corresponds to A -action robust policy Tg- (as per Definition 2), with:
A [(0')
where J(0") Vopoo())omoo- is the Jacobian matrix of the mean action (a(() with respect to 0, evaluated at 0
Consider stochastic policy (alss) parameterized by 0, assumed be Gaussian policy with mean action (8)) and fixed covariance matrix Σ:
[a(a = (a: (o(s) E).
(10)
Let E R be perturbation such that ||||| < &. Performing first-order Taylor expansion around 0":
(g)+ mo- (s) + J(0")e O(]ee)3). O((e)))
(11)
Define the action perturbation:
8, pe (s:) no- (st) =
(12)
Then,
||&.1| s + {(]0)) +
(13)
From Definition 1, for ||||| E:
E Be
Σ Er(r(s), =r-
(14)
Noting that shifting the mean of the Gaussian policy is equivalent to shifting the action variable, we have:
T st) = TT- (ar - 6e 5c(8)).
(15)
Thus, the expected cumulative reward under the perturbed policy becomes:
Eggp.eraye
Σ >rr(s(,
Σ (()())))))
(16)
rrr(st, +
(17)
14

---

Published as a conference paper at ICLR 2025
Since Bemp,awto is EEt= ('r(st, ac)l = = it follows that:
E Bawppant
Σ >'r(s a + +z) = =r"
(18)
This implies that for all || s A with A" |]( + O(E2), the expected cumulative reward remains r", satisfying the first condition of a A -action robust policy as per Definition
Moreover, since there exists with |||]| & such that Es,a 7 ("r(st; < r-, the corre- sponding ór satisfies || > A", and thus the expected cumulative reward decreases below r' for perturbations larger than A fulfilling the second condition.
A.1 DISCUSSION ON THE BOUNDS OF THE JACOBIAN
In the proof of Proposition 1, we derive a bound on the action robustmess parameter A in terms of the Jacobian [|[(0)) and the flatness parameter E:
A ((00))+++(0
We acknowledge that in practical applications, especially when using deep neural networks for pol-
icy representations, it is challenging guarantee that the Jacobian [0] is bounded. Deep neural networks can have complex architectures and nonlinear activation functions that may lead to large gradients and, consequently, large Jacobian norms.
However, several common practices during neural network training help control the magnitude of the Jacobian:
Weight Regularization: Techniques such as Le regularization (weight decay) penalize large weights, which indirectly constrains the Jacobian. By limiting the magnitude of the weights, the sensitivity of the network outputs to changes in the inputs and parameters is reduced.
Applying gradient clipping during optimization prevents the gradients from becoming excessively large. This helps maintain the stability of parameter updates and controls the growth of the Jacobian norm.
Gradient Clipping:
Using activation functions with bounded derivatives (e.g.. hyperbolic tangent, sigmoid) limits the rate of change of the network outputs with respect to the inputs and parameters, thus contributing to bounding the Jacobian.
Bounded Activation Functions:
Normalization Techniques: Methods like batch normalization and layer normalization help stabilize the learning process and control the scale of activations and gradients, affecting the Jacobian norm.
Regarding the SAM objective, SAM promotes convergence flatter regions in the loss landscape by considering adversarial perturbations of the network parameters during optimization with the objective function Equation 3.
where (()) is the loss function, and defines the size of the neighborhood around the current pa- rameters By optimizing for the worst-case loss within a neighborhood, SAM discourages sharp minima with high curvature, which are associated with large Hessian norms. This process inherently suppresses large gradients and promotes smoother variations of the loss with respect to parameter changes. While SAM primarily targets the curvature of the loss landscape (second-order informa- tion), this effect also influences the Jacobian the network outputs with respect to the parameters (first-order information).
Our empirical results demonstrate that SAM+PPO leads to policies that are more robust to pertur- bations in both the parameters and the environment dynamics. This supports the practical relevance of our theoretical findings, suggesting that SAM effectively contributes to controlling the Jacobian norm in practice.
While it may not be possible to guarantee a bounded Jacobian everywhere in the network, the combi- nation of SAM and common training practices helps in maintaining the Jacobian within reasonable
15

---

Published as a conference paper at ICLR 2025
range. This enhances the robustness of the learned policies and aligns with the conditions assumed in our theoretical analysis.
B EXPERIMENTAL DETAILS
In this appendix, we provide detailed information about the hyperparameters and experimental set- tings used in our experiments.
B.I NETWORK ARCHITECTURE
For all agents, including PPO, SAM+PPO, and RNAC, we employ multi-layer perceptron (MLP) architecture for both the actor (policy network) and the critic (value network). The network consists of an input layer matching the state dimension of the environment, followed by three fully connected hidden layers, each with 64 neurons and Tanh activation functions. The output layer of the actor network produces the parameters of the action distribution (mean and log standard deviation for Gaussian policies), while the critic network outputs single scalar value representing the state value estimate. We apply orthogonal initialization to all layers to enhance training stability.
B.2 HYPERPARAMETER,
TRAINING SETTINGS
We use identical hyperparameters for PPO, SAM+PPO, and RNAC to ensure a fair comparison. The shared hyperparameters are as follows: the discount factor 7 is set to 0.99, the GAE parameter is 0.95, and the PPO clip parameter is 0.2. Both the actor and critic learning rates set to 3 x 10 with the Adam optimizer used for optimization. The batch size is 2048, and the mini-batch size is 64, with 10 PPO epochs per update (Kepcchs 10). We employ gradient clipping with maximum norm of 0.5, learning rate decay, and state normalization for all agents.
For SAM+PPO, we introduce the SAM optimizer with P parameters of 0.008 for HalfCheetah-v3 and Walker2d-v3, and 0.01 for Hopper-v3, which control the neighborhood size for sharpness-aware minimization. For RNAC, we adjust the uncertainty set to "IPM" model uncertainties in the transition dynamics and configure the number of next steps to 2, allowing the agent to consider multiple potential future states.
For comparison with existing robust reinforcement learning algorithms, we trained Robust Adver- sarial Reinforcement Learning (RARL) (Pinto et al., 2017) using the following hyperparameters specific to the Hopper-v3 environment. The batch size was set to 4000 to align with the original RARL implementation, and the mini-batch size remained at 64. The network's hidden width was increased to 100 neurons per layer to match the architecture used in RARL. Both the protagonist and adversary actors and critics were trained with learning rates of 0. 0.0003. The discount factor was set to 0.995, and the GAE parameter À was set to 0.97, consistent with the original RARL settings. The entropy coefficient was maintained at 0.0, and weight regularization was applied with a parameter of 0.00005. We the adversarial fraction to 0.25 to control the adversary's action strength and configured the number of iterations for both the protagonist and adversary to 1. All other settings, including gradient clipping, learning rate decay, state normalization, and orthogonal initialization, were kept consistent with PPO, SAM+PPO, and RNAC to ensure a fair comparison.
C DETAILS AND ANALYSIS OF SAM INTEGRATED WITH PPO
In this appendix, we provide a detailed discussion on the integration of SAM with RL, especially with PPO. We elaborate on the optimization process of the min-max objective, discuss the compu- tational overhead introduced by SAM, and analyze the sensitivity to the perturbation radius p.
C. OPTIMIZATION OF THE MIN-MAX OBJECTIVE
The SAM optimization modifies the standard PPO objective by introducing maximization parameter perturbations within an & norm ball of radius p. For convenience, we restate the objective presented in Equation from Section 3.3:
16

---

Published as a conference paper at ICLR 2025
min max L(O + c), Iep
(19)
where I(0) is the PPO loss function, and 9 represents the policy parameters.
C.I.1 LINKAGE TO ACTION ROBUST MDP
In PPO, the loss function L(0) is inherently tied to the optimization of the expected cumulative re- ward. Specifically, PPO aims maximize the expected cumulative discounted reward by optimizing the policy parameters 0. The PPO loss function can be expressed in relation to the reward function rs8, a) as follows:
00
L(0) -Esrpp
Σ
r'r(s( 0t)
(20)
where the negative sign indicates that minimizing [(0) equivalent to maximizing the expected cumulative reward. This formulation aligns with the reward-based learning paradigm of RL, where the objective is to find a policy that yields the highest expected cumulative reward across all possible state-action pairs.
With Equation 20, Equation 19 can be expressed as:
oc Σ
min max -Egnggartate Sel
Nrr(s)
(21)
which aligns with Equation in Section 5, incorporating the cumulative discounted reward into the expectation. Both formulations aim to optimize the policy parameters / by considering worst-case perturbations within a specified norm bound.
Based on Proposition Remark 1.1 of Section 4 and the corresponding proof in Appendix A, we derive that the SAM+PPO objective function aims to find the E-flat reward maxima as defined in Definition 1. Consequently, this leads to the formulation of A-action robust policy as defined in Definition 2. Therefore, the objective function aligns with Equation 2 of the main text:
max min Eguppane 1
7'r(s + +6i)
(22)
where 8, represents perturbations to the actions time t, and 3 is scaling factor analogous to P. This formulation corresponds to an Action Robust MDP, where the policy seeks to maximize the expected cumulative reward while minimizing the impact of worst-case perturbations to actions within a specified norm bound.
C.1.2 OPTIMIZATION PROCEDURE
Optimizing the min-max objective in Equation 19 involves integrating the SAM approach with the PPO framework. Below, we present the integrated algorithm and provide a detailed explanation of each step.
17

---

Published as a conference paper at ICLR 2025
Algorithm SAM Integrated with PPO
Require: Initial policy parameters Oo. perturbation radius learning rate clipping parameter
1: for each iteration do
2: Collect Trajectories: Use current policy to collect trajectories by interacting with the
environment.
3: Compute Advantages: Calculate advantages A, and returns based collected trajectories.
4: Compute PPO Loss:
L(O) = min (r: (ii()))
5: Compute Base Gradient: go VoT(0)
6:
Apply SAM Perturbation:
90
c" =
1190 l|
Compute Perturbed Loss: Epertutedd L(B + E*) Compute Perturbed Gradient: go+t- V [0 + +
8:
9:
Update Policy Parameters: 0 +- 0g0+c
10:
end for=0
1. Collect Trajectories: Te - Collect trajectories using Th
Using the current policy parameters A, we collect batch of trajectories by interacting with the environment. These trajectories essential for estimating the PPO loss.
2. Compute Advantages: Å = Calculate advantages based on collected trajectories
Advantages A, are computed to reduce variance in the policy gradient estimates, enhancing the stability of training.
3. Compute PPO Loss: L( E min (r, clpp)r,)) +))+
The PPO loss function L(0) incorporates clipping mechanism to prevent large updates, ensuring stable policy optimization.
4. Compute Base Gradient: ge VeL(O)
We compute the gradient of the PPO loss with respect to the policy parameters A. This gradient indicates the direction in which the loss increases most rapidly.
5. Apply SAM Perturbation: e" === 9e
P
While the inner maximization problem theoretically considers all possible directions within the p-ball around 0, choosing in the direction of the gradient go is first-order approxima- tion that efficiently captures the worst-case perturbation. This approximation rooted in the Taylor expansion of the loss function, where the gradient direction signifies the direction maximum loss increase. By perturbing in this direction, we effectively approximate the maximum increase in loss within the allowed perturbation magnitude p.
6. Compute Perturbed Loss: Lpentured
L(
An additional forward pass is conducted using the perturbed parameters 0 + E* to evaluate the loss under perturbed conditions. This step assesses how the loss behaves when the parameters are subjected to the worst-case perturbation identified in the previous step.
7. Compute Perturbed Gradient:
ge+e- V (0++)))
We compute the gradient of the perturbed loss with respect to the policy parameters. This gradient informs the parameter update step by indicating how to adjust @ to minimize the worst-case loss.
8. Update Policy Parameters: 0- ogge+e"
Where a is the learning rate. The policy parameters 9 are updated using the gradient of the perturbed loss, thereby incorporating the robustness introduced by SAM. This update step moves the parameters in the direction that minimizes the worst-case loss, enhancing the policy's robustness against adversarial perturbations.
18

---

Published as a conference paper at ICLR 2025
This optimization procedure effectively seeks parameters that minimize the worst-case loss within neighborhood radius around the current parameters 8, thereby enhancing the robustness of the policy.
C.2
COMPUTATIONAL OVERHEAD ANALYSTS
Integrating SAM with PPO introduces additional computational overhead due the extra forward and backward passes required to compute the perturbation and update the policy parameters using the perturbed loss.
C.2.1 THEORETICAL ANALYSIS
standard PPO, each optimization step involves one forward pass and one backward pass to com- pute the gradient VoL(0). With SAM, the optimization step involves:
First Forward and Backward Pass: Compute go VVL(0).
Second Forward Pass: Compute L(O + <*) with perturbed parameters.
Second Backward Pass: Compute go+c- 7oI((0+c")
This effectively doubles the number of forward and backward passes per optimization step. In Big O notation, the per-iteration computational complexity increases from ((N) for PPO to ((2N) for SAM+PPO, where N is the number of parameters.
C.2.2 EMPIRICAL ANALYSIS
We empirically measured the training time per iteration for both PPO and SAM+PPO. The results are summarized in Table 4.
Table 4: Comparison of agents Training Time (sec) per Iteration across Environments
Algorithm
HalfCheetah-v3
Hopper-v3
Walker2d-v3
PPO
1.22 1.5(x1.88)
0.13
0.2 (.24(x1.20)
SAM+PPO
0.23(x1.76)
The results show that SAM+PPO incurs approximately 61% increase in training time per iteration compared to PPO. Considering the significant gains in robustness, this overhead is a reasonable trade-off. While the per-iteration training and update times for SAM+PPO are higher compared to standard PPO, the overall training time presents a more favorable comparison. Specifically, as shown in Table 9, SAM+PPO achieves comparable slightly increased overall training times while converging faster and attaining rewards more efficiently. This indicates that SAM+PPO not only enhances robustness and generalization but also maintains efficient training dynamics.
C.3 SENSITIVITY TO HYPERPARAMETER p
The perturbation radius p introduced by SAM is eritical hyperparameter that controls the extent of parameter perturbations during optimization. Selecting an appropriate value for p is crucial for balancing robustness and training stability.
Large pt Encourages the optimizer to find flatter maxima, potentially enhancing robustness. How- ever, if p too large, it may lead training instability convergence issues.
Small p: May result in less robustness gain, as the perturbations are insufficient promote sig- nificant flatness in the parameter space.
We conducted experiments by varying p within practical range to assess its impact on performance.
The values tested were p E {0.001, 0.005, 0.01, 0.05, 0.1}. Figure illustrates the performance of SAM+PPO with different values of p on the Hopper-v3 environment.
19

---

Published as a conference paper at ICLR 2025
Action Robustness Evaluation
4000
PRO
SAM--PO((th0000)
Average Reward
3000
SAM-PPO(tthOO.000 0050)
SAM-PPO(thbO.00)
SAM- PPO(roo0..5)
2000
SAM-PPO(rho0.)
1000
0.0
0.1
0.2
0.3
0.4
0.5
Action Noise Standard Deviation
Figure 8: Sensitivity of SAM+PPO performance to the perturbation radius p in Hopper-v3 environ- ment.
Optimal Range: A value of / == 0.01 provided the best results of robustness and training stability.
Small (t At 0.001, 0.005, the robustness gains minimal compared standard PPO.
Large At 0.05, 0.1, training became unstable, and performance degraded significantly.
C.3.1 EFFECT PPO HYPERPARAMETERS
The integration of SAM into PPO modifies the optimization process by adding an outer loop that seeks parameters robust to perturbations within specified radius p. This outer loop operates in- dependently of PPO's s internal mechanisms. Since SAM's perturbation process focuses on finding flatter regions in the parameter space without altering PPO's update equations or loss function struc- ture, it does not inherently change the role or sensitivity of PPO's 's original hyperparameters.
The perturbation radius P is a pivotal hyperparameter SAM+PPO that directly influences the balance between robustness and training stability. Through our sensitivity analysis, we identified = 0.05 as the optimal value that offers significant robustness gains without compromising training stability. Additionally, integrating SAM does not necessitate retuning of PPO's existing hyperpa- rameters, making the combination both effective and practical.
D ADDITIONAL EXPERIMENTAL RESULTS
D.1 EXPERIMENTS ON DISCRETE ACTION ENVIRONMENTS
In the main text, our experiments focused on continuous control enviromments. To validate the appli- cability and reliability of our SAM-enhanced method in a broader context, we extended our experi- ments to discrete action environments provided by OpenAl Gym: CartPole-vl and Lunarl ander-v2. These environments allow us assess the performance of SAM+PPO in settings with discrete action spaces and different reward structures.
D.1.1 EXPERIMENTAL SETUP
Similar to the experiments in the main text, we compared the standard PPO algorithm with our SAM+PPO. Except for the SAM-specific parameter p, all hyperparameters were kept identical be- tween PPO and SAM+PPO to ensure fair comparison. The network architecture for both agents consisted of three fully connected layers with 64 neurons each, and the output layer used softmax activation function to produce probability distribution the learning rate to x 10-4, the discount factor y to 0.99, and used the Adam optimizer. The perturbation radius for SAM+PPO was set 0.05.
the possible discrete actions. We set
20

---

Published as a conference paper at ICLR 2025
In CartPole-vl and LunarLander-v2. there common elements that allow for changing the transition probabilities (e.g.. mass or friction adjustments), so we focused on evaluating robustness to action and reward perturbations.
D.1.2 ACTION ROBUSTNESS EVALUATION
To evaluate action robustness, we introduced action perturbations by adding zero-mean Gaussian noise with standard deviation da 0.2 to the logits of the policy network before the softmax activation during evaluation, similar to the action noise added the main text experiments.
Table 5: Action robustness evaluation across discrete action environments. The average return is reported over 100 evaluation episodes.
Algorithm
CartPole-vl
LunarLanderrvv2
Nominal
Perturbed
Nominal
Perturbed
PPO
500
464
(-36)
200
175 (-25)
SAM+PPO
500
481 ( 19)
200
188 (-12)
As shown in Table 5, both PPO and SAM+PPO achieve the maximum average return in the nom- inal setting for both environments. Under action perturbations, SAM+PPO consistently outper- forms PPO, exhibiting smaller performance degradation. In CartPole-v.. 1. PPO's 's average return de- creases by 36 points, whereas SAM+PPO's average return decreases by only 19 points. Similarly, in PPO's performance drops by 25 points, while SAM+PPO's performance drops by only 12 points. These results indicate that SAM+PPO enhances robustness action perturbations in discrete action environments.
D.1.3 REWARD ROBUSTNESS EVALUATION
For the reward robustness evaluation, we introduced zero-mean Gaussian noise with standard deviation of o, 0.1 to the rewards during training. After training, we evaluated the agents in the nominal environment without reward noise.
Table 6: Performance comparison of agents trained with and without reward noise (a,, = 0.1).
CartPole-vl
LunarLander-v2
Algorithm
Nominal
Noisy
Nominal
Noisy
PPO
500
432
(-68)
200
165
35)
SAM+PPO
500
458
42)
200
182
-18)
Table 6 shows that SAM+PPO is less sensitive to reward noise during training compared to PPO. In CartPole-vl, PPO's performance decreases by 68 points when trained with reward noise, while SAM+PPO's performance decreases by only 42 points. In JunarL.nner-vy drops by 35 points, whereas SAM+PPO's performance drops by 18 points. This suggests that the flatter reward maxima achieved by SAM+PPO contribute to better robustness against reward pertur- bations.
PPO's performance
D.2 SAM ENHANCED WITH OTHER POLICY GRADIENT ALGORITHMS
To investigate the applicability of our SAM-enhanced approach to other policy gradient algorithms, we conducted experiments by integrating SAM with Trust Region Policy Optimization (TRPO), resulting in TRP+
21

---

Published as a conference paper at ICLR 2025
D.2.1 EXPERIMENTAL SETUP
We used the same experimental setup as the main text, with hyperparameters and network archi- tectures kept consistent between TRPO and SAM+TRPO, except for the SAM-specific parameter p. The perturbation radius p for SAM+TRPO was set to 0.006.
Algorithms to be considered: We compared TRPO and SAM+TRPO.
Experiments to be done: We evaluated the robustness of the agents under action perturbations and transition probability perturbations (mass and friction variations) in the same manner as in the main text.
D.2.2 PERFORMANCE EVALUATION
We report the overall performance of TRPO and SAM+TRPO under various perturbation scenarios in Table 7.
Table 7: Performance comparison under various perturbations for TRPO and SAM+TRPO.
Perturbation
Metric
HalfCheetah-
Hopper-v3
Walker2d-v3
Nominal
Petturbed
Nominal
Pertured
Nominal
Perturbed
Action Noise a = 0.2
TRPO SAM-TRPO
4805
1502 [502(-3300) (-333)
3118
1452 (-1606)
1975
(-3
5502
3975 (5227)
3547
2313 ( 1234)
5097
2052(-3045)
Mass Scale Factor 1.2
TRPO
4837
3865 (-972)
$215
1556 (-1659)
4957
782 (-475))
SAM-TRPO
5562
5210
(2100(352)
3499
3508 (+9)
5205
5-84(++79)
Friction Coefficient
TRPO
4723
4774 (+51)
3075
1580 (-1495)
4996
4756 (-22))
SAM-TRPO
5562
5539 (-00--23)) 23)
3498
2728
770)
5073
5131(+-1)
values means the performance degradation from 'Nominal' each perturbations
The results in Table 7 show that SAM+TRPO outperforms standard TRPO in both nominal and perturbed settings across all environments and perturbation types. Under action noise, SAM+TRPO exhibits significantly less performance degradation compared TRPO. Similarly, under and friction perturbations, SAM+TRPO maintains higher returns, indicating enhanced robustness. In- tegrating SAM with TRPO introduces additional computational overhead due the extra gradient computations required for the SAM optimization step. However, the overall training time remains reasonable, and the robustness gains justify the additional computational cost. The positive results with SAM+TRPO suggest that our SAM-enhanced approach can be applied to other policy gradient and actor-critic algorithms, such as A2C and SAC, The key requirement is that the algorithm must be amenable to gradient-based optimization, allowing for the incorporation of the SAM perturbation step.
E COMPARISON WITH EXISTING ROBUST RL ALGORITHMS
We compare SAM+PPO with other robust RL, mainly focusing on the performance and computa- tional costs. To further validate the robustness of our approach, compare SAM+PPO with two state-of-the-art robust reinforcement learning algorithms: RARL and RNAC. We selected RARL. and RNAC for comparison because they represent prominent approaches in robust reinforcement learning and have demonstrated effectiveness in enhancing robustness to uncertainties.
E.1 ALGORITHM PRINCIPLES
RARL introduces an adversary during training that applies perturbations the environment (Pinto al., 2017). RNAC, on the other hand, leverages natural gradient methods achieve robustness un- der the worst-case distribution within a specified uncertainty set. (Zhou et al., 2024). By comparing with these algorithms, aim to highlight the effectiveness of our approach in achieving robustness while training solely in the nominal environment, without requiring adversarial training or explicit uncertainty modeling.
Existing RRL algorithms like RARL and RNAC improve robustness by training the agent under un- certainty sets adversarial conditions. RNAC uses distributional robustness, optimizing the policy
22

---

Published as a conference paper at ICLR 2025
to perform well under the worst-case distribution within specified uncertainty set. RARL involves training an adversarial agent alongside the protagonist, introducing perturbations during training to simulate worst-case scenarios. In contrast, SAM+PPO trains solely in the nominal environment without explicitly modeling uncertainties or adversaries. Despite this, SAM+PPO achieves compa- rable or superior robustness.
RRL algorithms require knowledge of the uncertainties or adversarial models during training. which may not always be available or accurate. SAM+PPO does not rely on such information. RRL algorithms often involve more complex training procedures, such as training adversarial agents or solving min-max optimization problems. SAM+PPO introduces minimal overhead by incorporating SAM into the optimization process. SAM+PPO is more readily applicable in real-world scenarios where modeling uncertainties is challenging. It enhances robustness without additional assumptions or modifications to the environment.
E.2 BETTER ROBUSTNESS BY SAM+PPO
Table summarizes the performance of each algorithm under specific perturbations Hopper-v3. SAM+PPO outperforms PPO and the RRL algorithms, demonstrating its effectiveness in enhanc- ing robustness without requiring specialized training conditions. We used PPO as an underlying algorithm in RARL, in which TRPO was used in the original paper, ensure consistency in our comparison and elaborate on how SAM bring benefits in comparison to the other approaches
Table 8: Performance comparison under the perturbations in Hopper-v3.
Algorithm
Nominal
Action Noise 0.2
Mass Scale 1.3
Friction coefficient 1.24
PPO
3169
1467 1702)
1400( 1769)
581( 2588)
RNAC
3467
1321 2146)
1411( 2056)
822( 2645)
RARL
3030
1645 1385)
1304[ 1726)
887( 2143)
SAM+PPO
3588
2312(
1276)
3513( 75)
1571( (571(-2017)
values means the performance degradation from 'Nominal' each perturbations
E.3 LOWER COMPUTATIONAL COSTS
We compared the computational efficiency of each algorithm by measuring the total training time and total training steps each algorithm required to reach an average reward of 3000 on the Hopper- v3 environment. All experiments were conducted on the same hardware setup, and each algorithm was evaluated every 5000 environment steps. As shown Table 9, SAM+PPO requires less training time than RARL and RNAC to converge while achieving better robustness. For the training time, SAM+PPO shows the comparative costs with PPO, with emphasis on even lower total steps than PPO. However, the early work, RARL, shows immense costs, which is x4.41 of PPO in time and x3.84 in the total counts of total steps. Although the recent work, RNAC, significantly reduces the costs, but SAM+PPO shows much smaller costs in training.
Table 9: Computational costs comparison achieve convergence (done with NVIDIA RTX 3090)
Algorithm Overall Training Time (sec) J Overall Iterations Counts
Overall (action selection) Step Counts
PPO
546.6
1790
4.30 x 10
RNAC
781.2(xx143)
6.40 x (x14xx)4
RARI.
2411.4(x4.41)
4288 x2.40)
1.65 x 10 (x3.84)
SAM-PPO
590.0( (100
2103( x 1.17)
3.65 10 (x0.85)
23


Reviews:
[{"review_id": "gpPIqcossX", "reviewer": "Reviewer_jacH", "summary": "The paper investigates the relationship between flat reward maxima in policy parameter space and the robustness of reinforcement learning (RL) agents. It claims that flatter reward maxima lead to more robust policies, particularly against action perturbations. The paper presents a theoretical proposition linking flat reward to action robustness and supports this claim through empirical experiments in MuJoCo environments (e.g., Hopper-v3, Walker2d-v3, HalfCheetah-v3). The authors demonstrate that an RL algorithm enhanced with Sharpness-Aware Minimization (SAM), called SAM+PPO, consistently outperforms standard PPO and a recent robust RL baseline (RNAC) in various robustness tests, including action noise, transition probability changes, and reward function variations. The paper also provides visualizations and quantitative measurements of reward surfaces, further confirming the link between flatness and robustness.", "questions": "- Do you have an intuition on why SAM doesn't perform better on Walker2d-v3 for high friction factor?\n- Have you tested SAM+PPO on non-MuJoCo environments to assess robustness in discrete action spaces or varying reward structures?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- This paper provides a formal link between flat reward surfaces and robustness in policy space. Proposition 1 establishes a clear theoretical foundation for the paper's main claim.\n- The authors comprehensively test SAM+PPO across multiple challenging environments and scenarios, including noisy actions and varying transition probabilities, to demonstrate robustness.\n- The authors compare SAM+PPO with RNAC, PPO, and RARL, which shows both performance and computational efficiency, which strengthens their findings.\n- The use of reward surface visualizations and flatness metrics strengthens the paper's argument by providing visual and quantitative evidence for the flatness achieved by SAM+PPO.", "weaknesses": "- While SAM is shown to be effective, the paper lacks a discussion of its potential limitations, such as computational overhead or sensitivity to hyperparameter tuning.\n- The justification for reward noise being added during training for reward function robustness evaluation could be clearer: The paper mentions this difference in methodology but could expand on why this is necessary for a valid evaluation.\n- I don't know if the preliminary experiment is best placed in the introduction, it feels a bit out of place for me.\n- typos 234 \"objeective\", 249 \" funciton\"", "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": "Jt5M0SwtsN", "reviewer": "Reviewer_RSs8", "summary": "The paper presents a study on using sharpness-aware regularization to obtain robust reinforcement learning policies. Drawing a theoretical connection between flatness in the reward, action and parameter space to action-robust RL, the authors present both a theoretical justification and experiments to show that the proposed method achieves good robustness properties.", "questions": "Is there a specific advantage to using PPO with SAM, or could any PG or even AC algorithm be used? It might be that the clipping approximation to the trust region synergizes well with the SAM objective? I think this is an optional extension to the paper.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The authors propose a simple yet intuitive approach for robust RL. I was somewhat surprised that this combination has apparently not been tried in the literature, but a brief literature survey has not brought up any similar algorithms. I actually think the authors are somewhat underselling their contributions here! While SAM has been used to train PPO before, the authors appropriately cite prior work here, previous papers have not drawn any connections to robust RL at all and the authors should feel entitled to proudly claim this connection as their connection! They do not merely provide theoretical backing, as far as I can tell, they make a connection that was wholly absent in cited work.\n\nThe theoretical statements are mostly correct as far as I can tell. See questions below however.", "weaknesses": "The main problem with the paper as it stands are writing problems and baseline comparisons.\n\nEspecially the beginning of the paper, abstract and introduction, suffer from very frequent grammar mistakes which make the paper much harder to read. I strongly encourage the author to revise the paper wrt to the writing.\n\nIn definition 1, I'm unsure if $\\epsilon$ is added to the policy, parameters or action? From the proof it seems this is a parameter perturbation, this should be stated directly. I think adding parentheses in the equation would already make this much clearer, as we have two nested subscripts here.\nIn addition, the state is sampled from the policy, which seems strange?\n\nAs the theoretical statement depends on the Jacobian of the policy network, which is not bounded anywhere, I'm slightly skeptical that the theoretical results are sufficient to practically guarantee robust RL. Does the SAM objective guarantee or incentivize a flat Jacobian?\n\nGiven the surprisingly (?) bad results of RNAC - it barely seems to outperform PPO - I think it would be appropriate to apply SAM+PPO in the same environments as used in the RNAC paper. As far as I can tell, the code is available, so this should be feasible within the rebuttal timeline? If not, I will not hold this against the authors. I think it is important to verify that used examples are not cherry-picked to make the presented algorithm look stronger. This is the higher priority comment in terms of baseline comparisons.\n\nI would encourage the authors to present some additional baselines. I acknowledge that more baselines is a somewhat lazy comment. However, given that there are several different formulations of robust RL, I believe it would be helpful to pick a variety of environments and algorithms presented with different robust formulations for comparison to understand how well the algorithm does in comparison to others. This doesn't have to be many or complex environments, just a larger variety of formalisms. This is a soft concern and not a large barrier to acceptance for me.\nBoth safe-control-gym [1] and Safety Gymnasium [2] provide a variety of tasks and implemented baselines to speed up experimentation.\n\n[1] https://github.com/utiasDSL/safe-control-gym\n[2] https://github.com/PKU-Alignment/safety-gymnasium", "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": "fgtJOstnJR", "reviewer": "Reviewer_uQqo", "summary": "This paper explores the impact of flat minima in reinforcement learning (RL), linking flatter reward surfaces to improved model robustness. The authors show that flatter rewards lead to more consistent actions despite parameter changes, enhancing robustness against variations in state transitions and reward functions. The authors show through extensive experiments to confirm that flatter rewards significantly bolster RL model performance across diverse scenarios.\n\n-------------------\nAfter the rebuttal: The authors have addressed some of my concerns. I raised the score.", "questions": "- Is the perturbation domain $\\rho$ in Eq.8 known to the agent? Probably the optimization of the objective in Eq.8 needs elaboration, and with pseudo-code.\n\n- Why in \"Nominal\"  SMA+PPO still has a higher reward, e.g. Table 1+2, Fig. 3,4. Similarly,  experiment in 5.2, why when action noise is small, i.e. even equal to 0, SAM+PPO still performs better than the others, because the objectives of PPO and SAM+PPO would converge to the same one? And in 5.3, SAM+PPO has a higher return, while with variation in Friction Coefficient shows mixed results. \n\n- Joint variation of friction and mass shows quite clear that SAM+PPO is performing better than baselines, except on Walker2d-v3 with a mixed result. Can the authors elaborate on why or provide ablation to explain the mixed performance of SAM+PPO?\n\n\n- The proof of proposition 1 is a bit not standard. The policy is sometimes referred as a distribution, but sometime used as a deterministic mapping. It needs revised.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- Provide a link of flat reward to action robustness. The authors show this through both theoretical results in section 4, and various experiment results. The motivation of having a robust objective is good. The theoretical result seems correct.\n\n- Positive experiment results showing the benefit of optimizing for a flat reward maxima. The authors show this through different experiment settings: variation to physics properties of the underlying MDP, and visualization of the reward surface.", "weaknesses": "- The performance of SAM + PPO is mixed in comparisons to the baselines, e.g. some visible ones at Fig 5.c, 4.b.\n\n- Ablations are not provided to understand how such an objective can bring benefits in comparisons to similar approaches, e.g. RNAC or robust RL.", "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": "HJS9ymeHuT", "reviewer": "Reviewer_jafF", "summary": "This work proposes a new method to ensure robustness in RL based on variations in the loss landscape: \"SAM\": Sharpness-Aware Minimization. By posing the policy optimization as a min/max objective with respect to perturbations in the parameter space, the authors show robustness to changes in reward and dynamics. A theoretical result is given, linking parameter and reward robustness, and a diverse set of experiments on 3 MuJoCo environments is provided.", "questions": "I'm really curious about \"flat rewards\" in general. Definitions 1 and 2 seem too strict at first glance (the equalities therein), so it is actually a bit surprising to me that they are even possible at all; however IIUC, Fig 6 does give evidence of this. I think that these definitions can be further elaborated on (do you have a toy example where it is easy to see in parameter or action space?) Realistically, what values of $\\epsilon$ do you think are reasonable? Something like $10^{-11}$ or $10^{-2}$? (I might've missed it somewhere, sorry.) If these are novel definitions not previously given in the literature, that can be stated as a contribution of the paper. I think it can spark future work in both theory and experimental directions.\n\nHere are some follow up questions/comments:\n\n- In Sec 5, how long are those agents trained for? Equal number of env steps for each? How were hparams tuned for each algo?\n- What is the agent's action scale for these environments (cf L337)? What do you do if the noise added is outside the action range?\n- Do you have any ideas about the sharp dropoff in Fig 3b for SAM PPO? it looks interesting, but I'm not sure what to make of it... is there some \"critical\" mass ratio? I.e., if we zoom in, how sharp is that transition, and have you averaged over enough random seeds?\n- you mention \"flatter reward maxima\" in L70. I think a formal definition or good visualization of this phenomenon early on would really improve the paper.\n- How does this work relate at all to other trust region methods like TRPO? How about e.g. [1]\n\n[1]: https://arxiv.org/abs/2103.06257\n\nTypos/minor\n\n- Fig 3 caption \"nomial\"\n- some missing +/- signs in Table 1 (in parens)\n- citations in sec 2 often have a missing leading space.\n- can you improve the visual in Fig 1? I think it's important but not quite capturing the essence. Maybe just to remove axes and grid and zoom in a bit: is there indeed a channel for the agent? It's hard to see\n- The introduction paragraphs have some grammatical issues. A cleanup/re-write here can help to crystallize the main message early on\n\nWith a rewrite to clean up the presentation, deeper explanation for SAM (i)-(iv), and perhaps a few more visualizations, this could be a really strong paper; but unfortunately I don't think it's quite there yet.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This use of robustness in policy parameter space seems to be fairly new\n- The experiments demonstrate a strong performance boost across a range of perturbations \n- The visualizations in Figure 6 offer an interesting insight into the optimizations produced by PPO vs SAM+PPO. The Hopper example is quite striking. Could you elaborate on the distinction and sharp dropoffs seen there?\n- Theory provides a potential link between flatness in parameter space and action robustness\n- Provided a solid comparison wrt computational overhead / sample complexity and wall time versus other algorithms\n- Figure 5 is quite nice, I think it should be emphasized\n\nOverall, the paper seems like a nice first step in the direction of understanding the relationship between robustness in reward, policy parameter, and dynamics spaces. The notion of \"flat rewards\" is an interesting one.", "weaknesses": "**Writing:**\n\n- Overall, I think the clarity of the paper can be enhanced with a re-write fixing grammar and overall structure:\n- It would be helpful for example, to also include some visualizations of the definitions 1 & 2.\n- Also, Proposition 1 and the following remarks are not very clear. As an example, for Prop1, if I understand correctly, the result would be better phrased as \"if $\\mathcal{E}$-flat, then $\\Delta$-robust, with $\\Delta \\leq ...$ otherwise the current phrasing is a bit confusing.\n\n\n**Discussion:**\n\n- The discussion of the main idea, \"SAM\" is lacking:\n- After it is introduced in Sec 3.3, the authors give a way to solve the optimization problem in Eq (3) by their steps (i)-(iv). However, (to me at least), it is not clear why this method is used. Is there prior work demonstrating the efficacy of this method? Are there experiments or maybe some minimal example illustrating the utility of this setup? E.g. why is $\\epsilon$ chosen to be in the direction of the previously computed gradient, if theoretically it should represent an arbitrary direction in the ball. \n- At the very least, can the authors provide some visual demonstration as to what is happening here in the loss landscape? Getting a better intuition would help to understand the core method of the paper. \n- Remark 1.1 seems to be a restatement of Prop 1 unless I am missing something. Could you please explain?\n- Remark 1.2 can be improved by using more technically accurate statements (i.e. what is meant by \"when a reward function slightly changes\")? What is meant by the \"direct [correspondence] to the changes of loss function in the supervised learning case\"? I think the latter is very unclear, and maybe even misleading.\n\n\n**Experiments:**\n- My only issue with the experiments (minor) is that you are missing RNAC in Table 2 (why?). Also why not compare against RARL? Missing explanation of the shaded regions in each figure caption.", "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 (Oral)