import unittest from types import SimpleNamespace import torch from algorithms.algorithm.HATTNetwork import HATTActor, HATTV2Actor from algorithms.algorithm.HATTV20Network import ( HATTV20Actor, initialize_full_actor_from_compact_reference, ) from algorithms.algorithm.HATTV30Network import ( HATTV30AActor, HATTV30BActor, HATTV30CActor, ) from algorithms.algorithm.HATTV31Network import HATTV31AActor, HATTV31BActor from algorithms.algorithm.HATTV32RNetwork import ( HATTV32RFullActor, HATTV32RNoMemberActor, HATTV32RNoTransitionActor, ) from algorithms.algorithm.HATTV33AANetwork import ( HATTV33AAActor, TypedRelationEdge, ) from algorithms.algorithm.HATTV34Network import HATTV34Actor class HATTNetworkTest(unittest.TestCase): def setUp(self): torch.manual_seed(11) self.num_agents = 3 self.num_tasks = 5 self.cap_dim = 2 self.args = SimpleNamespace( hidden_size=16, use_policy_active_masks=True, use_naive_recurrent_policy=False, use_recurrent_policy=False, recurrent_N=1, use_orthogonal=True, cap_dim=self.cap_dim, gat_out_dim=16, hatt_hidden_dim=16, hatt_num_heads=4, hatt_dropout=0.0, hatt_capability_scale=5.0, ) self.ego_dim = self.cap_dim + 7 self.other_dim = self.cap_dim + 7 self.task_dim = self.cap_dim + 4 self.actor = HATTActor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) self.obs = self._make_obs() self.available = torch.ones(self.num_agents, self.num_tasks) self.available[0, 2] = 0.0 self.rnn_states = torch.zeros(self.num_agents, 1, self.args.hidden_size) self.masks = torch.ones(self.num_agents, 1) def _make_obs(self): a, t, c = self.num_agents, self.num_tasks, self.cap_dim ego = torch.rand(1, a, self.ego_dim) others = torch.rand(1, a, a, self.other_dim) tasks = torch.rand(1, a, t, self.task_dim) coalition_index = 2 + c other_coalition_index = c + 4 living_index = c + 6 ids = torch.tensor([0, 2, 4]) ego[..., coalition_index] = ids.float() / (t - 1) others[..., other_coalition_index] = ids.view(1, 1, a).float() / (t - 1) others[..., living_index] = 1.0 for observer in range(a): others[:, observer, observer] = 0.0 return {"ego": ego, "others": others, "tasks": tasks} def _logits(self, obs, available): dist, _ = self.actor._distribution( obs, self.rnn_states, self.masks, available ) return dist.logits.reshape(1, self.num_agents, self.num_tasks) def test_task_permutation_equivariance_and_mask(self): base = self._logits(self.obs, self.available).detach() self.assertLess(float(base.exp()[0, 0, 2]), 1e-7) permutation = torch.tensor([2, 4, 0, 3, 1]) inverse = torch.empty_like(permutation) inverse[permutation] = torch.arange(self.num_tasks) permuted_obs = {k: v.clone() for k, v in self.obs.items()} permuted_obs["tasks"] = permuted_obs["tasks"][:, :, permutation] permuted_available = self.available[:, permutation] coalition_index = 2 + self.cap_dim other_coalition_index = self.cap_dim + 4 old = ( permuted_obs["ego"][..., coalition_index] * (self.num_tasks - 1) ).round().long() permuted_obs["ego"][..., coalition_index] = ( inverse[old].float() / (self.num_tasks - 1) ) old = ( permuted_obs["others"][..., other_coalition_index] * (self.num_tasks - 1) ).round().long() permuted_obs["others"][..., other_coalition_index] = ( inverse[old].float() / (self.num_tasks - 1) ) actual = self._logits(permuted_obs, permuted_available).detach() torch.testing.assert_close(actual, base[:, :, permutation], atol=2e-5, rtol=2e-5) def test_agent_permutation_equivariance(self): base = self._logits(self.obs, self.available).detach() permutation = torch.tensor([2, 0, 1]) permuted_obs = {k: v.clone() for k, v in self.obs.items()} permuted_obs["ego"] = permuted_obs["ego"][:, permutation] permuted_obs["tasks"] = permuted_obs["tasks"][:, permutation] permuted_obs["others"] = permuted_obs["others"][:, permutation][ :, :, permutation ] actual = self._logits( permuted_obs, self.available[permutation] ).detach() torch.testing.assert_close(actual, base[:, permutation], atol=2e-5, rtol=2e-5) def test_all_major_branches_receive_gradients(self): dist, _ = self.actor._distribution( self.obs, self.rnn_states, self.masks, self.available ) actions = torch.zeros(self.num_agents, 1, dtype=torch.long) (-dist.log_probs(actions).mean()).backward() gradients = [ self.actor.backbone.self_encoder.net[0].weight.grad, self.actor.backbone.other_encoder.net[0].weight.grad, self.actor.backbone.task_encoder.net[0].weight.grad, self.actor.backbone.aa_attention.query.weight.grad, self.actor.backbone.ta_attention.query.weight.grad, self.actor.backbone.fusion[0].weight.grad, self.actor.decoder.agent_query.weight.grad, ] self.assertTrue( all( grad is not None and torch.isfinite(grad).all() and grad.abs().sum() > 0 for grad in gradients ) ) class HATTV2NetworkTest(HATTNetworkTest): def setUp(self): super().setUp() torch.manual_seed(11) self.actor = HATTV2Actor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) def test_all_major_branches_receive_gradients(self): dist, _ = self.actor._distribution( self.obs, self.rnn_states, self.masks, self.available ) actions = torch.zeros(self.num_agents, 1, dtype=torch.long) (-dist.log_probs(actions).mean()).backward() gradients = [ self.actor.backbone.self_encoder.net[0].weight.grad, self.actor.backbone.other_encoder.net[0].weight.grad, self.actor.backbone.task_encoder.net[0].weight.grad, self.actor.backbone.member_encoder.net[0].weight.grad, self.actor.backbone.aa_attention.query.weight.grad, self.actor.backbone.agent_fusion[0].weight.grad, self.actor.backbone.task_fusion[0].weight.grad, self.actor.backbone.decision_fusion[0].weight.grad, self.actor.decoder.network[0].weight.grad, self.actor.decoder.network[-1].weight.grad, ] self.assertTrue( all( grad is not None and torch.isfinite(grad).all() and grad.abs().sum() > 0 for grad in gradients ) ) def test_assignment_moves_coalition_context_between_task_slots(self): obs = {k: v.clone() for k, v in self.obs.items()} other_coalition_index = self.cap_dim + 4 observer = 0 other_agent = 1 old_task = 2 new_task = 3 obs["others"][0, observer, other_agent, other_coalition_index] = ( old_task / (self.num_tasks - 1) ) _, before, _, _ = self.actor.backbone( obs, self.available.unsqueeze(0) ) obs["others"][0, observer, other_agent, other_coalition_index] = ( new_task / (self.num_tasks - 1) ) _, after, _, _ = self.actor.backbone( obs, self.available.unsqueeze(0) ) unchanged_tasks = [0, 1, 4] torch.testing.assert_close( after[0, observer, unchanged_tasks], before[0, observer, unchanged_tasks], ) self.assertGreater( float( (after[0, observer, old_task] - before[0, observer, old_task]) .abs() .sum() .detach() ), 1e-5, ) self.assertGreater( float( (after[0, observer, new_task] - before[0, observer, new_task]) .abs() .sum() .detach() ), 1e-5, ) def test_task_pooling_ignores_unavailable_candidates(self): before_context, _, _, _ = self.actor.backbone( self.obs, self.available.unsqueeze(0) ) before_logits = self._logits(self.obs, self.available).detach() changed_obs = {key: value.clone() for key, value in self.obs.items()} changed_obs["tasks"][0, 0, 2] += 100.0 after_context, _, _, _ = self.actor.backbone( changed_obs, self.available.unsqueeze(0) ) after_logits = self._logits(changed_obs, self.available).detach() torch.testing.assert_close( after_context[0, 0], before_context[0, 0] ) available = self.available[0].bool() torch.testing.assert_close( after_logits[0, 0, available], before_logits[0, 0, available] ) def test_non_spatial_tasks_do_not_use_relative_position(self): obs = {key: value.clone() for key, value in self.obs.items()} type_index = self.cap_dim + 3 obs["tasks"][..., 0, type_index] = 0.0 obs["tasks"][..., 1, type_index] = 2.0 / 3.0 obs["tasks"][..., 2, type_index] = 1.0 / 3.0 obs["tasks"][..., 3, type_index] = 1.0 _, _, edge_features, _ = self.actor.backbone( obs, self.available.unsqueeze(0) ) torch.testing.assert_close( edge_features[..., 0:2, 0:2], torch.zeros_like(edge_features[..., 0:2, 0:2]), ) torch.testing.assert_close( edge_features[..., 0:2, 3], torch.zeros_like(edge_features[..., 0:2, 3]), ) torch.testing.assert_close( edge_features[..., 2:4, 3], torch.ones_like(edge_features[..., 2:4, 3]), ) def test_initial_policy_has_nontrivial_finite_logits(self): available = torch.ones_like(self.available) dist, _ = self.actor._distribution( self.obs, self.rnn_states, self.masks, available ) self.assertTrue(torch.isfinite(dist.logits).all()) self.assertGreater(float(dist.logits.std().detach()), 1e-4) class HATTV20NetworkTest(HATTNetworkTest): def setUp(self): super().setUp() torch.manual_seed(11) self.actor = HATTV20Actor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) def test_full_schema_scale_restore_and_compact_projection(self): full_dim = 2 + self.cap_dim + 1 + 1 + 3 + 1 full_tasks = torch.zeros(1, self.num_agents, self.num_tasks, full_dim) full_tasks[..., 0] = 0.1 full_tasks[..., 1] = 0.2 full_tasks[..., 2] = 0.15 full_tasks[..., 2 + self.cap_dim] = 0.6 full_tasks[..., 3 + self.cap_dim] = 0.7 full_tasks[..., 7 + self.cap_dim] = 1.0 self.args.hatt_v20_task_input_mode = "restore_req_scale" restored_actor = HATTV20Actor( self.args, self.ego_dim, self.other_dim, full_dim, self.num_agents, self.num_tasks, ) restored = restored_actor.backbone._prepare_task_encoder_input(full_tasks) self.assertEqual(restored.shape[-1], full_dim) torch.testing.assert_close(restored[..., 2], torch.full_like(restored[..., 2], 0.75)) self.args.hatt_v20_task_input_mode = "compact_projection" projected_actor = HATTV20Actor( self.args, self.ego_dim, self.other_dim, full_dim, self.num_agents, self.num_tasks, ) projected = projected_actor.backbone._prepare_task_encoder_input(full_tasks) expected = torch.tensor([0.1, 0.2, 0.6, 0.7, 0.75, 1.0]) torch.testing.assert_close(projected[0, 0, 0], expected) self.assertEqual(projected_actor.backbone.task_encoder.net[0].in_features, 6) def test_compact_projection_supports_exact_actor_pairing(self): compact_args = SimpleNamespace(**vars(self.args)) compact_args.hatt_v20_task_input_mode = "native" torch.manual_seed(91) compact_actor = HATTV20Actor( compact_args, self.ego_dim, self.other_dim, 6, self.num_agents, self.num_tasks, ) projected_args = SimpleNamespace(**vars(self.args)) projected_args.hatt_v20_task_input_mode = "compact_projection" torch.manual_seed(91) projected_actor = HATTV20Actor( projected_args, self.ego_dim, self.other_dim, 2 + self.cap_dim + 1 + 1 + 3 + 1, self.num_agents, self.num_tasks, ) compact_state = compact_actor.state_dict() projected_state = projected_actor.state_dict() self.assertEqual(compact_state.keys(), projected_state.keys()) for key in compact_state: torch.testing.assert_close(compact_state[key], projected_state[key]) def test_full_actor_function_preserving_initialization(self): full_dim = 2 + self.cap_dim + 1 + 1 + 3 + 1 compact_args = SimpleNamespace(**vars(self.args)) compact_args.hatt_v20_task_input_mode = "native" torch.manual_seed(137) compact_actor = HATTV20Actor( compact_args, self.ego_dim, self.other_dim, 6, self.num_agents, self.num_tasks, ) full_args = SimpleNamespace(**vars(self.args)) full_args.hatt_v20_task_input_mode = "restore_req_scale" torch.manual_seed(999) full_actor = HATTV20Actor( full_args, self.ego_dim, self.other_dim, full_dim, self.num_agents, self.num_tasks, ) initialize_full_actor_from_compact_reference( full_actor, compact_actor, self.cap_dim ) compact_tasks = torch.randn( 2, self.num_agents, self.num_tasks, 6 ) full_tasks = torch.zeros( 2, self.num_agents, self.num_tasks, full_dim ) full_tasks[..., 0:2] = compact_tasks[..., 0:2] full_tasks[..., 2] = compact_tasks[..., 4] / 5.0 full_tasks[..., 2 + self.cap_dim] = compact_tasks[..., 2] full_tasks[..., 3 + self.cap_dim] = compact_tasks[..., 3] full_tasks[..., 7 + self.cap_dim] = compact_tasks[..., 5] compact_encoded = compact_actor.backbone.task_encoder(compact_tasks) full_prepared = full_actor.backbone._prepare_task_encoder_input(full_tasks) full_encoded = full_actor.backbone.task_encoder(full_prepared) # GEMM reduction order differs for 6-D and 12-D inputs even when the # added products are all zero. The construction is function-preserving # up to normal floating-point accumulation error. torch.testing.assert_close(compact_encoded, full_encoded, rtol=1e-6, atol=1e-6) full_first = full_actor.backbone.task_encoder.net[0].weight semantic_columns = {0, 1, 2, 2 + self.cap_dim, 3 + self.cap_dim, 7 + self.cap_dim} zero_columns = [idx for idx in range(full_dim) if idx not in semantic_columns] torch.testing.assert_close( full_first[:, zero_columns], torch.zeros_like(full_first[:, zero_columns]) ) projected_args = SimpleNamespace(**vars(self.args)) projected_args.hatt_v20_task_input_mode = "compact_projection" torch.manual_seed(137) projected_actor = HATTV20Actor( projected_args, self.ego_dim, self.other_dim, full_dim, self.num_agents, self.num_tasks, ) full_obs = { "ego": torch.rand(2, self.num_agents, self.ego_dim), "others": torch.rand( 2, self.num_agents, self.num_agents, self.other_dim ), "tasks": full_tasks, } available = torch.ones(2, self.num_agents, self.num_tasks) rnn_states = torch.zeros( 2 * self.num_agents, 1, self.args.hidden_size ) masks = torch.ones(2 * self.num_agents, 1) projected_dist, _ = projected_actor._distribution( full_obs, rnn_states, masks, available ) full_dist, _ = full_actor._distribution( full_obs, rnn_states, masks, available ) torch.testing.assert_close( projected_dist.logits, full_dist.logits, rtol=1e-6, atol=1e-6 ) def test_no_aa_attention_ablation_keeps_shape_and_zeroes_attention(self): self.args.hatt_use_aa_attention = False actor = HATTV20Actor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) agent, tasks, edges, diagnostics = actor.backbone( self.obs, self.available, return_diagnostics=True ) self.assertEqual(agent.shape, (1, self.num_agents, self.args.hatt_hidden_dim)) self.assertEqual(tasks.shape, (1, self.num_agents, self.num_tasks, self.args.hatt_hidden_dim)) self.assertEqual(edges.shape[:3], (1, self.num_agents, self.num_tasks)) self.assertEqual(float(diagnostics["aa_attention"].abs().sum()), 0.0) def test_no_coalition_context_ablation_zeroes_coalition_diagnostics(self): self.args.hatt_use_coalition_context = False actor = HATTV20Actor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) _, _, _, diagnostics = actor.backbone( self.obs, self.available, return_diagnostics=True ) self.assertEqual(float(diagnostics["coalition_context_norm"].sum()), 0.0) self.assertEqual(float(diagnostics["occupied_task_ratio"].sum()), 0.0) def test_all_major_branches_receive_gradients(self): dist, _ = self.actor._distribution( self.obs, self.rnn_states, self.masks, self.available ) actions = torch.zeros(self.num_agents, 1, dtype=torch.long) (-dist.log_probs(actions).mean()).backward() gradients = [ self.actor.backbone.self_encoder.net[0].weight.grad, self.actor.backbone.other_encoder.net[0].weight.grad, self.actor.backbone.task_encoder.net[0].weight.grad, self.actor.backbone.aa_attention.query.weight.grad, self.actor.backbone.agent_fusion[0].weight.grad, self.actor.backbone.task_fusion[0].weight.grad, self.actor.decoder.network[0].weight.grad, self.actor.decoder.network[-1].weight.grad, ] self.assertTrue( all( grad is not None and torch.isfinite(grad).all() and grad.abs().sum() > 0 for grad in gradients ) ) class HATTV32RNetworkTest(HATTNetworkTest): def setUp(self): super().setUp() # Explicit task types: Deliver, Explore, Idle, Resupply, Deliver. self.obs["tasks"][..., -1] = torch.tensor( [1.0 / 3.0, 0.0, 2.0 / 3.0, 1.0, 1.0 / 3.0] ) # Give observer 0 two real teammates in its current Deliver coalition. coalition_index = self.cap_dim + 4 self.obs["others"][0, 0, 1, coalition_index] = 0.0 self.obs["others"][0, 0, 2, coalition_index] = 0.0 torch.manual_seed(11) self.actor = HATTV32RFullActor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) def test_agent_permutation_equivariance(self): # scatter_add may sum coalition members in a different floating-point # order after relabeling; the policy remains equivariant numerically. base = self._logits(self.obs, self.available).detach() permutation = torch.tensor([2, 0, 1]) permuted_obs = {key: value.clone() for key, value in self.obs.items()} permuted_obs["ego"] = permuted_obs["ego"][:, permutation] permuted_obs["tasks"] = permuted_obs["tasks"][:, permutation] permuted_obs["others"] = permuted_obs["others"][:, permutation][ :, :, permutation ] actual = self._logits( permuted_obs, self.available[permutation] ).detach() torch.testing.assert_close( actual, base[:, permutation], atol=2e-3, rtol=2e-3 ) def test_all_major_branches_receive_gradients(self): self.test_new_relation_branches_receive_gradients() def test_transition_semantics_ignore_control_target_and_stay(self): _, _, _, diagnostics = self.actor.backbone( self.obs, self.available.unsqueeze(0), return_diagnostics=True ) target_norm = diagnostics["candidate_target_delta_norm"] source_norm = diagnostics["candidate_source_delta_norm"] # Explore/Idle/Resupply never receive a fictitious target coalition. torch.testing.assert_close( target_norm[..., 1:4], torch.zeros_like(target_norm[..., 1:4]) ) # Observer 0 currently occupies task 0, so staying has no leave effect. self.assertEqual(float(source_norm[0, 0, 0].detach()), 0.0) self.assertGreater(float(source_norm[0, 0, 4].detach()), 1e-6) def test_both_ablation_paths_zero_only_their_diagnostics(self): no_transition = HATTV32RNoTransitionActor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) no_member = HATTV32RNoMemberActor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) _, _, _, transition_diag = no_transition.backbone( self.obs, self.available.unsqueeze(0), return_diagnostics=True ) _, _, _, member_diag = no_member.backbone( self.obs, self.available.unsqueeze(0), return_diagnostics=True ) self.assertEqual( float( transition_diag["task_conditioned_member_context_norm"] .sum() .detach() ) > 0.0, True, ) self.assertEqual( float(transition_diag["transition_context_norm"].sum().detach()), 0.0 ) self.assertEqual( float( member_diag["task_conditioned_member_context_norm"] .sum() .detach() ), 0.0, ) def test_new_relation_branches_receive_gradients(self): dist, _ = self.actor._distribution( self.obs, self.rnn_states, self.masks, self.available ) actions = torch.zeros(self.num_agents, 1, dtype=torch.long) (-dist.log_probs(actions).mean()).backward() gradients = [ self.actor.backbone.transition_projection[0].weight.grad, self.actor.backbone.member_attention.self_query.weight.grad, self.actor.backbone.member_projection[0].weight.grad, ] self.assertTrue( all( grad is not None and torch.isfinite(grad).all() and grad.abs().sum() > 0 for grad in gradients ) ) class HATTV33AANetworkTest(HATTNetworkTest): def setUp(self): super().setUp() self.obs["tasks"][..., -1] = torch.tensor( [1.0 / 3.0, 0.0, 2.0 / 3.0, 1.0, 1.0 / 3.0] ) coalition_index = self.cap_dim + 4 # Observer 0 sees two members in Deliver candidate 0. self.obs["others"][0, 0, 1, coalition_index] = 0.0 self.obs["others"][0, 0, 2, coalition_index] = 0.0 self.aa_args = SimpleNamespace( **vars(self.args), hatt_candidate_aa_mode="learned", hatt_shared_task_obs=False, ) torch.manual_seed(11) self.actor = HATTV33AAActor( self.aa_args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) def _actor_with_mode(self, mode): args = SimpleNamespace(**vars(self.aa_args)) args.hatt_candidate_aa_mode = mode actor = HATTV33AAActor( args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) actor.load_state_dict(self.actor.state_dict()) return actor def test_aa_and_at_use_the_same_typed_edge_operator(self): self.assertIsInstance(self.actor.backbone.aa_relation, TypedRelationEdge) self.assertIsInstance(self.actor.backbone.at_relation, TypedRelationEdge) self.assertIsNot( self.actor.backbone.aa_relation.query.weight, self.actor.backbone.at_relation.query.weight, ) def test_membership_routes_aa_only_to_deliver_candidates(self): _, _, _, diagnostics = self.actor.backbone( self.obs, self.available.unsqueeze(0), return_diagnostics=True ) counts = diagnostics["candidate_aa_member_count"] self.assertEqual(float(counts[0, 0, 0]), 2.0) self.assertEqual(float(counts[0, 0, 1:4].sum()), 0.0) self.assertGreater( float(diagnostics["candidate_aa_context_norm"][0, 0, 0].detach()), 0.0, ) def test_uniform_and_zero_are_checkpoint_compatible_interventions(self): uniform = self._actor_with_mode("uniform") zero = self._actor_with_mode("zero") _, _, _, uniform_diag = uniform.backbone( self.obs, self.available.unsqueeze(0), return_diagnostics=True ) valid = self.obs["others"].abs().sum(dim=-1) > 0 uniform_gates = uniform_diag["aa_relation_gate"] torch.testing.assert_close( uniform_gates.masked_select(valid.unsqueeze(-1)), torch.full_like(uniform_gates.masked_select(valid.unsqueeze(-1)), 0.5), ) learned_logits = self._logits(self.obs, self.available).detach() zero_dist, _ = zero._distribution( self.obs, self.rnn_states, self.masks, self.available ) zero_logits = zero_dist.logits.reshape_as(learned_logits).detach() self.assertGreater(float((learned_logits - zero_logits).abs().max()), 1e-6) def test_all_major_branches_receive_gradients(self): dist, _ = self.actor._distribution( self.obs, self.rnn_states, self.masks, self.available ) actions = torch.zeros(self.num_agents, 1, dtype=torch.long) (-dist.log_probs(actions).mean()).backward() gradients = [ self.actor.backbone.aa_relation.query.weight.grad, self.actor.backbone.aa_relation.key.weight.grad, self.actor.backbone.aa_relation.value.weight.grad, self.actor.backbone.at_relation.query.weight.grad, self.actor.backbone.at_relation.key.weight.grad, self.actor.backbone.at_relation.value.weight.grad, self.actor.backbone.agent_fusion[0].weight.grad, self.actor.backbone.candidate_fusion[0].weight.grad, self.actor.decoder.network[0].weight.grad, ] self.assertTrue( all( grad is not None and torch.isfinite(grad).all() and grad.abs().sum() > 0 for grad in gradients ) ) class HATTV34NetworkTest(HATTNetworkTest): def setUp(self): super().setUp() self.obs["tasks"][..., -1] = torch.tensor( [1.0 / 3.0, 0.0, 2.0 / 3.0, 1.0, 1.0 / 3.0] ) coalition_index = self.cap_dim + 4 self.obs["others"][0, 0, 1, coalition_index] = 0.0 self.obs["others"][0, 0, 2, coalition_index] = 0.0 self.aa_args = SimpleNamespace( **vars(self.args), hatt_candidate_aa_mode="learned", hatt_shared_task_obs=False, hatt_aa_task_context=True, coalition_task_feature_mode="compact", ) torch.manual_seed(11) self.actor = HATTV34Actor( self.aa_args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) def test_task_context_toggles_aa_edge_dim(self): self.assertEqual( self.actor.backbone.aa_relation.edge_bias[0].in_features, 7 ) off_args = SimpleNamespace(**vars(self.aa_args)) off_args.hatt_aa_task_context = False off = HATTV34Actor( off_args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) self.assertEqual(off.backbone.aa_relation.edge_bias[0].in_features, 3) def test_task_context_off_reproduces_v33aa(self): off_args = SimpleNamespace(**vars(self.aa_args)) off_args.hatt_aa_task_context = False torch.manual_seed(11) off = HATTV34Actor( off_args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) v33_args = SimpleNamespace(**vars(self.aa_args)) torch.manual_seed(11) v33 = HATTV33AAActor( v33_args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) _, ctx33, _, _ = v33.backbone( self.obs, self.available.unsqueeze(0), return_diagnostics=True ) _, ctx34, _, _ = off.backbone( self.obs, self.available.unsqueeze(0), return_diagnostics=True ) torch.testing.assert_close(ctx34, ctx33) def test_all_major_branches_receive_gradients(self): dist, _ = self.actor._distribution( self.obs, self.rnn_states, self.masks, self.available ) actions = torch.zeros(self.num_agents, 1, dtype=torch.long) (-dist.log_probs(actions).mean()).backward() gradients = [ self.actor.backbone.aa_relation.query.weight.grad, self.actor.backbone.aa_relation.key.weight.grad, self.actor.backbone.aa_relation.value.weight.grad, self.actor.backbone.aa_relation.edge_bias[0].weight.grad, self.actor.backbone.at_relation.query.weight.grad, self.actor.backbone.agent_fusion[0].weight.grad, self.actor.backbone.candidate_fusion[0].weight.grad, self.actor.decoder.network[0].weight.grad, ] self.assertTrue( all( grad is not None and torch.isfinite(grad).all() and grad.abs().sum() > 0 for grad in gradients ) ) class HATTV20ReconfigurationGateTest(unittest.TestCase): def _make_obs(self): return HATTNetworkTest._make_obs(self) def setUp(self): HATTNetworkTest.setUp(self) torch.manual_seed(11) gate_args = SimpleNamespace( **vars(self.args), reconfiguration_mode="learned", reconfiguration_gate_init=0.8, reconfiguration_gate_floor=0.02, ) self.actor = HATTV20Actor( gate_args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) def test_gate_forms_valid_masked_distribution(self): available = self.available.clone() # Agent 0 currently occupies task 0; make it unavailable to force release. available[0, 0] = 0.0 dist, _ = self.actor._distribution( self.obs, self.rnn_states, self.masks, available ) probabilities = dist.probs.reshape( 1, self.num_agents, self.num_tasks ) self.assertTrue(torch.allclose(probabilities.sum(-1), torch.ones(1, self.num_agents))) self.assertLess(float(probabilities[0, 0, 0].detach()), 1e-7) gate = self.actor.last_diagnostics[ "reconfiguration_gate_probability" ] self.assertTrue(torch.all((gate >= 0.02) & (gate <= 1.0))) self.assertAlmostEqual(float(gate[0, 0, 0]), 1.0, places=6) def test_ppo_action_loss_reaches_gate_parameters(self): dist, _ = self.actor._distribution( self.obs, self.rnn_states, self.masks, self.available ) current = torch.tensor([[0], [2], [4]], dtype=torch.long) alternatives = (current + 1) % self.num_tasks loss = -dist.log_probs(alternatives).mean() loss.backward() gate_grad = self.actor.reconfiguration_gate[-1].weight.grad self.assertIsNotNone(gate_grad) self.assertTrue(torch.isfinite(gate_grad).all()) self.assertGreater(float(gate_grad.abs().sum()), 0.0) def test_explicit_context_changes_only_the_gate_input(self): context_args = SimpleNamespace( **vars(self.args), reconfiguration_mode="learned", reconfiguration_gate_init=0.8, reconfiguration_gate_floor=0.02, reconfiguration_context_features=True, reconfiguration_context_dim=9, ) actor = HATTV20Actor( context_args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) obs = {k: value.clone() for k, value in self.obs.items()} obs["reconfiguration"] = torch.zeros(1, self.num_agents, 9) dist, _ = actor._distribution( obs, self.rnn_states, self.masks, self.available ) self.assertEqual( actor.reconfiguration_gate[0].in_features, actor.hidden_size + context_args.hatt_hidden_dim + 2 + 9, ) self.assertEqual( dist.probs.reshape(1, self.num_agents, self.num_tasks).shape, (1, self.num_agents, self.num_tasks), ) def test_switch_parameterization_makes_gate_equal_switch_probability(self): switch_args = SimpleNamespace( **vars(self.args), reconfiguration_mode="learned", reconfiguration_gate_init=0.3, reconfiguration_gate_floor=0.0, reconfiguration_gate_parameterization="switch", ) actor = HATTV20Actor( switch_args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) dist, _ = actor._distribution( self.obs, self.rnn_states, self.masks, self.available ) probabilities = dist.probs.reshape(1, self.num_agents, self.num_tasks) current = torch.tensor([0, 2, 4]) current_probability = probabilities[0, torch.arange(self.num_agents), current] gate = actor.last_diagnostics["reconfiguration_gate_probability"][0, :, 0] torch.testing.assert_close(current_probability, 1.0 - gate) torch.testing.assert_close( actor.last_diagnostics["reconfiguration_switch_probability"][0, :, 0], gate, ) def test_event_override_opens_gate_only_on_event(self): event_args = SimpleNamespace( **vars(self.args), reconfiguration_mode="learned", reconfiguration_gate_init=0.8, reconfiguration_gate_floor=0.02, reconfiguration_context_features=True, reconfiguration_context_dim=9, reconfiguration_gate_override="event", ) actor = HATTV20Actor( event_args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) obs = {k: value.clone() for k, value in self.obs.items()} obs["reconfiguration"] = torch.zeros(1, self.num_agents, 9) obs["reconfiguration"][0, 0, 0] = 1.0 actor._distribution(obs, self.rnn_states, self.masks, self.available) gate = actor.last_diagnostics["reconfiguration_gate_probability"][0, :, 0] torch.testing.assert_close(gate, torch.tensor([1.0, 0.0, 0.0])) def test_hierarchical_gate_joint_log_prob_is_replayable(self): hierarchical_args = SimpleNamespace( **vars(self.args), reconfiguration_mode="learned", reconfiguration_gate_init=0.05, reconfiguration_gate_floor=0.0, reconfiguration_gate_parameterization="hierarchical", reconfiguration_gate_override="event_or_learned", reconfiguration_context_features=True, reconfiguration_context_dim=11, ) actor = HATTV20Actor( hierarchical_args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) obs = {k: value.clone() for k, value in self.obs.items()} obs["reconfiguration"] = torch.zeros(1, self.num_agents, 11) obs["reconfiguration"][..., 10] = 1.0 actions, old_log_probs, _ = actor( obs, self.rnn_states, self.masks, self.available, deterministic=False, ) gate_actions = actor.last_gate_actions.clone() new_log_probs, entropy = actor.evaluate_actions( obs, self.rnn_states, actions, self.masks, self.available, torch.ones(self.num_agents, 1), gate_actions, ) torch.testing.assert_close(gate_actions, torch.ones_like(gate_actions)) torch.testing.assert_close(new_log_probs, old_log_probs) self.assertTrue(torch.isfinite(entropy)) self.assertTrue(torch.isfinite(actor.last_gate_regularization)) class HATTV30ANetworkTest(HATTV20NetworkTest): def setUp(self): super().setUp() torch.manual_seed(11) self.actor = HATTV30AActor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) class HATTV30BNetworkTest(HATTV20NetworkTest): def setUp(self): super().setUp() torch.manual_seed(11) self.actor = HATTV30BActor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) def test_residual_features_receive_gradients(self): dist, _ = self.actor._distribution( self.obs, self.rnn_states, self.masks, self.available ) actions = torch.zeros(self.num_agents, 1, dtype=torch.long) (-dist.log_probs(actions).mean()).backward() self.assertIsNotNone(self.actor.decoder.network[0].weight.grad) self.assertTrue(torch.isfinite(self.actor.decoder.network[0].weight.grad).all()) class HATTV30CNetworkTest(HATTV30BNetworkTest): def setUp(self): super().setUp() torch.manual_seed(11) self.actor = HATTV30CActor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) def test_dual_channel_interaction_receives_gradients(self): obs = {k: v.clone() for k, v in self.obs.items()} other_coalition_index = self.cap_dim + 4 obs["others"][..., other_coalition_index] = 0.0 dist, _ = self.actor._distribution( obs, self.rnn_states, self.masks, self.available ) actions = torch.zeros(self.num_agents, 1, dtype=torch.long) (-dist.log_probs(actions).mean()).backward() dual = self.actor.backbone.dual_channel_interaction gradients = [ dual.self_channel_1.weight.grad, dual.self_channel_2.weight.grad, dual.task_channel_1.weight.grad, dual.task_channel_2.weight.grad, dual.member_channel_1_key.weight.grad, dual.member_channel_2_key.weight.grad, dual.member_channel_1_value.weight.grad, dual.member_channel_2_value.weight.grad, ] self.assertTrue( all( grad is not None and torch.isfinite(grad).all() and grad.abs().sum() > 0 for grad in gradients ) ) class HATTV31ANetworkTest(HATTV20NetworkTest): def setUp(self): super().setUp() torch.manual_seed(11) self.actor = HATTV31AActor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) def test_typed_residual_gate_and_mask(self): obs = {k: v.clone() for k, v in self.obs.items()} zero_req_task = 1 obs["tasks"][..., zero_req_task, 2 : 2 + self.cap_dim] = 0.0 _, candidate, diagnostics = self.actor.backbone( obs, self.available.unsqueeze(0) ) residual_context = candidate["extras"][0] self.assertLess( float(residual_context[..., zero_req_task, :].abs().sum().detach()), 1e-7, ) self.assertGreater(float(diagnostics["residual_gate"].detach()), 0.0) def test_capability_residuals_use_normalized_scale(self): c = self.cap_dim obs = {k: torch.zeros_like(v) for k, v in self.obs.items()} living_index = c + 6 other_coalition_index = c + 4 obs["others"][..., living_index] = 0.0 obs["others"][0, 0, 0, living_index] = 1.0 obs["others"][0, 0, 0, other_coalition_index] = 0.0 obs["tasks"][0, 0, 0, 2 : 2 + c] = torch.tensor([0.6, 0.0]) obs["others"][0, 0, 0, 2 : 2 + c] = torch.tensor([1.0, 0.0]) obs["ego"][0, 0, 2 : 2 + c] = torch.tensor([1.5, 0.0]) _, candidate, _ = self.actor.backbone(obs, torch.ones(1, self.num_agents, self.num_tasks)) residual_context = candidate["extras"][0] residuals = self.actor.backbone._capability_residuals( obs["ego"], obs["others"], obs["tasks"], self.actor.backbone._coalition_indices(obs["others"][..., other_coalition_index]), (obs["others"][..., living_index] > 0.5).float(), ) torch.testing.assert_close( residuals[0, 0, 0], torch.tensor([0.3, 0.0, 0.0, 0.0]), atol=1e-6, rtol=1e-6, ) self.assertLess(float(residual_context[0, 0, 1:].abs().max().detach()), 1e-7) obs["others"][0, 0, 0, 2 : 2 + c] = torch.tensor([2.5, 0.0]) obs["ego"][0, 0, 2 : 2 + c] = torch.tensor([2.0, 0.0]) residuals = self.actor.backbone._capability_residuals( obs["ego"], obs["others"], obs["tasks"], self.actor.backbone._coalition_indices(obs["others"][..., other_coalition_index]), (obs["others"][..., living_index] > 0.5).float(), ) torch.testing.assert_close( residuals[0, 0, 0], torch.tensor([0.1, 0.0, 0.3, 0.0]), atol=1e-6, rtol=1e-6, ) def test_zero_residual_produces_zero_context(self): obs = {k: torch.zeros_like(v) for k, v in self.obs.items()} obs["others"][..., self.cap_dim + 6] = 1.0 _, candidate, diagnostics = self.actor.backbone( obs, torch.ones(1, self.num_agents, self.num_tasks) ) residual_context = candidate["extras"][0] self.assertLess(float(residual_context.abs().max().detach()), 1e-7) self.assertLess(float(diagnostics["nonzero_residual_ratio"].max()), 1e-7) def test_residual_gate_receives_gradients(self): dist, _ = self.actor._distribution( self.obs, self.rnn_states, self.masks, self.available ) actions = torch.zeros(self.num_agents, 1, dtype=torch.long) (-dist.log_probs(actions).mean()).backward() self.assertIsNotNone(self.actor.backbone.residual_gate.grad) self.assertTrue(torch.isfinite(self.actor.backbone.residual_gate.grad).all()) self.assertGreater(float(self.actor.backbone.residual_gate.grad.abs()), 0.0) class HATTV31BNetworkTest(HATTV31ANetworkTest): def setUp(self): super().setUp() torch.manual_seed(11) self.actor = HATTV31BActor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) def test_dual_channel_gate_receives_gradients(self): obs = {k: v.clone() for k, v in self.obs.items()} other_coalition_index = self.cap_dim + 4 obs["others"][..., other_coalition_index] = 0.0 dist, _ = self.actor._distribution( obs, self.rnn_states, self.masks, self.available ) actions = torch.zeros(self.num_agents, 1, dtype=torch.long) (-dist.log_probs(actions).mean()).backward() self.assertIsNotNone(self.actor.backbone.dual_channel_gate.grad) self.assertTrue(torch.isfinite(self.actor.backbone.dual_channel_gate.grad).all()) self.assertGreater(float(self.actor.backbone.dual_channel_gate.grad.abs()), 0.0) self.assertIsNotNone( self.actor.backbone.dual_channel_interaction.member_channel_1_key.weight.grad ) def test_empty_coalition_dual_channel_context_is_zero(self): obs = {k: v.clone() for k, v in self.obs.items()} obs["others"].zero_() _, candidate, diagnostics = self.actor.backbone( obs, torch.ones(1, self.num_agents, self.num_tasks) ) dual_channel_context = candidate["extras"][1] self.assertLess(float(dual_channel_context.abs().max().detach()), 1e-7) self.assertLess(float(diagnostics["occupied_task_mask"].max().detach()), 1e-7) from algorithms.algorithm.HATTV32Network import ( HATTV32AActor, HATTV32BActor, HATTV32CActor, HATTV32DActor, ) class HATTV32ANetworkTest(HATTV31ANetworkTest): def setUp(self): super().setUp() torch.manual_seed(11) self.actor = HATTV32AActor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) class HATTV32BNetworkTest(HATTV32ANetworkTest): def setUp(self): super().setUp() torch.manual_seed(11) self.actor = HATTV32BActor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) def test_true_signed_interaction_receives_gradients(self): obs = {k: v.clone() for k, v in self.obs.items()} other_coalition_index = self.cap_dim + 4 obs["others"][..., other_coalition_index] = 0.0 dist, _ = self.actor._distribution( obs, self.rnn_states, self.masks, self.available ) actions = torch.zeros(self.num_agents, 1, dtype=torch.long) (-dist.log_probs(actions).mean()).backward() signed = self.actor.backbone.true_signed_interaction self.assertIsNotNone(signed.rel_query_self.weight.grad) self.assertIsNotNone(signed.pol_query_self.weight.grad) self.assertTrue(torch.isfinite(signed.rel_query_self.weight.grad).all()) self.assertTrue(torch.isfinite(signed.pol_query_self.weight.grad).all()) def test_polarity_in_range(self): obs = {k: v.clone() for k, v in self.obs.items()} other_coalition_index = self.cap_dim + 4 obs["others"][..., other_coalition_index] = 0.0 self.actor._distribution(obs, self.rnn_states, self.masks, self.available) polarity = self.actor.last_diagnostics["signed_polarity_mean"] self.assertTrue((polarity.abs() <= 1.0 + 1e-6).all()) class HATTV32CNetworkTest(HATTV32BNetworkTest): def setUp(self): super().setUp() torch.manual_seed(11) self.actor = HATTV32CActor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) def test_source_leave_delta_nonzero_on_switch(self): obs = {k: v.clone() for k, v in self.obs.items()} other_coalition_index = self.cap_dim + 4 coalition_index = 2 + self.cap_dim # All agents assigned to task 0; candidate task 1 is a switch. obs["others"][..., other_coalition_index] = 0.0 obs["ego"][..., coalition_index] = 0.0 self.actor._distribution(obs, self.rnn_states, self.masks, self.available) leave_norm = self.actor.last_diagnostics[ "source_leave_delta_norm" ] # Task 0 is stay -> leave delta should be zero; task 1 is switch -> nonzero. self.assertLess(float(leave_norm[0, 0, 0].abs()), 1e-6) self.assertGreater(float(leave_norm[0, 0, 1].abs()), 1e-6) class HATTV32DActorNetworkTest(HATTV32CNetworkTest): def setUp(self): super().setUp() torch.manual_seed(11) self.actor = HATTV32DActor( self.args, self.ego_dim, self.other_dim, self.task_dim, self.num_agents, self.num_tasks, ) def test_utility_head_receives_gradients(self): dist, _ = self.actor._distribution( self.obs, self.rnn_states, self.masks, self.available ) actions = torch.zeros(self.num_agents, 1, dtype=torch.long) (-dist.log_probs(actions).mean()).backward() self.assertIsNotNone(self.actor.backbone.utility_head.network[0].weight.grad) self.assertTrue( torch.isfinite( self.actor.backbone.utility_head.network[0].weight.grad ).all() ) def test_utility_loss_finite(self): loss = self.actor.get_utility_loss(self.obs, self.available) self.assertTrue(torch.isfinite(loss)) self.assertGreaterEqual(float(loss), 0.0) if __name__ == "__main__": unittest.main()