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

Use `TrainState` in evaluator #996

Closed sash-a closed 2 months ago

sash-a commented 8 months ago

As discussed in #994 we should now be able to use the MetricsWrapper and the TrainState in the evaluator instead of manually recording the episode return and creating a special EvalState.