""" Commander AI Lab — State Encoder ═════════════════════════════════ Converts raw decision snapshots (JSONL) into fixed-size tensor vectors suitable for neural network input. Pipeline: 1. Load card embeddings from mtg-embeddings (768-dim, ~32K cards) 2. For each decision snapshot: a. Encode global features (life, mana, turn, phase, etc.) b. For each zone per player: mean-pool card embeddings c. Encode playstyle as one-hot vector d. Concatenate into single fixed-length vector Output dimensions: global_features (29) + zone_embeddings (6144) + playstyle (4) = 6177 """ import json import logging import os from pathlib import Path from typing import Dict, List, Optional, Tuple import numpy as np logger = logging.getLogger("ml.encoder") # Import scope config import sys sys.path.insert(0, str(Path(__file__).parent.parent.parent)) from ml.config.scope import ( STATE_DIMS, PLAYSTYLE_TO_IDX, PHASE_TO_IDX, Playstyle, GamePhase, GAME_SCOPE, ) class CardEmbeddingIndex: """ Loads and caches card embeddings from the minimaxir/mtg-embeddings dataset. Provides name → 768-dim vector lookup. Uses the same NPZ cache as coach/embeddings.py for consistency. """ def __init__(self, embeddings_dir: str = None): self.embeddings_dir = embeddings_dir or str( Path(__file__).parent.parent.parent / "embeddings" ) self._name_to_idx: Dict[str, int] = {} self._embeddings: Optional[np.ndarray] = None self._names: List[str] = [] self._loaded = False self._zero_vec = np.zeros(STATE_DIMS.card_embedding_dim, dtype=np.float32) def load(self) -> bool: """Load embeddings from NPZ cache. Searches multiple known locations and handles two NPZ formats: - Coach format: keys 'vectors', 'names' (saved by coach/embeddings.py) - ML format: keys 'embeddings', 'names' (legacy) If no NPZ is found anywhere, auto-downloads from HuggingFace via the coach embedding index and converts to NPZ. """ project_root = Path(__file__).parent.parent.parent candidates = [ str(project_root / "data" / "mtg-embeddings.npz"), str(project_root / "data" / "mtg_embeddings.npz"), os.path.join(self.embeddings_dir, "mtg_embeddings.npz"), os.path.join(self.embeddings_dir, "mtg-embeddings.npz"), str(project_root / "embeddings" / "mtg_embeddings.npz"), str(project_root / "embeddings" / "mtg-embeddings.npz"), ] npz_path = None for candidate in candidates: if os.path.exists(candidate): npz_path = candidate break # Auto-download if nothing found if not npz_path: logger.info("Embeddings NPZ not found — attempting auto-download...") npz_path = self._auto_download(project_root) if not npz_path: return False try: data = np.load(npz_path, allow_pickle=True) # Handle both NPZ formats: # Coach format has 'vectors', ML legacy has 'embeddings' if "vectors" in data: self._embeddings = data["vectors"].astype(np.float32) elif "embeddings" in data: self._embeddings = data["embeddings"].astype(np.float32) else: logger.error("NPZ at %s has no 'vectors' or 'embeddings' key. Keys: %s", npz_path, list(data.keys())) return False self._names = list(data["names"]) self._name_to_idx = { name.lower(): i for i, name in enumerate(self._names) } self._loaded = True logger.info( "Loaded %d card embeddings (%d-dim) from %s", len(self._names), self._embeddings.shape[1], npz_path, ) return True except Exception as e: logger.error("Failed to load embeddings: %s", e) return False @staticmethod def _auto_download(project_root: Path) -> Optional[str]: """Download embeddings from HuggingFace via the coach module. Uses coach/embeddings.py MTGEmbeddingIndex which handles downloading the parquet from HuggingFace and converting to NPZ. Returns the NPZ path on success, or None on failure. """ try: import sys pr = str(project_root) if pr not in sys.path: sys.path.insert(0, pr) from coach.embeddings import MTGEmbeddingIndex from coach.config import EMBEDDINGS_NPZ idx = MTGEmbeddingIndex() idx.download_parquet() idx.convert_parquet_to_npz() npz_str = str(EMBEDDINGS_NPZ) if os.path.exists(npz_str): logger.info("Auto-downloaded embeddings to %s", npz_str) return npz_str logger.error("Auto-download completed but NPZ not found at %s", npz_str) return None except Exception as e: logger.error("Auto-download failed: %s", e) return None def get_embedding(self, card_name: str) -> np.ndarray: """Get embedding vector for a card name. Returns zero vector if not found.""" if not self._loaded: return self._zero_vec idx = self._name_to_idx.get(card_name.lower()) if idx is not None: return self._embeddings[idx] # Try fuzzy match: strip trailing set codes, articles clean = card_name.lower().split("|")[0].strip() idx = self._name_to_idx.get(clean) if idx is not None: return self._embeddings[idx] return self._zero_vec def mean_pool_zone(self, card_names: List[str]) -> np.ndarray: """Mean-pool embeddings for a list of card names (a game zone).""" if not card_names: return self._zero_vec vectors = [self.get_embedding(name) for name in card_names] return np.mean(vectors, axis=0).astype(np.float32) @property def is_loaded(self) -> bool: return self._loaded class StateEncoder: """ Encodes a decision snapshot JSON dict into a fixed-size numpy vector. Input: a dict from the JSONL file (one decision snapshot) Output: np.ndarray of shape (6177,) Components: [0:29] - Global features (life, mana, turn, phase, etc.) [29:6173] - Zone embeddings (4 zones × 2 players × 768) [6173:6177] - Playstyle one-hot (4 dims) Global feature layout per player block (14 features × 2 players + 1 turn = 29): idx 0 life (normalised /40) idx 1 cmdr_dmg (normalised /21) idx 2 mana (normalised /20, clamped) idx 3 cmdr_tax (normalised /10, clamped) idx 4 hand_size (normalised /15, clamped) idx 5 graveyard_size (normalised /100, clamped) idx 6 creatures (normalised /30, clamped) idx 7 total_power (normalised /100, clamped) └─ real Forge value when present; falls back to creatures*3 heuristic idx 8 is_active (binary 0/1) idx 9-12 phase one-hot (4 dims: main_1/combat/main_2/end) idx 13 lands (normalised /15, clamped) --- repeat for player 1 at idx 14-27 --- idx 28 turn (normalised /max_turns) """ def __init__(self, card_index: CardEmbeddingIndex): self.card_index = card_index self.dim = STATE_DIMS def encode(self, decision: dict, playstyle: str = "midrange") -> np.ndarray: """ Encode a single decision snapshot into a state vector. Args: decision: Dict loaded from JSONL line playstyle: Deck archetype string (aggro/control/midrange/combo) Returns: np.ndarray of shape (total_state_dim,) """ players = decision.get("players", []) if len(players) < 2: while len(players) < 2: players.append(self._empty_player(len(players))) # 1. Global features global_vec = self._encode_global(decision, players) # 2. Zone embeddings zone_vec = self._encode_zones(players) # 3. Playstyle one-hot style_vec = self._encode_playstyle(playstyle) # Concatenate state = np.concatenate([global_vec, zone_vec, style_vec]) assert state.shape[0] == self.dim.total_state_dim, ( f"State dim mismatch: {state.shape[0]} != {self.dim.total_state_dim}" ) return state.astype(np.float32) def _encode_global(self, decision: dict, players: list) -> np.ndarray: """Encode global scalar features. Shape: (29,) Per-player block (14 features each, 2 players = 28) + 1 turn feature. Index 7 priority for total_power: 1. p['total_power'] — real Forge value from Step 4 / _build_encoder_snapshot() 2. p['creatures'] * 3 — legacy heuristic for old JSONL snapshots """ features = [] for p in players[:2]: # idx 0: life life = p.get("life", 40) features.append(life / 40.0) # idx 1: commander damage taken cmdr_dmg = p.get("cmdr_dmg", 0) features.append(cmdr_dmg / 21.0) # idx 2: mana available mana = p.get("mana", 0) features.append(min(mana / 20.0, 1.0)) # idx 3: commander tax cmdr_tax = p.get("cmdr_tax", 0) features.append(min(cmdr_tax / 10.0, 1.0)) # idx 4: cards in hand hand_size = len(p.get("hand", [])) features.append(min(hand_size / 15.0, 1.0)) # idx 5: graveyard size grave_size = len(p.get("graveyard", [])) features.append(min(grave_size / 100.0, 1.0)) # idx 6: creatures on battlefield creatures = p.get("creatures", 0) features.append(min(creatures / 30.0, 1.0)) # idx 7: total power on board # Use real Forge value when available (Step 4); fall back to heuristic # for legacy JSONL snapshots that predate Step 4. raw_power = p.get("total_power", 0) total_power = raw_power if raw_power > 0 else creatures * 3 features.append(min(total_power / 100.0, 1.0)) # idx 8: is active player is_active = 1.0 if p.get("seat") == decision.get("active_seat") else 0.0 features.append(is_active) # idx 9-12: phase one-hot (4 dims) phase_str = decision.get("phase", "main_1") phase_oh = [0.0] * 4 try: phase_enum = GamePhase(phase_str) phase_oh[PHASE_TO_IDX[phase_enum]] = 1.0 except (ValueError, KeyError): phase_oh[0] = 1.0 # default to main_1 features.extend(phase_oh) # idx 13: land count lands = p.get("lands", 0) features.append(min(lands / 15.0, 1.0)) # idx 28: turn number (shared, 1 feature) turn = decision.get("turn", 1) features.append(min(turn / GAME_SCOPE.max_turns, 1.0)) return np.array(features, dtype=np.float32) def _encode_zones(self, players: list) -> np.ndarray: """ Encode zone contents as mean-pooled card embeddings. Shape: (4 zones × 2 players × 768) = (6144,) Zone embedding boundary: each zone slot is exactly ``zone_pool_dim`` (768) floats. Empty zones produce a zero vector of that size. This fixed-size contract is relied upon by every downstream consumer (PolicyNetwork, PPOTrainer, dataset builder). See ``StateDimensions`` in ``ml/config/scope.py`` for the canonical definition. """ zone_keys = ["hand", "battlefield", "graveyard", "command_zone"] vectors = [] for p in players[:2]: for zone_key in zone_keys: cards = p.get(zone_key, []) pooled = self.card_index.mean_pool_zone(cards) vectors.append(pooled) return np.concatenate(vectors) def _encode_playstyle(self, playstyle_str: str) -> np.ndarray: """One-hot encode the playstyle. Shape: (4,)""" vec = np.zeros(self.dim.playstyle_dim, dtype=np.float32) try: style = Playstyle(playstyle_str.lower()) vec[PLAYSTYLE_TO_IDX[style]] = 1.0 except (ValueError, KeyError): vec[PLAYSTYLE_TO_IDX[Playstyle.MIDRANGE]] = 1.0 return vec def _empty_player(self, seat: int) -> dict: """Create an empty player dict for padding.""" return { "seat": seat, "life": 40, "cmdr_dmg": 0, "mana": 0, "cmdr_tax": 0, "creatures": 0, "total_power": 0, "total_toughness": 0, "artifacts": 0, "enchantments": 0, "lands": 0, "hand": [], "battlefield": [], "graveyard": [], "command_zone": [], } def encode_decisions_file( jsonl_path: str, card_index: CardEmbeddingIndex, max_samples: int = None, ) -> Tuple[np.ndarray, List[dict]]: """ Encode all decisions from a JSONL file into a state matrix. Args: jsonl_path: Path to ml-decisions-*.jsonl file card_index: Loaded card embedding index max_samples: Maximum number of samples to encode (None = all) Returns: (states, raw_decisions) where: states: np.ndarray of shape (N, 6177) raw_decisions: list of original JSON dicts (for action labels) """ encoder = StateEncoder(card_index) states = [] raw = [] with open(jsonl_path, "r") as f: for i, line in enumerate(f): if max_samples and i >= max_samples: break line = line.strip() if not line: continue try: decision = json.loads(line) playstyle = decision.get("archetype", "midrange") state_vec = encoder.encode(decision, playstyle) states.append(state_vec) raw.append(decision) except (json.JSONDecodeError, Exception) as e: logger.warning("Skipping line %d: %s", i, e) continue if not states: return np.empty((0, STATE_DIMS.total_state_dim), dtype=np.float32), [] return np.stack(states), raw