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

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

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

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

Mastering Diverse Domains through World Models

Really Fast End-to-End Jax RL Implementations

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

RL Environments in JAX 🌍

Multi-Agent Reinforcement Learning with JAX

Evolution Strategies in JAX 🦎

♟️ Vectorized RL game environments in JAX

Neural MMO - A Massively Multiagent Environment for Artificial Intelligence Research

jax-triton contains integrations between JAX and OpenAI Triton

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

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

A framework for multi-agent reinforcement learning.

Hardware-Accelerated Reinforcement Learning Algorithms in pure Jax!

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

Pytorch Lightning Distributed Accelerators using Ray

Efficient baselines for autocurricula in JAX.

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

Accelerated minigrid environments with JAX

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

Train self-modifying neural networks with neuromodulated plasticity

Contains JAX implementation of algorithms for inverse reinforcement learning

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

Highly scalable 2D JAX physics engine.

An Open-Ended Agentic Simulator

Neuroevolution Benchmark in JAX 🦕