⚡ electrical engineering

Parallel Differentiable Reachability for Learning and Planning with Certified Neural Dynamics and Controllers

This paper introduces a parallelizable, differentiable reachability framework in JAX that combines Taylor-model flowpipes with CROWN-style bound propagation to enable certified training and gradient-based planning for neural network dynamics and controllers under uncertainty.

Original authors: Keyi Shen, Glen Chou

Published 2026-05-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Keyi Shen, Glen Chou

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 perform a delicate task, like pushing a T-shaped object across a table or flying a drone through a narrow gap. You use a "brain" (a neural network) to help the robot figure out how to move. This brain is incredibly good at guessing what will happen next based on past experience.

However, there's a big problem: What if the robot is slightly wrong? Maybe the floor is slippery, or a sensor misread the distance by a tiny bit. In the real world, these small errors can snowball. A tiny mistake in step one can lead to a crash in step ten.

Traditionally, engineers have two choices:

Trust the robot completely: This is fast and works well most of the time, but if something goes wrong, the robot might crash because it didn't account for the "what ifs."
Check the math formally: This guarantees safety by calculating every possible outcome, but it's so slow and complicated that it can't be used while the robot is actually moving or learning.

This paper introduces a new tool called DiffReach that acts like a "super-fast safety calculator" that can run in real-time. Here is how it works, using some everyday analogies:

1. The "Shadow" Analogy (Reachability)

Imagine the robot is a person walking in a foggy forest. You don't know exactly where they are, only that they are somewhere inside a small circle around a known point.

Old Way: To be safe, you might draw a giant, fuzzy cloud around the person that gets bigger and bigger with every step, eventually covering the whole forest. This is safe, but it's too scary to use because it says "they might hit a tree" even if they are far away.
The Paper's Way: DiffReach draws a tight, flexible shadow around the robot. As the robot moves, this shadow stretches and twists to follow the robot's path exactly, but it never gets bigger than necessary. It tells you, "If the robot starts here, it could end up anywhere inside this specific shape, but nowhere else."

2. The "Assembly Line" Analogy (Parallel & Differentiable)

Usually, calculating these shadows is like a single person doing a math problem on a piece of paper. It takes forever.

The Innovation: The authors built this tool using JAX (a powerful computing engine). Imagine instead of one person, you have a factory with thousands of workers (GPUs) doing the math simultaneously.
Differentiable: This is the magic part. Usually, safety checks are like a "stop sign"—you check, and if it's unsafe, you stop. You can't use the stop sign to teach the robot how to drive better. DiffReach is like a smooth, sliding ramp. Because the math is "differentiable," the robot can look at the shadow, see where it's getting too wide, and immediately adjust its brain (neural network) to make the shadow tighter. It turns safety from a "stop sign" into a "steering wheel."

3. The Two-Part Engine

The paper combines two different mathematical techniques into one seamless machine:

