Uncertainty-Aware Predictive Safety Filters for Probabilistic Neural Network Dynamics

Abstract: Predictive safety filters (PSFs) leverage model predictive control to enforce constraint satisfaction during deep reinforcement learning (RL) exploration, yet their reliance on first-principles models or Gaussian processes limits scalability and broader applicability. Meanwhile, model-based RL (MBRL) methods routinely employ probabilistic ensemble (PE) neural networks to capture complex, high-dimensional dynamics from data with minimal prior knowledge. However, existing attempts to integrate PEs into PSFs lack rigorous uncertainty quantification. We introduce the Uncertainty-Aware Predictive Safety Filter (UPSi), a PSF that provides rigorous safety predictions using PE dynamics models by formulating future outcomes as reachable sets. UPSi introduces an explicit certainty constraint that prevents model exploitation and integrates seamlessly into common MBRL frameworks. We evaluate UPSi within Dyna-style MBRL on standard safe RL benchmarks and report substantial improvements in exploration safety over prior neural network PSFs while maintaining performance on par with standard MBRL. UPSi bridges the gap between the scalability and generality of modern MBRL and the safety guarantees of predictive safety filters.

• The paper introduces UPSi, an uncertainty-aware safety filter that formalizes model certainty using probabilistic neural networks for safe exploration.
• It employs probabilistic ensemble models to separate epistemic and aleatoric uncertainties and overapproximate reachable sets via ellipsoidal methods.
• Experimental results demonstrate lower constraint violations and comparable return performance to baseline MBPO, underscoring its practical applicability.

Uncertainty-Aware Predictive Safety Filters for Probabilistic Neural Network Dynamics (UPSi)

Introduction

This work presents Uncertainty-Aware Predictive Safety Filters (UPSi), a formal mechanism for safe exploration in deep reinforcement learning (RL) using data-driven probabilistic neural network models. The development addresses a key limitation of current Predictive Safety Filters (PSFs): a reliance on first-principles or Gaussian process models for robust reachable set computation, which impedes their scalability to high-dimensional systems. UPSi integrates with state-of-the-art Model-Based Policy Optimization (MBPO) while enforcing rigorous safety constraints using scalable probabilistic ensemble (PE) models and explicit in-distribution evaluation via a certainty constraint.

Figure 1: UPSi architecture within Dyna-style MBRL. The filter leverages a neural network ensemble for dynamics, uncertainty estimation, and constraint satisfaction.

The UPSi framework is significant in that it achieves robust reachable set computation, formalizes the notion of model certainty to prevent model exploitation, and maintains sample efficiency and asymptotic performance of modern MBRL algorithms while yielding lower constraint violations.

Background and Motivation

Model-based RL approaches often employ PE neural networks due to their expressivity and data efficiency on high-dimensional, nonlinear tasks. PSFs can provide strong formal guarantees by computing sets of states that are reachable under all admissible uncertainties, ensuring constraint satisfaction through model predictive control (MPC). Current approaches integrating PE models into PSFs, however, are heuristic: they typically overestimate region safety due to unaddressed epistemic uncertainty, leading to constraint violations and model exploitation.

UPSi addresses these deficits by leveraging recent uncertainty quantification results for PE models that separate epistemic and aleatoric uncertainties, and by formalizing a "certain set"—regions of state/action space where model predictions are known to be reliable.

Methodology

Certain Set Construction and Uncertainty Quantification

Within UPSi, the "certain set" $C$ is defined via the trace of the Kalman gain, which quantifies the ratio of propagated (aleatoric) variance to total predictive variance, encapsulating model epistemic uncertainty. Only states and actions within $C$ are considered valid for control planning, ensuring that predictions during filtering remain within the model's trust region.

Rigorous Overapproximation of Reachable Sets

