Link CleanRL's PPO Implementation

[vwxyzjn]

What does this PR do?

Hello, this is an awesome repo! We really love using Flax at CleanRL to seed up our implementations. I was wondering if you'd be interested in linking our PPO implementation, which runs faster (1850% speedup) by using EnvPool, a library for fast parallelizing environments, jitting the entire rollout through EnvPool's XLA interface, storing data more efficiently, and jax.scan. Thanks for your consideration!

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other
  checks if that's the case).
• This change is discussed in a Github issue/
  discussion (please add a
  link).
• The documentation and docstrings adhere to the
  documentation guidelines.
• This change includes necessary high-coverage tests.
  (No quality testing = no merge!)

[marcvanzee]

Nice! @jheek

[codecov-commenter]

Codecov Report

Merging #2780 (a621fbe) into main (62617db) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##             main    #2780   +/-   ##
=======================================
  Coverage   81.22%   81.22%           
=======================================
  Files          53       53           
  Lines        5636     5636           
=======================================
  Hits         4578     4578           
  Misses       1058     1058           

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[jheek]

Looks great!