A High-Performance JAX Framework for Programmatic Reinforcement Learning

Open source codebase for the s-FEAST algorithm described in "Online tree-based planning for active spacecraft fault estimation and collision avoidance."

Code accompanying the NeurIPS 2025 paper "Sequential Monte Carlo for Policy Optimization in Continuous POMDPs".

Partially Observable Benchmarks in JAX

POMDP Planner with Hybrid Dynamics

A LaTeX template for a basic DFG (Deutsche Forschungsgemeinschaft, German Research Foundation) grant proposal.

Optimal feedback control + interventions, in JAX.

GPU-accelerated MDP solvers in Python with JAX

Diffilqrax is an open-source Python library for differentiable iLQR and LQR in JAX

The COntextual INference (COIN) model

Switching linear dynamical systems (SLDS) models in JAX

BlackJAX is a Bayesian Inference library designed for ease of use, speed and modularity.

Companion code in JAX for the paper Parallel Iterated Extended and Sigma-Point Kalman Smoothers.

This is a community-driven collection of resources around amortized inference.

Python/TensorFlow implementations of the parallel control methods

A generic library for linear and non-linear Gaussian smoothing problems. The code leverages JAX and implements several linearization algorithms, both in a sequential and parallel fashion, as well a…

This is a collection of code samples aimed at illustrating temporal parallelization methods for sequential data.

Models of Sequential Decision-Making

Fast Bayesian Inference in Python. 50+ conjugate models with vectorized updates, sufficient statistic helpers, built-in plotting, and SciPy integration.

A wrapper API for working with numpyro models.

Probabilistic programming with programmable variational inference.

Mitigating Partial Observability in Sequential Decision Processes via the Lambda Discrepancy