feat: wire DecisionExporter into Forge pipeline + fix reward call signatures

- Add _export_forge_episode() adapter: translates ForgeEpisode transitions
  back into DecisionExporter.record_step() calls so Forge-sourced wins flow
  into the distillation JSONL pipeline (closes Phase 1 for real games)
- ppo_consumer() now accepts an optional `exporter: DecisionExporter` arg;
  begin_episode / record_step / end_episode are called around each episode
  as it is dequeued
- run_forge_pipeline() forwards the exporter kwarg down to ppo_consumer and
  calls exporter.flush() after the pipeline completes, returning the output
  path in the stats dict under "exporter_file"
- Fix snapshots_to_episode(): compute_terminal_reward() was called with
  non-existent kwargs `turns=` and `life_remaining=`; compute_intermediate_reward()
  was called with non-existent kwargs `life_delta=`, `cards_drawn=`,
  `creatures_played=`.  Both now use the correct signatures from reward.py
  (prev_state/curr_state/action_taken for intermediate; won/draw/timeout for
  terminal).  Forge snapshots carry `players` arrays so we reconstruct minimal
  prev/curr dicts for the intermediate path.

Ref: docs/CLOSED-LOOP-DISTILLATION.md Phase 1
This commit is contained in:
KoalaTrapLord committed 2026-04-03 17:09:05 -04:00
1 parent 9338ac881b
commit 98e1092ebc
1 file changed
+153 -10
+153 -10
View File
@@ -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-<id>.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)