Taylor Models (The Analytical Part): For parts of the robot we understand with physics (like a drone's aerodynamics), they use a method that predicts the path like a smooth curve.
CROWN (The Neural Part): For the "brain" parts (the neural networks), they use a technique that checks the limits of the AI's guesses.
The Unification: They created a universal language (a "Unified TM-CROWN Interface") so these two different methods can talk to each other without losing precision. It's like translating a conversation between a physicist and an AI expert without losing any details in the translation.

4. What They Actually Did (The Results)

The authors didn't just build the tool; they used it to do two specific things:

Certified Training: They taught the robot's brain to be "safety-conscious." Instead of just learning to hit a target, the brain learned to keep its "shadow" small. If the shadow gets too big (meaning the robot is unsure), the training penalizes it. The result? Robots that are just as good at their jobs but much safer when things go wrong.
Safe Planning (MPC): They used the tool to plan paths in real-time.
- In Simulation: They tested this on a robot arm and a drone swarm (up to 72 dimensions of movement!). The tool was 100 times faster than previous methods while keeping the safety bounds tight.
- In the Real World: They tested it on a real robot arm pushing a T-shaped object. The robot successfully planned its moves while avoiding obstacles, even when the environment was slightly unpredictable.

The Bottom Line

This paper presents a way to make robots both smart and safe at the same time. By making the safety check fast enough to run on a video game card (GPU) and flexible enough to be part of the learning process, they allow robots to learn from experience without forgetting to watch out for the "what ifs."

What they did NOT claim:

They did not claim this works for medical surgery or autonomous driving on public roads yet.
They did not claim it solves all safety problems (they admit it still struggles with very long time horizons or complex physical contacts like bouncing balls).
They focused strictly on robotics tasks like pushing objects and flying drones.

Technical Summary: Parallel Differentiable Reachability for Learning and Planning with Certified Neural Dynamics and Controllers

Problem Statement

Neural network (NN) dynamics models and control policies have demonstrated strong empirical performance in robotics, particularly in complex environments where analytical modeling is infeasible. However, providing robust safety guarantees for closed-loop NN systems under uncertainty remains a significant challenge. Classical reachability analysis offers formal guarantees via reachable-set over-approximations but often scales poorly, struggles with NN components, and produces overly conservative bounds. Conversely, modern neural network verification (NNV) methods like CROWN scale well to large networks but are primarily designed for static verification, lack support for continuous-time dynamics, and are generally non-differentiable.

Consequently, reachability analysis is typically restricted to post-hoc certification, creating costly cycles of non-certified training, failed verification, and retraining. There is a critical need for a framework that integrates certified reachability directly into robot learning and online planning pipelines, supporting both continuous- and discrete-time systems with analytical and NN-based components while remaining differentiable and efficient.

Methodology

The authors propose DiffReach, a parallelizable, differentiable reachability framework implemented in JAX. The core innovation is a unified representation that combines Taylor-model (TM) flowpipe construction for analytical dynamics with CROWN-style linear bound propagation for neural components.

1. Unified Representation (Taylor-CROWN Interface)

To enable GPU-batched computation and automatic differentiation, the framework adopts a fixed-shape Taylor model representation. Instead of variable-size polynomials, reachable sets are represented as:
$F(z) = c + Az \oplus I \quad \text{or} \quad F(z) = c + Az + B \tau z \oplus I$
where $z$ augments the state with time, $c$ is a constant, $A$ and $B$ are matrices capturing affine and quasi-quadratic dependencies, and $I$ is an interval remainder. This representation preserves affine dependencies without intermediate interval relaxation.

Neural Components: The framework leverages CROWN (via jax_verify) to compute affine bounds for neural networks. A key contribution is enforcing a shared-slope constraint in CROWN, ensuring the lower and upper bounds share the same linear term. This allows neural network outputs to be represented as linear TMs, maintaining compatibility with the analytical dynamics pipeline.
Continuous-Time (CT) Systems: For CT systems with NN feedback control, the framework alternates between:
1. Controller Certification: Using the TM-CROWN interface to certify the neural controller's output given the current state's TM enclosure.
2. Dynamics Propagation: Propagating the system dynamics (analytical or neural ODE) using TM flowpipes under a zero-order hold policy for the control input.
Discrete-Time (DT) Systems: Reachability reduces to stepwise propagation where the TM-CROWN interface certifies the one-step map at each time step.

2. DiffReach-Robotics Applications

Building on the DiffReach engine, the authors develop two robotics-focused methods:

Certified Training: A training method for neural dynamics and controllers that augments standard prediction or tracking losses with a reachability-aware regularization term. This term penalizes the volume growth of certified reachable sets under bounded uncertainty. To stabilize training, a curriculum learning strategy is employed, progressively increasing the rollout horizon and decreasing the perturbation radius.
Reachability-Aware MPC: A sampling-based Model Predictive Control (MPC) scheme that integrates certified reachable sets into online planning. The planner optimizes a cost function comprising the task objective and a penalty for unsafe trajectories or large reachable sets. Parallelization allows efficient evaluation of many candidate trajectories, with gradient-based refinement applied to top candidates to further reduce conservativeness.

Key Contributions

Parallel Differentiable Reachability Framework: A unified JAX-based representation combining Taylor-model flowpipes and CROWN-style bound propagation for both continuous- and discrete-time systems with analytical and NN-based dynamics/controllers.
Certified Training Method: A learning approach that uses differentiable reachable-set objectives to improve the robustness of neural components under uncertainty, encouraging "reachability-friendly" models.
Reachability-Aware MPC: An online planning scheme that integrates certified reachable sets into sampling-based MPC with gradient-based refinement, enabling verified-by-construction planning.
Experimental Validation: Comprehensive evaluation on non-prehensile manipulation (T-pushing) and quadrotor tasks, including hardware experiments and high-dimensional evaluations (up to 72D), demonstrating practical online planning with certified guarantees.

Results

The framework was evaluated on an RTX 4090 GPU across several dimensions:

Efficiency and Tightness: Compared to state-of-the-art DT (CROWN) and CT (immrax) baselines, DiffReach achieves comparable reachable-set tightness while being approximately 100× faster than CROWN for DT systems. For CT systems, it produces substantially tighter reachable sets than immrax, which frequently suffers from numerical overflow and infinite-volume tubes.
Scalability: The framework scales effectively to higher-dimensional systems, including a 6-quadrotor swarm (72D state) and a 10-joint planar arm (40D state). While memory scales as $O(n^2)$ , runtime scales favorably due to GPU-parallel matrix operations.
Certified Training: Models trained with reachability regularization achieved nominal performance comparable to standard training but produced substantially tighter reachable sets. These sets closely matched sampled lower-bound volumes, whereas non-certified models yielded loose, less verifiable predictions.
Robust Planning: In MPC experiments under bounded disturbances, the reachability-aware planner maintained high success rates and satisfied constraints, whereas vanilla MPC with non-certified models failed significantly under larger disturbances (e.g., unintended contact in T-pushing tasks).
Hardware Deployment: The system was deployed on a Kuka iiwa platform for T-pushing, running online at approximately 10 Hz (5 Hz with refinement), demonstrating the feasibility of integrating differentiable reachability into real-time robotic control.

Significance and Claims

The paper claims that DiffReach bridges the gap between formal verification and robot learning by providing a differentiable, GPU-accelerated reachability primitive. Unlike previous methods restricted to offline verification, this framework allows reachable-set objectives to directly inform the training of neural dynamics and controllers and the execution of online planning.

The authors emphasize that their approach enables "verified-by-construction" learning and planning, reducing the reliance on costly post-hoc certification cycles. By unifying Taylor models and CROWN within a single computation graph, the method preserves affine dependencies and supports both continuous and discrete-time systems, addressing a critical limitation in current NN-in-the-loop analysis tools.

The work concludes by acknowledging current limitations, such as challenges with long-horizon conservatism and hybrid/contact-rich dynamics, but posits that this framework motivates future research into scalable differentiable reachability and tighter integration between formal verification, robot learning, and control.