diff --git a/ml/training/forge_episode_generator.py b/ml/training/forge_episode_generator.py index 561fd2f..539d210 100644 --- a/ml/training/forge_episode_generator.py +++ b/ml/training/forge_episode_generator.py @@ -34,6 +34,7 @@ from ml.training.reward import ( compute_intermediate_reward, compute_terminal_reward, ) +from ml.training.decision_exporter import DecisionExporter, ExporterConfig logger = logging.getLogger("ml.forge_episode_generator") @@ -51,6 +52,8 @@ class ForgeTransition: log_prob: float value: float done: bool + # Raw snapshot retained so the exporter can reconstruct game-state context + snapshot: Optional[Dict[str, Any]] = None @dataclass @@ -61,6 +64,8 @@ class ForgeEpisode: winner: Optional[int] = None total_turns: int = 0 deck_names: List[str] = field(default_factory=list) + agent_seat: int = 0 + playstyle: str = "midrange" @property def length(self) -> int: @@ -70,6 +75,10 @@ class ForgeEpisode: def total_reward(self) -> float: return sum(t.reward for t in self.transitions) + @property + def won(self) -> bool: + return self.winner == self.agent_seat + # ══════════════════════════════════════════════════════════ # JSONL Parser — reads Forge ML decision snapshots @@ -95,9 +104,44 @@ def parse_forge_jsonl(path: Path) -> List[Dict[str, Any]]: return snapshots +def _make_minimal_state(snap: Dict[str, Any]) -> Dict[str, Any]: + """Build a minimal prev/curr state dict from a Forge snapshot. + + compute_intermediate_reward() expects a dict with a 'players' list + where each player has 'life_total', 'cards_in_hand', + 'creatures_on_battlefield', etc. Forge snapshots carry a 'players' + array directly; we forward it as-is and fall back to scalar fields + when the array is absent. + """ + if "players" in snap: + return {"players": snap["players"]} + + # Fallback: reconstruct a two-player view from flat snapshot fields + agent_seat = snap.get("active_seat", 0) + opp_seat = 1 - agent_seat + players = [None, None] + + def _player(life_key: str, hand_key: str, creatures_key: str) -> Dict: + return { + "life_total": snap.get(life_key, 40), + "cards_in_hand": snap.get(hand_key, 0), + "creatures_on_battlefield": snap.get(creatures_key, 0), + "cards_in_graveyard": 0, + "mana_available": snap.get("mana_available", 0), + "commander_damage_taken": 0, + "total_power_on_board": 0, + } + + players[agent_seat] = _player("agent_life", "agent_hand_size", "agent_creatures") + players[opp_seat] = _player("opp_life", "opp_hand_size", "opp_creatures") + return {"players": players} + + def snapshots_to_episode( snapshots: List[Dict[str, Any]], reward_config: Optional[RewardConfig] = None, + agent_seat: int = 0, + playstyle: str = "midrange", ) -> ForgeEpisode: """Convert a sequence of Forge snapshots into a ForgeEpisode. @@ -111,6 +155,8 @@ def snapshots_to_episode( episode = ForgeEpisode( game_id=first.get("game_id", "unknown"), deck_names=first.get("deck_names", []), + agent_seat=agent_seat, + playstyle=playstyle, ) rc = reward_config or RewardConfig() @@ -130,20 +176,26 @@ def snapshots_to_episode( except (ValueError, KeyError): action_idx = ACTION_TO_IDX[MacroAction.PASS] - # Reward + # Reward — use correct reward.py signatures is_done = i == len(snapshots) - 1 if is_done and "outcome" in snap: + outcome = snap["outcome"] reward = compute_terminal_reward( - won=snap["outcome"].get("won", False), - turns=snap.get("turn", 1), - life_remaining=snap["outcome"].get("life_remaining", 0), + won=outcome.get("won", False), + draw=outcome.get("draw", False), + timeout=outcome.get("timeout", False), config=rc, ) else: + # Build minimal prev/curr dicts for intermediate reward + prev_snap = snapshots[i - 1] if i > 0 else snap + prev_state = _make_minimal_state(prev_snap) + curr_state = _make_minimal_state(snap) reward = compute_intermediate_reward( - life_delta=snap.get("life_delta", 0), - cards_drawn=snap.get("cards_drawn", 0), - creatures_played=snap.get("creatures_played", 0), + prev_state=prev_state, + curr_state=curr_state, + action_taken=action_label, + agent_seat=agent_seat, config=rc, ) @@ -154,6 +206,7 @@ def snapshots_to_episode( log_prob=snap.get("log_prob", 0.0), value=snap.get("value", 0.0), done=is_done, + snapshot=snap, )) episode.total_turns = snapshots[-1].get("turn", len(snapshots)) @@ -163,6 +216,61 @@ def snapshots_to_episode( return episode +# ══════════════════════════════════════════════════════════ +# DecisionExporter adapter for Forge episodes +# ══════════════════════════════════════════════════════════ + +def _export_forge_episode(episode: ForgeEpisode, exporter: DecisionExporter) -> None: + """Feed a completed ForgeEpisode into the DecisionExporter. + + Translates each ForgeTransition back into the game_state_snapshot + format expected by DecisionExporter.record_step(), then finalises + the episode so the exporter can apply its win/length gates. + + The raw snapshot stored on each ForgeTransition is used directly + when available; otherwise a minimal dict is synthesised from the + transition's action index and reward so the schema is always valid. + """ + exporter.begin_episode( + agent_seat=episode.agent_seat, + playstyle=episode.playstyle, + ) + + for t in episode.transitions: + if t.snapshot is not None: + game_state = _make_minimal_state(t.snapshot) + # Carry through turn/phase/active_seat if present + game_state["turn"] = t.snapshot.get("turn", 1) + game_state["phase"] = t.snapshot.get("phase", "main_1") + game_state["active_seat"] = t.snapshot.get("active_seat", episode.agent_seat) + else: + # Minimal fallback — exporter only needs players list for schema + game_state = { + "turn": 1, + "phase": "main_1", + "active_seat": episode.agent_seat, + "players": [ + {"life_total": 40, "cards_in_hand": 0, "creatures_on_battlefield": 0, + "cards_in_graveyard": 0, "mana_available": 0, + "commander_damage_taken": 0, "total_power_on_board": 0, "commander_tax": 0}, + {"life_total": 40, "cards_in_hand": 0, "creatures_on_battlefield": 0, + "cards_in_graveyard": 0, "mana_available": 0, + "commander_damage_taken": 0, "total_power_on_board": 0, "commander_tax": 0}, + ], + } + + exporter.record_step( + game_state_snapshot=game_state, + action_idx=t.action, + reward=t.reward, + ) + + exporter.end_episode( + won=episode.won, + episode_return=episode.total_reward, + ) + + # ══════════════════════════════════════════════════════════ # Async Episode Queue (producer-consumer) # ══════════════════════════════════════════════════════════ @@ -271,6 +379,7 @@ async def ppo_consumer( queue: ForgeEpisodeQueue, buffer_size: int = 64, on_batch_ready=None, + exporter: Optional[DecisionExporter] = None, ) -> int: """Consumer coroutine: pulls episodes from the queue and batches them for PPO training updates. @@ -279,6 +388,11 @@ async def ppo_consumer( queue: ForgeEpisodeQueue to consume from buffer_size: number of episodes per training batch on_batch_ready: callback(episodes) when a batch is full + exporter: Optional DecisionExporter — when provided, every + consumed episode is fed into the exporter so winning + Forge games accumulate in the distillation JSONL pool. + Call exporter.flush() after the pipeline completes to + write the JSONL file. Returns total episodes consumed. """ @@ -292,6 +406,14 @@ async def ppo_consumer( break continue + # Feed into distillation exporter (win-gate applied inside exporter) + if exporter is not None: + try: + _export_forge_episode(episode, exporter) + except Exception as exc: + logger.warning("[CONSUMER] Exporter error for episode %s: %s", + episode.game_id, exc) + batch.append(episode) total_consumed += 1 @@ -322,12 +444,24 @@ async def run_forge_pipeline( max_queue: int = 256, on_batch_ready=None, reward_config: Optional[RewardConfig] = None, -) -> Dict[str, int]: + exporter: Optional[DecisionExporter] = None, +) -> Dict[str, Any]: """Run the full Forge producer → PPO consumer pipeline. The Forge Java batch must be started separately (via forge_runner.py or MultiThreadBatchRunner). This function handles the Python side: reading JSONL output and feeding episodes into the PPO trainer. + + If an exporter is provided, winning Forge episodes are recorded and + flushed to a PPO JSONL file at the end of the pipeline run. The + output path is returned in the stats dict under "exporter_file". + + Example:: + + from ml.training.decision_exporter import DecisionExporter, ExporterConfig + exporter = DecisionExporter(ExporterConfig(model_version="gen-1")) + stats = await run_forge_pipeline(num_games=128, exporter=exporter) + print(stats["exporter_file"]) # data/results/ml-decisions-ppo-.jsonl """ queue = ForgeEpisodeQueue(max_size=max_queue) @@ -336,16 +470,25 @@ async def run_forge_pipeline( reward_config=reward_config) ) consumer_task = asyncio.create_task( - ppo_consumer(queue, buffer_size, on_batch_ready) + ppo_consumer(queue, buffer_size, on_batch_ready, exporter=exporter) ) produced = await producer_task queue.close() consumed = await consumer_task - stats = { + # Flush exporter to disk and capture output path + exporter_file = None + if exporter is not None: + exporter_file = exporter.flush() + if exporter_file: + logger.info("[PIPELINE] Distillation export: %s stats=%s", + exporter_file, exporter.stats) + + stats: Dict[str, Any] = { "produced": produced, "consumed": consumed, + "exporter_file": exporter_file, **queue.stats, } logger.info("[PIPELINE] Complete: %s", stats)