🤖 machine learning

SHAPO: Sharpness-Aware Policy Optimization for Safe Exploration

This paper introduces SHAPO, a safe exploration method for reinforcement learning that leverages epistemic uncertainty by applying sharpness-aware policy updates to bias learning toward conservative behavior in under-explored regions, thereby improving both safety and task performance across continuous-control tasks.

Original authors: Kaustubh Mani, Yann Pequignot, Vincent Mai, Liam Paull

Published 2026-06-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Kaustubh Mani, Yann Pequignot, Vincent Mai, Liam Paull

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to walk across a room filled with invisible traps. The robot has never been there before, so it doesn't know where the traps are. This is the core problem of Safe Exploration in Artificial Intelligence: how do you teach an agent to learn without accidentally falling into a "trap" (a catastrophic failure) while it's still figuring things out?

The paper introduces a new method called SHAPO (Sharpness-Aware Policy Optimization) to solve this. Here is the concept broken down into simple analogies.

1. The Problem: The "Overconfident" Student

In standard AI training, the robot (the "actor") looks at the data it has collected so far and says, "I think I know the best way to move." It calculates a path based on its current knowledge.

However, in areas the robot hasn't visited much (high uncertainty), its knowledge is shaky. It might think a dangerous cliff is a safe path because it hasn't seen the edge yet. Standard training methods often trust these shaky estimates too much, leading the robot to take risky shortcuts that result in crashes.

2. The Solution: The "Paranoid" Optimist

SHAPO changes how the robot learns by introducing a healthy dose of pessimism. Instead of asking, "What is the best I can do right now?", SHAPO asks, "What is the worst that could happen if I'm slightly wrong about my knowledge?"

Think of it like a hiker walking through fog:

Standard AI: "The path looks clear, so I'll run fast."
SHAPO: "The path looks clear, but the fog is thick. If I'm wrong, I might fall off a cliff. So, I will walk slowly and carefully, assuming the ground might be slippery."

3. How It Works: The "Wobbly Table" Analogy

The paper uses a concept called Sharpness-Aware Optimization. Imagine you are trying to balance a ball on top of a hill.

