env_discrete.py
局部网格观测到 HATT 输入的适配
原文件 · 206 行 · 8844 字节
envs/gymrt2a/env_discrete.py
SHA-256: 2b941a90ba5d840e4c454ebc47c403827c2cf15118b808bbb1ce8ce4a7e9a630
1"""Project adapter for the upstream GyMRT2A grid environment.23The upstream source is vendored unchanged under ``upstream/gymrt2a``. This4module only translates its spatial observations and flattened EPyMARL actions5to the interface used by this repository's vector environment and buffers.6"""78from __future__ import annotations910import random11import sys12from pathlib import Path1314import gym15import numpy as np16from gym import spaces171819UPSTREAM_ROOT = Path(__file__).resolve().parent / "upstream"20if str(UPSTREAM_ROOT) not in sys.path:21 sys.path.insert(0, str(UPSTREAM_ROOT))2223from gymrt2a.env import GridWorldEnv # noqa: E402242526class DiscreteActionEnv:27 """Expose GyMRT2A through the project's multi-agent environment contract."""2829 def __init__(self, config):30 self.config = config31 self.num_agent = int(config.gymrt2a_num_agents)32 self.max_steps = int(config.max_steps)33 self._seed = int(config.seed)34 self._episode_step = 03536 raw = GridWorldEnv(37 sz=(int(config.gymrt2a_map_size),) * 2,38 n_agents=self.num_agent,39 n_obj=list(config.gymrt2a_objects),40 view_range=int(config.gymrt2a_view_range),41 comm_range=int(config.gymrt2a_comm_range),42 comm=int(config.gymrt2a_comm_dim),43 render=bool(config.use_render),44 max_obj_lvl=len(config.gymrt2a_objects),45 one_hot_obj_lvl=True,46 obj_lvl_rwd_exp=float(config.gymrt2a_reward_exponent),47 Lsec=int(config.gymrt2a_sectors),48 action_grid=True,49 share_intention=config.gymrt2a_share_intention,50 respawn=config.gymrt2a_respawn,51 )52 self.env = raw53 sample_obs = self.env.reset()[0]54 self.channels, self.grid_side, _ = sample_obs.shape55 self.num_candidates = self.grid_side * self.grid_side56 self.candidate_dim = self.channels + 257 self.ego_dim = 358 self.other_dim = 459 self.channel_info = self.env.get_channel_info()60 self.object_channels = tuple(self.channel_info["objects"])6162 action_n = self.num_candidates * max(1, int(config.gymrt2a_comm_dim))63 if action_n != self.num_candidates * max(1, int(config.gymrt2a_comm_dim)):64 raise ValueError("GyMRT2A action grid and candidate observation do not align.")65 if int(config.gymrt2a_comm_dim) > 1:66 raise NotImplementedError(67 "The first migration supports navigation/intention sharing with silent communication only."68 )6970 obs_space = spaces.Dict({71 "ego": spaces.Box(-np.inf, np.inf, (self.ego_dim,), dtype=np.float32),72 "tasks": spaces.Box(73 -np.inf, np.inf, (self.num_candidates, self.candidate_dim), dtype=np.float3274 ),75 "others": spaces.Box(76 -np.inf, np.inf, (self.num_agent, self.other_dim), dtype=np.float3277 ),78 })79 self.observation_space = [obs_space for _ in range(self.num_agent)]80 self.action_space = [spaces.Discrete(action_n) for _ in range(self.num_agent)]8182 map_cells = int(config.gymrt2a_map_size) ** 283 state_space = spaces.Dict({84 "agents": spaces.Box(-np.inf, np.inf, (self.num_agent, 3), dtype=np.float32),85 "tasks": spaces.Box(86 -np.inf, np.inf, (map_cells, len(config.gymrt2a_objects) + 2), dtype=np.float3287 ),88 "global": spaces.Box(-np.inf, np.inf, (2,), dtype=np.float32),89 "focal": spaces.Box(0.0, 1.0, (self.num_agent, 1), dtype=np.float32),90 })91 self.share_observation_space = [state_space for _ in range(self.num_agent)]9293 def seed(self, seed):94 self._seed = int(seed)95 return [self._seed]9697 def set_reset_mode(self, mode):98 # GyMRT2A always samples a fresh scenario. Accept the project's eval hook.99 del mode100101 def _seed_episode(self, episode):102 seed = self._seed + max(0, int(episode))103 np.random.seed(seed)104 random.seed(seed)105106 def _format_obs(self, raw_obs):107 rows, cols = np.meshgrid(108 np.linspace(-1.0, 1.0, self.grid_side, dtype=np.float32),109 np.linspace(-1.0, 1.0, self.grid_side, dtype=np.float32),110 indexing="ij",111 )112 coordinates = np.stack([rows, cols], axis=-1).reshape(self.num_candidates, 2)113 robot_channel = int(self.channel_info["robots"])114 formatted = {"ego": [], "tasks": [], "others": []}115 for obs in raw_obs:116 obs = np.asarray(obs, dtype=np.float32)117 candidate_values = obs.reshape(self.channels, -1).T118 formatted["tasks"].append(np.concatenate([coordinates, candidate_values], axis=-1))119120 center = self.grid_side // 2121 formatted["ego"].append(122 np.asarray([self._episode_step / max(1, self.max_steps), 0.0, 1.0], dtype=np.float32)123 )124 others = np.zeros((self.num_agent, self.other_dim), dtype=np.float32)125 visible = np.argwhere(obs[robot_channel] > 0.5)126 for slot, (row, col) in enumerate(visible[: self.num_agent]):127 others[slot] = (128 (row - center) / max(1, center),129 (col - center) / max(1, center),130 1.0,131 1.0,132 )133 formatted["others"].append(others)134 return {key: np.asarray(value, dtype=np.float32) for key, value in formatted.items()}135136 def _available_actions(self, raw_obs):137 """复现官方控制器：视野内有物体时，只允许选择物体所在格点。"""138 available = np.ones((self.num_agent, self.action_space[0].n), dtype=np.float32)139 for agent_index, obs in enumerate(raw_obs):140 object_grid = np.asarray(obs)[list(self.object_channels)].any(axis=0).reshape(-1)141 if object_grid.any():142 available[agent_index] = object_grid.astype(np.float32)143 return available144145 def reset(self, episode=0):146 self._seed_episode(episode)147 self._episode_step = 0148 raw_obs = self.env.reset()149 obs = self._format_obs(raw_obs)150 active = np.ones((self.num_agent, 1), dtype=np.float32)151 available = self._available_actions(raw_obs)152 return obs, active, available153154 def step(self, actions):155 flat_actions = np.asarray(actions).reshape(self.num_agent).astype(np.int64)156 comm_n = max(1, int(self.config.gymrt2a_comm_dim))157 shaped_actions = []158 for action in flat_actions:159 cell, comm = divmod(int(action), comm_n)160 row, col = divmod(cell, self.grid_side)161 center = self.grid_side // 2162 shaped_actions.append(((row - center, col - center), comm))163 raw_obs, rewards, dones, info = self.env.step(tuple(shaped_actions))164 self._episode_step += 1165 truncated = self._episode_step >= self.max_steps166 done = bool(all(dones) or truncated)167 info = dict(info)168 info.update({169 "end": self._episode_step if done else 0,170 "success": bool(all(dones)),171 "truncated": truncated and not all(dones),172 "objects_collected": int(sum(self.env.obj_pickup_count)),173 "objects_by_level": list(self.env.obj_pickup_count),174 })175 reward = np.asarray(rewards, dtype=np.float32).reshape(self.num_agent, 1)176 done_array = np.full(self.num_agent, done, dtype=bool)177 active = np.ones((self.num_agent, 1), dtype=np.float32)178 bad = np.full((self.num_agent, 1), 0.0 if truncated else 1.0, dtype=np.float32)179 available = self._available_actions(raw_obs)180 return self._format_obs(raw_obs), reward, done_array, info, active, bad, available181182 def get_global_state(self):183 size = float(self.env.sz[0])184 agents = np.zeros((self.num_agent, 3), dtype=np.float32)185 for index, robot in enumerate(self.env.map.robots):186 agents[index] = (robot.pos[0] / size, robot.pos[1] / size, 1.0)187 task_dim = len(self.config.gymrt2a_objects) + 2188 tasks = np.zeros((int(size) ** 2, task_dim), dtype=np.float32)189 for obj in self.env.map.objects:190 index = int(obj.pos[0] * size + obj.pos[1])191 tasks[index, :2] = (obj.pos[0] / size, obj.pos[1] / size)192 tasks[index, 1 + int(obj.lvl)] = 1.0193 return {194 "agents": agents,195 "tasks": tasks,196 "global": np.asarray(197 [self._episode_step / max(1, self.max_steps), len(self.env.map.objects) / max(1, sum(self.config.gymrt2a_objects))],198 dtype=np.float32,199 ),200 }201202 def close(self):203 self.env.close()204205 def render(self, mode="human"):206 return self.env.render(mode=mode)