From 0f6faa5fe44c8065338e0cdc7eb152b840bcdce2 Mon Sep 17 00:00:00 2001
From: KoalaTrapLord <265133236+KoalaTrapLord@users.noreply.github.com>
Date: Sun, 22 Mar 2026 00:06:55 +0000
Subject: [PATCH] feat(ml): integrate DecisionExporter into self-play rollouts
 (#65)

---
 ml/training/self_play.py | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/ml/training/self_play.py b/ml/training/self_play.py
index 476f531..f8195c8 100644
--- a/ml/training/self_play.py
+++ b/ml/training/self_play.py
@@ -38,6 +38,7 @@ from ml.config.scope import (
     Playstyle, PLAYSTYLE_TO_IDX, STATE_DIMS, GamePhase, PHASE_TO_IDX,
 )
 from ml.training.reward import compute_intermediate_reward, compute_terminal_reward, RewardConfig
+from ml.training.decision_exporter import DecisionExporter
 
 logger = logging.getLogger("ml.self_play")
 
@@ -384,6 +385,7 @@ def run_self_play_episode(
     playstyle: str = "midrange",
     max_steps: int = 50,
     reward_config: RewardConfig = None,
+    exporter: DecisionExporter = None,
 ) -> List[dict]:
     """Run a single self-play episode and return the trajectory.
 
@@ -395,6 +397,7 @@ def run_self_play_episode(
         playstyle: Playstyle for state encoding
         max_steps: Maximum decision steps per episode
         reward_config: Reward configuration
+        exporter: Optional DecisionExporter to record steps for distillation
 
     Returns:
         List of step dicts: [{state, action, reward, value, log_prob, done}, ...]
@@ -458,6 +461,12 @@ def run_self_play_episode(
                 "log_prob": log_prob_scalar,
                 "done": True,
             })
+            if exporter is not None:
+                exporter.record_step(
+                    game_state_snapshot=prev_snapshot,
+                    action_idx=action_idx,
+                    reward=reward,
+                )
             break
 
         # Opponent's turn
@@ -483,6 +492,12 @@ def run_self_play_episode(
                 "log_prob": log_prob_scalar,
                 "done": True,
             })
+            if exporter is not None:
+                exporter.record_step(
+                    game_state_snapshot=prev_snapshot,
+                    action_idx=action_idx,
+                    reward=reward,
+                )
             break
 
         # Intermediate reward
@@ -499,6 +514,12 @@ def run_self_play_episode(
             "log_prob": log_prob_scalar,
             "done": False,
         })
+        if exporter is not None:
+            exporter.record_step(
+                game_state_snapshot=prev_snapshot,
+                action_idx=action_idx,
+                reward=reward,
+            )
 
         # Advance turn
         state = advance_turn(state)
@@ -513,6 +534,7 @@ def collect_rollouts(
     num_episodes: int = 64,
     playstyle: str = "midrange",
     reward_config: RewardConfig = None,
+    exporter: DecisionExporter = None,
 ) -> dict:
     """Collect multiple self-play episodes into a rollout buffer.
 
@@ -523,6 +545,7 @@ def collect_rollouts(
         num_episodes: Number of episodes to collect
         playstyle: Agent playstyle
         reward_config: Reward configuration
+        exporter: Optional DecisionExporter to record decisions for distillation
 
     Returns:
         Collection statistics dict
@@ -538,12 +561,16 @@ def collect_rollouts(
         # Alternate agent seat for balanced training
         agent_seat = ep % 2
 
+        if exporter is not None:
+            exporter.begin_episode(agent_seat=agent_seat, playstyle=playstyle)
+
         trajectory = run_self_play_episode(
             agent_model=agent_model,
             opponent_policy=opponent_policy,
             agent_seat=agent_seat,
             playstyle=playstyle,
             reward_config=reward_config,
+            exporter=exporter,
         )
 
         for step in trajectory:
@@ -569,6 +596,11 @@ def collect_rollouts(
             else:
                 total_draws += 1
 
+        if exporter is not None:
+            won = trajectory and trajectory[-1]["done"] and trajectory[-1]["reward"] > 0.5
+            ep_return = sum(s["reward"] for s in trajectory)
+            exporter.end_episode(won=won, episode_return=ep_return)
+
     # Compute GAE after collection
     buffer.compute_returns_and_advantages(last_value=0.0)