A JAX-based research framework for differentiable and parallelizable acoustic simulations, on CPU, GPUs and TPUs

A JAX-based research framework for writing differentiable numerical simulators with arbitrary discretizations

Learning in infinite dimension with neural operators.

🕹️ A diverse suite of scalable reinforcement learning environments in JAX

Official reinforcement learning environment for demand response and load shaping

Datasets with baselines for offline multi-agent reinforcement learning 🤖

Level-based Foraging (LBF): A multi-agent environment for RL

🦁 A library of multi-agent reinforcement learning systems and components

FAST iteration of MARL research ideas: A starting point for Multi-Agent Reinforcement Learning

Callable PyTrees and filtered transforms => neural networks in JAX. https://docs.kidger.site/equinox/