Hacker News Board Hacker News, browsed like a classic message board
Live data from Hacker News

4000x Speedup in Reinforcement Learning with Jax

chrislu.page

Post reply on HN
Print view View original

21–30 of 32 posts

Re: 4000x Speedup in Reinforcement Learning with Jax

#21

Post by canucker2016 »

percentcer wrote: post #6

Reminds me of this evergreen tweet from ryg: https://mobile.twitter.com/rygorous/status/12712968344392826... if you made something 2x faster, you might have done something smart if you made something 100x faster, you definitely just stopped doing something stupid

I recall the Programming Pearls article from the Sept. 1984 issue of the Communications of the ACM journal, which compared various algorithms to determine the max sum of a continuous subarray of an array.

The article showed how a linear, O(N), algorithm running on a lowly 8-bit CPU can beat a cubic algorithm, O(N^3), running on a Cray supercomputer, when N is sufficiently large.

see https://www.cs.rpi.edu/~moorthy/Courses/CSCI2300/p865-bentle...

Re: 4000x Speedup in Reinforcement Learning with Jax

#22

Post by thatguymike »

percentcer wrote: post #6

Reminds me of this evergreen tweet from ryg: https://mobile.twitter.com/rygorous/status/12712968344392826... if you made something 2x faster, you might have done something smart if you made something 100x faster, you definitely just stopped doing something stupid

Or, you started doing something stupid!

Re: 4000x Speedup in Reinforcement Learning with Jax

#23

Post by Ar-Curunir »

squeaky-clean wrote: post #8
hyperbovine wrote: post #7

Earlier quoted context omitted.

Meh. This tweet is a lot less clever than it seems. Shave a factor of n off the complexity of your your algorithm, as happens regularly in CS and informatics, and have all the 1000x speedups you want.

If you shave a factor of n off of your algorithm, it usually isn't the same algorithm anymore. That's what they mean, the previous algorithm choice was "stupid" and you've stopped doing something stupid.

That's like saying, everybody struggling to solve SAT problems is just being stupid; just prove P = NP and solve the damn thing!

Re: 4000x Speedup in Reinforcement Learning with Jax

#24

Post by __turbobrew__ »

percentcer wrote: post #6

Reminds me of this evergreen tweet from ryg: https://mobile.twitter.com/rygorous/status/12712968344392826... if you made something 2x faster, you might have done something smart if you made something 100x faster, you definitely just stopped doing something stupid

if you made something 100x faster you don’t understand how mmap works

Re: 4000x Speedup in Reinforcement Learning with Jax

#26

Post by Inufu »

sillysaurusx wrote: post #5

It's a little disingenuous to say that the 4000x speedup is due to Jax. I'm a huge Jax fanboy (one of the biggest) but the speedup here is thanks to running the simulation environment on a GPU. But as much as I love Jax, it's still extraordinarily difficult to implement even simple environments purely on a GPU. My long-term ambition is to replicate OpenAI's Dota 2 reinforcement learning work, since it's one of the mo…

AlphaZero did not run game logic on TPUs (neither chess nor other games), implementing it in C++ is more than fast enough and much simpler.

TPUs were used for neural network inference and training, but game logic as well as MCTS was on the CPU using C++.

JAX is awesome though, I use it for all my neural network stuff!

Re: 4000x Speedup in Reinforcement Learning with Jax

#27

Post by erwincoumans »

schizo89 wrote: post #3
_hark wrote: post #2

jax.vmap() is all you need?

Not only vectorization, but the plethora of environments written in jax. Hopefully someone will port MuJoCo to jax soon

There is Brax, a differentiable physics simulator written in Jax. It includes Gym tasks such as Ant, Humanoid and more: https://github.com/google/brax It is not full MuJoCo but a good base to add more features. Aside from position based dynamics (xpbd) it features motion in generalized coordinates using the same accurate robot dynamics algorithms as MuJoCo and TDS (Tiny Differentiable Simulator).

Re: 4000x Speedup in Reinforcement Learning with Jax

#28

Post by sillysaurusx »

Inufu wrote: post #26
sillysaurusx wrote: post #5

It's a little disingenuous to say that the 4000x speedup is due to Jax. I'm a huge Jax fanboy (one of the biggest) but the speedup here is thanks to running the simulation environment on a GPU. But as much as I love Jax, it's still extraordinarily difficult to implement even simple environments purely on a GPU. My long-term ambition is to replicate OpenAI's Dota 2 reinforcement learning work, since it's one of the mo…

AlphaZero did not run game logic on TPUs (neither chess nor other games), implementing it in C++ is more than fast enough and much simpler. TPUs were used for neural network inference and training, but game logic as well as MCTS was on the CPU using C++. JAX is awesome though, I use it for all my neural network stuff!

According to the AlphaZero paper (https://arxiv.org/pdf/1712.01815.pdf) they ran game logic on TPUs:

> Training proceeded for 700,000 steps (mini-batches of size 4,096) starting from randomly initialised parameters, using 5,000 first-generation TPUs to generate self-play games and 64 second-generation TPUs to train the neural networks. Further details of the training procedure are provided in the Methods.

Re: 4000x Speedup in Reinforcement Learning with Jax

#29

Post by luchris429 »

sillysaurusx wrote: post #5

It's a little disingenuous to say that the 4000x speedup is due to Jax. I'm a huge Jax fanboy (one of the biggest) but the speedup here is thanks to running the simulation environment on a GPU. But as much as I love Jax, it's still extraordinarily difficult to implement even simple environments purely on a GPU. My long-term ambition is to replicate OpenAI's Dota 2 reinforcement learning work, since it's one of the mo…

Author here! I didn't realize this got posted on HN. While indeed we do get a speedup by putting the environments on the GPU, most of the speedup seems to come from the ability to easily parallelize RL training with Jax.

While there is work on putting RL environments on accelerators, the main speedup from this work comes from also training many RL agents in parallel. This is largely because the neural networks we use in RL are relatively small and thus don't utilize the GPU very efficiently.

While this was always possible to do, Jax makes it way easier because we just need to call `jax.vmap` to get it to work.

Post reply on HN