A "Sharp" Hill: If the hill is a sharp peak, the ball is very unstable. A tiny nudge (a small change in the robot's brain) will make it roll off quickly. This represents high uncertainty.
A "Flat" Hill: If the hill is a wide, flat plateau, the ball is stable. You can nudge it a little, and it stays put. This represents low uncertainty (you are confident).

SHAPO looks at the "landscape" of the robot's decisions. If the robot is making a decision on a "sharp" part of the landscape (where it is uncertain), SHAPO says, "This decision is too risky. Let's adjust our learning to be more conservative."

4. The Magic Trick: The "What-If" Step

Here is the clever part of the algorithm, explained simply:

The Nudge: Before the robot updates its brain, SHAPO gives its current settings a tiny, artificial "nudge" in the direction that would make it perform worse. It asks, "If I were slightly wrong, how bad would this be?"
The Reaction: The robot then calculates the learning step based on this "worse" scenario.
The Result:
- If the robot was planning to take a risky action (one that could lead to a crash), the "worse" scenario looks very scary. This makes the robot update its brain strongly to avoid that action in the future.
- If the robot was planning a safe action, the "worse" scenario isn't that bad. The robot makes a smaller, more gentle update.

In short: SHAPO makes the robot pay extra attention to the rare, dangerous mistakes and ignore the safe, boring ones. It effectively says, "Don't just learn from what worked; learn hard from what could have gone wrong."

5. The Outcome: A Better Safety Net

The paper tested this on various tasks, like a robot navigating a maze with hidden hazards or a robot running without falling over.

The Result: Robots using SHAPO learned faster and, crucially, crashed much less often than robots using standard methods.
The Balance: They didn't just become super cautious and refuse to move. They found a better balance: they were safe enough to explore, but confident enough to get the job done.

Summary

SHAPO is like a safety instructor for AI. Instead of letting the AI guess its way through the unknown, it forces the AI to imagine the worst-case scenario for every move it considers. By doing this, the AI learns to be cautious in dangerous, unfamiliar areas while still being efficient in safe ones, ensuring it doesn't crash before it's ready.

Technical Summary: SHAPO - Sharpness-Aware Policy Optimization for Safe Exploration

Problem Statement

Deploying reinforcement learning (RL) agents in safety-critical domains requires safe exploration, a process where agents gather informative experience while avoiding catastrophic failures. A fundamental challenge arises during early training when agents visit states with high epistemic uncertainty (uncertainty about model parameters due to scarce on-policy data). In these regions, policy predictions are unreliable, and exploratory actions may lead to unsafe outcomes.

While prior work has largely addressed uncertainty on the critic side (e.g., using ensembles, dropout, or distributional methods) to quantify return variability, there is a gap in methods for policy updates that explicitly account for the actor's parameter uncertainty. Maintaining an explicit posterior over policy parameters is often infeasible for deep on-policy RL, making actor-side risk aversion difficult to implement. Consequently, standard policy updates often rely on point estimates of gradients, which can be overconfident in regions of high uncertainty.

Methodology: Sharpness-Aware Policy Optimization (SHAPO)

The authors propose SHAPO, a policy update rule that adapts the Sharpness-Aware Minimization (SAM) framework (Foret et al., 2020) to the context of safe exploration. Instead of optimizing the expected return directly, SHAPO seeks a policy that maximizes the worst-case expected return within a local neighborhood of the current policy parameters.

Core Mechanism

SHAPO modifies the standard policy gradient update by evaluating gradients at a perturbed parameter rather than the current parameters. The update process involves two steps:

Inner Minimization (Pessimistic Perturbation): The algorithm finds a perturbation $\epsilon_{Down}$ within a trust region (defined by the Fisher Information Matrix) that minimizes the Lagrangian objective $L_\lambda^{\pi_0}$ . This identifies a "worst-case" policy $\tilde{\pi}$ in the local neighborhood that performs worse than the current policy $\pi_0$ .
Outer Maximization (Update): The policy is updated using the gradient of the objective evaluated at this perturbed parameter $\theta_0 + \epsilon_{Down}$ , rather than at $\theta_0$ .

Theoretical Interpretation: Pessimism via Epistemic Uncertainty

The paper reinterprets this sharpness-aware adjustment through the lens of epistemic uncertainty:

Uncertainty Modeling: The current policy parameters $\theta_0$ are modeled as a random variable following a multivariate normal distribution $N(\theta_0, \frac{1}{n}F^{-1})$ , where $F$ is the Fisher Information Matrix and $n$ represents the effective sample size.
Pessimistic Bias: The perturbation $\epsilon_{Down}$ corresponds to the most likely parameter value that falls within the lower tail (e.g., the $\alpha$ -quantile) of the expected return distribution. By optimizing based on the gradient at this pessimistic point, SHAPO implicitly reweighs policy gradients.
Gradient Re-weighting: Analytical results on a simplified Gaussian policy show that SHAPO amplifies the gradient magnitude for rare, unsafe actions (negative advantage) while tempering the gradient for rare, safe actions (positive advantage). This biases learning toward conservative behavior in under-explored regions, effectively penalizing actions that could lead to catastrophic outcomes if the current policy estimate is slightly off.

Implementation

SHAPO is designed to be compatible with various on-policy algorithms (e.g., TRPO, PPO, CPO, CRPO). It utilizes the Fisher Information Matrix to define the neighborhood of policies, ensuring the perturbation respects the geometry of the policy manifold. The method introduces a hyperparameter $\delta_{Down}$ to control the level of pessimism, which can be kept constant or scheduled based on the number of collected samples.

Key Contributions

Novel Update Rule: Introduction of SHAPO, which integrates sharpness-aware optimization into the actor's policy update to explicitly handle epistemic uncertainty.
Theoretical Insight: A formal demonstration that SHAPO's gradient adjustment is equivalent to optimizing a pessimistic objective under a parameter uncertainty model, leading to a systematic re-weighting of gradients that prioritizes safety in uncertain regions.
Empirical Validation: Extensive experiments showing that SHAPO expands the safety-efficiency Pareto frontier across multiple baselines (TRPO, CPO, PIDLag, CRPO, SauteRL) and environments (Safety Gym, MuJoCo).
Ablation Studies: Evidence that:
- Perturbations defined in the Fisher metric outperform Euclidean counterparts.
- Applying sharpness-aware optimization to the actor is more critical for safe exploration than applying it to the critic.
- SHAPO reduces the frequency of heavy-tailed episodic cost distributions (catastrophic failures).

Experimental Results

The authors evaluated SHAPO on Safety Gym (PointGoal1, PointButton1) and MuJoCo (Ant-v4, Walker2d-v4) environments.

Safety and Efficiency: Across all benchmarks, SHAPO consistently improved either safety (lower cost/failure rates), task performance (higher return), or both, significantly expanding the Pareto frontiers of the baseline algorithms.
Tail Risk Reduction: SHAPO significantly suppressed heavy tails in episodic cost distributions (e.g., reducing the 95th percentile of costs in Ant-v4), indicating a reduction in rare, catastrophic failures.
Robustness: The method demonstrated robustness to hyperparameter choices, particularly when $\delta_{Down}$ was treated as a tunable constant rather than a dynamically scheduled value based on uncertain sample size estimates.
Unconstrained RL: In unconstrained environments with termination conditions (Ant, Walker2d), SHAPO reduced premature terminations, improving sample efficiency. In environments without such tail risks (HalfCheetah), the benefits were limited, confirming that the method specifically targets tail-risk scenarios.

Significance and Claims

The paper claims that SHAPO addresses a critical gap in safe exploration by shifting the focus from merely enforcing constraints on the final policy to shaping the updates themselves to respect uncertainty. By treating the "sharpness" of the policy surrogate (its sensitivity to parameter perturbations) as a proxy for epistemic uncertainty, SHAPO creates a closed loop:

Safer updates lead to safer trajectories.
Safer trajectories generate cleaner data in risky regions.
Cleaner data reduces epistemic uncertainty.
Reduced uncertainty naturally relaxes conservatism as learning progresses.

The authors position SHAPO not as a replacement for existing safe RL algorithms, but as a modular, orthogonal improvement that enhances the safety and reliability of standard on-policy updates without requiring complex Bayesian inference or explicit uncertainty quantification modules.