Creating the environment

from torchrl.envs.utils import check_env_specs
from momentum_env.momentum_env import Momentum
from torchrl.envs.libs.pettingzoo import PettingZooWrapper

env = Momentum()
map = {
        "team0": ['player1', 'player2', 'player3'],
        "team1": ['player4', 'player5', 'player6']
}


env = PettingZooWrapper(env=env, return_state=False, use_mask=False, group_map=map)

env.action_keys

#Output: [('team0', 'action'), ('team1', 'action')]