train: remove policy target-mask leakage and harden runtime settings

Author: Dieg0Code

src/model/system.py

@@ -101,13 +101,9 @@ def _common_step(
         batch: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
     ) -> dict[str, torch.Tensor]:
         boards, target_pis, target_vs = batch
-        legal_mask = (target_pis > 0).to(dtype=boards.dtype)
-        has_legal_support = torch.sum(legal_mask, dim=1, keepdim=True) > 0
-        if not bool(torch.all(has_legal_support).item()):
-            # Defensive fallback: keep logits finite even if a malformed target row is all zeros.
-            legal_mask = torch.where(has_legal_support, legal_mask, torch.ones_like(legal_mask))
-
-        pi_logits, v_pred = self(boards, action_mask=legal_mask)
+        # Training must not see a target-derived action mask, otherwise policy loss can
+        # become artificially easy (label leakage) when targets are sparse/one-hot.
+        pi_logits, v_pred = self(boards)
 
         loss_v = functional.mse_loss(v_pred.view(-1), target_vs.view(-1))
         log_probs = functional.log_softmax(pi_logits, dim=1)

tests/test_training_step_numerics.py

@@ -102,7 +102,7 @@ def test_forward_passes_action_mask_to_inner_model(self) -> None:
             self.assertIsInstance(action_mask_obj, torch.Tensor)
             self.assertTrue(torch.equal(action_mask_obj, mask))
 
-    def test_common_step_uses_policy_support_as_action_mask(self) -> None:
+    def test_common_step_does_not_pass_target_derived_action_mask(self) -> None:
         system = AtaxxZero(
             learning_rate=1e-3,
             d_model=64,
@@ -121,15 +121,11 @@ def test_common_step_uses_policy_support_as_action_mask(self) -> None:
 
         with patch.object(system.model, "forward", return_value=(pi_logits, v_pred)) as forward_spy:
             _ = system._common_step((boards, target_pis, target_vs))
-            _, kwargs = forward_spy.call_args
+            args, kwargs = forward_spy.call_args
 
-        self.assertIn("action_mask", kwargs)
-        action_mask = kwargs["action_mask"]
-        self.assertIsInstance(action_mask, torch.Tensor)
-        self.assertEqual(action_mask.shape, target_pis.shape)
-        self.assertEqual(float(action_mask[0, 7].item()), 1.0)
-        self.assertEqual(float(torch.sum(action_mask[0]).item()), 1.0)
-        self.assertTrue(torch.all(action_mask[1] == 1.0).item())
+        self.assertEqual(len(args), 1)
+        self.assertIsInstance(args[0], torch.Tensor)
+        self.assertEqual(kwargs.get("action_mask"), None)
 
     def test_common_step_applies_value_loss_coefficient(self) -> None:
         system = AtaxxZero(

train.py

@@ -16,11 +16,9 @@
 src = root / "src"
 if str(src) not in sys.path:
     sys.path.insert(0, str(src))
-
 if TYPE_CHECKING:
     from data.replay_buffer import ReplayBuffer
     from model.system import AtaxxZero
-
 from training.bootstrap import generate_imitation_data  # noqa: E402
 from training.callbacks import OptimizerStateTransfer  # noqa: E402
 from training.checkpointing import (  # noqa: E402
@@ -73,7 +71,7 @@ def _build_train_loader(buffer: ReplayBuffer, device: str) -> DataLoader[object]
             batch_size=cfg_int("batch_size"),
             shuffle=True,
             num_workers=cfg_int("num_workers"),
-            persistent_workers=True,
+            persistent_workers=cfg_bool("persistent_workers"),
             pin_memory=(device == "cuda"),
             prefetch_factor=2,
         )
@@ -99,7 +97,7 @@ def _build_val_loader(buffer: ReplayBuffer, device: str) -> DataLoader[object] |
             batch_size=cfg_int("batch_size"),
             shuffle=False,
             num_workers=cfg_int("num_workers"),
-            persistent_workers=True,
+            persistent_workers=cfg_bool("persistent_workers"),
             pin_memory=(device == "cuda"),
             prefetch_factor=2,
         )
@@ -494,6 +492,8 @@ def main() -> None:
                     fail_on_error=cfg_bool("fail_on_hf_upload_error"),
                 )
             except Exception as exc:
+                if cfg_bool("fail_on_hf_upload_error"):
+                    raise
                 log(f"HF upload wait failed: {exc}")
             hf_upload_executor.shutdown(wait=False, cancel_futures=True)
 if __name__ == "__main__":