"""Project adapter for the upstream GyMRT2A grid environment. The upstream source is vendored unchanged under ``upstream/gymrt2a``. This module only translates its spatial observations and flattened EPyMARL actions to the interface used by this repository's vector environment and buffers. """ from __future__ import annotations import random import sys from pathlib import Path import gym import numpy as np from gym import spaces UPSTREAM_ROOT = Path(__file__).resolve().parent / "upstream" if str(UPSTREAM_ROOT) not in sys.path: sys.path.insert(0, str(UPSTREAM_ROOT)) from gymrt2a.env import GridWorldEnv # noqa: E402 class DiscreteActionEnv: """Expose GyMRT2A through the project's multi-agent environment contract.""" def __init__(self, config): self.config = config self.num_agent = int(config.gymrt2a_num_agents) self.max_steps = int(config.max_steps) self._seed = int(config.seed) self._episode_step = 0 raw = GridWorldEnv( sz=(int(config.gymrt2a_map_size),) * 2, n_agents=self.num_agent, n_obj=list(config.gymrt2a_objects), view_range=int(config.gymrt2a_view_range), comm_range=int(config.gymrt2a_comm_range), comm=int(config.gymrt2a_comm_dim), render=bool(config.use_render), max_obj_lvl=len(config.gymrt2a_objects), one_hot_obj_lvl=True, obj_lvl_rwd_exp=float(config.gymrt2a_reward_exponent), Lsec=int(config.gymrt2a_sectors), action_grid=True, share_intention=config.gymrt2a_share_intention, respawn=config.gymrt2a_respawn, ) self.env = raw sample_obs = self.env.reset()[0] self.channels, self.grid_side, _ = sample_obs.shape self.num_candidates = self.grid_side * self.grid_side self.candidate_dim = self.channels + 2 self.ego_dim = 3 self.other_dim = 4 self.channel_info = self.env.get_channel_info() self.object_channels = tuple(self.channel_info["objects"]) action_n = self.num_candidates * max(1, int(config.gymrt2a_comm_dim)) if action_n != self.num_candidates * max(1, int(config.gymrt2a_comm_dim)): raise ValueError("GyMRT2A action grid and candidate observation do not align.") if int(config.gymrt2a_comm_dim) > 1: raise NotImplementedError( "The first migration supports navigation/intention sharing with silent communication only." ) obs_space = spaces.Dict({ "ego": spaces.Box(-np.inf, np.inf, (self.ego_dim,), dtype=np.float32), "tasks": spaces.Box( -np.inf, np.inf, (self.num_candidates, self.candidate_dim), dtype=np.float32 ), "others": spaces.Box( -np.inf, np.inf, (self.num_agent, self.other_dim), dtype=np.float32 ), }) self.observation_space = [obs_space for _ in range(self.num_agent)] self.action_space = [spaces.Discrete(action_n) for _ in range(self.num_agent)] map_cells = int(config.gymrt2a_map_size) ** 2 state_space = spaces.Dict({ "agents": spaces.Box(-np.inf, np.inf, (self.num_agent, 3), dtype=np.float32), "tasks": spaces.Box( -np.inf, np.inf, (map_cells, len(config.gymrt2a_objects) + 2), dtype=np.float32 ), "global": spaces.Box(-np.inf, np.inf, (2,), dtype=np.float32), "focal": spaces.Box(0.0, 1.0, (self.num_agent, 1), dtype=np.float32), }) self.share_observation_space = [state_space for _ in range(self.num_agent)] def seed(self, seed): self._seed = int(seed) return [self._seed] def set_reset_mode(self, mode): # GyMRT2A always samples a fresh scenario. Accept the project's eval hook. del mode def _seed_episode(self, episode): seed = self._seed + max(0, int(episode)) np.random.seed(seed) random.seed(seed) def _format_obs(self, raw_obs): rows, cols = np.meshgrid( np.linspace(-1.0, 1.0, self.grid_side, dtype=np.float32), np.linspace(-1.0, 1.0, self.grid_side, dtype=np.float32), indexing="ij", ) coordinates = np.stack([rows, cols], axis=-1).reshape(self.num_candidates, 2) robot_channel = int(self.channel_info["robots"]) formatted = {"ego": [], "tasks": [], "others": []} for obs in raw_obs: obs = np.asarray(obs, dtype=np.float32) candidate_values = obs.reshape(self.channels, -1).T formatted["tasks"].append(np.concatenate([coordinates, candidate_values], axis=-1)) center = self.grid_side // 2 formatted["ego"].append( np.asarray([self._episode_step / max(1, self.max_steps), 0.0, 1.0], dtype=np.float32) ) others = np.zeros((self.num_agent, self.other_dim), dtype=np.float32) visible = np.argwhere(obs[robot_channel] > 0.5) for slot, (row, col) in enumerate(visible[: self.num_agent]): others[slot] = ( (row - center) / max(1, center), (col - center) / max(1, center), 1.0, 1.0, ) formatted["others"].append(others) return {key: np.asarray(value, dtype=np.float32) for key, value in formatted.items()} def _available_actions(self, raw_obs): """复现官方控制器：视野内有物体时，只允许选择物体所在格点。""" available = np.ones((self.num_agent, self.action_space[0].n), dtype=np.float32) for agent_index, obs in enumerate(raw_obs): object_grid = np.asarray(obs)[list(self.object_channels)].any(axis=0).reshape(-1) if object_grid.any(): available[agent_index] = object_grid.astype(np.float32) return available def reset(self, episode=0): self._seed_episode(episode) self._episode_step = 0 raw_obs = self.env.reset() obs = self._format_obs(raw_obs) active = np.ones((self.num_agent, 1), dtype=np.float32) available = self._available_actions(raw_obs) return obs, active, available def step(self, actions): flat_actions = np.asarray(actions).reshape(self.num_agent).astype(np.int64) comm_n = max(1, int(self.config.gymrt2a_comm_dim)) shaped_actions = [] for action in flat_actions: cell, comm = divmod(int(action), comm_n) row, col = divmod(cell, self.grid_side) center = self.grid_side // 2 shaped_actions.append(((row - center, col - center), comm)) raw_obs, rewards, dones, info = self.env.step(tuple(shaped_actions)) self._episode_step += 1 truncated = self._episode_step >= self.max_steps done = bool(all(dones) or truncated) info = dict(info) info.update({ "end": self._episode_step if done else 0, "success": bool(all(dones)), "truncated": truncated and not all(dones), "objects_collected": int(sum(self.env.obj_pickup_count)), "objects_by_level": list(self.env.obj_pickup_count), }) reward = np.asarray(rewards, dtype=np.float32).reshape(self.num_agent, 1) done_array = np.full(self.num_agent, done, dtype=bool) active = np.ones((self.num_agent, 1), dtype=np.float32) bad = np.full((self.num_agent, 1), 0.0 if truncated else 1.0, dtype=np.float32) available = self._available_actions(raw_obs) return self._format_obs(raw_obs), reward, done_array, info, active, bad, available def get_global_state(self): size = float(self.env.sz[0]) agents = np.zeros((self.num_agent, 3), dtype=np.float32) for index, robot in enumerate(self.env.map.robots): agents[index] = (robot.pos[0] / size, robot.pos[1] / size, 1.0) task_dim = len(self.config.gymrt2a_objects) + 2 tasks = np.zeros((int(size) ** 2, task_dim), dtype=np.float32) for obj in self.env.map.objects: index = int(obj.pos[0] * size + obj.pos[1]) tasks[index, :2] = (obj.pos[0] / size, obj.pos[1] / size) tasks[index, 1 + int(obj.lvl)] = 1.0 return { "agents": agents, "tasks": tasks, "global": np.asarray( [self._episode_step / max(1, self.max_steps), len(self.env.map.objects) / max(1, sum(self.config.gymrt2a_objects))], dtype=np.float32, ), } def close(self): self.env.close() def render(self, mode="human"): return self.env.render(mode=mode)