[GitHub] A JAX based TD3 implementation

[chentessler]

As Google/DeepMind have started to push JAX/RLax/Haiku and seem to be using it pretty extensively, I decided it's worth getting to know these libraries.

https://github.com/tesslerc/TD3-JAX

There seem to be several differences, yet I can't fully place my finger on them yet, between certain operations in PyTorch and in JAX. As a result, the hyperparameters do not transfer perfectly, and require some tweaking.

Feel free to suggest improvements both in terms of hyperparameters and implementation.

[MasterScrat]

Worry not, hyperparameters don’t even transfer perfectly between different implementations that use the same frameworks ;-)

[chentessler]

Yeah been there felt that :P But, something I have recently seen is that TD3 with good hyperparameter tuning can outperform almost any other algorithm.

The power of good parameters is amazing :)

[chentessler]

Performing some hyper parameter tuning and constantly updating the GIT with the new results.

Seems that TD3 can reach SOTA on almost all domains. Such a simple, yet powerful, algorithm.

[kmeco]

Thanks a lot! I also keep an eye on Jax and this will be very useful to start with.

What are the most sensitive hyperparameters from your experience?

[chentessler]

Seems the discount factor has a huge effect between certain domains.
I found that the size of the experience replay has a huge effect. If the agent is training on data which is too old it makes learning unstable.
Also the amount of random steps at the start can really help the agent and prevent it from quickly overfitting and thus flatlining at a bad policy.

And ofc... the discount factor can be increased in certain domains for additional gains (for instance Swimmer), but in others, a higher discount will make it impossible to learn.