test_hatt_network.py
排列、梯度、任务上下文、schema 初始化检查
原文件 · 1368 行 · 51589 字节
tests/test_hatt_network.py
SHA-256: 31ae3158a7b6fe8fcdfbcfa72b311b2213bddc51db364c8b682c2f7492104e9b
1import unittest2from types import SimpleNamespace34import torch56from algorithms.algorithm.HATTNetwork import HATTActor, HATTV2Actor7from algorithms.algorithm.HATTV20Network import (8 HATTV20Actor,9 initialize_full_actor_from_compact_reference,10)11from algorithms.algorithm.HATTV30Network import (12 HATTV30AActor,13 HATTV30BActor,14 HATTV30CActor,15)16from algorithms.algorithm.HATTV31Network import HATTV31AActor, HATTV31BActor17from algorithms.algorithm.HATTV32RNetwork import (18 HATTV32RFullActor,19 HATTV32RNoMemberActor,20 HATTV32RNoTransitionActor,21)22from algorithms.algorithm.HATTV33AANetwork import (23 HATTV33AAActor,24 TypedRelationEdge,25)26from algorithms.algorithm.HATTV34Network import HATTV34Actor272829class HATTNetworkTest(unittest.TestCase):30 def setUp(self):31 torch.manual_seed(11)32 self.num_agents = 333 self.num_tasks = 534 self.cap_dim = 235 self.args = SimpleNamespace(36 hidden_size=16,37 use_policy_active_masks=True,38 use_naive_recurrent_policy=False,39 use_recurrent_policy=False,40 recurrent_N=1,41 use_orthogonal=True,42 cap_dim=self.cap_dim,43 gat_out_dim=16,44 hatt_hidden_dim=16,45 hatt_num_heads=4,46 hatt_dropout=0.0,47 hatt_capability_scale=5.0,48 )49 self.ego_dim = self.cap_dim + 750 self.other_dim = self.cap_dim + 751 self.task_dim = self.cap_dim + 452 self.actor = HATTActor(53 self.args,54 self.ego_dim,55 self.other_dim,56 self.task_dim,57 self.num_agents,58 self.num_tasks,59 )60 self.obs = self._make_obs()61 self.available = torch.ones(self.num_agents, self.num_tasks)62 self.available[0, 2] = 0.063 self.rnn_states = torch.zeros(self.num_agents, 1, self.args.hidden_size)64 self.masks = torch.ones(self.num_agents, 1)6566 def _make_obs(self):67 a, t, c = self.num_agents, self.num_tasks, self.cap_dim68 ego = torch.rand(1, a, self.ego_dim)69 others = torch.rand(1, a, a, self.other_dim)70 tasks = torch.rand(1, a, t, self.task_dim)7172 coalition_index = 2 + c73 other_coalition_index = c + 474 living_index = c + 675 ids = torch.tensor([0, 2, 4])76 ego[..., coalition_index] = ids.float() / (t - 1)77 others[..., other_coalition_index] = ids.view(1, 1, a).float() / (t - 1)78 others[..., living_index] = 1.079 for observer in range(a):80 others[:, observer, observer] = 0.081 return {"ego": ego, "others": others, "tasks": tasks}8283 def _logits(self, obs, available):84 dist, _ = self.actor._distribution(85 obs, self.rnn_states, self.masks, available86 )87 return dist.logits.reshape(1, self.num_agents, self.num_tasks)8889 def test_task_permutation_equivariance_and_mask(self):90 base = self._logits(self.obs, self.available).detach()91 self.assertLess(float(base.exp()[0, 0, 2]), 1e-7)9293 permutation = torch.tensor([2, 4, 0, 3, 1])94 inverse = torch.empty_like(permutation)95 inverse[permutation] = torch.arange(self.num_tasks)96 permuted_obs = {k: v.clone() for k, v in self.obs.items()}97 permuted_obs["tasks"] = permuted_obs["tasks"][:, :, permutation]98 permuted_available = self.available[:, permutation]99100 coalition_index = 2 + self.cap_dim101 other_coalition_index = self.cap_dim + 4102 old = (103 permuted_obs["ego"][..., coalition_index] * (self.num_tasks - 1)104 ).round().long()105 permuted_obs["ego"][..., coalition_index] = (106 inverse[old].float() / (self.num_tasks - 1)107 )108 old = (109 permuted_obs["others"][..., other_coalition_index]110 * (self.num_tasks - 1)111 ).round().long()112 permuted_obs["others"][..., other_coalition_index] = (113 inverse[old].float() / (self.num_tasks - 1)114 )115116 actual = self._logits(permuted_obs, permuted_available).detach()117 torch.testing.assert_close(actual, base[:, :, permutation], atol=2e-5, rtol=2e-5)118119 def test_agent_permutation_equivariance(self):120 base = self._logits(self.obs, self.available).detach()121 permutation = torch.tensor([2, 0, 1])122 permuted_obs = {k: v.clone() for k, v in self.obs.items()}123 permuted_obs["ego"] = permuted_obs["ego"][:, permutation]124 permuted_obs["tasks"] = permuted_obs["tasks"][:, permutation]125 permuted_obs["others"] = permuted_obs["others"][:, permutation][126 :, :, permutation127 ]128 actual = self._logits(129 permuted_obs, self.available[permutation]130 ).detach()131 torch.testing.assert_close(actual, base[:, permutation], atol=2e-5, rtol=2e-5)132133 def test_all_major_branches_receive_gradients(self):134 dist, _ = self.actor._distribution(135 self.obs, self.rnn_states, self.masks, self.available136 )137 actions = torch.zeros(self.num_agents, 1, dtype=torch.long)138 (-dist.log_probs(actions).mean()).backward()139 gradients = [140 self.actor.backbone.self_encoder.net[0].weight.grad,141 self.actor.backbone.other_encoder.net[0].weight.grad,142 self.actor.backbone.task_encoder.net[0].weight.grad,143 self.actor.backbone.aa_attention.query.weight.grad,144 self.actor.backbone.ta_attention.query.weight.grad,145 self.actor.backbone.fusion[0].weight.grad,146 self.actor.decoder.agent_query.weight.grad,147 ]148 self.assertTrue(149 all(150 grad is not None151 and torch.isfinite(grad).all()152 and grad.abs().sum() > 0153 for grad in gradients154 )155 )156157158class HATTV2NetworkTest(HATTNetworkTest):159 def setUp(self):160 super().setUp()161 torch.manual_seed(11)162 self.actor = HATTV2Actor(163 self.args,164 self.ego_dim,165 self.other_dim,166 self.task_dim,167 self.num_agents,168 self.num_tasks,169 )170171 def test_all_major_branches_receive_gradients(self):172 dist, _ = self.actor._distribution(173 self.obs, self.rnn_states, self.masks, self.available174 )175 actions = torch.zeros(self.num_agents, 1, dtype=torch.long)176 (-dist.log_probs(actions).mean()).backward()177 gradients = [178 self.actor.backbone.self_encoder.net[0].weight.grad,179 self.actor.backbone.other_encoder.net[0].weight.grad,180 self.actor.backbone.task_encoder.net[0].weight.grad,181 self.actor.backbone.member_encoder.net[0].weight.grad,182 self.actor.backbone.aa_attention.query.weight.grad,183 self.actor.backbone.agent_fusion[0].weight.grad,184 self.actor.backbone.task_fusion[0].weight.grad,185 self.actor.backbone.decision_fusion[0].weight.grad,186 self.actor.decoder.network[0].weight.grad,187 self.actor.decoder.network[-1].weight.grad,188 ]189 self.assertTrue(190 all(191 grad is not None192 and torch.isfinite(grad).all()193 and grad.abs().sum() > 0194 for grad in gradients195 )196 )197198 def test_assignment_moves_coalition_context_between_task_slots(self):199 obs = {k: v.clone() for k, v in self.obs.items()}200 other_coalition_index = self.cap_dim + 4201 observer = 0202 other_agent = 1203 old_task = 2204 new_task = 3205 obs["others"][0, observer, other_agent, other_coalition_index] = (206 old_task / (self.num_tasks - 1)207 )208 _, before, _, _ = self.actor.backbone(209 obs, self.available.unsqueeze(0)210 )211 obs["others"][0, observer, other_agent, other_coalition_index] = (212 new_task / (self.num_tasks - 1)213 )214 _, after, _, _ = self.actor.backbone(215 obs, self.available.unsqueeze(0)216 )217218 unchanged_tasks = [0, 1, 4]219 torch.testing.assert_close(220 after[0, observer, unchanged_tasks],221 before[0, observer, unchanged_tasks],222 )223 self.assertGreater(224 float(225 (after[0, observer, old_task] - before[0, observer, old_task])226 .abs()227 .sum()228 .detach()229 ),230 1e-5,231 )232 self.assertGreater(233 float(234 (after[0, observer, new_task] - before[0, observer, new_task])235 .abs()236 .sum()237 .detach()238 ),239 1e-5,240 )241242 def test_task_pooling_ignores_unavailable_candidates(self):243 before_context, _, _, _ = self.actor.backbone(244 self.obs, self.available.unsqueeze(0)245 )246 before_logits = self._logits(self.obs, self.available).detach()247248 changed_obs = {key: value.clone() for key, value in self.obs.items()}249 changed_obs["tasks"][0, 0, 2] += 100.0250 after_context, _, _, _ = self.actor.backbone(251 changed_obs, self.available.unsqueeze(0)252 )253 after_logits = self._logits(changed_obs, self.available).detach()254255 torch.testing.assert_close(256 after_context[0, 0], before_context[0, 0]257 )258 available = self.available[0].bool()259 torch.testing.assert_close(260 after_logits[0, 0, available], before_logits[0, 0, available]261 )262263 def test_non_spatial_tasks_do_not_use_relative_position(self):264 obs = {key: value.clone() for key, value in self.obs.items()}265 type_index = self.cap_dim + 3266 obs["tasks"][..., 0, type_index] = 0.0267 obs["tasks"][..., 1, type_index] = 2.0 / 3.0268 obs["tasks"][..., 2, type_index] = 1.0 / 3.0269 obs["tasks"][..., 3, type_index] = 1.0270271 _, _, edge_features, _ = self.actor.backbone(272 obs, self.available.unsqueeze(0)273 )274 torch.testing.assert_close(275 edge_features[..., 0:2, 0:2],276 torch.zeros_like(edge_features[..., 0:2, 0:2]),277 )278 torch.testing.assert_close(279 edge_features[..., 0:2, 3],280 torch.zeros_like(edge_features[..., 0:2, 3]),281 )282 torch.testing.assert_close(283 edge_features[..., 2:4, 3],284 torch.ones_like(edge_features[..., 2:4, 3]),285 )286287 def test_initial_policy_has_nontrivial_finite_logits(self):288 available = torch.ones_like(self.available)289 dist, _ = self.actor._distribution(290 self.obs, self.rnn_states, self.masks, available291 )292 self.assertTrue(torch.isfinite(dist.logits).all())293 self.assertGreater(float(dist.logits.std().detach()), 1e-4)294295296class HATTV20NetworkTest(HATTNetworkTest):297 def setUp(self):298 super().setUp()299 torch.manual_seed(11)300 self.actor = HATTV20Actor(301 self.args,302 self.ego_dim,303 self.other_dim,304 self.task_dim,305 self.num_agents,306 self.num_tasks,307 )308309 def test_full_schema_scale_restore_and_compact_projection(self):310 full_dim = 2 + self.cap_dim + 1 + 1 + 3 + 1311 full_tasks = torch.zeros(1, self.num_agents, self.num_tasks, full_dim)312 full_tasks[..., 0] = 0.1313 full_tasks[..., 1] = 0.2314 full_tasks[..., 2] = 0.15315 full_tasks[..., 2 + self.cap_dim] = 0.6316 full_tasks[..., 3 + self.cap_dim] = 0.7317 full_tasks[..., 7 + self.cap_dim] = 1.0318319 self.args.hatt_v20_task_input_mode = "restore_req_scale"320 restored_actor = HATTV20Actor(321 self.args,322 self.ego_dim,323 self.other_dim,324 full_dim,325 self.num_agents,326 self.num_tasks,327 )328 restored = restored_actor.backbone._prepare_task_encoder_input(full_tasks)329 self.assertEqual(restored.shape[-1], full_dim)330 torch.testing.assert_close(restored[..., 2], torch.full_like(restored[..., 2], 0.75))331332 self.args.hatt_v20_task_input_mode = "compact_projection"333 projected_actor = HATTV20Actor(334 self.args,335 self.ego_dim,336 self.other_dim,337 full_dim,338 self.num_agents,339 self.num_tasks,340 )341 projected = projected_actor.backbone._prepare_task_encoder_input(full_tasks)342 expected = torch.tensor([0.1, 0.2, 0.6, 0.7, 0.75, 1.0])343 torch.testing.assert_close(projected[0, 0, 0], expected)344 self.assertEqual(projected_actor.backbone.task_encoder.net[0].in_features, 6)345346 def test_compact_projection_supports_exact_actor_pairing(self):347 compact_args = SimpleNamespace(**vars(self.args))348 compact_args.hatt_v20_task_input_mode = "native"349 torch.manual_seed(91)350 compact_actor = HATTV20Actor(351 compact_args,352 self.ego_dim,353 self.other_dim,354 6,355 self.num_agents,356 self.num_tasks,357 )358359 projected_args = SimpleNamespace(**vars(self.args))360 projected_args.hatt_v20_task_input_mode = "compact_projection"361 torch.manual_seed(91)362 projected_actor = HATTV20Actor(363 projected_args,364 self.ego_dim,365 self.other_dim,366 2 + self.cap_dim + 1 + 1 + 3 + 1,367 self.num_agents,368 self.num_tasks,369 )370371 compact_state = compact_actor.state_dict()372 projected_state = projected_actor.state_dict()373 self.assertEqual(compact_state.keys(), projected_state.keys())374 for key in compact_state:375 torch.testing.assert_close(compact_state[key], projected_state[key])376377 def test_full_actor_function_preserving_initialization(self):378 full_dim = 2 + self.cap_dim + 1 + 1 + 3 + 1379 compact_args = SimpleNamespace(**vars(self.args))380 compact_args.hatt_v20_task_input_mode = "native"381 torch.manual_seed(137)382 compact_actor = HATTV20Actor(383 compact_args,384 self.ego_dim,385 self.other_dim,386 6,387 self.num_agents,388 self.num_tasks,389 )390391 full_args = SimpleNamespace(**vars(self.args))392 full_args.hatt_v20_task_input_mode = "restore_req_scale"393 torch.manual_seed(999)394 full_actor = HATTV20Actor(395 full_args,396 self.ego_dim,397 self.other_dim,398 full_dim,399 self.num_agents,400 self.num_tasks,401 )402 initialize_full_actor_from_compact_reference(403 full_actor, compact_actor, self.cap_dim404 )405406 compact_tasks = torch.randn(407 2, self.num_agents, self.num_tasks, 6408 )409 full_tasks = torch.zeros(410 2, self.num_agents, self.num_tasks, full_dim411 )412 full_tasks[..., 0:2] = compact_tasks[..., 0:2]413 full_tasks[..., 2] = compact_tasks[..., 4] / 5.0414 full_tasks[..., 2 + self.cap_dim] = compact_tasks[..., 2]415 full_tasks[..., 3 + self.cap_dim] = compact_tasks[..., 3]416 full_tasks[..., 7 + self.cap_dim] = compact_tasks[..., 5]417418 compact_encoded = compact_actor.backbone.task_encoder(compact_tasks)419 full_prepared = full_actor.backbone._prepare_task_encoder_input(full_tasks)420 full_encoded = full_actor.backbone.task_encoder(full_prepared)421 # GEMM reduction order differs for 6-D and 12-D inputs even when the422 # added products are all zero. The construction is function-preserving423 # up to normal floating-point accumulation error.424 torch.testing.assert_close(compact_encoded, full_encoded, rtol=1e-6, atol=1e-6)425426 full_first = full_actor.backbone.task_encoder.net[0].weight427 semantic_columns = {0, 1, 2, 2 + self.cap_dim, 3 + self.cap_dim, 7 + self.cap_dim}428 zero_columns = [idx for idx in range(full_dim) if idx not in semantic_columns]429 torch.testing.assert_close(430 full_first[:, zero_columns], torch.zeros_like(full_first[:, zero_columns])431 )432433 projected_args = SimpleNamespace(**vars(self.args))434 projected_args.hatt_v20_task_input_mode = "compact_projection"435 torch.manual_seed(137)436 projected_actor = HATTV20Actor(437 projected_args,438 self.ego_dim,439 self.other_dim,440 full_dim,441 self.num_agents,442 self.num_tasks,443 )444 full_obs = {445 "ego": torch.rand(2, self.num_agents, self.ego_dim),446 "others": torch.rand(447 2, self.num_agents, self.num_agents, self.other_dim448 ),449 "tasks": full_tasks,450 }451 available = torch.ones(2, self.num_agents, self.num_tasks)452 rnn_states = torch.zeros(453 2 * self.num_agents, 1, self.args.hidden_size454 )455 masks = torch.ones(2 * self.num_agents, 1)456 projected_dist, _ = projected_actor._distribution(457 full_obs, rnn_states, masks, available458 )459 full_dist, _ = full_actor._distribution(460 full_obs, rnn_states, masks, available461 )462 torch.testing.assert_close(463 projected_dist.logits, full_dist.logits, rtol=1e-6, atol=1e-6464 )465466 def test_no_aa_attention_ablation_keeps_shape_and_zeroes_attention(self):467 self.args.hatt_use_aa_attention = False468 actor = HATTV20Actor(469 self.args,470 self.ego_dim,471 self.other_dim,472 self.task_dim,473 self.num_agents,474 self.num_tasks,475 )476 agent, tasks, edges, diagnostics = actor.backbone(477 self.obs, self.available, return_diagnostics=True478 )479 self.assertEqual(agent.shape, (1, self.num_agents, self.args.hatt_hidden_dim))480 self.assertEqual(tasks.shape, (1, self.num_agents, self.num_tasks, self.args.hatt_hidden_dim))481 self.assertEqual(edges.shape[:3], (1, self.num_agents, self.num_tasks))482 self.assertEqual(float(diagnostics["aa_attention"].abs().sum()), 0.0)483484 def test_no_coalition_context_ablation_zeroes_coalition_diagnostics(self):485 self.args.hatt_use_coalition_context = False486 actor = HATTV20Actor(487 self.args,488 self.ego_dim,489 self.other_dim,490 self.task_dim,491 self.num_agents,492 self.num_tasks,493 )494 _, _, _, diagnostics = actor.backbone(495 self.obs, self.available, return_diagnostics=True496 )497 self.assertEqual(float(diagnostics["coalition_context_norm"].sum()), 0.0)498 self.assertEqual(float(diagnostics["occupied_task_ratio"].sum()), 0.0)499500 def test_all_major_branches_receive_gradients(self):501 dist, _ = self.actor._distribution(502 self.obs, self.rnn_states, self.masks, self.available503 )504 actions = torch.zeros(self.num_agents, 1, dtype=torch.long)505 (-dist.log_probs(actions).mean()).backward()506 gradients = [507 self.actor.backbone.self_encoder.net[0].weight.grad,508 self.actor.backbone.other_encoder.net[0].weight.grad,509 self.actor.backbone.task_encoder.net[0].weight.grad,510 self.actor.backbone.aa_attention.query.weight.grad,511 self.actor.backbone.agent_fusion[0].weight.grad,512 self.actor.backbone.task_fusion[0].weight.grad,513 self.actor.decoder.network[0].weight.grad,514 self.actor.decoder.network[-1].weight.grad,515 ]516 self.assertTrue(517 all(518 grad is not None519 and torch.isfinite(grad).all()520 and grad.abs().sum() > 0521 for grad in gradients522 )523 )524525526class HATTV32RNetworkTest(HATTNetworkTest):527 def setUp(self):528 super().setUp()529 # Explicit task types: Deliver, Explore, Idle, Resupply, Deliver.530 self.obs["tasks"][..., -1] = torch.tensor(531 [1.0 / 3.0, 0.0, 2.0 / 3.0, 1.0, 1.0 / 3.0]532 )533 # Give observer 0 two real teammates in its current Deliver coalition.534 coalition_index = self.cap_dim + 4535 self.obs["others"][0, 0, 1, coalition_index] = 0.0536 self.obs["others"][0, 0, 2, coalition_index] = 0.0537 torch.manual_seed(11)538 self.actor = HATTV32RFullActor(539 self.args,540 self.ego_dim,541 self.other_dim,542 self.task_dim,543 self.num_agents,544 self.num_tasks,545 )546547 def test_agent_permutation_equivariance(self):548 # scatter_add may sum coalition members in a different floating-point549 # order after relabeling; the policy remains equivariant numerically.550 base = self._logits(self.obs, self.available).detach()551 permutation = torch.tensor([2, 0, 1])552 permuted_obs = {key: value.clone() for key, value in self.obs.items()}553 permuted_obs["ego"] = permuted_obs["ego"][:, permutation]554 permuted_obs["tasks"] = permuted_obs["tasks"][:, permutation]555 permuted_obs["others"] = permuted_obs["others"][:, permutation][556 :, :, permutation557 ]558 actual = self._logits(559 permuted_obs, self.available[permutation]560 ).detach()561 torch.testing.assert_close(562 actual, base[:, permutation], atol=2e-3, rtol=2e-3563 )564565 def test_all_major_branches_receive_gradients(self):566 self.test_new_relation_branches_receive_gradients()567568 def test_transition_semantics_ignore_control_target_and_stay(self):569 _, _, _, diagnostics = self.actor.backbone(570 self.obs, self.available.unsqueeze(0), return_diagnostics=True571 )572 target_norm = diagnostics["candidate_target_delta_norm"]573 source_norm = diagnostics["candidate_source_delta_norm"]574 # Explore/Idle/Resupply never receive a fictitious target coalition.575 torch.testing.assert_close(576 target_norm[..., 1:4], torch.zeros_like(target_norm[..., 1:4])577 )578 # Observer 0 currently occupies task 0, so staying has no leave effect.579 self.assertEqual(float(source_norm[0, 0, 0].detach()), 0.0)580 self.assertGreater(float(source_norm[0, 0, 4].detach()), 1e-6)581582 def test_both_ablation_paths_zero_only_their_diagnostics(self):583 no_transition = HATTV32RNoTransitionActor(584 self.args,585 self.ego_dim,586 self.other_dim,587 self.task_dim,588 self.num_agents,589 self.num_tasks,590 )591 no_member = HATTV32RNoMemberActor(592 self.args,593 self.ego_dim,594 self.other_dim,595 self.task_dim,596 self.num_agents,597 self.num_tasks,598 )599 _, _, _, transition_diag = no_transition.backbone(600 self.obs, self.available.unsqueeze(0), return_diagnostics=True601 )602 _, _, _, member_diag = no_member.backbone(603 self.obs, self.available.unsqueeze(0), return_diagnostics=True604 )605 self.assertEqual(606 float(607 transition_diag["task_conditioned_member_context_norm"]608 .sum()609 .detach()610 )611 > 0.0,612 True,613 )614 self.assertEqual(615 float(transition_diag["transition_context_norm"].sum().detach()), 0.0616 )617 self.assertEqual(618 float(619 member_diag["task_conditioned_member_context_norm"]620 .sum()621 .detach()622 ),623 0.0,624 )625626 def test_new_relation_branches_receive_gradients(self):627 dist, _ = self.actor._distribution(628 self.obs, self.rnn_states, self.masks, self.available629 )630 actions = torch.zeros(self.num_agents, 1, dtype=torch.long)631 (-dist.log_probs(actions).mean()).backward()632 gradients = [633 self.actor.backbone.transition_projection[0].weight.grad,634 self.actor.backbone.member_attention.self_query.weight.grad,635 self.actor.backbone.member_projection[0].weight.grad,636 ]637 self.assertTrue(638 all(639 grad is not None640 and torch.isfinite(grad).all()641 and grad.abs().sum() > 0642 for grad in gradients643 )644 )645646647class HATTV33AANetworkTest(HATTNetworkTest):648 def setUp(self):649 super().setUp()650 self.obs["tasks"][..., -1] = torch.tensor(651 [1.0 / 3.0, 0.0, 2.0 / 3.0, 1.0, 1.0 / 3.0]652 )653 coalition_index = self.cap_dim + 4654 # Observer 0 sees two members in Deliver candidate 0.655 self.obs["others"][0, 0, 1, coalition_index] = 0.0656 self.obs["others"][0, 0, 2, coalition_index] = 0.0657 self.aa_args = SimpleNamespace(658 **vars(self.args),659 hatt_candidate_aa_mode="learned",660 hatt_shared_task_obs=False,661 )662 torch.manual_seed(11)663 self.actor = HATTV33AAActor(664 self.aa_args,665 self.ego_dim,666 self.other_dim,667 self.task_dim,668 self.num_agents,669 self.num_tasks,670 )671672 def _actor_with_mode(self, mode):673 args = SimpleNamespace(**vars(self.aa_args))674 args.hatt_candidate_aa_mode = mode675 actor = HATTV33AAActor(676 args,677 self.ego_dim,678 self.other_dim,679 self.task_dim,680 self.num_agents,681 self.num_tasks,682 )683 actor.load_state_dict(self.actor.state_dict())684 return actor685686 def test_aa_and_at_use_the_same_typed_edge_operator(self):687 self.assertIsInstance(self.actor.backbone.aa_relation, TypedRelationEdge)688 self.assertIsInstance(self.actor.backbone.at_relation, TypedRelationEdge)689 self.assertIsNot(690 self.actor.backbone.aa_relation.query.weight,691 self.actor.backbone.at_relation.query.weight,692 )693694 def test_membership_routes_aa_only_to_deliver_candidates(self):695 _, _, _, diagnostics = self.actor.backbone(696 self.obs, self.available.unsqueeze(0), return_diagnostics=True697 )698 counts = diagnostics["candidate_aa_member_count"]699 self.assertEqual(float(counts[0, 0, 0]), 2.0)700 self.assertEqual(float(counts[0, 0, 1:4].sum()), 0.0)701 self.assertGreater(702 float(diagnostics["candidate_aa_context_norm"][0, 0, 0].detach()),703 0.0,704 )705706 def test_uniform_and_zero_are_checkpoint_compatible_interventions(self):707 uniform = self._actor_with_mode("uniform")708 zero = self._actor_with_mode("zero")709 _, _, _, uniform_diag = uniform.backbone(710 self.obs, self.available.unsqueeze(0), return_diagnostics=True711 )712 valid = self.obs["others"].abs().sum(dim=-1) > 0713 uniform_gates = uniform_diag["aa_relation_gate"]714 torch.testing.assert_close(715 uniform_gates.masked_select(valid.unsqueeze(-1)),716 torch.full_like(uniform_gates.masked_select(valid.unsqueeze(-1)), 0.5),717 )718 learned_logits = self._logits(self.obs, self.available).detach()719 zero_dist, _ = zero._distribution(720 self.obs, self.rnn_states, self.masks, self.available721 )722 zero_logits = zero_dist.logits.reshape_as(learned_logits).detach()723 self.assertGreater(float((learned_logits - zero_logits).abs().max()), 1e-6)724725 def test_all_major_branches_receive_gradients(self):726 dist, _ = self.actor._distribution(727 self.obs, self.rnn_states, self.masks, self.available728 )729 actions = torch.zeros(self.num_agents, 1, dtype=torch.long)730 (-dist.log_probs(actions).mean()).backward()731 gradients = [732 self.actor.backbone.aa_relation.query.weight.grad,733 self.actor.backbone.aa_relation.key.weight.grad,734 self.actor.backbone.aa_relation.value.weight.grad,735 self.actor.backbone.at_relation.query.weight.grad,736 self.actor.backbone.at_relation.key.weight.grad,737 self.actor.backbone.at_relation.value.weight.grad,738 self.actor.backbone.agent_fusion[0].weight.grad,739 self.actor.backbone.candidate_fusion[0].weight.grad,740 self.actor.decoder.network[0].weight.grad,741 ]742 self.assertTrue(743 all(744 grad is not None745 and torch.isfinite(grad).all()746 and grad.abs().sum() > 0747 for grad in gradients748 )749 )750751752class HATTV34NetworkTest(HATTNetworkTest):753 def setUp(self):754 super().setUp()755 self.obs["tasks"][..., -1] = torch.tensor(756 [1.0 / 3.0, 0.0, 2.0 / 3.0, 1.0, 1.0 / 3.0]757 )758 coalition_index = self.cap_dim + 4759 self.obs["others"][0, 0, 1, coalition_index] = 0.0760 self.obs["others"][0, 0, 2, coalition_index] = 0.0761 self.aa_args = SimpleNamespace(762 **vars(self.args),763 hatt_candidate_aa_mode="learned",764 hatt_shared_task_obs=False,765 hatt_aa_task_context=True,766 coalition_task_feature_mode="compact",767 )768 torch.manual_seed(11)769 self.actor = HATTV34Actor(770 self.aa_args,771 self.ego_dim,772 self.other_dim,773 self.task_dim,774 self.num_agents,775 self.num_tasks,776 )777778 def test_task_context_toggles_aa_edge_dim(self):779 self.assertEqual(780 self.actor.backbone.aa_relation.edge_bias[0].in_features, 7781 )782 off_args = SimpleNamespace(**vars(self.aa_args))783 off_args.hatt_aa_task_context = False784 off = HATTV34Actor(785 off_args,786 self.ego_dim,787 self.other_dim,788 self.task_dim,789 self.num_agents,790 self.num_tasks,791 )792 self.assertEqual(off.backbone.aa_relation.edge_bias[0].in_features, 3)793794 def test_task_context_off_reproduces_v33aa(self):795 off_args = SimpleNamespace(**vars(self.aa_args))796 off_args.hatt_aa_task_context = False797 torch.manual_seed(11)798 off = HATTV34Actor(799 off_args,800 self.ego_dim,801 self.other_dim,802 self.task_dim,803 self.num_agents,804 self.num_tasks,805 )806 v33_args = SimpleNamespace(**vars(self.aa_args))807 torch.manual_seed(11)808 v33 = HATTV33AAActor(809 v33_args,810 self.ego_dim,811 self.other_dim,812 self.task_dim,813 self.num_agents,814 self.num_tasks,815 )816 _, ctx33, _, _ = v33.backbone(817 self.obs, self.available.unsqueeze(0), return_diagnostics=True818 )819 _, ctx34, _, _ = off.backbone(820 self.obs, self.available.unsqueeze(0), return_diagnostics=True821 )822 torch.testing.assert_close(ctx34, ctx33)823824 def test_all_major_branches_receive_gradients(self):825 dist, _ = self.actor._distribution(826 self.obs, self.rnn_states, self.masks, self.available827 )828 actions = torch.zeros(self.num_agents, 1, dtype=torch.long)829 (-dist.log_probs(actions).mean()).backward()830 gradients = [831 self.actor.backbone.aa_relation.query.weight.grad,832 self.actor.backbone.aa_relation.key.weight.grad,833 self.actor.backbone.aa_relation.value.weight.grad,834 self.actor.backbone.aa_relation.edge_bias[0].weight.grad,835 self.actor.backbone.at_relation.query.weight.grad,836 self.actor.backbone.agent_fusion[0].weight.grad,837 self.actor.backbone.candidate_fusion[0].weight.grad,838 self.actor.decoder.network[0].weight.grad,839 ]840 self.assertTrue(841 all(842 grad is not None843 and torch.isfinite(grad).all()844 and grad.abs().sum() > 0845 for grad in gradients846 )847 )848849850class HATTV20ReconfigurationGateTest(unittest.TestCase):851 def _make_obs(self):852 return HATTNetworkTest._make_obs(self)853854 def setUp(self):855 HATTNetworkTest.setUp(self)856 torch.manual_seed(11)857 gate_args = SimpleNamespace(858 **vars(self.args),859 reconfiguration_mode="learned",860 reconfiguration_gate_init=0.8,861 reconfiguration_gate_floor=0.02,862 )863 self.actor = HATTV20Actor(864 gate_args,865 self.ego_dim,866 self.other_dim,867 self.task_dim,868 self.num_agents,869 self.num_tasks,870 )871872 def test_gate_forms_valid_masked_distribution(self):873 available = self.available.clone()874 # Agent 0 currently occupies task 0; make it unavailable to force release.875 available[0, 0] = 0.0876 dist, _ = self.actor._distribution(877 self.obs, self.rnn_states, self.masks, available878 )879 probabilities = dist.probs.reshape(880 1, self.num_agents, self.num_tasks881 )882 self.assertTrue(torch.allclose(probabilities.sum(-1), torch.ones(1, self.num_agents)))883 self.assertLess(float(probabilities[0, 0, 0].detach()), 1e-7)884 gate = self.actor.last_diagnostics[885 "reconfiguration_gate_probability"886 ]887 self.assertTrue(torch.all((gate >= 0.02) & (gate <= 1.0)))888 self.assertAlmostEqual(float(gate[0, 0, 0]), 1.0, places=6)889890 def test_ppo_action_loss_reaches_gate_parameters(self):891 dist, _ = self.actor._distribution(892 self.obs, self.rnn_states, self.masks, self.available893 )894 current = torch.tensor([[0], [2], [4]], dtype=torch.long)895 alternatives = (current + 1) % self.num_tasks896 loss = -dist.log_probs(alternatives).mean()897 loss.backward()898 gate_grad = self.actor.reconfiguration_gate[-1].weight.grad899 self.assertIsNotNone(gate_grad)900 self.assertTrue(torch.isfinite(gate_grad).all())901 self.assertGreater(float(gate_grad.abs().sum()), 0.0)902903 def test_explicit_context_changes_only_the_gate_input(self):904 context_args = SimpleNamespace(905 **vars(self.args),906 reconfiguration_mode="learned",907 reconfiguration_gate_init=0.8,908 reconfiguration_gate_floor=0.02,909 reconfiguration_context_features=True,910 reconfiguration_context_dim=9,911 )912 actor = HATTV20Actor(913 context_args,914 self.ego_dim,915 self.other_dim,916 self.task_dim,917 self.num_agents,918 self.num_tasks,919 )920 obs = {k: value.clone() for k, value in self.obs.items()}921 obs["reconfiguration"] = torch.zeros(1, self.num_agents, 9)922 dist, _ = actor._distribution(923 obs, self.rnn_states, self.masks, self.available924 )925926 self.assertEqual(927 actor.reconfiguration_gate[0].in_features,928 actor.hidden_size + context_args.hatt_hidden_dim + 2 + 9,929 )930 self.assertEqual(931 dist.probs.reshape(1, self.num_agents, self.num_tasks).shape,932 (1, self.num_agents, self.num_tasks),933 )934935 def test_switch_parameterization_makes_gate_equal_switch_probability(self):936 switch_args = SimpleNamespace(937 **vars(self.args),938 reconfiguration_mode="learned",939 reconfiguration_gate_init=0.3,940 reconfiguration_gate_floor=0.0,941 reconfiguration_gate_parameterization="switch",942 )943 actor = HATTV20Actor(944 switch_args,945 self.ego_dim,946 self.other_dim,947 self.task_dim,948 self.num_agents,949 self.num_tasks,950 )951 dist, _ = actor._distribution(952 self.obs, self.rnn_states, self.masks, self.available953 )954 probabilities = dist.probs.reshape(1, self.num_agents, self.num_tasks)955 current = torch.tensor([0, 2, 4])956 current_probability = probabilities[0, torch.arange(self.num_agents), current]957 gate = actor.last_diagnostics["reconfiguration_gate_probability"][0, :, 0]958 torch.testing.assert_close(current_probability, 1.0 - gate)959 torch.testing.assert_close(960 actor.last_diagnostics["reconfiguration_switch_probability"][0, :, 0],961 gate,962 )963964 def test_event_override_opens_gate_only_on_event(self):965 event_args = SimpleNamespace(966 **vars(self.args),967 reconfiguration_mode="learned",968 reconfiguration_gate_init=0.8,969 reconfiguration_gate_floor=0.02,970 reconfiguration_context_features=True,971 reconfiguration_context_dim=9,972 reconfiguration_gate_override="event",973 )974 actor = HATTV20Actor(975 event_args,976 self.ego_dim,977 self.other_dim,978 self.task_dim,979 self.num_agents,980 self.num_tasks,981 )982 obs = {k: value.clone() for k, value in self.obs.items()}983 obs["reconfiguration"] = torch.zeros(1, self.num_agents, 9)984 obs["reconfiguration"][0, 0, 0] = 1.0985 actor._distribution(obs, self.rnn_states, self.masks, self.available)986 gate = actor.last_diagnostics["reconfiguration_gate_probability"][0, :, 0]987 torch.testing.assert_close(gate, torch.tensor([1.0, 0.0, 0.0]))988989 def test_hierarchical_gate_joint_log_prob_is_replayable(self):990 hierarchical_args = SimpleNamespace(991 **vars(self.args),992 reconfiguration_mode="learned",993 reconfiguration_gate_init=0.05,994 reconfiguration_gate_floor=0.0,995 reconfiguration_gate_parameterization="hierarchical",996 reconfiguration_gate_override="event_or_learned",997 reconfiguration_context_features=True,998 reconfiguration_context_dim=11,999 )1000 actor = HATTV20Actor(1001 hierarchical_args,1002 self.ego_dim,1003 self.other_dim,1004 self.task_dim,1005 self.num_agents,1006 self.num_tasks,1007 )1008 obs = {k: value.clone() for k, value in self.obs.items()}1009 obs["reconfiguration"] = torch.zeros(1, self.num_agents, 11)1010 obs["reconfiguration"][..., 10] = 1.01011 actions, old_log_probs, _ = actor(1012 obs,1013 self.rnn_states,1014 self.masks,1015 self.available,1016 deterministic=False,1017 )1018 gate_actions = actor.last_gate_actions.clone()1019 new_log_probs, entropy = actor.evaluate_actions(1020 obs,1021 self.rnn_states,1022 actions,1023 self.masks,1024 self.available,1025 torch.ones(self.num_agents, 1),1026 gate_actions,1027 )1028 torch.testing.assert_close(gate_actions, torch.ones_like(gate_actions))1029 torch.testing.assert_close(new_log_probs, old_log_probs)1030 self.assertTrue(torch.isfinite(entropy))1031 self.assertTrue(torch.isfinite(actor.last_gate_regularization))103210331034class HATTV30ANetworkTest(HATTV20NetworkTest):1035 def setUp(self):1036 super().setUp()1037 torch.manual_seed(11)1038 self.actor = HATTV30AActor(1039 self.args,1040 self.ego_dim,1041 self.other_dim,1042 self.task_dim,1043 self.num_agents,1044 self.num_tasks,1045 )104610471048class HATTV30BNetworkTest(HATTV20NetworkTest):1049 def setUp(self):1050 super().setUp()1051 torch.manual_seed(11)1052 self.actor = HATTV30BActor(1053 self.args,1054 self.ego_dim,1055 self.other_dim,1056 self.task_dim,1057 self.num_agents,1058 self.num_tasks,1059 )10601061 def test_residual_features_receive_gradients(self):1062 dist, _ = self.actor._distribution(1063 self.obs, self.rnn_states, self.masks, self.available1064 )1065 actions = torch.zeros(self.num_agents, 1, dtype=torch.long)1066 (-dist.log_probs(actions).mean()).backward()1067 self.assertIsNotNone(self.actor.decoder.network[0].weight.grad)1068 self.assertTrue(torch.isfinite(self.actor.decoder.network[0].weight.grad).all())106910701071class HATTV30CNetworkTest(HATTV30BNetworkTest):1072 def setUp(self):1073 super().setUp()1074 torch.manual_seed(11)1075 self.actor = HATTV30CActor(1076 self.args,1077 self.ego_dim,1078 self.other_dim,1079 self.task_dim,1080 self.num_agents,1081 self.num_tasks,1082 )10831084 def test_dual_channel_interaction_receives_gradients(self):1085 obs = {k: v.clone() for k, v in self.obs.items()}1086 other_coalition_index = self.cap_dim + 41087 obs["others"][..., other_coalition_index] = 0.01088 dist, _ = self.actor._distribution(1089 obs, self.rnn_states, self.masks, self.available1090 )1091 actions = torch.zeros(self.num_agents, 1, dtype=torch.long)1092 (-dist.log_probs(actions).mean()).backward()1093 dual = self.actor.backbone.dual_channel_interaction1094 gradients = [1095 dual.self_channel_1.weight.grad,1096 dual.self_channel_2.weight.grad,1097 dual.task_channel_1.weight.grad,1098 dual.task_channel_2.weight.grad,1099 dual.member_channel_1_key.weight.grad,1100 dual.member_channel_2_key.weight.grad,1101 dual.member_channel_1_value.weight.grad,1102 dual.member_channel_2_value.weight.grad,1103 ]1104 self.assertTrue(1105 all(1106 grad is not None1107 and torch.isfinite(grad).all()1108 and grad.abs().sum() > 01109 for grad in gradients1110 )1111 )111211131114class HATTV31ANetworkTest(HATTV20NetworkTest):1115 def setUp(self):1116 super().setUp()1117 torch.manual_seed(11)1118 self.actor = HATTV31AActor(1119 self.args,1120 self.ego_dim,1121 self.other_dim,1122 self.task_dim,1123 self.num_agents,1124 self.num_tasks,1125 )11261127 def test_typed_residual_gate_and_mask(self):1128 obs = {k: v.clone() for k, v in self.obs.items()}1129 zero_req_task = 11130 obs["tasks"][..., zero_req_task, 2 : 2 + self.cap_dim] = 0.01131 _, candidate, diagnostics = self.actor.backbone(1132 obs, self.available.unsqueeze(0)1133 )1134 residual_context = candidate["extras"][0]1135 self.assertLess(1136 float(residual_context[..., zero_req_task, :].abs().sum().detach()),1137 1e-7,1138 )1139 self.assertGreater(float(diagnostics["residual_gate"].detach()), 0.0)11401141 def test_capability_residuals_use_normalized_scale(self):1142 c = self.cap_dim1143 obs = {k: torch.zeros_like(v) for k, v in self.obs.items()}1144 living_index = c + 61145 other_coalition_index = c + 41146 obs["others"][..., living_index] = 0.01147 obs["others"][0, 0, 0, living_index] = 1.01148 obs["others"][0, 0, 0, other_coalition_index] = 0.01149 obs["tasks"][0, 0, 0, 2 : 2 + c] = torch.tensor([0.6, 0.0])11501151 obs["others"][0, 0, 0, 2 : 2 + c] = torch.tensor([1.0, 0.0])1152 obs["ego"][0, 0, 2 : 2 + c] = torch.tensor([1.5, 0.0])1153 _, candidate, _ = self.actor.backbone(obs, torch.ones(1, self.num_agents, self.num_tasks))1154 residual_context = candidate["extras"][0]1155 residuals = self.actor.backbone._capability_residuals(1156 obs["ego"],1157 obs["others"],1158 obs["tasks"],1159 self.actor.backbone._coalition_indices(obs["others"][..., other_coalition_index]),1160 (obs["others"][..., living_index] > 0.5).float(),1161 )1162 torch.testing.assert_close(1163 residuals[0, 0, 0],1164 torch.tensor([0.3, 0.0, 0.0, 0.0]),1165 atol=1e-6,1166 rtol=1e-6,1167 )1168 self.assertLess(float(residual_context[0, 0, 1:].abs().max().detach()), 1e-7)11691170 obs["others"][0, 0, 0, 2 : 2 + c] = torch.tensor([2.5, 0.0])1171 obs["ego"][0, 0, 2 : 2 + c] = torch.tensor([2.0, 0.0])1172 residuals = self.actor.backbone._capability_residuals(1173 obs["ego"],1174 obs["others"],1175 obs["tasks"],1176 self.actor.backbone._coalition_indices(obs["others"][..., other_coalition_index]),1177 (obs["others"][..., living_index] > 0.5).float(),1178 )1179 torch.testing.assert_close(1180 residuals[0, 0, 0],1181 torch.tensor([0.1, 0.0, 0.3, 0.0]),1182 atol=1e-6,1183 rtol=1e-6,1184 )11851186 def test_zero_residual_produces_zero_context(self):1187 obs = {k: torch.zeros_like(v) for k, v in self.obs.items()}1188 obs["others"][..., self.cap_dim + 6] = 1.01189 _, candidate, diagnostics = self.actor.backbone(1190 obs, torch.ones(1, self.num_agents, self.num_tasks)1191 )1192 residual_context = candidate["extras"][0]1193 self.assertLess(float(residual_context.abs().max().detach()), 1e-7)1194 self.assertLess(float(diagnostics["nonzero_residual_ratio"].max()), 1e-7)11951196 def test_residual_gate_receives_gradients(self):1197 dist, _ = self.actor._distribution(1198 self.obs, self.rnn_states, self.masks, self.available1199 )1200 actions = torch.zeros(self.num_agents, 1, dtype=torch.long)1201 (-dist.log_probs(actions).mean()).backward()1202 self.assertIsNotNone(self.actor.backbone.residual_gate.grad)1203 self.assertTrue(torch.isfinite(self.actor.backbone.residual_gate.grad).all())1204 self.assertGreater(float(self.actor.backbone.residual_gate.grad.abs()), 0.0)120512061207class HATTV31BNetworkTest(HATTV31ANetworkTest):1208 def setUp(self):1209 super().setUp()1210 torch.manual_seed(11)1211 self.actor = HATTV31BActor(1212 self.args,1213 self.ego_dim,1214 self.other_dim,1215 self.task_dim,1216 self.num_agents,1217 self.num_tasks,1218 )12191220 def test_dual_channel_gate_receives_gradients(self):1221 obs = {k: v.clone() for k, v in self.obs.items()}1222 other_coalition_index = self.cap_dim + 41223 obs["others"][..., other_coalition_index] = 0.01224 dist, _ = self.actor._distribution(1225 obs, self.rnn_states, self.masks, self.available1226 )1227 actions = torch.zeros(self.num_agents, 1, dtype=torch.long)1228 (-dist.log_probs(actions).mean()).backward()1229 self.assertIsNotNone(self.actor.backbone.dual_channel_gate.grad)1230 self.assertTrue(torch.isfinite(self.actor.backbone.dual_channel_gate.grad).all())1231 self.assertGreater(float(self.actor.backbone.dual_channel_gate.grad.abs()), 0.0)1232 self.assertIsNotNone(1233 self.actor.backbone.dual_channel_interaction.member_channel_1_key.weight.grad1234 )12351236 def test_empty_coalition_dual_channel_context_is_zero(self):1237 obs = {k: v.clone() for k, v in self.obs.items()}1238 obs["others"].zero_()1239 _, candidate, diagnostics = self.actor.backbone(1240 obs, torch.ones(1, self.num_agents, self.num_tasks)1241 )1242 dual_channel_context = candidate["extras"][1]1243 self.assertLess(float(dual_channel_context.abs().max().detach()), 1e-7)1244 self.assertLess(float(diagnostics["occupied_task_mask"].max().detach()), 1e-7)124512461247from algorithms.algorithm.HATTV32Network import (1248 HATTV32AActor,1249 HATTV32BActor,1250 HATTV32CActor,1251 HATTV32DActor,1252)125312541255class HATTV32ANetworkTest(HATTV31ANetworkTest):1256 def setUp(self):1257 super().setUp()1258 torch.manual_seed(11)1259 self.actor = HATTV32AActor(1260 self.args,1261 self.ego_dim,1262 self.other_dim,1263 self.task_dim,1264 self.num_agents,1265 self.num_tasks,1266 )126712681269class HATTV32BNetworkTest(HATTV32ANetworkTest):1270 def setUp(self):1271 super().setUp()1272 torch.manual_seed(11)1273 self.actor = HATTV32BActor(1274 self.args,1275 self.ego_dim,1276 self.other_dim,1277 self.task_dim,1278 self.num_agents,1279 self.num_tasks,1280 )12811282 def test_true_signed_interaction_receives_gradients(self):1283 obs = {k: v.clone() for k, v in self.obs.items()}1284 other_coalition_index = self.cap_dim + 41285 obs["others"][..., other_coalition_index] = 0.01286 dist, _ = self.actor._distribution(1287 obs, self.rnn_states, self.masks, self.available1288 )1289 actions = torch.zeros(self.num_agents, 1, dtype=torch.long)1290 (-dist.log_probs(actions).mean()).backward()1291 signed = self.actor.backbone.true_signed_interaction1292 self.assertIsNotNone(signed.rel_query_self.weight.grad)1293 self.assertIsNotNone(signed.pol_query_self.weight.grad)1294 self.assertTrue(torch.isfinite(signed.rel_query_self.weight.grad).all())1295 self.assertTrue(torch.isfinite(signed.pol_query_self.weight.grad).all())12961297 def test_polarity_in_range(self):1298 obs = {k: v.clone() for k, v in self.obs.items()}1299 other_coalition_index = self.cap_dim + 41300 obs["others"][..., other_coalition_index] = 0.01301 self.actor._distribution(obs, self.rnn_states, self.masks, self.available)1302 polarity = self.actor.last_diagnostics["signed_polarity_mean"]1303 self.assertTrue((polarity.abs() <= 1.0 + 1e-6).all())130413051306class HATTV32CNetworkTest(HATTV32BNetworkTest):1307 def setUp(self):1308 super().setUp()1309 torch.manual_seed(11)1310 self.actor = HATTV32CActor(1311 self.args,1312 self.ego_dim,1313 self.other_dim,1314 self.task_dim,1315 self.num_agents,1316 self.num_tasks,1317 )13181319 def test_source_leave_delta_nonzero_on_switch(self):1320 obs = {k: v.clone() for k, v in self.obs.items()}1321 other_coalition_index = self.cap_dim + 41322 coalition_index = 2 + self.cap_dim1323 # All agents assigned to task 0; candidate task 1 is a switch.1324 obs["others"][..., other_coalition_index] = 0.01325 obs["ego"][..., coalition_index] = 0.01326 self.actor._distribution(obs, self.rnn_states, self.masks, self.available)1327 leave_norm = self.actor.last_diagnostics[1328 "source_leave_delta_norm"1329 ]1330 # Task 0 is stay -> leave delta should be zero; task 1 is switch -> nonzero.1331 self.assertLess(float(leave_norm[0, 0, 0].abs()), 1e-6)1332 self.assertGreater(float(leave_norm[0, 0, 1].abs()), 1e-6)133313341335class HATTV32DActorNetworkTest(HATTV32CNetworkTest):1336 def setUp(self):1337 super().setUp()1338 torch.manual_seed(11)1339 self.actor = HATTV32DActor(1340 self.args,1341 self.ego_dim,1342 self.other_dim,1343 self.task_dim,1344 self.num_agents,1345 self.num_tasks,1346 )13471348 def test_utility_head_receives_gradients(self):1349 dist, _ = self.actor._distribution(1350 self.obs, self.rnn_states, self.masks, self.available1351 )1352 actions = torch.zeros(self.num_agents, 1, dtype=torch.long)1353 (-dist.log_probs(actions).mean()).backward()1354 self.assertIsNotNone(self.actor.backbone.utility_head.network[0].weight.grad)1355 self.assertTrue(1356 torch.isfinite(1357 self.actor.backbone.utility_head.network[0].weight.grad1358 ).all()1359 )13601361 def test_utility_loss_finite(self):1362 loss = self.actor.get_utility_loss(self.obs, self.available)1363 self.assertTrue(torch.isfinite(loss))1364 self.assertGreaterEqual(float(loss), 0.0)136513661367if __name__ == "__main__":1368 unittest.main()