qmix.py
QMIX/VDN trainer
原文件 · 261 行 · 11524 字节
algorithms/qmix/qmix.py
SHA-256: 2937b005cf3701e106a8d34d3169a2ca6f1da7afb5e3f427e463b19d3ea7eea8
1"""Local recurrent QMIX/VDN trainer without the historical offpolicy package."""23import copy45import torch6import torch.nn.functional as F78from algorithms.qmix.algorithm.q_mixer import QMixer9from algorithms.vdn.algorithm.vdn_mixer import VDNMixer101112def _to_tensor_tree(value, device):13 if isinstance(value, dict):14 return {key: _to_tensor_tree(item, device) for key, item in value.items()}15 return torch.as_tensor(value, dtype=torch.float32, device=device)161718def _time(tree, index):19 if isinstance(tree, dict):20 return {key: value[:, index] for key, value in tree.items()}21 return tree[:, index]222324def _time_slice(tree, start, stop=None):25 if isinstance(tree, dict):26 return {key: _time_slice(value, start, stop) for key, value in tree.items()}27 return tree[:, start:stop]282930class QMix:31 def __init__(self, args, policy, state_dim, num_agents, device):32 self.args = args33 self.policy = policy34 self.target_policy = copy.deepcopy(policy).to(device)35 self.device = device36 self.num_agents = int(num_agents)37 trainer_name = getattr(args, "trainer_algorithm_name", None)38 self.vdn = str(trainer_name or args.algorithm_name).lower() == "vdn"39 if self.vdn:40 self.mixer = VDNMixer().to(device)41 else:42 self.mixer = QMixer(43 state_dim, num_agents,44 args.qmix_mixer_hidden_dim,45 args.qmix_hypernet_hidden_dim,46 ).to(device)47 self.target_mixer = copy.deepcopy(self.mixer).to(device)48 parameters = list(self.policy.q_network.parameters()) + list(self.mixer.parameters())49 if args.qmix_optimizer == "rmsprop":50 self.optimizer = torch.optim.RMSprop(51 parameters,52 lr=args.lr,53 alpha=args.qmix_rmsprop_alpha,54 eps=args.opti_eps,55 weight_decay=args.weight_decay,56 )57 else:58 self.optimizer = torch.optim.Adam(59 parameters,60 lr=args.lr,61 eps=args.opti_eps,62 weight_decay=args.weight_decay,63 )64 self.parameters = parameters65 self.update_count = 06667 def prep_training(self):68 self.policy.train()69 self.mixer.train()7071 def prep_rollout(self):72 self.policy.eval()73 self.mixer.eval()7475 def train_batch(self, batch):76 self.prep_training()77 batch = _to_tensor_tree(batch, self.device)78 observations = batch["obs"]79 states = batch["states"]80 available = batch["available_actions"]81 actions = batch["actions"].long()82 batch_size, transitions = actions.shape[:2]83 hidden = torch.zeros(84 batch_size, self.num_agents, self.policy.hidden_size, device=self.device85 )86 target_hidden = torch.zeros_like(hidden)87 burn_in = min(int(self.args.qmix_burn_in), transitions)88 recurrent = bool(self.args.qmix_use_recurrent_agent)89 if burn_in > 0 and recurrent:90 prefix_obs = _time_slice(observations, 0, burn_in)91 prefix_available = available[:, :burn_in]92 prefix_alive = batch["alive_masks"][:, :burn_in]93 if self.policy.use_hatt:94 chunk_size = int(self.args.qmix_hatt_sequence_chunk_size)95 with torch.no_grad():96 _, hidden = self.policy.q_network.forward_sequence(97 prefix_obs, hidden, prefix_available, chunk_size,98 alive_masks=prefix_alive,99 )100 _, target_hidden = self.target_policy.q_network.forward_sequence(101 prefix_obs, target_hidden, prefix_available, chunk_size,102 alive_masks=prefix_alive,103 )104 else:105 with torch.no_grad():106 for step in range(burn_in):107 hidden = hidden * prefix_alive[:, step]108 target_hidden = target_hidden * prefix_alive[:, step]109 _, hidden = self.policy.q_network.forward_step(110 _time(prefix_obs, step), hidden, prefix_available[:, step]111 )112 _, target_hidden = self.target_policy.q_network.forward_step(113 _time(prefix_obs, step), target_hidden,114 prefix_available[:, step]115 )116 hidden = hidden.detach()117 target_hidden = target_hidden.detach()118119 # Burn-in establishes recurrent context without retaining its graph. The120 # same leading samples are dropped for feedforward runs to keep their121 # number of supervised transitions comparable.122 observations = _time_slice(observations, burn_in)123 states = states[:, burn_in:]124 available = available[:, burn_in:]125 for key in (126 "actions", "rewards", "dones", "held_actions", "discounts", "valid"127 ):128 batch[key] = batch[key][:, burn_in:]129 for key in ("active_masks", "alive_masks"):130 batch[key] = batch[key][:, burn_in:]131 actions = batch["actions"].long()132 transitions = actions.shape[1]133 if self.policy.use_hatt:134 chunk_size = int(self.args.qmix_hatt_sequence_chunk_size)135 online_values, hidden = self.policy.q_network.forward_sequence(136 observations, hidden, available, chunk_size,137 alive_masks=batch["alive_masks"],138 )139 with torch.no_grad():140 target_values, target_hidden = (141 self.target_policy.q_network.forward_sequence(142 observations, target_hidden, available, chunk_size,143 alive_masks=batch["alive_masks"],144 )145 )146 else:147 online_steps, target_steps = [], []148 for step in range(transitions + 1):149 hidden = hidden * batch["alive_masks"][:, step]150 target_hidden = target_hidden * batch["alive_masks"][:, step]151 q_values, hidden = self.policy.q_network.forward_step(152 _time(observations, step), hidden, available[:, step]153 )154 with torch.no_grad():155 target_q, target_hidden = self.target_policy.q_network.forward_step(156 _time(observations, step), target_hidden, available[:, step]157 )158 online_steps.append(q_values)159 target_steps.append(target_q)160 online_values = torch.stack(online_steps, dim=1)161 target_values = torch.stack(target_steps, dim=1)162 chosen = online_values[:, :-1].gather(-1, actions.unsqueeze(-1)).squeeze(-1)163 next_online = online_values[:, 1:].masked_fill(164 available[:, 1:] <= 0, torch.finfo(online_values.dtype).min165 )166 next_target = target_values[:, 1:].masked_fill(167 available[:, 1:] <= 0, torch.finfo(target_values.dtype).min168 )169 if self.args.qmix_double_q:170 greedy_actions = next_online.argmax(dim=-1)171 else:172 greedy_actions = next_target.argmax(dim=-1)173 transition_mode = str(self.args.qmix_transition_mode)174 if transition_mode == "physical_legacy":175 next_actions = greedy_actions176 current_alive = torch.ones_like(batch["alive_masks"][:, :-1, :, 0])177 next_alive = torch.ones_like(batch["alive_masks"][:, 1:, :, 0])178 else:179 next_active = batch["active_masks"][:, 1:, :, 0] > 0.5180 next_alive = batch["alive_masks"][:, 1:, :, 0]181 fixed_actions = batch["held_actions"].long()182 next_actions = torch.where(183 next_active & (next_alive > 0.5), greedy_actions, fixed_actions184 )185 current_alive = batch["alive_masks"][:, :-1, :, 0]186 next_chosen = target_values[:, 1:].gather(187 -1, next_actions.unsqueeze(-1)188 ).squeeze(-1)189 chosen = chosen * current_alive190 next_chosen = next_chosen * next_alive191 joint_q = self.mixer(chosen, states[:, :-1])192 with torch.no_grad():193 target_joint_q = self.target_mixer(next_chosen, states[:, 1:])194 targets = batch["rewards"] + batch["discounts"] * (195 1.0 - batch["dones"]196 ) * target_joint_q197 valid = batch["valid"].clone()198 td_error = joint_q - targets199 max_abs_q = torch.maximum(joint_q.detach().abs().max(), targets.abs().max())200 threshold = float(self.args.qmix_max_abs_q)201 if not torch.isfinite(td_error).all() or not torch.isfinite(max_abs_q):202 raise FloatingPointError("QMIX produced a non-finite Q value or TD error")203 if threshold > 0.0 and float(max_abs_q) > threshold:204 raise FloatingPointError(205 f"QMIX Q magnitude {float(max_abs_q):.3g} exceeded "206 f"qmix_max_abs_q={threshold:.3g}"207 )208 if self.args.use_huber_loss:209 losses = F.smooth_l1_loss(joint_q, targets, reduction="none")210 else:211 losses = td_error.square()212 loss = (losses * valid).sum() / valid.sum().clamp_min(1.0)213 self.optimizer.zero_grad(set_to_none=True)214 loss.backward()215 grad_norm = torch.nn.utils.clip_grad_norm_(216 self.parameters, self.args.max_grad_norm217 )218 self.optimizer.step()219 self.update_count += 1220 info = {221 "loss": float(loss.detach()),222 "grad_norm": float(grad_norm),223 "Q_tot": float((joint_q * valid).sum().detach() / valid.sum().clamp_min(1.0)),224 "td_error": float((td_error.abs() * valid).sum().detach() / valid.sum().clamp_min(1.0)),225 "max_abs_q": float(max_abs_q),226 }227 if transition_mode != "physical_legacy":228 constrained = (batch["active_masks"][:, 1:, :, 0] <= 0.5) & (229 next_alive > 0.5230 )231 denom = constrained.float().sum().clamp_min(1.0)232 info["inactive_argmax_mismatch"] = float(233 (((greedy_actions != batch["held_actions"].long()) & constrained)234 .float().sum() / denom).detach()235 )236 with torch.no_grad():237 unconstrained_utilities = target_values[:, 1:].gather(238 -1, greedy_actions.unsqueeze(-1)239 ).squeeze(-1) * next_alive240 unconstrained_joint = self.target_mixer(241 unconstrained_utilities, states[:, 1:]242 )243 info["constrained_target_gap"] = float(244 ((unconstrained_joint - target_joint_q).abs() * valid).sum()245 / valid.sum().clamp_min(1.0)246 )247 return info248249 def hard_target_updates(self):250 self.target_policy.q_network.load_state_dict(self.policy.q_network.state_dict())251 self.target_mixer.load_state_dict(self.mixer.state_dict())252253 def state_dict(self):254 return {255 "policy": self.policy.get_actor_checkpoint(),256 "mixer": self.mixer.state_dict(),257 "target_policy": self.target_policy.q_network.state_dict(),258 "target_mixer": self.target_mixer.state_dict(),259 "optimizer": self.optimizer.state_dict(),260 "update_count": self.update_count,261 }