instadeepai / Mava

🦁 A research-friendly codebase for fast experimentation of multi-agent reinforcement learning in JAX
Apache License 2.0
709 stars 83 forks source link

feat: upgrade recurrent networks to new flax api #957

Closed RuanJohn closed 10 months ago

RuanJohn commented 10 months ago

What?

Upgrade the version of flax supported by Mava to the latest one.

Extra:

Raw becnhmarking results on nepuine Benchmarking summary on notion