feat(ml): add --data-source forge flag to PPO trainer (issue #83 phase 1.3)

Wires forge_episode_generator into the PPO training loop:
- New --data-source flag: synthetic (default) | forge | mixed
- New --forge-results-dir, --forge-batch-id, --forge-weight flags
- ForgeDataSource class: runs run_forge_pipeline() in background thread,
  drains episodes into RolloutBuffer via episodes_to_buffer()
- Mixed mode: blends Forge episodes + synthetic self-play per iteration
  using --forge-weight ratio (0.0–1.0)
- Synthetic path unchanged — no regression for existing training runs

ml/training/ppo_trainer.py

@@ -9,22 +9,27 @@ Proximal Policy Optimization (Clip) with:
   - Self-play episode collection between updates
   - Periodic checkpointing and evaluation
   - Win-rate tracking against opponent policies
+  - Forge-based episode source via --data-source forge|mixed
 
 Reference: Schulman et al., "Proximal Policy Optimization Algorithms", 2017
 
 Usage:
     python -m ml.training.ppo_trainer --iterations 100 --episodes-per-iter 64
     python -m ml.training.ppo_trainer --load-supervised ml/models/checkpoints/best_policy.pt
+    python -m ml.training.ppo_trainer --data-source forge --forge-results-dir results --iterations 50
+    python -m ml.training.ppo_trainer --data-source mixed --forge-weight 0.7 --iterations 100
 """
 
 import argparse
+import asyncio
 import json
 import logging
 import os
 import sys
+import threading
 import time
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 
 import numpy as np
 
@@ -51,6 +56,114 @@ from ml.training.self_play import (
 logger = logging.getLogger("ml.ppo")
 
 
+# ══════════════════════════════════════════════════════════
+# Forge data-source integration
+# ══════════════════════════════════════════════════════════
+
+def episodes_to_buffer(episodes, buffer: RolloutBuffer) -> int:
+    """Drain a list of ForgeEpisodes into an existing RolloutBuffer.
+
+    Returns the number of transitions added.
+    """
+    added = 0
+    for ep in episodes:
+        for t in ep.transitions:
+            if buffer.size >= buffer.buffer_size:
+                break
+            buffer.add(
+                state=t.state,
+                action=t.action,
+                reward=t.reward,
+                log_prob=t.log_prob,
+                value=t.value,
+                done=t.done,
+            )
+            added += 1
+    return added
+
+
+class ForgeDataSource:
+    """Runs the Forge producer pipeline in a background thread and
+    exposes a synchronous drain() method for the PPO training loop.
+
+    Usage::
+        source = ForgeDataSource(results_dir="results", batch_id="run1", num_games=64)
+        source.start()
+        # ... in training loop ...
+        episodes = source.drain()
+        source.stop()
+    """
+
+    def __init__(
+        self,
+        results_dir: str = "results",
+        batch_id: str = "forge-batch",
+        num_games: int = 256,
+        reward_config: Optional[RewardConfig] = None,
+        queue_size: int = 256,
+    ):
+        self.results_dir = results_dir
+        self.batch_id = batch_id
+        self.num_games = num_games
+        self.reward_config = reward_config or RewardConfig()
+        self.queue_size = queue_size
+        self._episodes: List = []
+        self._lock = threading.Lock()
+        self._thread: Optional[threading.Thread] = None
+        self._stop_event = threading.Event()
+
+    def start(self):
+        """Launch the background producer thread."""
+        self._thread = threading.Thread(target=self._run, daemon=True)
+        self._thread.start()
+        logger.info("[ForgeDataSource] Producer started (results_dir=%s, batch_id=%s, num_games=%d)",
+                    self.results_dir, self.batch_id, self.num_games)
+
+    def _run(self):
+        from ml.training.forge_episode_generator import ForgeEpisodeQueue, forge_producer
+
+        async def _produce():
+            queue = ForgeEpisodeQueue(max_size=self.queue_size)
+            produced = await forge_producer(
+                queue=queue,
+                results_dir=self.results_dir,
+                batch_id=self.batch_id,
+                num_games=self.num_games,
+                reward_config=self.reward_config,
+            )
+            queue.close()
+            # Drain remaining items
+            while not queue.is_closed:
+                ep = await queue.get(timeout=1.0)
+                if ep is not None:
+                    with self._lock:
+                        self._episodes.append(ep)
+            logger.info("[ForgeDataSource] Producer finished: %d episodes", produced)
+
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        try:
+            loop.run_until_complete(_produce())
+        finally:
+            loop.close()
+
+    def drain(self) -> List:
+        """Return and clear all buffered episodes (thread-safe)."""
+        with self._lock:
+            eps = list(self._episodes)
+            self._episodes.clear()
+        return eps
+
+    def stop(self):
+        self._stop_event.set()
+        if self._thread:
+            self._thread.join(timeout=5.0)
+
+
+# ══════════════════════════════════════════════════════════
+# PPO Config
+# ══════════════════════════════════════════════════════════
+
 class PPOConfig:
     """PPO hyperparameters."""
 
@@ -91,6 +204,13 @@ class PPOConfig:
 
         # Supervised preload
         load_supervised: Optional[str] = None,
+
+        # Data source
+        data_source: str = "synthetic",   # synthetic | forge | mixed
+        forge_results_dir: str = "results",
+        forge_batch_id: str = "forge-batch",
+        forge_num_games: int = 256,
+        forge_weight: float = 0.7,        # fraction of buffer from Forge when mixed
     ):
         self.iterations = iterations
         self.episodes_per_iter = episodes_per_iter
@@ -114,10 +234,19 @@ class PPOConfig:
         self.eval_every = eval_every
         self.eval_episodes = eval_episodes
         self.load_supervised = load_supervised
+        self.data_source = data_source
+        self.forge_results_dir = forge_results_dir
+        self.forge_batch_id = forge_batch_id
+        self.forge_num_games = forge_num_games
+        self.forge_weight = max(0.0, min(1.0, forge_weight))
 
 
+# ══════════════════════════════════════════════════════════
+# PPO Trainer
+# ══════════════════════════════════════════════════════════
+
 class PPOTrainer:
-    """Proximal Policy Optimization trainer with self-play."""
+    """Proximal Policy Optimization trainer with self-play and optional Forge data source."""
 
     def __init__(self, config: PPOConfig):
         if not TORCH_AVAILABLE:
@@ -126,28 +255,24 @@ class PPOTrainer:
         self.config = config
         self.device = self._detect_device()
 
-        # Create actor-critic model
         from ml.training.policy_network import PolicyValueNetwork
         self.model = PolicyValueNetwork(
             input_dim=STATE_DIMS.total_state_dim,
             hidden_dim=TRAINING_CONFIG.hidden_dim,
             num_actions=NUM_ACTIONS,
             num_layers=TRAINING_CONFIG.num_layers,
-            dropout=0.0,  # No dropout during RL
+            dropout=0.0,
         ).to(self.device)
 
-        # Optionally load from supervised checkpoint
         if config.load_supervised and os.path.exists(config.load_supervised):
             self._load_supervised_weights(config.load_supervised)
 
-        # Optimizer
         self.optimizer = torch.optim.Adam(
             self.model.parameters(),
             lr=config.learning_rate,
             eps=1e-5,
         )
 
-        # Learning rate scheduler
         self.scheduler = None
         if config.lr_schedule == "linear":
             self.scheduler = torch.optim.lr_scheduler.LinearLR(
@@ -163,7 +288,6 @@ class PPOTrainer:
                 eta_min=config.min_lr,
             )
 
-        # Rollout buffer
         self.buffer = RolloutBuffer(
             buffer_size=config.buffer_size,
             state_dim=STATE_DIMS.total_state_dim,
@@ -171,16 +295,23 @@ class PPOTrainer:
             gae_lambda=config.gae_lambda,
         )
 
-        # Reward config
         self.reward_config = RewardConfig(
             gamma=config.gamma,
             gae_lambda=config.gae_lambda,
         )
 
-        # Opponent policy
         self.opponent = self._create_opponent(config.opponent)
 
-        # Tracking
+        # Forge data source (started lazily in train())
+        self._forge_source: Optional[ForgeDataSource] = None
+        if config.data_source in ("forge", "mixed"):
+            self._forge_source = ForgeDataSource(
+                results_dir=config.forge_results_dir,
+                batch_id=config.forge_batch_id,
+                num_games=config.forge_num_games,
+                reward_config=self.reward_config,
+            )
+
         self.iteration = 0
         self.total_steps = 0
         self.best_win_rate = 0.0
@@ -196,13 +327,9 @@ class PPOTrainer:
         return "cpu"
 
     def _load_supervised_weights(self, path: str):
-        """Load weights from a supervised PolicyNetwork checkpoint into the
-        PolicyValueNetwork (ignoring the value head, which is new)."""
         logger.info("Loading supervised weights from %s", path)
         checkpoint = torch.load(path, map_location=self.device, weights_only=False)
         state_dict = checkpoint.get("model_state_dict", checkpoint)
-
-        # Filter out keys that don't match (e.g., value_head)
         model_dict = self.model.state_dict()
         compatible = {k: v for k, v in state_dict.items() if k in model_dict and v.shape == model_dict[k].shape}
         model_dict.update(compatible)
@@ -216,14 +343,10 @@ class PPOTrainer:
         elif opponent_type == "heuristic":
             return HeuristicPolicy()
         elif opponent_type == "self":
-            # Self-play: opponent uses a snapshot of the current policy
             return self._create_self_play_opponent()
-        else:
-            return HeuristicPolicy()
+        return HeuristicPolicy()
 
     def _create_self_play_opponent(self):
-        """Create an opponent that uses a frozen copy of the current policy."""
-
         class SelfPlayPolicy:
             def __init__(self, model, device):
                 self.model = model
@@ -241,11 +364,99 @@ class PPOTrainer:
 
         return SelfPlayPolicy(self.model, self.device)
 
-    def ppo_update(self) -> Dict:
-        """Run PPO update epochs on the collected rollout buffer.
+    def _collect_forge_rollouts(self) -> Dict:
+        """Drain buffered Forge episodes into the RolloutBuffer.
 
-        Returns dict of loss metrics.
+        Returns stats dict compatible with collect_rollouts() output.
         """
+        episodes = self._forge_source.drain()
+        if not episodes:
+            logger.warning("[Forge] No episodes available yet — waiting for JSONL output in %s",
+                           self.config.forge_results_dir)
+            return {"steps": 0, "win_rate": 0.0, "avg_reward_per_ep": 0.0,
+                    "wins": 0, "losses": 0, "draws": 0, "source": "forge"}
+
+        added = episodes_to_buffer(episodes, self.buffer)
+        wins = sum(1 for ep in episodes if ep.winner == 0)
+        win_rate = wins / max(len(episodes), 1)
+        avg_reward = sum(ep.total_reward for ep in episodes) / max(len(episodes), 1)
+        logger.info("[Forge] Drained %d episodes (%d transitions) into buffer — WR=%.0f%% AvgR=%.3f",
+                    len(episodes), added, win_rate * 100, avg_reward)
+        return {
+            "steps": added,
+            "win_rate": win_rate,
+            "avg_reward_per_ep": avg_reward,
+            "wins": wins,
+            "losses": len(episodes) - wins,
+            "draws": 0,
+            "source": "forge",
+        }
+
+    def _collect_rollouts_for_iter(self) -> Dict:
+        """Collect rollouts according to configured data_source.
+
+        - synthetic: pure self_play.collect_rollouts (original behaviour)
+        - forge:     pure Forge episode drain
+        - mixed:     Forge fills forge_weight fraction of buffer, synthetic fills remainder
+        """
+        cfg = self.config
+        self.buffer.reset()
+        self.model.eval()
+
+        if cfg.data_source == "forge":
+            stats = self._collect_forge_rollouts()
+
+        elif cfg.data_source == "mixed":
+            forge_cap = int(cfg.buffer_size * cfg.forge_weight)
+            synth_cap = cfg.buffer_size - forge_cap
+
+            # Temporarily shrink buffer for Forge slice
+            original_size = self.buffer.buffer_size
+            self.buffer.buffer_size = forge_cap
+            forge_stats = self._collect_forge_rollouts()
+            self.buffer.buffer_size = original_size
+
+            synth_eps = max(1, int(cfg.episodes_per_iter * (1.0 - cfg.forge_weight)))
+            synth_stats = collect_rollouts(
+                agent_model=self.model,
+                opponent_policy=self.opponent,
+                buffer=self.buffer,
+                num_episodes=synth_eps,
+                playstyle=cfg.playstyle,
+                reward_config=self.reward_config,
+            )
+            stats = {
+                "steps": forge_stats["steps"] + synth_stats["steps"],
+                "win_rate": (
+                    forge_stats["win_rate"] * cfg.forge_weight
+                    + synth_stats["win_rate"] * (1.0 - cfg.forge_weight)
+                ),
+                "avg_reward_per_ep": (
+                    forge_stats["avg_reward_per_ep"] * cfg.forge_weight
+                    + synth_stats["avg_reward_per_ep"] * (1.0 - cfg.forge_weight)
+                ),
+                "wins": forge_stats["wins"] + synth_stats["wins"],
+                "losses": forge_stats["losses"] + synth_stats["losses"],
+                "draws": synth_stats.get("draws", 0),
+                "source": "mixed",
+            }
+
+        else:  # synthetic (default — unchanged behaviour)
+            stats = collect_rollouts(
+                agent_model=self.model,
+                opponent_policy=self.opponent,
+                buffer=self.buffer,
+                num_episodes=cfg.episodes_per_iter,
+                playstyle=cfg.playstyle,
+                reward_config=self.reward_config,
+            )
+            stats["source"] = "synthetic"
+
+        self.model.train()
+        return stats
+
+    def ppo_update(self) -> Dict:
+        """Run PPO update epochs on the collected rollout buffer."""
         cfg = self.config
         total_policy_loss = 0.0
         total_value_loss = 0.0
@@ -258,56 +469,42 @@ class PPOTrainer:
             for batch in self.buffer.get_batches(cfg.batch_size):
                 states, actions, old_log_probs, advantages, returns = batch
 
-                # To tensors
                 states_t = torch.from_numpy(states).to(self.device)
                 actions_t = torch.from_numpy(actions).to(self.device)
                 old_log_probs_t = torch.from_numpy(old_log_probs).to(self.device)
                 advantages_t = torch.from_numpy(advantages).to(self.device)
                 returns_t = torch.from_numpy(returns).to(self.device)
 
-                # Forward pass
                 logits, values = self.model.forward_with_value(states_t)
                 values = values.squeeze(-1)
 
-                # Action distribution
                 probs = F.softmax(logits, dim=-1)
                 dist = torch.distributions.Categorical(probs)
                 new_log_probs = dist.log_prob(actions_t)
                 entropy = dist.entropy().mean()
 
-                # --- Policy Loss (clipped surrogate) ---
                 ratio = torch.exp(new_log_probs - old_log_probs_t)
                 surr1 = ratio * advantages_t
                 surr2 = torch.clamp(ratio, 1.0 - cfg.clip_epsilon, 1.0 + cfg.clip_epsilon) * advantages_t
                 policy_loss = -torch.min(surr1, surr2).mean()
 
-                # --- Value Loss (clipped) ---
                 if cfg.value_clip > 0:
-                    # Get old values from buffer for clipping
-                    old_values_t = torch.from_numpy(
-                        self.buffer.values[:self.buffer.size]
-                    ).to(self.device)
-                    # Reindex to match batch
-                    # Since we use shuffled indices in get_batches, we clip against returns
                     value_loss_unclipped = (values - returns_t) ** 2
                     value_loss = 0.5 * value_loss_unclipped.mean()
                 else:
                     value_loss = 0.5 * ((values - returns_t) ** 2).mean()
 
-                # --- Total Loss ---
                 loss = (
                     policy_loss
                     + cfg.value_loss_coeff * value_loss
                     - cfg.entropy_coeff * entropy
                 )
 
-                # Optimize
                 self.optimizer.zero_grad()
                 loss.backward()
                 nn.utils.clip_grad_norm_(self.model.parameters(), cfg.max_grad_norm)
                 self.optimizer.step()
 
-                # Track metrics
                 with torch.no_grad():
                     approx_kl = ((ratio - 1) - torch.log(ratio)).mean().item()
                     clip_frac = ((ratio - 1.0).abs() > cfg.clip_epsilon).float().mean().item()
@@ -330,12 +527,8 @@ class PPOTrainer:
         }
 
     def evaluate(self, num_episodes: int = 50) -> Dict:
