wbrenton

nanax

Minimal implementations of various deep learning architectures and training procedures in JAX (Flax)

cleanba

CleanRL's implementation of DeepMind's Podracer Sebulba Architecture for Distributed DRL

cleanrl

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

dm-haiku

JAX-based neural network library

ent-reg-marl

jax_sebulba

🪐 The Sebulba architecture to scale reinforcement learning on Cloud TPUs in JAX

lm-human-preference-details

[lm-human-preference-details](https://github.com/vwxyzjn/lm-human-preference-details) ported to JAX

open_spiel

OpenSpiel is a collection of environments and algorithms for research in general reinforcement learning and search/planning in games.

openrlbenchmark

quick_tpu

Git clone this repo to develop JAX on TPU