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

Evolution Strategies in JAX 🦎

RL Environments in JAX 🌍

🚀 Efficient implementations for emerging model architectures

ALIEN is a CUDA-powered artificial life simulation program.

jax-triton contains integrations between JAX and OpenAI Triton

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

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

Awesome Open-ended AI

Highly scalable 2D JAX physics engine.

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

Multi-Agent Reinforcement Learning with JAX

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

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

Your favourite classical machine learning algos on the GPU/TPU

The simplest, fastest repository for training/finetuning medium-sized GPTs.

Hardware-Accelerated Reinforcement Learning Algorithms in pure Jax!

Accelerated minigrid environments with JAX

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

Mastering Diverse Domains through World Models

High-quality single file implementation of Deep Reinforcement Learning algorithms with research-friendly features (PPO, DQN, C51, DDPG, TD3, SAC, PPG)

A framework for multi-agent reinforcement learning.

♟️ Vectorized RL game environments in JAX

Simple JAX Graphics Library.

Really Fast End-to-End Jax RL Implementations

Neural MMO - A Massively Multiagent Environment for Artificial Intelligence Research

Efficient baselines for autocurricula in JAX.

Contains JAX implementation of algorithms for inverse reinforcement learning

An Open-Ended Agentic Simulator