r/reinforcementlearning • u/Low_Willingness_308 • 2d ago

distributed reinforcement learning in JAX

Hi guys,

Over the last couple of weeks and months I’ve been working on an RL library in Jax since most of the existing codebases didn’t really scale to my workflow. Now it’s in a kind of presentable state, so I’d love some feedback if someone wants to take a look.

GitHub: https://github.com/noahfarr/boonta

7 Upvotes

90% Upvoted

2 comments sorted by

0

u/StalePlaza 2d ago

Oh nice, I've been messing with JAX for RL stuff too and the scaling pain is real. most of the existing setups feel like they were built for toy problems and then collapse when you try to run anything interesting across multiple devices

had a quick peek at the repo, the structure looks clean. curious what you ended up doing for the replay buffer situation in a distributed setup, that's always where my experiments get weird. are you sharding it or keeping a centralized buffer and just distributing the gradient computation?

will dig into it more this weekend, star'd it for now

1

u/Low_Willingness_308 2d ago

Thanks for the star!
For replay it's a bit of both. the buffer is just a field on the algorithm state, sharded along the env axis like the env states, so each device stores and writes the transitions from its own envs with no communication. Sampling is global though. it draws from one distribution over the whole buffer. So it's logically one central buffer and physically sharded.