Highly scalable 2D JAX physics engine.

Simple JAX Graphics Library.

Reinforcement learning on general 2D physics environments in JAX. ICLR 2025 Oral.

The AI Scientist: Towards Fully Automated Open-Ended Scientific Discovery 🧑‍🔬

Your favourite classical machine learning algos on the GPU/TPU

An Open-Ended Agentic Simulator

Drop-in environment replacements that make your RL algorithm train faster.

Awesome Open-ended AI

Research Papers and Code Repository on the Integration of Evolutionary Algorithms and Reinforcement Learning

Code for Discovering Preference Optimization Algorithms with and for Large Language Models

Code for Discovering Preference Optimization Algorithms with and for Large Language Models

Accelerated minigrid environments with JAX

Hardware-Accelerated Reinforcement Learning Algorithms in pure Jax!

POPGym Library in JAX

(Crafter + NetHack) in JAX. ICML 2024 Spotlight.

🚀 Efficient implementations of state-of-the-art linear attention models

jax-triton contains integrations between JAX and OpenAI Triton

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

JAX-accelerated Meta-Reinforcement Learning Environments Inspired by XLand and MiniGrid 🏎️

Scalable Opponent Shaping Experiments in JAX

Contains JAX implementation of algorithms for inverse reinforcement learning

Efficient baselines for autocurricula in JAX.

Benchmarking RL for POMDPs in Pure JAX [Code for "Structured State Space Models for In-Context Reinforcement Learning" (NeurIPS 2023)]

Baselines for gymnax 🤖

Analysis & Experiment Configurations for NEB

🦁 A research-friendly codebase for fast experimentation of multi-agent reinforcement learning in JAX

Sample Complexity of Model-Free Opponent Shaping

Implementation of ReLU to the Rescue: Improve Your On-Policy Actor-Critic with Positive Advantages