[Python] AttributeError: Can't get attribute 'CustomActionMaskedEnvironment.observation_space'...

I have this very basic custom parallel multi agent environment written in PettingZoo.

import functools
import random
from copy import copy

import numpy as np
from gymnasium.spaces import Discrete, MultiDiscrete

from pettingzoo import ParallelEnv


class CustomActionMaskedEnvironment(ParallelEnv):
"""The metadata holds environment constants.

The "name" metadata allows the environment to be pretty printed.
"""

metadata = {
"name": "masked_parallel_env",
}

def __init__(self):
"""The init method takes in environment arguments.

Should define the following attributes:
- escape x and y coordinates
- guard x and y coordinates
- prisoner x and y coordinates
- timestamp
- possible_agents

Note: as of v1.18.1, the action_spaces and observation_spaces attributes are deprecated.
Spaces should be defined in the action_space() and observation_space() methods.
If these methods are not overridden, spaces will be inferred from self.observation_spaces/action_spaces, raising a warning.

These attributes should not be changed after initialization.
"""
self.escape_y = None
self.escape_x = None
self.guard_y = None
self.guard_x = None
self.prisoner_y = None
self.prisoner_x = None
self.timestep = None
self.possible_agents = ["prisoner", "guard"]

def reset(self, seed=None, options=None):
"""Reset set the environment to a starting point.

It needs to initialize the following attributes:
- agents
- timestamp
- prisoner x and y coordinates
- guard x and y coordinates
- escape x and y coordinates
- observation
- infos

And must set up the environment so that render(), step(), and observe() can be called without issues.
"""
self.agents = copy(self.possible_agents)
self.timestep = 0

self.prisoner_x = 0
self.prisoner_y = 0

self.guard_x = 6
self.guard_y = 6

self.escape_x = random.randint(2, 5)
self.escape_y = random.randint(2, 5)

observation = (
self.prisoner_x + 7 * self.prisoner_y,
self.guard_x + 7 * self.guard_y,
self.escape_x + 7 * self.escape_y,
)

observations = {
"prisoner": {
"observation": observation,
"action_mask": [0, 1, 1, 0]
},
"guard": {
"observation": observation,
"action_mask": [1, 0, 0, 1]
},
}

# Get dummy infos. Necessary for proper parallel_to_aec conversion
infos = {
a: {

} for a in self.agents
}

return observations, infos

def step(self, actions):
"""Takes in an action for the current agent (specified by agent_selection).

Needs to update:
- prisoner x and y coordinates
- guard x and y coordinates
- terminations
- truncations
- rewards
- timestamp
- infos

And any internal state used by observe() or render()
"""
# Execute actions
prisoner_action = actions["prisoner"]
guard_action = actions["guard"]

if prisoner_action == 0 and self.prisoner_x > 0:
self.prisoner_x -= 1
elif prisoner_action == 1 and self.prisoner_x < 6:
self.prisoner_x += 1
elif prisoner_action == 2 and self.prisoner_y > 0:
self.prisoner_y -= 1
elif prisoner_action == 3 and self.prisoner_y < 6:
self.prisoner_y += 1

if guard_action == 0 and self.guard_x > 0:
self.guard_x -= 1
elif guard_action == 1 and self.guard_x < 6:
self.guard_x += 1
elif guard_action == 2 and self.guard_y > 0:
self.guard_y -= 1
elif guard_action == 3 and self.guard_y < 6:
self.guard_y += 1

# Generate action masks
prisoner_action_mask = np.ones(4, dtype=np.int8)
if self.prisoner_x == 0:
prisoner_action_mask[0] = 0 # Block left movement
elif self.prisoner_x == 6:
prisoner_action_mask[1] = 0 # Block right movement
if self.prisoner_y == 0:
prisoner_action_mask[2] = 0 # Block down movement
elif self.prisoner_y == 6:
prisoner_action_mask[3] = 0 # Block up movement

guard_action_mask = np.ones(4, dtype=np.int8)
if self.guard_x == 0:
guard_action_mask[0] = 0
elif self.guard_x == 6:
guard_action_mask[1] = 0
if self.guard_y == 0:
guard_action_mask[2] = 0
elif self.guard_y == 6:
guard_action_mask[3] = 0

# Action mask to prevent guard from going over escape cell
if self.guard_x - 1 == self.escape_x:
guard_action_mask[0] = 0
elif self.guard_x + 1 == self.escape_x:
guard_action_mask[1] = 0
if self.guard_y - 1 == self.escape_y:
guard_action_mask[2] = 0
elif self.guard_y + 1 == self.escape_y:
guard_action_mask[3] = 0

