4000x Speedup in Reinforcement Learning with Jax

Re:4000x Speedup in Reinforcement Learning with Jax by xyzzy4747 on Thursday April 06, 2023 @09:41PM (#35477020)

How does this compare with PyTorch / Tensorflow / etc.? Obviously doing heavy data processing on the GPU will have a large speedup compared to a single thread on the CPU.

[ Parent | on HN ]

Re:4000x Speedup in Reinforcement Learning with Jax by m00x on Thursday April 06, 2023 @11:29PM (#35477698)

I don't get that sense at all.

You could do the same with tensorflow and pytorch, but in my experience, with more difficulty since they're more opinionated about how you should do your operations.

JAX is definitely easier to do things that aren't on rails.

[ Parent | on HN ]

Re:4000x Speedup in Reinforcement Learning with Jax by luchris429 on Saturday April 08, 2023 @11:42AM (#35494376)

The baseline we are comparing to is standard RL training that is widely used in academia. The technique mentioned in the blog post is not widely used amongst researchers.

The reason we write about Jax is that doing this technique is really hard in PyTorch / Tensorflow. This is because:

1. Jax has vmap. (PyTorch does now too, but it is far more recent).

2. There are RL environments that others have written in pure Jax (see the blog post for four different repos of RL environments)

3. As m00x hints to, Jax replicates Numpy's API. This makes it way easier to use for non-neural network programming (e.g. RL environments).

[ Parent | on HN ]