Spaces:
Sleeping
Sleeping
File size: 9,190 Bytes
3f7c971 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 |
import math
import warnings
import magent
import numpy as np
from gym.spaces import Box, Discrete
from gym.utils import seeding
from pettingzoo import AECEnv
from pettingzoo.magent.render import Renderer
from pettingzoo.utils import agent_selector, wrappers
from pettingzoo.utils.env import ParallelEnv
def make_env(raw_env):
def env_fn(**kwargs):
env = raw_env(**kwargs)
env = wrappers.AssertOutOfBoundsWrapper(env)
env = wrappers.OrderEnforcingWrapper(env)
return env
return env_fn
class magent_parallel_env(ParallelEnv):
def __init__(self, env, active_handles, names, map_size, max_cycles, reward_range, minimap_mode, extra_features):
self.map_size = map_size
self.max_cycles = max_cycles
self.minimap_mode = minimap_mode
self.extra_features = extra_features
self.env = env
self.handles = active_handles
self._all_handles = self.env.get_handles()
env.reset()
self.generate_map()
self.team_sizes = team_sizes = [env.get_num(handle) for handle in self.handles]
self.agents = [f"{names[j]}_{i}" for j in range(len(team_sizes)) for i in range(team_sizes[j])]
self.possible_agents = self.agents[:]
num_actions = [env.get_action_space(handle)[0] for handle in self.handles]
action_spaces_list = [Discrete(num_actions[j]) for j in range(len(team_sizes)) for i in range(team_sizes[j])]
# may change depending on environment config? Not sure.
team_obs_shapes = self._calc_obs_shapes()
state_shape = self._calc_state_shape()
observation_space_list = [Box(low=0., high=2., shape=team_obs_shapes[j], dtype=np.float32) for j in range(len(team_sizes)) for i in range(team_sizes[j])]
self.state_space = Box(low=0., high=2., shape=state_shape, dtype=np.float32)
reward_low, reward_high = reward_range
if extra_features:
for space in observation_space_list:
idx = space.shape[2] - 3 if minimap_mode else space.shape[2] - 1
space.low[:, :, idx] = reward_low
space.high[:, :, idx] = reward_high
idx_state = self.state_space.shape[2] - 3 if minimap_mode else self.state_space.shape[2] - 1
self.state_space.low[:, :, idx_state] = reward_low
self.state_space.high[:, :, idx_state] = reward_high
self.action_spaces = {agent: space for agent, space in zip(self.agents, action_spaces_list)}
self.observation_spaces = {agent: space for agent, space in zip(self.agents, observation_space_list)}
self._zero_obs = {agent: np.zeros_like(space.low) for agent, space in self.observation_spaces.items()}
self.base_state = np.zeros(self.state_space.shape, dtype='float32')
walls = self.env._get_walls_info()
wall_x, wall_y = zip(*walls)
self.base_state[wall_x, wall_y, 0] = 1
self._renderer = None
self.frames = 0
def observation_space(self, agent):
return self.observation_spaces[agent]
def action_space(self, agent):
return self.action_spaces[agent]
def seed(self, seed=None):
if seed is None:
seed = seeding.create_seed(seed, max_bytes=4)
self.env.set_seed(seed)
def _calc_obs_shapes(self):
view_spaces = [self.env.get_view_space(handle) for handle in self.handles]
feature_spaces = [self.env.get_feature_space(handle) for handle in self.handles]
assert all(len(tup) == 3 for tup in view_spaces)
assert all(len(tup) == 1 for tup in feature_spaces)
feat_size = [[fs[0]] for fs in feature_spaces]
for feature_space in feat_size:
if not self.extra_features:
feature_space[0] = 2 if self.minimap_mode else 0
obs_spaces = [(view_space[:2] + (view_space[2] + feature_space[0],)) for view_space, feature_space in zip(view_spaces, feat_size)]
return obs_spaces
def _calc_state_shape(self):
feature_spaces = [
self.env.get_feature_space(handle) for handle in self._all_handles
]
self._minimap_features = 2 if self.minimap_mode else 0
# map channel and agent pair channel. Remove global agent position when minimap mode and extra features
state_depth = (
(max(feature_spaces)[0] - self._minimap_features) * self.extra_features
+ 1
+ len(self._all_handles) * 2
)
return (self.map_size, self.map_size, state_depth)
def render(self, mode="human"):
if self._renderer is None:
self._renderer = Renderer(self.env, self.map_size, mode)
assert mode == self._renderer.mode, "mode must be consistent across render calls"
return self._renderer.render(mode)
def close(self):
if self._renderer is not None:
self._renderer.close()
self._renderer = None
def reset(self):
self.agents = self.possible_agents[:]
self.env.reset()
self.frames = 0
self.all_dones = {agent: False for agent in self.possible_agents}
self.generate_map()
return self._observe_all()
def _observe_all(self):
observes = [None] * self.max_num_agents
for handle in self.handles:
ids = self.env.get_agent_id(handle)
view, features = self.env.get_observation(handle)
if self.minimap_mode and not self.extra_features:
features = features[:, -2:]
if self.minimap_mode or self.extra_features:
feat_reshape = np.expand_dims(np.expand_dims(features, 1), 1)
feat_img = np.tile(feat_reshape, (1, view.shape[1], view.shape[2], 1))
fin_obs = np.concatenate([view, feat_img], axis=-1)
else:
fin_obs = np.copy(view)
for id, obs in zip(ids, fin_obs):
observes[id] = obs
ret_agents = set(self.agents)
return {agent: obs if obs is not None else self._zero_obs[agent] for agent, obs in zip(self.possible_agents, observes) if agent in ret_agents}
def _all_rewards(self):
rewards = np.zeros(self.max_num_agents)
for handle in self.handles:
ids = self.env.get_agent_id(handle)
rewards[ids] = self.env.get_reward(handle)
ret_agents = set(self.agents)
return {agent: float(rew) for agent, rew in zip(self.possible_agents, rewards) if agent in ret_agents}
def _all_dones(self, step_done=False):
dones = np.ones(self.max_num_agents, dtype=bool)
if not step_done:
for handle in self.handles:
ids = self.env.get_agent_id(handle)
dones[ids] = ~self.env.get_alive(handle)
ret_agents = set(self.agents)
return {agent: bool(done) for agent, done in zip(self.possible_agents, dones) if agent in ret_agents}
def state(self):
'''
Returns an observation of the global environment
'''
state = np.copy(self.base_state)
for handle in self._all_handles:
view, features = self.env.get_observation(handle)
pos = self.env.get_pos(handle)
pos_x, pos_y = zip(*pos)
state[pos_x, pos_y, 1 + handle.value * 2] = 1
state[pos_x, pos_y, 2 + handle.value * 2] = view[:, view.shape[1] // 2, view.shape[2] // 2, 2]
if self.extra_features:
add_zeros = np.zeros(
(
features.shape[0],
state.shape[2]
- (1 + len(self.team_sizes) * 2 + features.shape[1] - self._minimap_features),
)
)
rewards = features[:, -1 - self._minimap_features]
actions = features[:, :-1 - self._minimap_features]
actions = np.concatenate((actions, add_zeros), axis=1)
rewards = rewards.reshape(len(rewards), 1)
state_features = np.hstack((actions, rewards))
state[pos_x, pos_y, 1 + len(self.team_sizes) * 2:] = state_features
return state
def step(self, all_actions):
action_list = [0] * self.max_num_agents
for i, agent in enumerate(self.possible_agents):
if agent in all_actions:
action_list[i] = all_actions[agent]
all_actions = np.asarray(action_list, dtype=np.int32)
start_point = 0
for i in range(len(self.handles)):
size = self.team_sizes[i]
self.env.set_action(self.handles[i], all_actions[start_point:(start_point + size)])
start_point += size
self.frames += 1
done = self.env.step() or self.frames >= self.max_cycles
all_infos = {agent: {} for agent in self.agents}
all_dones = self._all_dones(done)
all_rewards = self._all_rewards()
all_observes = self._observe_all()
self.all_dones = all_dones
self.env.clear_dead()
self.agents = [agent for agent in self.agents if not self.all_dones[agent]]
return all_observes, all_rewards, all_dones, all_infos
|