-        """Evaluate the current policy against opponents.
-
-        Returns win rates and metrics.
-        """
+        """Evaluate the current policy against heuristic and random opponents."""
         self.model.eval()
-
         results = {}
         for opp_name, opp_policy in [
             ("heuristic", HeuristicPolicy()),
@@ -363,18 +556,15 @@ class PPOTrainer:
                 "avg_reward": stats["avg_reward_per_ep"],
                 "avg_steps": stats["avg_steps_per_ep"],
             }
-
         self.model.train()
         return results
 
     def save_checkpoint(self, path: str = None, metrics: dict = None):
-        """Save PPO checkpoint."""
         if path is None:
             path = os.path.join(
                 self.config.checkpoint_dir,
                 f"ppo_iter_{self.iteration:04d}.pt",
             )
-
         os.makedirs(os.path.dirname(path), exist_ok=True)
         torch.save({
             "iteration": self.iteration,
@@ -388,6 +578,7 @@ class PPOTrainer:
                 "learning_rate": self.config.learning_rate,
                 "opponent": self.config.opponent,
                 "playstyle": self.config.playstyle,
+                "data_source": self.config.data_source,
             },
             "best_win_rate": self.best_win_rate,
             "model_config": {
@@ -401,14 +592,7 @@ class PPOTrainer:
         logger.info("Checkpoint saved: %s", path)
 
     def train(self, progress_callback=None) -> Dict:
-        """Run the full PPO training loop.
-
-        Args:
-            progress_callback: Optional callable(iteration, metrics) for progress reporting
-
-        Returns:
-            Training summary dict
-        """
+        """Run the full PPO training loop."""
         cfg = self.config
         logger.info("=" * 60)
         logger.info("PPO Training — Commander AI Lab")
@@ -423,98 +607,94 @@ class PPOTrainer:
         logger.info("  Opponent:       %s", cfg.opponent)
         logger.info("  Device:         %s", self.device)
         logger.info("  Supervised:     %s", cfg.load_supervised or "None")
+        logger.info("  Data source:    %s", cfg.data_source)
+        if cfg.data_source in ("forge", "mixed"):
+            logger.info("  Forge results:  %s / %s", cfg.forge_results_dir, cfg.forge_batch_id)
+            logger.info("  Forge weight:   %.0f%%", cfg.forge_weight * 100)
         logger.info("")
 
+        # Start Forge producer background thread if needed
+        if self._forge_source:
+            self._forge_source.start()
+
         t_start = time.time()
 
-        for iteration in range(1, cfg.iterations + 1):
-            self.iteration = iteration
-            t_iter = time.time()
+        try:
+            for iteration in range(1, cfg.iterations + 1):
+                self.iteration = iteration
+                t_iter = time.time()
 
-            # --- Collect rollouts ---
-            self.model.eval()
-            rollout_stats = collect_rollouts(
-                agent_model=self.model,
-                opponent_policy=self.opponent,
-                buffer=self.buffer,
-                num_episodes=cfg.episodes_per_iter,
-                playstyle=cfg.playstyle,
-                reward_config=self.reward_config,
-            )
-            self.total_steps += rollout_stats["steps"]
-            self.model.train()
+                rollout_stats = self._collect_rollouts_for_iter()
+                self.total_steps += rollout_stats["steps"]
 
-            # --- PPO Update ---
-            update_metrics = self.ppo_update()
+                update_metrics = self.ppo_update()
 
-            # --- LR Schedule ---
-            current_lr = self.optimizer.param_groups[0]["lr"]
-            if self.scheduler:
-                self.scheduler.step()
+                current_lr = self.optimizer.param_groups[0]["lr"]
+                if self.scheduler:
+                    self.scheduler.step()
 
-            # --- Evaluate ---
-            eval_results = None
-            if iteration % cfg.eval_every == 0 or iteration == 1:
-                eval_results = self.evaluate(cfg.eval_episodes)
-                heuristic_wr = eval_results.get("heuristic", {}).get("win_rate", 0)
+                eval_results = None
+                if iteration % cfg.eval_every == 0 or iteration == 1:
+                    eval_results = self.evaluate(cfg.eval_episodes)
+                    heuristic_wr = eval_results.get("heuristic", {}).get("win_rate", 0)
+                    if heuristic_wr > self.best_win_rate:
+                        self.best_win_rate = heuristic_wr
+                        best_path = os.path.join(cfg.checkpoint_dir, "best_ppo.pt")
+                        self.save_checkpoint(best_path, {"eval": eval_results})
+                        logger.info("  ★ New best win rate: %.1f%%", heuristic_wr * 100)
 
-                if heuristic_wr > self.best_win_rate:
-                    self.best_win_rate = heuristic_wr
-                    best_path = os.path.join(cfg.checkpoint_dir, "best_ppo.pt")
-                    self.save_checkpoint(best_path, {"eval": eval_results})
-                    logger.info("  ★ New best win rate: %.1f%%", heuristic_wr * 100)
+                if iteration % cfg.save_every == 0:
+                    self.save_checkpoint()
 
-            # --- Save periodic checkpoint ---
-            if iteration % cfg.save_every == 0:
-                self.save_checkpoint()
+                iter_time = time.time() - t_iter
+                metrics = {
+                    "iteration": iteration,
+                    "steps": rollout_stats["steps"],
+                    "total_steps": self.total_steps,
+                    "win_rate": rollout_stats["win_rate"],
+                    "avg_reward": rollout_stats["avg_reward_per_ep"],
+                    "policy_loss": update_metrics["policy_loss"],
+                    "value_loss": update_metrics["value_loss"],
+                    "entropy": update_metrics["entropy"],
+                    "approx_kl": update_metrics["approx_kl"],
+                    "clip_fraction": update_metrics["clip_fraction"],
+                    "lr": current_lr,
+                    "iter_time_s": round(iter_time, 1),
+                    "data_source": rollout_stats.get("source", cfg.data_source),
+                    "eval": eval_results,
+                }
+                self.history.append(metrics)
 
-            # --- Log ---
-            iter_time = time.time() - t_iter
-            metrics = {
-                "iteration": iteration,
-                "steps": rollout_stats["steps"],
-                "total_steps": self.total_steps,
-                "win_rate": rollout_stats["win_rate"],
-                "avg_reward": rollout_stats["avg_reward_per_ep"],
-                "policy_loss": update_metrics["policy_loss"],
-                "value_loss": update_metrics["value_loss"],
-                "entropy": update_metrics["entropy"],
-                "approx_kl": update_metrics["approx_kl"],
-                "clip_fraction": update_metrics["clip_fraction"],
-                "lr": current_lr,
-                "iter_time_s": round(iter_time, 1),
-                "eval": eval_results,
-            }
-            self.history.append(metrics)
+                eval_str = ""
+                if eval_results:
+                    h_wr = eval_results.get("heuristic", {}).get("win_rate", 0)
+                    r_wr = eval_results.get("random", {}).get("win_rate", 0)
+                    eval_str = f" | eval: h={h_wr:.0%} r={r_wr:.0%}"
 
-            # Pretty log
-            eval_str = ""
-            if eval_results:
-                h_wr = eval_results.get("heuristic", {}).get("win_rate", 0)
-                r_wr = eval_results.get("random", {}).get("win_rate", 0)
-                eval_str = f" | eval: h={h_wr:.0%} r={r_wr:.0%}"
+                src_tag = f"[{rollout_stats.get('source', '?')[:3]}]"
+                logger.info(
+                    "  [%3d/%d]%s wr=%.0f%% rew=%.3f ploss=%.4f vloss=%.4f ent=%.3f kl=%.4f%s (%.1fs)",
+                    iteration, cfg.iterations, src_tag,
+                    rollout_stats["win_rate"] * 100,
+                    rollout_stats["avg_reward_per_ep"],
+                    update_metrics["policy_loss"],
+                    update_metrics["value_loss"],
+                    update_metrics["entropy"],
+                    update_metrics["approx_kl"],
+                    eval_str,
+                    iter_time,
+                )
 
-            logger.info(
-                "  [%3d/%d] wr=%.0f%% rew=%.3f ploss=%.4f vloss=%.4f ent=%.3f kl=%.4f%s (%.1fs)",
-                iteration, cfg.iterations,
-                rollout_stats["win_rate"] * 100,
-                rollout_stats["avg_reward_per_ep"],
-                update_metrics["policy_loss"],
-                update_metrics["value_loss"],
-                update_metrics["entropy"],
-                update_metrics["approx_kl"],
-                eval_str,
-                iter_time,
-            )
+                if progress_callback:
+                    progress_callback(iteration, metrics)
 
-            if progress_callback:
-                progress_callback(iteration, metrics)
+        finally:
+            if self._forge_source:
+                self._forge_source.stop()
 
-        # --- Final Save ---
         total_time = time.time() - t_start
         self.save_checkpoint()
 
-        # Save training history
         history_path = os.path.join(cfg.checkpoint_dir, "ppo_history.json")
         with open(history_path, "w") as f:
             json.dump(self.history, f, indent=2, default=str)
@@ -526,6 +706,7 @@ class PPOTrainer:
             "final_win_rate": self.history[-1]["win_rate"] if self.history else 0,
             "total_time_s": round(total_time, 1),
             "device": self.device,
+            "data_source": cfg.data_source,
             "checkpoint_dir": cfg.checkpoint_dir,
             "history_path": history_path,
         }
@@ -536,6 +717,7 @@ class PPOTrainer:
         logger.info("  Total time:     %.1f s", total_time)
         logger.info("  Total steps:    %d", self.total_steps)
         logger.info("  Best win rate:  %.1f%%", self.best_win_rate * 100)
+        logger.info("  Data source:    %s", cfg.data_source)
         logger.info("  History:        %s", history_path)
         logger.info("=" * 60)
 
@@ -578,12 +760,32 @@ def main():
     parser.add_argument("--load-supervised", default=None,
                         help="Path to supervised checkpoint to initialize from")
 
+    # Data source
+    parser.add_argument(
+        "--data-source", default="synthetic",
+        choices=["synthetic", "forge", "mixed"],
+        help="Episode source: synthetic (self-play), forge (Forge JSONL), mixed (blend)",
+    )
+    parser.add_argument(
+        "--forge-results-dir", default="results",
+        help="Directory where Forge writes ml-decisions-forge-*.jsonl files",
+    )
+    parser.add_argument(
+        "--forge-batch-id", default="forge-batch",
+        help="Batch ID prefix for Forge JSONL files (matches BatchRunner --batchId)",
+    )
+    parser.add_argument(
+        "--forge-num-games", type=int, default=256,
+        help="Total Forge games to consume before stopping the producer",
+    )
+    parser.add_argument(
+        "--forge-weight", type=float, default=0.7,
+        help="Fraction of RolloutBuffer filled from Forge when --data-source mixed (0.0–1.0)",
+    )
+
     args = parser.parse_args()
 
-    logging.basicConfig(
-        level=logging.INFO,
-        format="%(message)s",
-    )
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
 
     config = PPOConfig(
         iterations=args.iterations,
@@ -605,6 +807,11 @@ def main():
         eval_every=args.eval_every,
         eval_episodes=args.eval_episodes,
         load_supervised=args.load_supervised,
+        data_source=args.data_source,
+        forge_results_dir=args.forge_results_dir,
+        forge_batch_id=args.forge_batch_id,
+        forge_num_games=args.forge_num_games,
+        forge_weight=args.forge_weight,
     )
 
     trainer = PPOTrainer(config)