Oct 25, 2021

JAXDL: JAX (Flax) Deep Learning Library

There are many implementations of many machine learning algorithms out there. However, most are complex in order to squeeze out the last dime in performance. JAXDL aims at providing clean and simple implementations of state-of-the-art algorithms. For now, the Soft-Actor-Critic (SAC) and Twin Delayed DDPG (TD3) reinforcement learning algorithms are implemented. Future implementations will include graph neural networks (GNNs), transformers, and more. Check it out on GitHub here.