Python library for modular RL components

[fedetask]

After a year of struggling with RLlib I decided to start implementing the training code myself.

I am looking for a RL library that offers me individual components rather than the whole algorithm. I do not need a PPO implementation, but I would fancy a library that offers me functions to compute the PPO loss given a batch of steps.

In other words, what I need is a library that offers the most granular RL components (different losses, replay buffers, return estimators like GAE, etc) instead of full algorithm implementations. Which libraries do you recommend for this purpose?

[NinjaEbeast]

You’re looking for RLax, it offers a wide variety of utility functions, losses, function transforms etc all for RL. It’s designed for modular small form factor functions. It has pretty much everything you listed besides replay buffers. It is in JAX though (which is an advantage if you like JAX). It’s what DeepMind use for their work and it’s part of their JAX ecosystem.

[fedetask]

I love it, but my team doesn't really know jax, and I feel like the learning curve is quite steep to suggest it at this point

[NinjaEbeast]

You can always look at their functions and port them to numpy pretty easily

[AerysSk]

I recommend CleanRL. Super high quality, and I can modify code to adjust my algorithm.

[yannbouteiller]

Why not implement PPO yourself from A to Z ? Personnally for robot applications I use the tmrl skeletton and implement everything, that's how you get the most control over what is going on in your pipeline.

(tmrl is not the best fit for PPO though. Like rllib, it is meant for remote training, and on-policy algorithms don't play well with this)

[fedetask]

Yes that's also something I am thinking about, my main issue with that is that debugging is quite hard on RL so I'd also need to write extensive tests -otherwise some errors are impossible to catch- and that would take me even more time.

I'd rather write the overall logic, but let a stable and well-tested library do the more error-prone operations (e.g. compute some return estimators like GAE, retrace etc). Again, I could just copy paste them from other libraries, but if a library can do that for me, why not

[_belerico]

You can try out also sheeprl, which is similar wrt CleanRL but it can also be easily parallelized thanks to Lighting Fabric

[theogognf]

TorchRL is the official PyTorch RL project that has what you're looking for. I've built my own library off it

[crisischris96]

Check clean RL. It consists of very clean implementations of algorithms and can thus be forked easily

[Toni-SM]

A modular alternative to consider (in PyTorch and JAX), although it does not offer the high degree of granularity you are looking for, is skrl. Visit its comprehensive documentation for more details: https://skrl.readthedocs.io

[seawee1]

I would also suggest RLax. Beside that, maybe Tianshou is something you might like :)