Feat: Custom global state for jumanji envs

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: Custom global state for jumanji envs #1046

Closed SimonDuToit closed 7 months ago

SimonDuToit commented 7 months ago

What?

The jumanji MA wrapper has been given a get_global_state method to allow jumanji wrappers to specify their own global state. The GlobalStateWrapper has been removed and the generic global state based on concatenating agent observations has been moved to the generic jumanji MA wrapper. The connector wrapper, which implements its own global state, has also been modified accordingly.

Why?

To allow jumanji wrappers to implement a custom global state.

SimonDuToit commented 7 months ago

I've implemented all the suggested changes.