QMIX from RLlib on Flatland


I m looking to experiment QMIX algorithm proposed in the RLlib framework on grids generated for the baselines. I am trying to adapt the baseline repo and wrap the agents into groups as required by the QMIX implementation of RLlib (https://docs.ray.io/en/master/rllib-env.html#grouping-agents). However I am not sure how I could introduce the grouping in an env class (I guess that’s where I should do the grouping).

Did anyone already attempt to use this grouping of agents and could give me hints ?
Have a good day