delivcoal_env.py
空间网格、legacy20/full27、路径意图适配
原文件 · 421 行 · 18828 字节
third_party/marl_mrt2a/marl_mrt2a/marl/src/envs/delivcoal_env.py
SHA-256: a2c36aac21fe51d816ba164e49bce8d0ea5a4a3493e57f1d6996cca0b9a72a6c
1"""让 PyMARL 的 MARL-MRT2A 策略可以在 DelivCoal 中运行。"""23from __future__ import annotations45import importlib.util6import sys7from pathlib import Path8from types import SimpleNamespace910import numpy as np11import yaml1213from .multiagentenv import MultiAgentEnv141516REPO_ROOT = Path(__file__).resolve().parents[6]17DEFAULT_CONFIG = REPO_ROOT / "experiments/reproduction/delivcoal_formal_base.yaml"181920def _load_discrete_env():21 """用单独的模块名加载 DelivCoal，避免和第三方代码中的 env 包重名。"""22 package_name = "gat_ppo_delivcoal"23 if package_name not in sys.modules:24 package_dir = REPO_ROOT / "envs/delivcoal"25 spec = importlib.util.spec_from_file_location(26 package_name,27 package_dir / "__init__.py",28 submodule_search_locations=[str(package_dir)],29 )30 module = importlib.util.module_from_spec(spec)31 sys.modules[package_name] = module32 spec.loader.exec_module(module)33 module_name = f"{package_name}.env_discrete"34 if module_name not in sys.modules:35 spec = importlib.util.spec_from_file_location(36 module_name, REPO_ROOT / "envs/delivcoal/env_discrete.py"37 )38 module = importlib.util.module_from_spec(spec)39 sys.modules[module_name] = module40 spec.loader.exec_module(module)41 return sys.modules[module_name].DiscreteActionEnv424344def _load_config(path):45 path = Path(path).expanduser().resolve()46 payload = yaml.safe_load(path.read_text(encoding="utf-8")) or {}47 rewards = Path(payload.get("rewards_config_file", "rewards.yaml"))48 if not rewards.is_absolute():49 rewards = REPO_ROOT / rewards50 payload["rewards_config_file"] = str(rewards)51 return SimpleNamespace(**payload)525354class DelivCoalWrapper(MultiAgentEnv):55 """把 DelivCoal 的任务选择转换成 MARL-MRT2A 使用的二维动作网格。5657 每个可选任务占用一个网格位置。网络选择位置后，再转换回 DelivCoal 任务编号。58 """5960 def __init__(self, key="delivcoal-formal", config_path=str(DEFAULT_CONFIG),61 grid_size=21, seed=0, time_limit=None,62 observation_layout="legacy20", share_intention="none", **kwargs):63 del kwargs, key64 if int(grid_size) % 2 != 1:65 raise ValueError("grid_size must be odd for MARL-MRT2A target mapping")66 self.grid_size = int(grid_size)67 self.half = self.grid_size // 268 self._seed = int(seed)69 self.observation_layout = str(observation_layout)70 if self.observation_layout not in ("legacy20", "full27"):71 raise ValueError("observation_layout must be legacy20 or full27")72 self.share_intention = str(share_intention)73 if self.share_intention not in ("none", "path"):74 raise ValueError("share_intention must be none or path")75 if self.share_intention == "path" and self.observation_layout != "full27":76 raise ValueError("path intention requires observation_layout=full27")77 self._test_mode = False78 self._episode = 079 self._config = _load_config(config_path)80 self.episode_limit = int(time_limit or self._config.max_steps)81 self._config.max_steps = self.episode_limit82 self._config.episode_length = self.episode_limit83 self._config.episode_time_limit = self.episode_limit84 env_class = _load_discrete_env()85 self._env = env_class(self._config)86 self._env.seed(self._seed)87 self.n_agents = self._env.num_agent88 self.n_actions = self.grid_size * self.grid_size89 # legacy20 使用 20 个通道；full27 另外保存完整任务字段和路径信息。90 self.n_channels = 20 if self.observation_layout == "legacy20" else 2791 self._obs = None92 self._state = None93 self._active = None94 self._available_tasks = None95 self._cell_to_task = None96 self._task_to_cell = None97 self._latest_info = {}98 self._sent_total = 0.099 self._return_total = 0.0100101 def _nearest_free(self, desired, occupied):102 """目标格已占用时，从近到远寻找一个空格。"""103 x0, y0 = desired104 if desired not in occupied:105 return desired106 for radius in range(1, self.grid_size):107 candidates = []108 for dx in range(-radius, radius + 1):109 for dy in range(-radius, radius + 1):110 if max(abs(dx), abs(dy)) != radius:111 continue112 x = int(np.clip(x0 + dx, 0, self.grid_size - 1))113 y = int(np.clip(y0 + dy, 0, self.grid_size - 1))114 candidates.append((abs(dx) + abs(dy), x, y))115 for _, x, y in sorted(set(candidates)):116 if (x, y) not in occupied:117 return (x, y)118 raise RuntimeError("spatial action grid has no free cell")119120 def _desired_cell(self, agent_id, task_id, task_row):121 task_pos = np.asarray(task_row[:2], dtype=np.float32)122 agent_pos = np.asarray(self._obs_dict["ego"][agent_id, :2])123 relative = np.clip(task_pos - agent_pos, -1.0, 1.0)124 desired = np.rint(relative * self.half).astype(np.int64) + self.half125 # Idle 没有位置，固定放在网格中心。126 if task_id == self._env.env.idle_task_id:127 desired[:] = self.half128 return tuple(int(value) for value in desired)129130 def _build_mapping(self):131 """为每个智能体建立任务编号和网格编号的双向映射。"""132 self._cell_to_task = []133 self._task_to_cell = []134 for agent_id in range(self.n_agents):135 cell_to_task = {}136 task_to_cell = {}137 available_ids = np.flatnonzero(self._available_tasks[agent_id] > 0.5)138 # 先放配送和探索任务，再给 Idle、Resupply 分配剩余位置。139 special = {self._env.env.idle_task_id, self._env.env.resupply_task_id}140 ordered = sorted(available_ids, key=lambda task: (task in special, int(task)))141 for task_id in ordered:142 desired = self._desired_cell(143 agent_id, int(task_id), self._obs_dict["tasks"][agent_id, task_id]144 )145 cell = self._nearest_free(desired, cell_to_task)146 flat = int(np.ravel_multi_index(cell, (self.grid_size, self.grid_size)))147 # cell_to_task 用于执行动作，task_to_cell 用于生成观测和动作 mask。148 cell_to_task[flat] = int(task_id)149 task_to_cell[int(task_id)] = cell150 self._cell_to_task.append(cell_to_task)151 self._task_to_cell.append(task_to_cell)152153 def _task_components(self, task):154 """从 compact 或 full 任务观测中取出需求、货量、期限和探索信息。"""155 cap_dim = int(self._config.cap_dim)156 mode = getattr(self._config, "coalition_task_feature_mode", "compact")157 requirements = np.zeros(4, dtype=np.float32)158 if mode == "full":159 values = np.asarray(task[2:2 + cap_dim], dtype=np.float32)160 requirements[:min(4, cap_dim)] = values[:4]161 demand = float(task[2 + cap_dim])162 deadline = float(task[3 + cap_dim])163 exploration = np.asarray(task[4 + cap_dim:7 + cap_dim])164 else:165 demand = float(task[2])166 deadline = float(task[3])167 requirements[0] = float(task[4])168 exploration = np.asarray(task[2:5])169 return requirements, demand, deadline, exploration170171 def _rasterize_path_intentions(self, focal_agent_id, grid):172 """把其他智能体的路径画到第 26 个观测通道。"""173 if self.share_intention != "path":174 return175 focal = self._env.env.agents[focal_agent_id]176 focal_pos = np.asarray(focal.position, dtype=np.float32) / float(177 self._config.map_size178 )179 for other_id, agent in enumerate(self._env.env.agents):180 if other_id == focal_agent_id or not agent.living or not agent.path:181 continue182 points = np.asarray(agent.path, dtype=np.float32)183 if points.ndim != 2 or points.shape[1] != 2:184 continue185 points = points / float(self._config.map_size)186 sampled = []187 for start, end in zip(points[:-1], points[1:]):188 cell_span = np.max(np.abs(end - start)) * self.half189 count = max(2, int(np.ceil(cell_span)) + 1)190 sampled.extend(np.linspace(start, end, count, endpoint=False))191 sampled.append(points[-1])192 total = max(1, len(sampled) - 1)193 for index, point in enumerate(sampled):194 relative = np.clip(point - focal_pos, -1.0, 1.0)195 cell = np.rint(relative * self.half).astype(np.int64) + self.half196 x, y = int(cell[0]), int(cell[1])197 value = 1.0 - 0.5 * index / total198 grid[26, x, y] = max(grid[26, x, y], value)199200 def _agent_grid(self, agent_id):201 """把一个智能体看到的任务和队友写入以自身为中心的栅格。"""202 grid = np.zeros(203 (self.n_channels, self.grid_size, self.grid_size), dtype=np.float32204 )205 grid[0, self.half, self.half] = 1.0206 delivery_end = self._env.env.delivery_task_num207 explore_start = self._env.env.explore_task_start208 explore_end = self._env.env.explore_task_end209 idle_id = self._env.env.idle_task_id210 for task_id, cell in self._task_to_cell[agent_id].items():211 # 任务类型和任务数值字段写入同一个网格位置的不同通道。212 x, y = cell213 if task_id < delivery_end:214 grid[1, x, y] = 1.0215 elif explore_start <= task_id < explore_end:216 grid[2, x, y] = 1.0217 else:218 grid[3, x, y] = 1.0 if task_id == idle_id else 0.5219 task = self._obs_dict["tasks"][agent_id, task_id]220 if self.observation_layout == "legacy20":221 grid[4:7, x, y] = task[2:5]222 else:223 requirements, demand, deadline, exploration = (224 self._task_components(task)225 )226 grid[4:8, x, y] = requirements227 grid[8, x, y] = demand228 grid[9, x, y] = deadline229 grid[10:13, x, y] = exploration230 ego_pos = self._obs_dict["ego"][agent_id, :2]231 # 其他智能体按相对位置写入队友通道，重合位置累计人数并保留最大属性。232 for other_id, other in enumerate(self._obs_dict["others"][agent_id]):233 if other_id == agent_id or other[-1] <= 0.5:234 continue235 relative = np.clip(other[:2] - ego_pos, -1.0, 1.0)236 cell = np.rint(relative * self.half).astype(np.int64) + self.half237 x, y = (int(cell[0]), int(cell[1]))238 if self.observation_layout == "legacy20":239 grid[7, x, y] = min(1.0, grid[7, x, y] + 1.0 / self.n_agents)240 grid[8, x, y] = max(grid[8, x, y], float(other[6]))241 grid[9, x, y] = max(grid[9, x, y], float(other[7]))242 grid[16:20, x, y] = np.maximum(243 grid[16:20, x, y], other[2:6]244 )245 else:246 grid[13, x, y] = min(247 1.0, grid[13, x, y] + 1.0 / self.n_agents248 )249 grid[14, x, y] = max(grid[14, x, y], float(other[6]))250 grid[15, x, y] = max(grid[15, x, y], float(other[7]))251 grid[22:26, x, y] = np.maximum(252 grid[22:26, x, y], other[2:6]253 )254 coalition = int(self._env.env.agents[agent_id].coalition)255 # current_target 通道标出当前正在执行的任务位置。256 current = self._task_to_cell[agent_id].get(coalition)257 if current is not None:258 target_channel = 10 if self.observation_layout == "legacy20" else 16259 grid[target_channel, current[0], current[1]] = 1.0260 if self.observation_layout == "legacy20":261 # 与当前智能体自身有关的字段只写在中心格。262 grid[11, self.half, self.half] = float(self._active[agent_id, 0])263 grid[12:16, self.half, self.half] = self._obs_dict["ego"][agent_id, 2:6]264 else:265 grid[17, self.half, self.half] = float(self._active[agent_id, 0])266 grid[18:22, self.half, self.half] = self._obs_dict["ego"][agent_id, 2:6]267 self._rasterize_path_intentions(agent_id, grid)268 return grid269270 def _refresh(self, obs, active, available):271 """环境推进后重新生成任务到网格的映射、局部观测和全局状态。"""272 self._obs_dict = obs273 self._active = np.asarray(active, dtype=np.float32)274 self._available_tasks = np.asarray(available, dtype=np.float32)275 self._build_mapping()276 self._obs = [self._agent_grid(agent_id) for agent_id in range(self.n_agents)]277 self._state = np.mean(np.stack(self._obs), axis=0).astype(np.float32)278279 def _fallback_task(self, agent_id):280 """动作无效或智能体未激活时，选择可继续执行的任务。"""281 coalition = int(self._env.env.agents[agent_id].coalition)282 if coalition in self._task_to_cell[agent_id]:283 return coalition284 idle = int(self._env.env.idle_task_id)285 if self._available_tasks[agent_id, idle] > 0.5:286 return idle287 return int(np.flatnonzero(self._available_tasks[agent_id] > 0.5)[0])288289 def step(self, actions):290 # 将网格动作转换成任务编号。未激活智能体继续当前任务或选择 Idle。291 decision_mask = self._active.astype(np.uint8, copy=True)292 task_actions = []293 for agent_id, raw_action in enumerate(actions):294 flat = int(raw_action)295 task_id = self._cell_to_task[agent_id].get(flat)296 if task_id is None or self._active[agent_id, 0] <= 0.5:297 # 网格位置没有任务，或智能体本步无需决策时，使用安全的继续动作。298 task_id = self._fallback_task(agent_id)299 task_actions.append(task_id)300 obs, rewards, dones, info, active, bad_mask, available = self._env.step(301 task_actions302 )303 self._sent_total += float(info.get("sent", 0.0))304 team_reward = float(np.sum(rewards))305 # PyMARL 接收团队总奖励，环境的逐智能体奖励在这里求和。306 self._return_total += team_reward307 self._refresh(obs, active, available)308 terminated = bool(np.all(dones))309 delivery = [task for task in info["tasks"] if int(task["type"]) == 1]310 completed = sum(bool(task["completed"]) for task in delivery)311 sent_goal = max(1e-8, float(info.get("sent_goal", 0.0)))312 episode_limit = bool(terminated and np.any(np.asarray(bad_mask) < 0.5))313 public_info = {314 # runner 依赖这些字段记录目标更新、episode 结束原因和任务指标。315 "robot_info": self._robot_info(),316 "action_exec": decision_mask,317 "episode_limit": episode_limit,318 "success": float(bool(delivery) and completed == len(delivery)),319 "task_completion_ratio": float(completed / max(1, len(delivery))),320 "sent_percentage": float(self._sent_total / sent_goal),321 "coalition_changes": float(322 info.get("reconfiguration", {}).get("total_coalition_changes", 0)323 ),324 }325 self._latest_info = public_info326 return team_reward, terminated, public_info327328 def _robot_info(self):329 """返回每个智能体在网格中心以及当前目标所在的格子。"""330 result = []331 center = (self.half, self.half)332 for agent_id, agent in enumerate(self._env.env.agents):333 cell = self._task_to_cell[agent_id].get(int(agent.coalition))334 result.append((center, cell))335 return result336337 def reset(self):338 # episode 编号传给 DelivCoal，用于环境重置模式和课程切换。339 obs, active, available = self._env.reset(self._episode)340 self._episode += 1341 self._sent_total = 0.0342 self._return_total = 0.0343 self._latest_info = {}344 self._refresh(obs, active, available)345 return self.get_obs(), self.get_state()346347 def get_obs(self):348 return self._obs349350 def get_obs_agent(self, agent_id):351 return self._obs[agent_id]352353 def get_obs_size(self):354 return (self.n_channels, self.grid_size, self.grid_size)355356 def get_state(self):357 return self._state358359 def get_state_size(self):360 return (self.n_channels, self.grid_size, self.grid_size)361362 def get_avail_actions(self):363 # PyMARL 需要网格长度的 mask，这里只开放已经分配给任务的格子。364 result = []365 for agent_id in range(self.n_agents):366 mask = np.zeros(self.n_actions, dtype=np.int32)367 if self._active[agent_id, 0] > 0.5:368 mask[list(self._cell_to_task[agent_id])] = 1369 else:370 task_id = self._fallback_task(agent_id)371 cell = self._task_to_cell[agent_id][task_id]372 mask[np.ravel_multi_index(cell, (self.grid_size, self.grid_size))] = 1373 result.append(mask)374 return result375376 def get_avail_agent_actions(self, agent_id):377 return self.get_avail_actions()[agent_id]378379 def get_total_actions(self):380 return self.n_actions381382 def set_mode(self, test_mode):383 self._test_mode = bool(test_mode)384 return True385386 def seed(self):387 return self._seed388389 def close(self):390 self._env.close()391392 def render(self):393 return None394395 def save_replay(self):396 return None397398 def get_stats(self):399 return {}400401 def get_channel_info(self):402 if self.observation_layout == "full27":403 return {404 "task_type": [1, 2, 3],405 "task_requirements": [4, 5, 6, 7],406 "task_demand": 8,407 "task_deadline": 9,408 "exploration_geometry": [10, 11, 12],409 "team": [13, 14, 15],410 "current_target": 16,411 "active": 17,412 "ego_capabilities": [18, 19, 20, 21],413 "team_capabilities": [22, 23, 24, 25],414 "path_intention": 26,415 }416 return {417 "task_type": [1, 2, 3],418 "task_attributes": [4, 5, 6],419 "team": [7, 8, 9],420 "current_target": 10,421 }