Custom PettingZoo Environment with Dictionary Spaces and PyTorch

Question

I am currently attempting to build a multi-agent model to player a game that I have made. My first thought was to convert it into an environment using PettingZoo, and I am using PyTorch to write the model itself.

Each agent in the map has 4 actions, each of which can take different values. Briefly, we have

• Forward: Float between -1 and +1
• Right: Float between -1 and +1
• Rotation Rate: Float between -pi and +pi
• Firing: Boolean

The observations are as follows. I used a LiDAR style spoke system for the detections, inspired from this PyTorch tutorial.

For each spoke:

• A set of distances as a proportion between 0 and 1
• The health of any players which have been detected
• Whether the player which was intersected with was on the same team or not

Additionally, the player's own health, and whether the player can fire in this timestep or not.

Within the environment code, I wrote the following.

        obs_space = Dict({
            "distances": Box(low=0, high=1, shape=(8,), dtype=np.float32),
            "healths": Box(low=0, high=100, shape=(8,), dtype='i'),
            "alliances": Box(low=-1, high=1, shape=(8,), dtype='i'),
            "player_health": Discrete(101, start=0),
            "can_fire": MultiBinary(1),
        })

        self.observation_spaces = {a: obs_space for a in self.agents}

        action_space = Dict({
            "forward": Box(low=-1, high=1, shape=(1,), dtype=np.float32),
            "right": Box(low=-1, high=1, shape=(1,), dtype=np.float32),
            "rot_rate": Box(low=-pi, high=pi, shape=(1,), dtype=np.float32),
            "firing": MultiBinary(1),
        })

        self.action_spaces = {a: action_space for a in self.agents}

I ran this code

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

env = Momentum()
env = PettingZooWrapper(env=env, return_state=False, use_mask=False)

check_env_specs(env)

And this was the error I got Screenshot of the error message..

I would really appreciate any way of solving this problem. Thanks!