wiki-tcg/backend/train_nn.py

import os
import random
import uuid
import numpy as np
from collections import deque
from dotenv import load_dotenv
load_dotenv()

from card import compute_deck_type
from ai import AIPersonality, choose_cards, choose_plan
from game import PlayerState, GameState, action_play_card, action_sacrifice, action_end_turn
from simulate import get_simulation_cards, _make_instances, MAX_TURNS
from nn import NeuralNet, NeuralPlayer

NN_WEIGHTS_PATH = os.path.join(os.path.dirname(__file__), "nn_weights.json")

P1 = "p1"
P2 = "p2"

FIXED_PERSONALITIES = [p for p in AIPersonality if p != AIPersonality.ARBITRARY]


# ==================== Game runner ====================

def _build_player(pid: str, name: str, cards: list, difficulty: int, personality: AIPersonality) -> PlayerState:
  deck = choose_cards(cards, difficulty, personality)
  instances = _make_instances(deck)
  random.shuffle(instances)
  p = PlayerState(
    user_id=pid, username=name,
    deck_type=compute_deck_type(deck) or "Balanced",
    deck=instances,
  )
  return p


def run_episode(
  p1_state: PlayerState,
  p2_state: PlayerState,
  p1_ctrl,   # (player, opponent) -> MovePlan
  p2_ctrl,   # (player, opponent) -> MovePlan
) -> str | None:
  """Returns winner_id (P1 or P2) or None on timeout."""
  p1_state.increment_energy_cap()
  p2_state.increment_energy_cap()
  p1_state.refill_energy()
  p1_state.draw_to_full()

  state = GameState(
    game_id=str(uuid.uuid4()),
    players={P1: p1_state, P2: p2_state},
    player_order=[P1, P2],
    active_player_id=P1,
    phase="main",
    turn=1,
  )
  ctrls = {P1: p1_ctrl, P2: p2_ctrl}

  for _ in range(MAX_TURNS):
    if state.result:
      break
    active_id = state.active_player_id
    player   = state.players[active_id]
    opponent = state.players[state.opponent_id(active_id)]

    plan = ctrls[active_id](player, opponent)

    for slot in plan.sacrifice_slots:
      if player.board[slot] is not None:
        action_sacrifice(state, slot)

    plays = list(plan.plays)
    random.shuffle(plays)
    for card, slot in plays:
      hand_idx = next((i for i, c in enumerate(player.hand) if c is card), None)
      if hand_idx is None or player.board[slot] is not None or card.cost > player.energy:
        continue
      action_play_card(state, hand_idx, slot)

    action_end_turn(state)

  return state.result.winner_id if state.result else None


# ==================== Training loop ====================

