587 lines
20 KiB
Python
587 lines
20 KiB
Python
"""
|
||
Commander AI Lab — Self-Play Episode Generator
|
||
════════════════════════════════════════════════
|
||
Generates synthetic 1v1 Commander game episodes for PPO training.
|
||
|
||
Since we can't run real Forge games at training speed, this module
|
||
simulates simplified game episodes using:
|
||
1. Random initial states (life, hand size, board)
|
||
2. Turn-by-turn state evolution with heuristic transitions
|
||
3. The agent's policy to select macro-actions at each decision point
|
||
4. An opponent policy (learned, random, or heuristic)
|
||
|
||
The episode generator produces trajectories of:
|
||
(state, action, reward, value, log_prob, done)
|
||
that feed directly into the RolloutBuffer for PPO training.
|
||
|
||
For real game integration (Phase 9+), the Java sim would call the
|
||
policy server and feed back actual game states — this synthetic
|
||
version allows training to begin immediately.
|
||
"""
|
||
|
||
import logging
|
||
import random
|
||
from dataclasses import dataclass, field
|
||
from typing import Dict, List, Optional, Tuple
|
||
|
||
import numpy as np
|
||
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
project_root = str(Path(__file__).parent.parent.parent)
|
||
if project_root not in sys.path:
|
||
sys.path.insert(0, project_root)
|
||
|
||
from ml.config.scope import (
|
||
NUM_ACTIONS, MacroAction, ACTION_TO_IDX, IDX_TO_ACTION,
|
||
Playstyle, PLAYSTYLE_TO_IDX, STATE_DIMS, GamePhase, PHASE_TO_IDX,
|
||
)
|
||
from ml.training.reward import compute_intermediate_reward, compute_terminal_reward, RewardConfig
|
||
|
||
logger = logging.getLogger("ml.self_play")
|
||
|
||
|
||
@dataclass
|
||
class PlayerState:
|
||
"""Simplified player state for synthetic games."""
|
||
life_total: int = 40
|
||
commander_damage_taken: int = 0
|
||
mana_available: int = 0
|
||
commander_tax: int = 0
|
||
cards_in_hand: int = 7
|
||
cards_in_graveyard: int = 0
|
||
creatures_on_battlefield: int = 0
|
||
total_power_on_board: int = 0
|
||
commander_in_play: bool = False
|
||
|
||
def to_dict(self) -> dict:
|
||
return {
|
||
"life_total": self.life_total,
|
||
"commander_damage_taken": self.commander_damage_taken,
|
||
"mana_available": self.mana_available,
|
||
"commander_tax": self.commander_tax,
|
||
"cards_in_hand": self.cards_in_hand,
|
||
"cards_in_graveyard": self.cards_in_graveyard,
|
||
"creatures_on_battlefield": self.creatures_on_battlefield,
|
||
"total_power_on_board": self.total_power_on_board,
|
||
}
|
||
|
||
|
||
@dataclass
|
||
class GameState:
|
||
"""Full synthetic game state."""
|
||
turn: int = 1
|
||
phase: str = "main_1"
|
||
active_seat: int = 0
|
||
players: list = field(default_factory=lambda: [PlayerState(), PlayerState()])
|
||
game_over: bool = False
|
||
winner: Optional[int] = None
|
||
draw: bool = False
|
||
timeout: bool = False
|
||
|
||
def to_snapshot(self) -> dict:
|
||
"""Convert to DecisionSnapshot-compatible dict."""
|
||
return {
|
||
"turn": self.turn,
|
||
"phase": self.phase,
|
||
"active_seat": self.active_seat,
|
||
"players": [p.to_dict() for p in self.players],
|
||
}
|
||
|
||
|
||
def create_random_initial_state() -> GameState:
|
||
"""Create a randomized starting game state (turn 1-3)."""
|
||
state = GameState()
|
||
|
||
# Random starting turn (simulate some early game development)
|
||
start_turn = random.randint(1, 3)
|
||
state.turn = start_turn
|
||
|
||
for i, p in enumerate(state.players):
|
||
p.life_total = 40
|
||
p.mana_available = start_turn + random.randint(0, 1) # Ramp variance
|
||
p.cards_in_hand = max(2, 7 - start_turn + random.randint(-1, 1))
|
||
|
||
# Some early board presence
|
||
if start_turn >= 2:
|
||
p.creatures_on_battlefield = random.randint(0, 2)
|
||
p.total_power_on_board = p.creatures_on_battlefield * random.randint(1, 3)
|
||
p.cards_in_graveyard = random.randint(0, 2)
|
||
|
||
state.active_seat = 0 # Agent always starts as seat 0
|
||
return state
|
||
|
||
|
||
def apply_action(
|
||
state: GameState,
|
||
action_idx: int,
|
||
player_seat: int,
|
||
) -> GameState:
|
||
"""Apply a macro-action to the game state, returning the new state.
|
||
|
||
This is a simplified heuristic simulation — not a real MTG rules engine.
|
||
It models approximate state transitions for training purposes.
|
||
"""
|
||
player = state.players[player_seat]
|
||
opponent = state.players[1 - player_seat]
|
||
action = IDX_TO_ACTION[action_idx]
|
||
|
||
# Decrement mana for actions that cost mana
|
||
mana_costs = {
|
||
MacroAction.CAST_CREATURE: random.randint(2, 5),
|
||
MacroAction.CAST_REMOVAL: random.randint(2, 4),
|
||
MacroAction.CAST_DRAW: random.randint(1, 3),
|
||
MacroAction.CAST_RAMP: random.randint(1, 3),
|
||
MacroAction.CAST_COMMANDER: max(3, 3 + player.commander_tax),
|
||
MacroAction.ATTACK_OPPONENT: 0,
|
||
MacroAction.HOLD_MANA: 0,
|
||
MacroAction.PASS: 0,
|
||
}
|
||
|
||
cost = mana_costs[action]
|
||
|
||
# Check if player has enough mana — if not, treat as pass
|
||
if cost > player.mana_available:
|
||
action = MacroAction.PASS
|
||
|
||
if action == MacroAction.CAST_CREATURE:
|
||
player.mana_available -= cost
|
||
player.cards_in_hand = max(0, player.cards_in_hand - 1)
|
||
player.creatures_on_battlefield += 1
|
||
power = random.randint(1, 4)
|
||
player.total_power_on_board += power
|
||
|
||
elif action == MacroAction.CAST_REMOVAL:
|
||
player.mana_available -= cost
|
||
player.cards_in_hand = max(0, player.cards_in_hand - 1)
|
||
if opponent.creatures_on_battlefield > 0:
|
||
opponent.creatures_on_battlefield -= 1
|
||
removed_power = random.randint(1, 4)
|
||
opponent.total_power_on_board = max(
|
||
0, opponent.total_power_on_board - removed_power
|
||
)
|
||
opponent.cards_in_graveyard += 1
|
||
|
||
elif action == MacroAction.CAST_DRAW:
|
||
player.mana_available -= cost
|
||
player.cards_in_hand = max(0, player.cards_in_hand - 1) # Cast the spell
|
||
cards_drawn = random.randint(1, 3)
|
||
player.cards_in_hand += cards_drawn
|
||
|
||
elif action == MacroAction.CAST_RAMP:
|
||
player.mana_available -= cost
|
||
player.cards_in_hand = max(0, player.cards_in_hand - 1)
|
||
# Ramp gives future mana — represented as a small immediate bonus
|
||
# (In the next turn, mana naturally increases)
|
||
|
||
elif action == MacroAction.CAST_COMMANDER:
|
||
if not player.commander_in_play:
|
||
player.mana_available -= cost
|
||
player.commander_in_play = True
|
||
player.creatures_on_battlefield += 1
|
||
power = random.randint(3, 6)
|
||
player.total_power_on_board += power
|
||
|
||
elif action == MacroAction.ATTACK_OPPONENT:
|
||
if player.creatures_on_battlefield > 0:
|
||
# Calculate combat damage
|
||
damage = player.total_power_on_board
|
||
# Opponent may block with some creatures
|
||
blocked_power = 0
|
||
if opponent.creatures_on_battlefield > 0:
|
||
blockers = random.randint(
|
||
0, min(opponent.creatures_on_battlefield,
|
||
player.creatures_on_battlefield)
|
||
)
|
||
blocked_power = blockers * random.randint(1, 3)
|
||
|
||
# Some creatures die in combat
|
||
attacking_deaths = random.randint(0, min(1, blockers))
|
||
blocking_deaths = random.randint(0, blockers)
|
||
player.creatures_on_battlefield = max(
|
||
0, player.creatures_on_battlefield - attacking_deaths
|
||
)
|
||
opponent.creatures_on_battlefield = max(
|
||
0, opponent.creatures_on_battlefield - blocking_deaths
|
||
)
|
||
|
||
player.total_power_on_board = max(
|
||
0, player.total_power_on_board - attacking_deaths * 2
|
||
)
|
||
opponent.total_power_on_board = max(
|
||
0, opponent.total_power_on_board - blocking_deaths * 2
|
||
)
|
||
opponent.cards_in_graveyard += blocking_deaths
|
||
player.cards_in_graveyard += attacking_deaths
|
||
|
||
net_damage = max(0, damage - blocked_power)
|
||
opponent.life_total -= net_damage
|
||
|
||
# Commander damage
|
||
if player.commander_in_play and net_damage > 0:
|
||
cmd_dmg = random.randint(1, min(net_damage, 5))
|
||
opponent.commander_damage_taken += cmd_dmg
|
||
|
||
elif action == MacroAction.HOLD_MANA:
|
||
pass # Intentionally hold mana — no state change
|
||
|
||
elif action == MacroAction.PASS:
|
||
pass # End step — no state change
|
||
|
||
# Check win conditions
|
||
if opponent.life_total <= 0:
|
||
state.game_over = True
|
||
state.winner = player_seat
|
||
elif opponent.commander_damage_taken >= 21:
|
||
state.game_over = True
|
||
state.winner = player_seat
|
||
|
||
return state
|
||
|
||
|
||
def advance_turn(state: GameState) -> GameState:
|
||
"""Advance to the next turn in the game.
|
||
|
||
Both players:
|
||
- Get mana = turn number (simplified)
|
||
- Draw a card
|
||
"""
|
||
state.turn += 1
|
||
|
||
# Check for timeout
|
||
if state.turn > 25:
|
||
state.game_over = True
|
||
state.timeout = True
|
||
# Whoever has more life wins, or draw
|
||
if state.players[0].life_total > state.players[1].life_total:
|
||
state.winner = 0
|
||
elif state.players[1].life_total > state.players[0].life_total:
|
||
state.winner = 1
|
||
else:
|
||
state.draw = True
|
||
return state
|
||
|
||
for p in state.players:
|
||
p.mana_available = min(state.turn + random.randint(0, 2), 15)
|
||
p.cards_in_hand += 1 # Draw for turn
|
||
|
||
# Alternate phase
|
||
state.phase = random.choice(["main_1", "main_2"])
|
||
return state
|
||
|
||
|
||
def encode_state_simple(
|
||
state: GameState,
|
||
agent_seat: int,
|
||
playstyle: str = "midrange",
|
||
) -> np.ndarray:
|
||
"""Encode a GameState into a flat feature vector for the policy network.
|
||
|
||
Uses actual state dimensions from scope.py to match the supervised
|
||
training format. Zones are filled with zero embeddings (no real cards).
|
||
"""
|
||
agent = state.players[agent_seat]
|
||
opp = state.players[1 - agent_seat]
|
||
|
||
# --- Global Features (29 dims) ---
|
||
features = []
|
||
|
||
# Per-player scalars (agent first, then opponent) × 14 each = 28
|
||
for p, is_active in [(agent, agent_seat == state.active_seat),
|
||
(opp, (1 - agent_seat) == state.active_seat)]:
|
||
features.extend([
|
||
p.life_total / 40.0,
|
||
p.commander_damage_taken / 21.0,
|
||
min(p.mana_available, 20) / 20.0,
|
||
min(p.commander_tax, 10) / 10.0,
|
||
min(p.cards_in_hand, 15) / 15.0,
|
||
min(p.cards_in_graveyard, 100) / 100.0,
|
||
min(p.creatures_on_battlefield, 30) / 30.0,
|
||
min(p.total_power_on_board, 100) / 100.0,
|
||
float(is_active),
|
||
state.turn / 25.0,
|
||
])
|
||
# Phase one-hot (4 dims)
|
||
phase_idx = {"main_1": 0, "combat": 1, "main_2": 2, "end": 3}.get(
|
||
state.phase, 0
|
||
)
|
||
phase_vec = [0.0] * 4
|
||
phase_vec[phase_idx] = 1.0
|
||
features.extend(phase_vec)
|
||
|
||
# Turn number (1 dim)
|
||
features.append(state.turn / 25.0)
|
||
|
||
# --- Zone Embeddings (6144 dims) — zeros for synthetic ---
|
||
zone_dims = STATE_DIMS.total_zone_dim # 6144
|
||
features.extend([0.0] * zone_dims)
|
||
|
||
# --- Playstyle One-Hot (4 dims) ---
|
||
ps_idx = {"aggro": 0, "control": 1, "midrange": 2, "combo": 3}.get(
|
||
playstyle, 2
|
||
)
|
||
ps_vec = [0.0] * 4
|
||
ps_vec[ps_idx] = 1.0
|
||
features.extend(ps_vec)
|
||
|
||
return np.array(features, dtype=np.float32)
|
||
|
||
|
||
class RandomPolicy:
|
||
"""Uniform random action selection (baseline opponent)."""
|
||
|
||
def select_action(self, state_vec: np.ndarray) -> Tuple[int, float]:
|
||
"""Returns (action_index, log_prob)."""
|
||
action = random.randint(0, NUM_ACTIONS - 1)
|
||
log_prob = -np.log(NUM_ACTIONS)
|
||
return action, log_prob
|
||
|
||
|
||
class HeuristicPolicy:
|
||
"""Simple heuristic opponent that makes reasonable-ish decisions."""
|
||
|
||
def select_action(self, state_vec: np.ndarray) -> Tuple[int, float]:
|
||
"""Returns (action_index, log_prob)."""
|
||
# Extract key features from the state vector
|
||
# agent_life = state_vec[0] * 40 (index 0 in the flat vector)
|
||
# agent_mana = state_vec[2] * 20 (index 2)
|
||
# agent_creatures = state_vec[6] * 30 (index 6)
|
||
# agent_hand = state_vec[4] * 15 (index 4)
|
||
|
||
mana = state_vec[2] * 20
|
||
creatures = state_vec[6] * 30
|
||
hand = state_vec[4] * 15
|
||
|
||
# Simple decision rules
|
||
weights = np.ones(NUM_ACTIONS) * 0.5
|
||
|
||
if mana >= 2:
|
||
weights[ACTION_TO_IDX[MacroAction.CAST_CREATURE]] = 3.0
|
||
if mana >= 2 and hand <= 3:
|
||
weights[ACTION_TO_IDX[MacroAction.CAST_DRAW]] = 2.5
|
||
if mana >= 1 and mana <= 4:
|
||
weights[ACTION_TO_IDX[MacroAction.CAST_RAMP]] = 2.0
|
||
if creatures >= 2:
|
||
weights[ACTION_TO_IDX[MacroAction.ATTACK_OPPONENT]] = 4.0
|
||
if mana >= 5:
|
||
weights[ACTION_TO_IDX[MacroAction.CAST_COMMANDER]] = 2.0
|
||
if mana < 2:
|
||
weights[ACTION_TO_IDX[MacroAction.PASS]] = 3.0
|
||
|
||
# Normalize to probabilities
|
||
probs = weights / weights.sum()
|
||
action = np.random.choice(NUM_ACTIONS, p=probs)
|
||
log_prob = np.log(probs[action] + 1e-8)
|
||
|
||
return int(action), float(log_prob)
|
||
|
||
|
||
def run_self_play_episode(
|
||
agent_model,
|
||
opponent_policy,
|
||
agent_seat: int = 0,
|
||
playstyle: str = "midrange",
|
||
max_steps: int = 50,
|
||
reward_config: RewardConfig = None,
|
||
) -> List[dict]:
|
||
"""Run a single self-play episode and return the trajectory.
|
||
|
||
Args:
|
||
agent_model: PolicyValueNetwork with forward_with_value method.
|
||
If None, uses random policy for the agent too.
|
||
opponent_policy: Policy object with select_action(state_vec) method
|
||
agent_seat: Agent's player seat (0 or 1)
|
||
playstyle: Playstyle for state encoding
|
||
max_steps: Maximum decision steps per episode
|
||
reward_config: Reward configuration
|
||
|
||
Returns:
|
||
List of step dicts: [{state, action, reward, value, log_prob, done}, ...]
|
||
"""
|
||
try:
|
||
import torch
|
||
import torch.nn.functional as F
|
||
except ImportError:
|
||
raise ImportError("PyTorch required for self-play")
|
||
|
||
state = create_random_initial_state()
|
||
trajectory = []
|
||
opp_seat = 1 - agent_seat
|
||
cfg = reward_config or RewardConfig()
|
||
|
||
for step in range(max_steps):
|
||
if state.game_over:
|
||
break
|
||
|
||
prev_snapshot = state.to_snapshot()
|
||
|
||
# Agent's turn
|
||
state_vec = encode_state_simple(state, agent_seat, playstyle)
|
||
|
||
if agent_model is not None:
|
||
# Detect device from model parameters (e.g. cuda:0)
|
||
_device = next(agent_model.parameters()).device
|
||
state_tensor = torch.from_numpy(state_vec).unsqueeze(0).to(_device)
|
||
with torch.no_grad():
|
||
logits, value = agent_model.forward_with_value(state_tensor)
|
||
probs = F.softmax(logits, dim=-1)
|
||
dist = torch.distributions.Categorical(probs)
|
||
action = dist.sample()
|
||
log_prob = dist.log_prob(action)
|
||
|
||
action_idx = action.item()
|
||
value_scalar = value.item()
|
||
log_prob_scalar = log_prob.item()
|
||
else:
|
||
action_idx = random.randint(0, NUM_ACTIONS - 1)
|
||
value_scalar = 0.0
|
||
log_prob_scalar = -np.log(NUM_ACTIONS)
|
||
|
||
# Apply agent's action
|
||
state = apply_action(state, action_idx, agent_seat)
|
||
|
||
if state.game_over:
|
||
curr_snapshot = state.to_snapshot()
|
||
reward = compute_intermediate_reward(
|
||
prev_snapshot, curr_snapshot, IDX_TO_ACTION[action_idx].value,
|
||
agent_seat, cfg
|
||
)
|
||
won = state.winner == agent_seat
|
||
reward += compute_terminal_reward(won, state.draw, state.timeout, cfg)
|
||
|
||
trajectory.append({
|
||
"state": state_vec,
|
||
"action": action_idx,
|
||
"reward": reward,
|
||
"value": value_scalar,
|
||
"log_prob": log_prob_scalar,
|
||
"done": True,
|
||
})
|
||
break
|
||
|
||
# Opponent's turn
|
||
opp_state_vec = encode_state_simple(state, opp_seat, "midrange")
|
||
opp_action, _ = opponent_policy.select_action(opp_state_vec)
|
||
state = apply_action(state, opp_action, opp_seat)
|
||
|
||
curr_snapshot = state.to_snapshot()
|
||
|
||
if state.game_over:
|
||
reward = compute_intermediate_reward(
|
||
prev_snapshot, curr_snapshot, IDX_TO_ACTION[action_idx].value,
|
||
agent_seat, cfg
|
||
)
|
||
won = state.winner == agent_seat
|
||
reward += compute_terminal_reward(won, state.draw, state.timeout, cfg)
|
||
|
||
trajectory.append({
|
||
"state": state_vec,
|
||
"action": action_idx,
|
||
"reward": reward,
|
||
"value": value_scalar,
|
||
"log_prob": log_prob_scalar,
|
||
"done": True,
|
||
})
|
||
break
|
||
|
||
# Intermediate reward
|
||
reward = compute_intermediate_reward(
|
||
prev_snapshot, curr_snapshot, IDX_TO_ACTION[action_idx].value,
|
||
agent_seat, cfg
|
||
)
|
||
|
||
trajectory.append({
|
||
"state": state_vec,
|
||
"action": action_idx,
|
||
"reward": reward,
|
||
"value": value_scalar,
|
||
"log_prob": log_prob_scalar,
|
||
"done": False,
|
||
})
|
||
|
||
# Advance turn
|
||
state = advance_turn(state)
|
||
|
||
return trajectory
|
||
|
||
|
||
def collect_rollouts(
|
||
agent_model,
|
||
opponent_policy,
|
||
buffer,
|
||
num_episodes: int = 64,
|
||
playstyle: str = "midrange",
|
||
reward_config: RewardConfig = None,
|
||
) -> dict:
|
||
"""Collect multiple self-play episodes into a rollout buffer.
|
||
|
||
Args:
|
||
agent_model: PolicyValueNetwork (or None for random agent)
|
||
opponent_policy: Opponent policy object
|
||
buffer: RolloutBuffer to fill
|
||
num_episodes: Number of episodes to collect
|
||
playstyle: Agent playstyle
|
||
reward_config: Reward configuration
|
||
|
||
Returns:
|
||
Collection statistics dict
|
||
"""
|
||
buffer.reset()
|
||
total_wins = 0
|
||
total_losses = 0
|
||
total_draws = 0
|
||
total_steps = 0
|
||
total_rewards = 0.0
|
||
|
||
for ep in range(num_episodes):
|
||
# Alternate agent seat for balanced training
|
||
agent_seat = ep % 2
|
||
|
||
trajectory = run_self_play_episode(
|
||
agent_model=agent_model,
|
||
opponent_policy=opponent_policy,
|
||
agent_seat=agent_seat,
|
||
playstyle=playstyle,
|
||
reward_config=reward_config,
|
||
)
|
||
|
||
for step in trajectory:
|
||
buffer.add(
|
||
state=step["state"],
|
||
action=step["action"],
|
||
reward=step["reward"],
|
||
value=step["value"],
|
||
log_prob=step["log_prob"],
|
||
done=step["done"],
|
||
)
|
||
total_rewards += step["reward"]
|
||
|
||
total_steps += len(trajectory)
|
||
|
||
# Track outcomes
|
||
if trajectory and trajectory[-1]["done"]:
|
||
final_reward = trajectory[-1]["reward"]
|
||
if final_reward > 0.5:
|
||
total_wins += 1
|
||
elif final_reward < -0.5:
|
||
total_losses += 1
|
||
else:
|
||
total_draws += 1
|
||
|
||
# Compute GAE after collection
|
||
buffer.compute_returns_and_advantages(last_value=0.0)
|
||
|
||
return {
|
||
"episodes": num_episodes,
|
||
"steps": total_steps,
|
||
"avg_steps_per_ep": total_steps / max(num_episodes, 1),
|
||
"total_reward": total_rewards,
|
||
"avg_reward_per_ep": total_rewards / max(num_episodes, 1),
|
||
"wins": total_wins,
|
||
"losses": total_losses,
|
||
"draws": total_draws,
|
||
"win_rate": total_wins / max(num_episodes, 1),
|
||
"buffer_size": buffer.size,
|
||
}
|