QMixPolicy.py
individual-Q 策略
原文件 · 255 行 · 10269 字节
algorithms/qmix/algorithm/QMixPolicy.py
SHA-256: 7266ead76a8d68641e9db739bb03f4122608d61589995b5b345b0b922042e014
1"""Shared recurrent individual-Q policies for QMIX, VDN and HATT-QMIX."""23import numpy as np4import torch5import torch.nn as nn67from algorithms.algorithm.HATTV20Network import (8 CandidateEdgeDecoder,9 HATTV20Backbone,10)11from utils.obs_utils import flat_obs_dim, flatten_obs_dict121314def _tensor(value, device):15 return value if isinstance(value, torch.Tensor) else torch.as_tensor(16 value, dtype=torch.float32, device=device17 )181920class RecurrentAgentQNetwork(nn.Module):21 """Parameter-shared DRQN applied independently to every agent."""2223 def __init__(self, obs_dim, action_dim, hidden_dim, use_recurrent=True):24 super().__init__()25 self.hidden_dim = int(hidden_dim)26 self.use_recurrent = bool(use_recurrent)27 self.encoder = nn.Sequential(28 nn.Linear(int(obs_dim), self.hidden_dim),29 nn.ReLU(),30 nn.LayerNorm(self.hidden_dim),31 )32 self.rnn = (33 nn.GRUCell(self.hidden_dim, self.hidden_dim)34 if self.use_recurrent else None35 )36 self.q_head = nn.Linear(self.hidden_dim, int(action_dim))3738 def forward_step(self, obs, hidden, available_actions=None):39 prefix = obs.shape[:-1]40 encoded = self.encoder(obs.reshape(-1, obs.shape[-1]))41 if self.use_recurrent:42 next_hidden = self.rnn(encoded, hidden.reshape(-1, self.hidden_dim))43 else:44 next_hidden = hidden.reshape(-1, self.hidden_dim)45 q_values = self.q_head(46 next_hidden if self.use_recurrent else encoded47 ).view(*prefix, -1)48 return q_values, next_hidden.view(*prefix, self.hidden_dim)495051class HATTAgentQNetwork(nn.Module):52 """HATT v2.0 relational encoder with a recurrent per-candidate Q decoder."""5354 def __init__(self, args, obs_space, num_agents, action_dim):55 super().__init__()56 spaces = obs_space.spaces57 hidden_dim = int(args.hatt_hidden_dim)58 self.hidden_dim = int(args.qmix_agent_hidden_dim)59 self.use_recurrent = bool(args.qmix_use_recurrent_agent)60 self.num_agents = int(num_agents)61 self.action_dim = int(action_dim)62 self.backbone = HATTV20Backbone(63 ego_dim=spaces["ego"].shape[-1],64 other_dim=spaces["others"].shape[-1],65 task_dim=spaces["tasks"].shape[-1],66 hidden_dim=hidden_dim,67 num_heads=int(args.hatt_num_heads),68 cap_dim=int(args.cap_dim),69 num_agents=self.num_agents,70 num_tasks=self.action_dim,71 dropout=float(args.hatt_dropout),72 shared_task_obs=bool(args.hatt_shared_task_obs),73 )74 self.agent_projection = nn.Sequential(75 nn.Linear(hidden_dim, self.hidden_dim),76 nn.ReLU(),77 nn.LayerNorm(self.hidden_dim),78 )79 self.rnn = (80 nn.GRUCell(self.hidden_dim, self.hidden_dim)81 if self.use_recurrent else None82 )83 self.decoder = CandidateEdgeDecoder(84 agent_dim=self.hidden_dim,85 task_dim=hidden_dim,86 hidden_dim=self.hidden_dim,87 )8889 def forward_step(self, obs, hidden, available_actions=None):90 agent_context, task_context, edges, _ = self.backbone(91 obs, available_actions, return_diagnostics=False92 )93 projected = self.agent_projection(agent_context)94 prefix = projected.shape[:-1]95 if self.use_recurrent:96 next_hidden = self.rnn(97 projected.reshape(-1, self.hidden_dim),98 hidden.reshape(-1, self.hidden_dim),99 ).view(*prefix, self.hidden_dim)100 decoder_context = next_hidden101 else:102 next_hidden = hidden103 decoder_context = projected104 q_values = self.decoder(decoder_context, task_context, edges)105 return q_values, next_hidden106107 @staticmethod108 def _flatten_time(tree, start, stop):109 """Flatten a contiguous [B, T, ...] slice to [B*T, ...]."""110111 if isinstance(tree, dict):112 return {113 key: HATTAgentQNetwork._flatten_time(value, start, stop)114 for key, value in tree.items()115 }116 chunk = tree[:, start:stop]117 return chunk.reshape(-1, *chunk.shape[2:])118119 def forward_sequence(120 self, obs, hidden, available_actions, chunk_size, alive_masks=None121 ):122 """Unroll the GRU while batching the memoryless HATT work over time.123124 HATT encoding and candidate decoding dominate the update cost and do not125 carry recurrent state. Processing a small time chunk in one call keeps126 the exact GRU recurrence while avoiding one set of tiny GPU kernels per127 transition. Chunking bounds the activation memory retained for BPTT.128 """129130 transitions = int(available_actions.shape[1])131 chunk_size = max(1, int(chunk_size))132 sequence_values = []133 for start in range(0, transitions, chunk_size):134 stop = min(transitions, start + chunk_size)135 width = stop - start136 chunk_obs = self._flatten_time(obs, start, stop)137 chunk_available = self._flatten_time(available_actions, start, stop)138 agent_context, task_context, edges, _ = self.backbone(139 chunk_obs, chunk_available, return_diagnostics=False140 )141 projected = self.agent_projection(agent_context).reshape(142 hidden.shape[0], width, self.num_agents, self.hidden_dim143 )144 if self.use_recurrent:145 recurrent = []146 for offset in range(width):147 if alive_masks is not None:148 hidden = hidden * alive_masks[:, start + offset]149 hidden = self.rnn(150 projected[:, offset].reshape(-1, self.hidden_dim),151 hidden.reshape(-1, self.hidden_dim),152 ).view_as(hidden)153 recurrent.append(hidden)154 decoder_context = torch.stack(recurrent, dim=1).reshape(155 -1, self.num_agents, self.hidden_dim156 )157 else:158 decoder_context = projected.reshape(159 -1, self.num_agents, self.hidden_dim160 )161 q_values = self.decoder(decoder_context, task_context, edges)162 sequence_values.append(163 q_values.reshape(hidden.shape[0], width, self.num_agents, -1)164 )165 return torch.cat(sequence_values, dim=1), hidden166167168class QMixPolicy(nn.Module):169 """Evaluation/checkpoint facade around an individual-Q network."""170171 def __init__(self, args, obs_space, action_space, num_agents, device):172 super().__init__()173 self.args = args174 self.device = device175 self.num_agents = int(num_agents)176 self.action_dim = int(action_space.n)177 self.hidden_size = int(args.qmix_agent_hidden_dim)178 trainer_name = getattr(args, "trainer_algorithm_name", None)179 self.use_hatt = str(trainer_name or args.algorithm_name).lower() == "hatt_qmix"180 if self.use_hatt:181 self.q_network = HATTAgentQNetwork(182 args, obs_space, self.num_agents, self.action_dim183 )184 else:185 flat_dim = flat_obs_dim(obs_space)186 self.q_network = RecurrentAgentQNetwork(187 flat_dim, self.action_dim, self.hidden_size,188 use_recurrent=bool(args.qmix_use_recurrent_agent),189 )190 self.to(device)191192 def _prepare_obs(self, obs):193 if self.use_hatt:194 result = {key: _tensor(value, self.device) for key, value in obs.items()}195 if result["ego"].ndim == 2:196 result = {key: value.unsqueeze(0) for key, value in result.items()}197 return result198 flat = _tensor(flatten_obs_dict(obs), self.device)199 if flat.ndim == 2:200 flat = flat.unsqueeze(0)201 return flat202203 def q_step(self, obs, hidden, available_actions):204 prepared = self._prepare_obs(obs)205 hidden = _tensor(hidden, self.device)206 if hidden.ndim == 2:207 hidden = hidden.unsqueeze(0)208 available = _tensor(available_actions, self.device)209 if available.ndim == 2:210 available = available.unsqueeze(0)211 return self.q_network.forward_step(prepared, hidden, available)212213 @torch.no_grad()214 def select_actions(self, obs, hidden, available_actions, epsilon=0.0):215 q_values, next_hidden = self.q_step(obs, hidden, available_actions)216 available = _tensor(available_actions, self.device)217 if available.ndim == 2:218 available = available.unsqueeze(0)219 masked_q = q_values.masked_fill(available <= 0, torch.finfo(q_values.dtype).min)220 greedy = masked_q.argmax(dim=-1)221 if epsilon > 0.0:222 random_scores = torch.rand_like(q_values).masked_fill(223 available <= 0, -1.0224 )225 random_actions = random_scores.argmax(dim=-1)226 choose_random = torch.rand_like(greedy.float()) < float(epsilon)227 actions = torch.where(choose_random, random_actions, greedy)228 else:229 actions = greedy230 return actions.unsqueeze(-1), next_hidden231232 @torch.no_grad()233 def act(self, obs, rnn_states, masks, available_actions=None, deterministic=True):234 del masks235 hidden = np.asarray(rnn_states, dtype=np.float32).reshape(236 1, self.num_agents, self.hidden_size237 )238 actions, next_hidden = self.select_actions(239 obs, hidden, available_actions, epsilon=0.0 if deterministic else 0.05240 )241 return actions.reshape(-1, 1), next_hidden.reshape(242 -1, 1, self.hidden_size243 )244245 def get_actor_checkpoint(self):246 return {247 "format": "qmix-v2",248 "algorithm_name": str(self.args.algorithm_name),249 "use_recurrent_agent": bool(self.args.qmix_use_recurrent_agent),250 "q_network": self.q_network.state_dict(),251 }252253 def load_actor_checkpoint(self, checkpoint):254 state = checkpoint.get("q_network", checkpoint)255 self.q_network.load_state_dict(state)