UPSi propagates the PE model's aleatoric uncertainty forward in time, constructing ellipsoidal overapproximations of the reachable sets using first-order Taylor expansions for the mean and covariance, and employing a provably conservative inflation to compensate for linearization errors. This yields, at each planning step, an ellipsoidal tube containing all possible future states under admissible noise realizations and model uncertainty, formalized by Theorem 1 in the paper.

Figure 2: Overapproximation of reachable sets by UPSi (blue ellipsoids) compared to ground truth (gray), illustrating the formal envelope around all possible future trajectories.

Certainty Constraint in PSF Formulation

An explicit certainty constraint is added to the standard robust MPC/PSF formulation to guarantee that every predicted state-control pair along a candidate trajectory remains inside the certain set $C^j$ for the current data regime, thereby preventing out-of-support extrapolation and model exploitation.

Terminal Set Expansion and Sampling-Based Implementation

To avoid the conservatism associated with small initialization regions, UPSi expands the admissible terminal set using statistical sampling and geometric enclosure methods (both polytopic and ellipsoidal), maintaining a balance between theoretical rigor and practical coverage for high-dimensional data.

Figure 3: Sampling-based terminal set expansion, where candidate safe end states are estimated and robustified using polytope and ellipsoidal methods.

Experimental Evaluation

UPSi is benchmarked on standard safe RL environments: Cartpole, Pendulum, and a high-dimensional Drone task, using MBPO as the base RL algorithm. The filtered agent enforces state and action constraints in the presence of stochastic process noise.

UPSi demonstrates:

• Substantially lower constraint violation rates during exploration versus the neural-network PSF baselines (e.g., XMPSC), particularly so in higher-dimensional environments where epistemic uncertainty dominates.
• Comparable asymptotic return performance to unconstrained MBPO, indicating that the conservativeness of the safety filter does not substantially hinder learning efficiency.

Filtering Behavior and Optimization Analysis

Comprehensive analysis of filter rates, infeasibility statistics, and optimization times demonstrates:

• UPSi filters (intervenes) more frequently than the XMPSC baseline, reflecting its more conservative and principled estimation of reachable uncertainty.
• For low-dimensional systems, the UPSi optimization can be solved efficiently. In high-dimensional settings, computational costs remain a challenge, and online feasibility is identified as a direction for future work.

  Figure 4: UPSi filter rates over time, illustrating frequent yet minimal corrections relative to baseline methods.

  

  Figure 5: Optimization time per environment step, contrasting the computational burden in low- and high-dimensional systems.

Implications and Future Directions

UPSi provides a critical contribution to safe RL with function approximators by scaling formal safety mechanisms to deep neural network models. It explicitly quantifies trust in model predictions, eliminates failures due to model exploitation, and provides a viable path for applying MPC-style safety guarantees in high-dimensional, data-driven settings.

Practically, this can accelerate safe RL deployment in robotics, autonomous vehicles, and any task with stringent safety constraints and scarce prior knowledge. Theoretically, UPSi closes a gap between uncertainty quantification in Bayesian neural networks and control-theoretic guarantees, and its modular architecture invites further advances in robust and computationally efficient safe filtering.

Key future directions include:

• Optimizing and accelerating the underlying robust MPC solvers to enable real-time deployment, especially in high-dimensional systems.
• Investigating tighter and potentially adaptive formulations for the certainty set, and automated calibration of uncertainty thresholds.
• Exploring hierarchical or hybrid architectures that combine UPSi with other shielded learning or formal synthesis techniques.

Conclusion

UPSi delivers a rigorous, uncertainty-aware safety filter for model-based RL with neural network ensemble dynamics that substantially improves safe exploration performance without sacrificing control quality. By introducing a principled mechanism for estimating and enforcing in-distribution prediction and reachable sets, UPSi advances the practical applicability of predictive safety filters to modern MBRL regimes, supporting both theoretical rigor and empirical utility.

Reference: "Uncertainty-Aware Predictive Safety Filters for Probabilistic Neural Network Dynamics" (2604.26836)