"""Shared recurrent individual-Q policies for QMIX, VDN and HATT-QMIX.""" import numpy as np import torch import torch.nn as nn from algorithms.algorithm.HATTV20Network import ( CandidateEdgeDecoder, HATTV20Backbone, ) from utils.obs_utils import flat_obs_dim, flatten_obs_dict def _tensor(value, device): return value if isinstance(value, torch.Tensor) else torch.as_tensor( value, dtype=torch.float32, device=device ) class RecurrentAgentQNetwork(nn.Module): """Parameter-shared DRQN applied independently to every agent.""" def __init__(self, obs_dim, action_dim, hidden_dim, use_recurrent=True): super().__init__() self.hidden_dim = int(hidden_dim) self.use_recurrent = bool(use_recurrent) self.encoder = nn.Sequential( nn.Linear(int(obs_dim), self.hidden_dim), nn.ReLU(), nn.LayerNorm(self.hidden_dim), ) self.rnn = ( nn.GRUCell(self.hidden_dim, self.hidden_dim) if self.use_recurrent else None ) self.q_head = nn.Linear(self.hidden_dim, int(action_dim)) def forward_step(self, obs, hidden, available_actions=None): prefix = obs.shape[:-1] encoded = self.encoder(obs.reshape(-1, obs.shape[-1])) if self.use_recurrent: next_hidden = self.rnn(encoded, hidden.reshape(-1, self.hidden_dim)) else: next_hidden = hidden.reshape(-1, self.hidden_dim) q_values = self.q_head( next_hidden if self.use_recurrent else encoded ).view(*prefix, -1) return q_values, next_hidden.view(*prefix, self.hidden_dim) class HATTAgentQNetwork(nn.Module): """HATT v2.0 relational encoder with a recurrent per-candidate Q decoder.""" def __init__(self, args, obs_space, num_agents, action_dim): super().__init__() spaces = obs_space.spaces hidden_dim = int(args.hatt_hidden_dim) self.hidden_dim = int(args.qmix_agent_hidden_dim) self.use_recurrent = bool(args.qmix_use_recurrent_agent) self.num_agents = int(num_agents) self.action_dim = int(action_dim) self.backbone = HATTV20Backbone( ego_dim=spaces["ego"].shape[-1], other_dim=spaces["others"].shape[-1], task_dim=spaces["tasks"].shape[-1], hidden_dim=hidden_dim, num_heads=int(args.hatt_num_heads), cap_dim=int(args.cap_dim), num_agents=self.num_agents, num_tasks=self.action_dim, dropout=float(args.hatt_dropout), shared_task_obs=bool(args.hatt_shared_task_obs), ) self.agent_projection = nn.Sequential( nn.Linear(hidden_dim, self.hidden_dim), nn.ReLU(), nn.LayerNorm(self.hidden_dim), ) self.rnn = ( nn.GRUCell(self.hidden_dim, self.hidden_dim) if self.use_recurrent else None ) self.decoder = CandidateEdgeDecoder( agent_dim=self.hidden_dim, task_dim=hidden_dim, hidden_dim=self.hidden_dim, ) def forward_step(self, obs, hidden, available_actions=None): agent_context, task_context, edges, _ = self.backbone( obs, available_actions, return_diagnostics=False ) projected = self.agent_projection(agent_context) prefix = projected.shape[:-1] if self.use_recurrent: next_hidden = self.rnn( projected.reshape(-1, self.hidden_dim), hidden.reshape(-1, self.hidden_dim), ).view(*prefix, self.hidden_dim) decoder_context = next_hidden else: next_hidden = hidden decoder_context = projected q_values = self.decoder(decoder_context, task_context, edges) return q_values, next_hidden @staticmethod def _flatten_time(tree, start, stop): """Flatten a contiguous [B, T, ...] slice to [B*T, ...].""" if isinstance(tree, dict): return { key: HATTAgentQNetwork._flatten_time(value, start, stop) for key, value in tree.items() } chunk = tree[:, start:stop] return chunk.reshape(-1, *chunk.shape[2:]) def forward_sequence( self, obs, hidden, available_actions, chunk_size, alive_masks=None ): """Unroll the GRU while batching the memoryless HATT work over time. HATT encoding and candidate decoding dominate the update cost and do not carry recurrent state. Processing a small time chunk in one call keeps the exact GRU recurrence while avoiding one set of tiny GPU kernels per transition. Chunking bounds the activation memory retained for BPTT. """ transitions = int(available_actions.shape[1]) chunk_size = max(1, int(chunk_size)) sequence_values = [] for start in range(0, transitions, chunk_size): stop = min(transitions, start + chunk_size) width = stop - start chunk_obs = self._flatten_time(obs, start, stop) chunk_available = self._flatten_time(available_actions, start, stop) agent_context, task_context, edges, _ = self.backbone( chunk_obs, chunk_available, return_diagnostics=False ) projected = self.agent_projection(agent_context).reshape( hidden.shape[0], width, self.num_agents, self.hidden_dim ) if self.use_recurrent: recurrent = [] for offset in range(width): if alive_masks is not None: hidden = hidden * alive_masks[:, start + offset] hidden = self.rnn( projected[:, offset].reshape(-1, self.hidden_dim), hidden.reshape(-1, self.hidden_dim), ).view_as(hidden) recurrent.append(hidden) decoder_context = torch.stack(recurrent, dim=1).reshape( -1, self.num_agents, self.hidden_dim ) else: decoder_context = projected.reshape( -1, self.num_agents, self.hidden_dim ) q_values = self.decoder(decoder_context, task_context, edges) sequence_values.append( q_values.reshape(hidden.shape[0], width, self.num_agents, -1) ) return torch.cat(sequence_values, dim=1), hidden class QMixPolicy(nn.Module): """Evaluation/checkpoint facade around an individual-Q network.""" def __init__(self, args, obs_space, action_space, num_agents, device): super().__init__() self.args = args self.device = device self.num_agents = int(num_agents) self.action_dim = int(action_space.n) self.hidden_size = int(args.qmix_agent_hidden_dim) trainer_name = getattr(args, "trainer_algorithm_name", None) self.use_hatt = str(trainer_name or args.algorithm_name).lower() == "hatt_qmix" if self.use_hatt: self.q_network = HATTAgentQNetwork( args, obs_space, self.num_agents, self.action_dim ) else: flat_dim = flat_obs_dim(obs_space) self.q_network = RecurrentAgentQNetwork( flat_dim, self.action_dim, self.hidden_size, use_recurrent=bool(args.qmix_use_recurrent_agent), ) self.to(device) def _prepare_obs(self, obs): if self.use_hatt: result = {key: _tensor(value, self.device) for key, value in obs.items()} if result["ego"].ndim == 2: result = {key: value.unsqueeze(0) for key, value in result.items()} return result flat = _tensor(flatten_obs_dict(obs), self.device) if flat.ndim == 2: flat = flat.unsqueeze(0) return flat def q_step(self, obs, hidden, available_actions): prepared = self._prepare_obs(obs) hidden = _tensor(hidden, self.device) if hidden.ndim == 2: hidden = hidden.unsqueeze(0) available = _tensor(available_actions, self.device) if available.ndim == 2: available = available.unsqueeze(0) return self.q_network.forward_step(prepared, hidden, available) @torch.no_grad() def select_actions(self, obs, hidden, available_actions, epsilon=0.0): q_values, next_hidden = self.q_step(obs, hidden, available_actions) available = _tensor(available_actions, self.device) if available.ndim == 2: available = available.unsqueeze(0) masked_q = q_values.masked_fill(available <= 0, torch.finfo(q_values.dtype).min) greedy = masked_q.argmax(dim=-1) if epsilon > 0.0: random_scores = torch.rand_like(q_values).masked_fill( available <= 0, -1.0 ) random_actions = random_scores.argmax(dim=-1) choose_random = torch.rand_like(greedy.float()) < float(epsilon) actions = torch.where(choose_random, random_actions, greedy) else: actions = greedy return actions.unsqueeze(-1), next_hidden @torch.no_grad() def act(self, obs, rnn_states, masks, available_actions=None, deterministic=True): del masks hidden = np.asarray(rnn_states, dtype=np.float32).reshape( 1, self.num_agents, self.hidden_size ) actions, next_hidden = self.select_actions( obs, hidden, available_actions, epsilon=0.0 if deterministic else 0.05 ) return actions.reshape(-1, 1), next_hidden.reshape( -1, 1, self.hidden_size ) def get_actor_checkpoint(self): return { "format": "qmix-v2", "algorithm_name": str(self.args.algorithm_name), "use_recurrent_agent": bool(self.args.qmix_use_recurrent_agent), "q_network": self.q_network.state_dict(), } def load_actor_checkpoint(self, checkpoint): state = checkpoint.get("q_network", checkpoint) self.q_network.load_state_dict(state)