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).
Neural differential equations are also easier with jax. sim2real may be easier with simulator where some of hard computations are replaced with neural approximations
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 most impactful (or at least most entertaining) use of RL. It would be more or less impossible to translate the game logic into Jax, short of transpiling C++ to Jax somehow. Which isn't a bad idea – someone should make that.
It should also be noted that there's a long history of RL being done on accelerators. AlphaZero's chess evaluations ran entirely on TPUs. Pytorch CUDA graphs also make it easier to implement this kind of thing nowadays, since (again, as much as I love Jax) some Pytorch constructs are simply easier to use than turning everything into a functional programming paradigm.
All that said, you should really try out Jax. The fact that you can calculate gradients w.r.t. any arbitrary function is just amazing, and you have complete control over what's JIT'ed into a GPU graph and what's not. It's a wonderful feeling compared to using Pytorch's accursed .backwards() accumulation scheme.
Can't wait for a framework that feels closer to pure arbitrary Python. Maybe AI can figure out how to do it.
> 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.
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.
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
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.
It's just an idiom and not meant to be taken so literally. Basically if you're getting a 100x improvement, you're jumping from one paradigm to another. And those are probably already solved algorithms / invented datastructures. If you're doing something novel and "clever" to speed up your algorithm for your circumstances, it's probably going to be under a 5x improvement.
Or to continue the SAT analogy, if you're doubling your SAT score (e.g. 700 to 1400) you're getting smarter. If your 100x'ing your SAT score, the best possible previous score you could have gotten is a 16. Which would mean you got about one and a half questions correct across both the reading and math SAT.
I think it’s fair to classify using O(n^2) sorting as stupid (if the input size is significant). In numerical computing, using a naive routine for computing eigenvalues in O(n^4) would equally be considered stupid, unless the input sizes are known to be small of course.
I don’t think it seems clever so much as pithy? I don’t think it’s necessarily about algorithms so much as it is a claim that the main way to be 1000x slower is by doing 999x times too much work or waiting.
An example that is basically unrelated to complexity theory is something like talking to a distant service but keeping a small number of requests because eg you were worried about load or didn’t notice you were waiting for acks or create a new tcp connection for each request or have a small sendbuf or somehow send way too fast and get rate limited and need to retry.
I believe the underlying point is that the bigger the N, the more obvious it probably is. If you’re in 1000x territory, there’s a strong likelihood that the improvement is so obvious that doing the opposite might be called “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.
From what I understand of Jax, it feels somewhat similar in flavor to Julia, but trying to live with the language constraints (and ecosystem benefits) of Python.
I wonder how Julia is placed for running reinforcement learning algorithms (efficiently) — particularly in cases when the “environment” is nicely wrapped in Python to fit some standardized interface.
I've done some RL experiments in Julia, and having all the in-between be fast was helpful; I saw significant speed increases. That said, Julia was probably just compensating for my own stupidity, because I was converting my environment objects into tensors over and over and over.
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.
It's almost like the author is claiming credit for creating Nvidia, when in fact he is just calling its APIs.
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.
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).
31 comments
[ 4.0 ms ] story [ 79.6 ms ] threadMy long-term ambition is to replicate OpenAI's Dota 2 reinforcement learning work, since it's one of the most impactful (or at least most entertaining) use of RL. It would be more or less impossible to translate the game logic into Jax, short of transpiling C++ to Jax somehow. Which isn't a bad idea – someone should make that.
It should also be noted that there's a long history of RL being done on accelerators. AlphaZero's chess evaluations ran entirely on TPUs. Pytorch CUDA graphs also make it easier to implement this kind of thing nowadays, since (again, as much as I love Jax) some Pytorch constructs are simply easier to use than turning everything into a functional programming paradigm.
All that said, you should really try out Jax. The fact that you can calculate gradients w.r.t. any arbitrary function is just amazing, and you have complete control over what's JIT'ed into a GPU graph and what's not. It's a wonderful feeling compared to using Pytorch's accursed .backwards() accumulation scheme.
Can't wait for a framework that feels closer to pure arbitrary Python. Maybe AI can figure out how to do it.
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!
> 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.
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.
Or to continue the SAT analogy, if you're doubling your SAT score (e.g. 700 to 1400) you're getting smarter. If your 100x'ing your SAT score, the best possible previous score you could have gotten is a 16. Which would mean you got about one and a half questions correct across both the reading and math SAT.
An example that is basically unrelated to complexity theory is something like talking to a distant service but keeping a small number of requests because eg you were worried about load or didn’t notice you were waiting for acks or create a new tcp connection for each request or have a small sendbuf or somehow send way too fast and get rate limited and need to retry.
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...
I wonder how Julia is placed for running reinforcement learning algorithms (efficiently) — particularly in cases when the “environment” is nicely wrapped in Python to fit some standardized interface.
It's almost like the author is claiming credit for creating Nvidia, when in fact he is just calling its APIs.
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.
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).