JAX is coming

[lowrollr]

I've decided to orient the bulk of my near-term future work in this project towards re-writing TurboZero in JAX.

Why JAX?

• A lot of open-source RL projects use JAX
• Allows easy integration with third-party JAX RL environments, through:
  • direct support via a wrapper included in TurboZero
  • the ability to easily write a custom wrapper for environments not officially supported
  • this is a big deal!
• thanks to jax.vmap, some of the scary vectorized code will be a lot more readable
• more code will be JIT compiled (MCTS in particular)

Perhaps the biggest issue with the current PyTorch implementation is that there's no easy way to support existing third-party environments (gym, etc.) in such a way that elegantly integrates with the vectorized algorithm implementations. However, there are a lot of JAX libraries with environments that inherently support vectorization, which makes integrating with these environments much more trivial. This lets TurboZero focus more on algorithms, utilities, and other infrastructure instead of custom environments.

I'm not sure exactly what I plan to do with the existing PyTorch implementation yet. Development of both versions in parallel seems unlikely, but in a perfect (backend-agnostic) world improvements to any training/evaluation infrastructure could be utilized by both implementations. In practice I suspect this would get ugly rather quickly.

It's all going down in the jax branch.

Coming soon!