# Check termination conditions
terminations = {
a: False for a in self.agents
}

rewards = {
a: 0 for a in self.agents
}

if self.prisoner_x == self.guard_x and self.prisoner_y == self.guard_y:
rewards = {
"prisoner": -1,
"guard": 1
}

terminations = {
a: True for a in self.agents
}

self.agents = []

elif self.prisoner_x == self.escape_x and self.prisoner_y == self.escape_y:
rewards = {
"prisoner": 1,
"guard": -1
}

terminations = {
a: True for a in self.agents
}

self.agents = []

# Check truncation conditions (overwrites termination conditions)
truncations = {
"prisoner": False,
"guard": False
}

if self.timestep > 100:
rewards = {
"prisoner": 0,
"guard": 0
}

truncations = {
"prisoner": True,
"guard": True
}

self.agents = []

self.timestep += 1

# Get observations
observation = (
self.prisoner_x + 7 * self.prisoner_y,
self.guard_x + 7 * self.guard_y,
self.escape_x + 7 * self.escape_y,
)

observations = {
"prisoner": {
"observation": observation,
"action_mask": prisoner_action_mask,
},
"guard": {
"observation": observation,
"action_mask": guard_action_mask
},
}

# Get dummy infos (not used in this example)
infos = {"prisoner": {}, "guard": {}}

return observations, rewards, terminations, truncations, infos

def render(self, mode='human'):
"""Renders the environment."""

grid = np.zeros((7, 7), dtype = object)
grid[self.prisoner_y, self.prisoner_x] = "P"
grid[self.guard_y, self.guard_x] = "G"
grid[self.escape_y, self.escape_x] = "E"
print(f"{grid} \n")

# Observation space should be defined here.
# lru_cache allows observation and action spaces to be memoized, reducing clock cycles required to get each agent's space.
# If your spaces change over time, remove this line (disable caching).
@functools.lru_cache(maxsize=None)
def observation_space(self, agent):
# gymnasium spaces are defined and documented here: https://gymnasium.farama.org/api/spaces/
return MultiDiscrete([7 * 7 - 1] * 3)

# Action space should be defined here.
# If your spaces change over time, remove this line (disable caching).
@functools.lru_cache(maxsize=None)
def action_space(self, agent):
return Discrete(4)


The test function pettingzoo.test.parallel_api_test passes. Now, I would like to train with RLlib.

from ray import tune
from ray.rllib.env import ParallelPettingZooEnv
from ray.tune.registry import register_env

env_name = "mpe"

# Register the environment
def env_creator(args):
return ParallelPettingZooEnv(CustomActionMaskedEnvironment())

register_env(env_name, env_creator)



This code, however, fails with the following traceback during some pickling event:

(raylet) Traceback (most recent call last):
File "python/ray/_raylet.pyx", line 1873, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 1981, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 1880, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 1820, in ray._raylet.execute_task.function_executor
File "..PATH../python3.10/site-packages/ray/_private/function_manager.py", line 689, in actor_method_executor
return method(__ray_actor, *args, **kwargs)
File "..PATH../python3.10/site-packages/ray/util/tracing/tracing_helper.py", line 463, in _resume_span
return method(self, *_args, **_kwargs)
File "..PATH../python3.10/site-packages/ray/rllib/algorithms/algorithm.py", line 471, in __init__
self._env_id, self.env_creator = self._get_env_id_and_creator(
File "..PATH../python3.10/site-packages/ray/rllib/algorithms/algorithm.py", line 3099, in _get_env_id_and_creator
return env_specifier, _global_registry.get(ENV_CREATOR, env_specifier)

File "..PATH../ray/tune/registry.py", line 275, in get
return pickle.loads(value)
`AttributeError: Can't get attribute 'CustomActionMaskedEnvironment.observation_space' on <module '__main__' from '/usr/local/lib/python3.10/dist-packages/ray/_private/workers/default_worker.py'>


from ray.rllib.algorithms.ppo import PPO

config = {
"env": env_name,
"env_config": {}, # config to pass to env class
"num_workers": 1, # Adjust as necessary
"framework": "torch", # or "tf" if using TensorFlow
"multiagent": {
"policies": {
"prisoner": (None, MultiDiscrete([7 * 7 - 1] * 3), Discrete(4), {}),
"guard": (None, MultiDiscrete([7 * 7 - 1] * 3), Discrete(4), {}),
},
"policy_mapping_fn": lambda agent_id: agent_id,
},
}

# Run the training
results = tune.run(PPO, config=config, stop={
"training_iteration": 5
})


Why?

Continue reading...