def __init__(self, env_name, device='cpu', frame_skip=4):
env = importlib.import_module(
'pettingzoo.atari.{}'.format(env_name)).parallel_env(
obs_type='grayscale_image')
env = MaxAndSkipMAALE(env, skip=frame_skip)
env = from_parallel(env)
env = resize_v0(env, 84, 84)
self._env = env
self.name = env_name
self.device = torch.device(device)
self.agents = self._env.possible_agents
self.subenvs = {
agent: SubEnv(agent, device, self.state_spaces[agent],
self.action_spaces[agent])
for agent in self.agents
}
def __init__(self, env_name, device='cpu', frame_skip=4):
def raw_env(**kwargs):
return BaseAtariEnv(game="surround",
num_players=1,
mode_num=2,
**kwargs)
env = parallel_wrapper_fn(
base_env_wrapper_fn(raw_env))(obs_type="grayscale_image")
env = MaxAndSkipMAALE(env, skip=frame_skip)
env = from_parallel(env)
env = recolor_observations(env)
env = resize_v0(env, 84, 84)
self._env = env
self.name = env_name
self.device = torch.device(device)
self.agents = self._env.possible_agents
self.subenvs = {
agent: SubEnv(agent, device, self.state_spaces[agent],
self.action_spaces[agent])
for agent in self.agents
}
def __call__(self, env, *args, **kwargs):
if isinstance(env, gym.Env):
if not self.gym_avaliable:
raise ValueError(
f"{self.wrapper_name} does not apply to gym environments, pettingzoo environments only"
)
wrap_class = getattr(gym_wrappers, self.wrapper_name)
return wrap_class(env, *args, **kwargs)
elif isinstance(env, AECEnv):
wrap_class = getattr(aec_wrappers, self.wrapper_name)
return wrap_class(env, *args, **kwargs)
elif isinstance(env, ParallelEnv):
wrap_class = getattr(parallel_wrappers, self.wrapper_name, None)
if wrap_class is not None:
return wrap_class(env, *args, **kwargs)
else:
wrap_class = getattr(aec_wrappers, self.wrapper_name)
return to_parallel(
wrap_class(from_parallel(env), *args, **kwargs))
else:
raise ValueError(
"environment passed to supersuit wrapper must either be a gym environment or a pettingzoo environment"
)