def train(
  n_episodes: int = 20_000,
  self_play_start: int = 5_000,
  self_play_max_frac: float = 0.4,
  lr: float = 1e-3,
  opp_difficulty: int = 10,
  temperature: float = 1.0,
  batch_size: int = 50,
  save_every: int = 5_000,
  save_path: str = NN_WEIGHTS_PATH,
) -> NeuralNet:
  cards = get_simulation_cards()

  if os.path.exists(save_path):
    print(f"Resuming from {save_path}")
    net = NeuralNet.load(save_path)
  else:
    print("Initializing new network")
    net = NeuralNet(seed=42)

  recent_outcomes: deque[int] = deque(maxlen=1000)  # rolling window for win rate display
  baseline = 0.0      # EMA of recent outcomes; subtracted before each update
  baseline_alpha = 0.99   # decay — roughly a 100-episode window

  batch_gw = [np.zeros_like(w) for w in net.weights]
  batch_gb = [np.zeros_like(b) for b in net.biases]
  batch_count = 0

  for episode in range(1, n_episodes + 1):
    # Ramp self-play fraction linearly from 0 to self_play_max_frac
    if episode >= self_play_start:
      progress = (episode - self_play_start) / max(1, n_episodes - self_play_start)
      self_play_prob = self_play_max_frac * progress
    else:
      self_play_prob = 0.0

    # Randomly decide who goes first (NN is always P1 for simplicity)
    nn_goes_first = random.random() < 0.5

    if random.random() < self_play_prob:
      # ---- Self-play ----
      nn1 = NeuralPlayer(net, training=True, temperature=temperature)
      nn2 = NeuralPlayer(net, training=True, temperature=temperature)

      p1_state = _build_player(P1, "NN1", cards, 10, AIPersonality.BALANCED)
      p2_state = _build_player(P2, "NN2", cards, 10, AIPersonality.BALANCED)

      if not nn_goes_first:
        p1_state, p2_state = p2_state, p1_state

      winner = run_episode(p1_state, p2_state, nn1.choose_plan, nn2.choose_plan)
      p1_outcome = 1.0 if winner == P1 else -1.0
      baseline = baseline_alpha * baseline + (1 - baseline_alpha) * p1_outcome

      for player_grads in [nn1.compute_grads(p1_outcome - baseline),
                  nn2.compute_grads(-p1_outcome - baseline)]:
        if player_grads is not None:
          gw, gb = player_grads
          for i in range(len(batch_gw)):
            batch_gw[i] += gw[i]
            batch_gb[i] += gb[i]
          batch_count += 1

    else:
      # ---- NN vs fixed opponent ----
      opp_personality = random.choice(FIXED_PERSONALITIES)
      nn_player = NeuralPlayer(net, training=True, temperature=temperature)
      opp_ctrl  = lambda p, o, pers=opp_personality, diff=opp_difficulty: choose_plan(p, o, pers, diff)

      if nn_goes_first:
        nn_id  = P1
        p1_state = _build_player(P1, "NN",  cards, 10, AIPersonality.BALANCED)
        p2_state = _build_player(P2, "OPP", cards, opp_difficulty, opp_personality)
        winner = run_episode(p1_state, p2_state, nn_player.choose_plan, opp_ctrl)
      else:
        nn_id  = P2
        p1_state = _build_player(P1, "OPP", cards, opp_difficulty, opp_personality)
        p2_state = _build_player(P2, "NN",  cards, 10, AIPersonality.BALANCED)
        winner = run_episode(p1_state, p2_state, opp_ctrl, nn_player.choose_plan)

      nn_outcome = 1.0 if winner == nn_id else -1.0
      player_grads = nn_player.compute_grads(nn_outcome - baseline)
      baseline = baseline_alpha * baseline + (1 - baseline_alpha) * nn_outcome

      if player_grads is not None:
        gw, gb = player_grads
        for i in range(len(batch_gw)):
          batch_gw[i] += gw[i]
          batch_gb[i] += gb[i]
        batch_count += 1

      recent_outcomes.append(1 if winner == nn_id else 0)

    if batch_count >= batch_size:
      for i in range(len(batch_gw)):
        batch_gw[i] /= batch_count
        batch_gb[i] /= batch_count
      net.adam_update(batch_gw, batch_gb, lr=lr)
      batch_gw = [np.zeros_like(w) for w in net.weights]
      batch_gb = [np.zeros_like(b) for b in net.biases]
      batch_count = 0

    if episode % 1000 == 0 or episode == n_episodes:
      wr = sum(recent_outcomes) / len(recent_outcomes) if recent_outcomes else 0.0
      print(f"[{episode:>6}/{n_episodes}]  win rate (last {len(recent_outcomes)}): {wr:.1%}  "
          f"self-play frac: {self_play_prob:.0%}", flush=True)

    if episode % save_every == 0:
      net.save(save_path)
      print(f"  → saved to {save_path}")

  net.save(save_path)
  wr = sum(recent_outcomes) / len(recent_outcomes) if recent_outcomes else 0.0
  print(f"Done. Final win rate (last {len(recent_outcomes)}): {wr:.1%}")
  return net


if __name__ == "__main__":
  train()