$\texttt{immrax}$: A Parallelizable and Differentiable Toolbox for Interval Analysis and Mixed Monotone Reachability in JAX

TL;DR

The paper presents immrax, a differentiable, parallelizable toolbox for interval analysis and mixed monotone reachability implemented as JAX function transforms, enabling efficient bound propagation over input intervals $[\underline{x},\overline{x}]$ with $\mathsf{F}=\underline{\mathsf{F}}\overline{\mathsf{F}}$. It introduces composable transforms (\(\text{natif}, \text{jacif}, \text{mjacif}\)) and a novel Jacobian-based mixed bound for embedding-based reachability, along with an embedding framework that propagates over-approximations via $\dot{\underline{x}}$ and $\dot{\overline{x}}$. The framework is demonstrated on two case studies: a nonlinear vehicle controlled by a neural network and a robust pendulum controller, showing GPU-accelerated reachability and AD-driven robust optimal control, respectively. By combining interval theory, differentiable programming, and parallel hardware, immrax advances practical, differentiable, and scalable certified analysis for learning-enabled control systems.

Abstract

We present an implementation of interval analysis and mixed monotone interval reachability analysis as function transforms in Python, fully composable with the computational framework JAX. The resulting toolbox inherits several key features from JAX, including computational efficiency through Just-In-Time Compilation, GPU acceleration for quick parallelized computations, and Automatic Differentiability. We demonstrate the toolbox's performance on several case studies, including a reachability problem on a vehicle model controlled by a neural network, and a robust closed-loop optimal control problem for a swinging pendulum.

Figures (3)

• Figure 1: The swinging trajectory of the pendulum embedding system induced by \ref{['eq:pendulumif']} controlled by the closed-loop control policy generated by the problem \ref{['eq:pendopt']} is visualized for various time instances. The angle $[\ul{\theta},\overline{\theta}]$ is represented as a wedge, where blue represents the interval of possible angles. The gray wedge represents the desired final state of the pendulum.
• Figure 2: The reachable set over-approximations computed by simulating the embedding system from SJ-AH-SC:23c using Euler integration are visualized in light blue. The initial set $[7.95,8.05]\times[6.95,7.05]\times[-\frac{2\pi}{3} - 0.01, -\frac{2\pi}{3} + 0.01]\times[1.99,2.01]$ is divided into different numbers of partitions. $100$ Monte Carlo trajectories are pictured in dark red. In all cases, the vehicle is certified to avoid the obstacle pictured in light red, with varying degrees of accuracy to the true reachable set.
• Figure 3: The swinging trajectory of the embedding system induced by \ref{['eq:pendulumif']} controlled by the closed-loop control policy generated by the optimization problem \ref{['eq:pendopt']} is plotted versus time. Left: The angle $[\ul{\theta},\overline{\theta}]$ vs. $t$ in seconds (blue), with the terminal set constraint $\theta\in[\pi-\frac{10\pi}{360},\pi+\frac{10\pi}{360}]$ (gray). Right: The angluar velocity $[\ul{\dot{\theta}},\overline{\dot{\theta}}]$ vs. $t$ in seconds, with the terminal set constraint $\dot{\theta}\in[-0.1,0.1]$ (gray).