Core Deep Reinforcement Learning algorithms using JAX for improved performance relative to PyTorch and TensorFlow. Control tasks rely on the DeepMind Control Suite; see repo for further setup if you don't have MuJoCo configured.

Current implementations

  • TD3
  • SAC
  • MPO
  • PPO
  • A2C/A3C


To test each algorithm on cartpole swingup:

python main_dm_control.py --max_timestep 100000
python main_dm_control.py --policy SAC --max_timesteps 100000

Jax Rl

JAX implementations of core Deep RL algorithms

Jax Rl Info

⭐ Stars 33
🔗 Source Code github.com
🕒 Last Update 4 months ago
🕒 Created a year ago
🐞 Open Issues 0
➗ Star-Issue Ratio Infinity
😎 Author henry-prior