JAX rewrite: 5k FPS → 1.4M FPS (280x speedup on Generals.io RL) ⚡
Six months ago I implemented a NumPy environment for generals.io and trained an agent that hit top 20 on human leaderboards. I reached 5k fps with that setup.
In the last couple of days I rewrote everything in JAX with help from opus4.5 (here we go again) and got 1.4M FPS on single H200, which is a 300x speedup!
I’m confident that with so much more fps going super-human is much more attainable!
For those interested in coding agents for games, here is the repo https://github.com/strakam/generals-bots
Lesson Learned
With current coding agents, writing fast JAX code is extremely easy. If you want rapid RL environments and quick experimental results, just do it in JAX. The speedup is absurd and the tools make it painless.
Environment is fully reproducible and easy to use. Check it out if you’re interested!
Happy to answer questions about the implementation or approach.
submitted by /u/shrekofspeed
[link] [comments]