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:
1 parent
9338ac881b
commit
98e1092ebc
1 file changed
+153
-10
@@ -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)
|
||||
|
||||
Reference in new issue
Block a user