Uncertainty-Aware Predictive Safety Filtering (UPSi)

       ██      ██ ██████  ███████ ██
       ██      ██ ██   ██ ██            
       ██      ██ ██████  ███████ ██      
       ██      ██ ██           ██ ██      
         ██████   ██      ███████ ██

This repository contains the code and additional materials of the paper "Uncertainty-Aware Predictive Safety Filters for Probabilistic Neural Network Dynamics", presented at RLC 2026, Montréal. The implementation is based on the code from "Reinforcement learning with ensemble model predictive safety certification", Gronauer et al., 2024. The paper and a video introducing our filter are available online:

Paper: https://arxiv.org/abs/2604.26836

Video: https://youtu.be/Qggju12rZ8U

More videos of our experiments can be found under the name tags env_initial.mp4 (start of training) and env_final.mp4 (end of training) in x_mpsc/materials/.

Installation

First, create and activate the conda environment from the provided environment.yml:

cd x_mpsc
conda env create -f environment.yml
conda activate upsi

Then, install the x_mpsc package in editable mode from the repository root (the directory containing setup.py):

pip install -e .

Note: The algorithm makes use of MPI which requires an OpenMPI installation. Furthermore, ffmpeg is required for logging videos.

Training

Each environment has a dedicated folder under experiments/ with four run scripts:

Script	Description
run_sac.py	Vanilla SAC (no safety filter)
run_mbpo.py	MBPO (no safety filter)
run_xmpsc.py	MBPO + XMPSC safety filter
run_upsi.py	MBPO + UPSi safety filter

To start a training run, navigate to the desired environment folder and execute the script:

cd experiments/cartpole
python run_upsi.py

This automatically spawns parallel processes via MPI (defaulting to the number of physical CPU cores). Training progress is logged to the console, to CSV files, and to Weights & Biases. Console verbosity can be controlled via the logging level (DEBUG > INFO > WARN > ERROR > DISABLED); the default is INFO.

Note: Training without MPI is possible but not recommended, as it will be significantly slower.

Logs and checkpoints are written to logs/experiments/<env>/.

The reachable_sets/ folder contains a standalone script for visualizing UPSi predictive tubes on the cartpole environment.

Structure of this repo

All experiment run scripts can be found in x_mpsc/experiments. Additionally, videos of the filter during training are contained in x_mpsc/materials. See below for a more detailed overview:

x_mpsc/
├── environment.yml                 Conda environment specification
├── setup.py                        Package installation script
│
├── experiments/                    Experiment run scripts
│   ├── cartpole/                   run_sac/run_mbpo/run_xmpsc/run_upsi
│   ├── drone/                      run_sac/run_mbpo/run_xmpsc/run_upsi
│   ├── pendulum/                   run_sac/run_mbpo/run_xmpsc/run_upsi
│   └── reachable_sets/             Reachable-set visualization
│
├── logs/                           Training logs & checkpoints
│
└── x_mpsc/                         Main package
    ├── benchmark.py                Benchmarking and training utilities
    │
    ├── algs/                       Algorithm implementations
    │   ├── terminal_set.py         Terminal set computation
    │   ├── sac/                    Soft Actor-Critic
    │   └── mbpo/                   Model-Based Policy Optimization
    │
    ├── common/                     Shared utilities
    │   ├── loggers.py              Logging (console, file, WandB)
    │   └── mpi_tools.py            MPI communication helpers
    │
    ├── envs/                       Gym environments
    │   ├── cartpole_extended_new/  Cartpole
    │   ├── drone/                  Drone
    │   └── simple_pendulum/        Pendulum
    │
    ├── models/                     Dynamical system models
    │   └── ensemble.py             PE ensemble implementation
    │
    └── mpsc/                       PSF implementation (UPSi / XMPSC)

Implemented Algorithms

• Model-based Policy Optimization (MBPO)
• Soft Actor-Critic (SAC)