385 lines
14 KiB
Python
385 lines
14 KiB
Python
import argparse
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from game import UltimateTicTacToe
|
|
from mcts import MCTS
|
|
from model import UltimateTicTacToeModel
|
|
from trainer import Trainer
|
|
|
|
|
|
DEFAULT_ARGS = {
|
|
"num_simulations": 100,
|
|
"numIters": 50,
|
|
"numEps": 20,
|
|
"epochs": 5,
|
|
"batch_size": 64,
|
|
"lr": 5e-4,
|
|
"weight_decay": 1e-4,
|
|
"replay_buffer_size": 50000,
|
|
"value_loss_weight": 1.0,
|
|
"grad_clip_norm": 5.0,
|
|
"checkpoint_path": "latest.pth",
|
|
"temperature_threshold": 10,
|
|
"root_dirichlet_alpha": 0.3,
|
|
"root_exploration_fraction": 0.25,
|
|
"arena_compare_games": 6,
|
|
"arena_accept_threshold": 0.55,
|
|
"arena_compare_simulations": 8,
|
|
}
|
|
|
|
|
|
def get_device(device_arg):
|
|
if device_arg:
|
|
return device_arg
|
|
return "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
|
|
def build_model(game, device):
|
|
return UltimateTicTacToeModel(
|
|
game.get_board_size(),
|
|
game.get_action_size(),
|
|
device,
|
|
)
|
|
|
|
|
|
def load_checkpoint(model, checkpoint_path, device, optimizer=None, required=True):
|
|
checkpoint = Path(checkpoint_path)
|
|
if not checkpoint.exists():
|
|
if required:
|
|
raise FileNotFoundError(f"Checkpoint not found: {checkpoint}")
|
|
return False
|
|
|
|
state = torch.load(checkpoint, map_location=device)
|
|
model.load_state_dict(state["state_dict"])
|
|
if optimizer is not None and "optimizer_state_dict" in state:
|
|
optimizer.load_state_dict(state["optimizer_state_dict"])
|
|
model.eval()
|
|
return True
|
|
|
|
|
|
def canonical_state(game, state, player):
|
|
board_data, active_board = state
|
|
return (game.get_canonical_board_data(board_data, player), active_board)
|
|
|
|
|
|
def apply_moves(game, moves):
|
|
state = game.get_init_board()
|
|
player = 1
|
|
for action in moves:
|
|
next_state = game.get_next_state(state, player, action, verify_move=True)
|
|
if next_state is False:
|
|
raise ValueError(f"Illegal move in sequence: {action}")
|
|
state, player = next_state
|
|
return state, player
|
|
|
|
|
|
def format_board(board_data):
|
|
symbols = {1: "X", -1: "O", 0: "."}
|
|
rows = []
|
|
for row in range(9):
|
|
cells = [symbols[int(board_data[row * 9 + col])] for col in range(9)]
|
|
groups = [" ".join(cells[idx:idx + 3]) for idx in (0, 3, 6)]
|
|
rows.append(" | ".join(groups))
|
|
if row in (2, 5):
|
|
rows.append("-" * 23)
|
|
return "\n".join(rows)
|
|
|
|
|
|
def top_policy_moves(policy, limit):
|
|
ranked = np.argsort(policy)[::-1][:limit]
|
|
return [(int(action), float(policy[action])) for action in ranked]
|
|
|
|
|
|
def parse_moves(text):
|
|
if not text:
|
|
return []
|
|
return [int(part.strip()) for part in text.split(",") if part.strip()]
|
|
|
|
|
|
def parse_action(text):
|
|
raw = text.strip().replace(",", " ").split()
|
|
if len(raw) == 1:
|
|
action = int(raw[0])
|
|
elif len(raw) == 2:
|
|
row, col = (int(value) for value in raw)
|
|
if not (0 <= row < 9 and 0 <= col < 9):
|
|
raise ValueError("Row and column must be in [0, 8].")
|
|
action = row * 9 + col
|
|
else:
|
|
raise ValueError("Enter either a flat move index or 'row col'.")
|
|
if not (0 <= action < 81):
|
|
raise ValueError("Move index must be in [0, 80].")
|
|
return action
|
|
|
|
|
|
def scalar_value(value):
|
|
return float(np.asarray(value).reshape(-1)[0])
|
|
|
|
|
|
def train_command(args):
|
|
device = get_device(args.device)
|
|
game = UltimateTicTacToe()
|
|
model = build_model(game, device)
|
|
|
|
train_args = dict(DEFAULT_ARGS)
|
|
train_args.update(
|
|
{
|
|
"num_simulations": args.num_simulations,
|
|
"numIters": args.num_iters,
|
|
"numEps": args.num_eps,
|
|
"epochs": args.epochs,
|
|
"batch_size": args.batch_size,
|
|
"lr": args.lr,
|
|
"weight_decay": args.weight_decay,
|
|
"replay_buffer_size": args.replay_buffer_size,
|
|
"value_loss_weight": args.value_loss_weight,
|
|
"grad_clip_norm": args.grad_clip_norm,
|
|
"checkpoint_path": args.checkpoint,
|
|
"temperature_threshold": args.temperature_threshold,
|
|
"root_dirichlet_alpha": args.root_dirichlet_alpha,
|
|
"root_exploration_fraction": args.root_exploration_fraction,
|
|
"arena_compare_games": args.arena_compare_games,
|
|
"arena_accept_threshold": args.arena_accept_threshold,
|
|
"arena_compare_simulations": args.arena_compare_simulations,
|
|
}
|
|
)
|
|
|
|
trainer = Trainer(game, model, train_args)
|
|
if args.resume:
|
|
load_checkpoint(model, args.checkpoint, device, optimizer=trainer.optimizer)
|
|
trainer.learn()
|
|
|
|
|
|
def eval_command(args):
|
|
device = get_device(args.device)
|
|
game = UltimateTicTacToe()
|
|
model = build_model(game, device)
|
|
load_checkpoint(model, args.checkpoint, device)
|
|
|
|
moves = parse_moves(args.moves)
|
|
state, player = apply_moves(game, moves)
|
|
current_state = canonical_state(game, state, player)
|
|
encoded = game.encode_state(current_state)
|
|
policy, value = model.predict(encoded)
|
|
legal_mask = np.array(game.get_valid_moves(state), dtype=np.float32)
|
|
policy = policy * legal_mask
|
|
if policy.sum() > 0:
|
|
policy = policy / policy.sum()
|
|
|
|
print("Board:")
|
|
print(format_board(state[0]))
|
|
print()
|
|
print(f"Side to move: {'X' if player == 1 else 'O'}")
|
|
print(f"Active small board: {state[1]}")
|
|
print(f"Model value: {scalar_value(value):.4f}")
|
|
print("Top policy moves:")
|
|
for action, prob in top_policy_moves(policy, args.top_k):
|
|
print(f" {action:2d} -> {prob:.4f}")
|
|
|
|
if args.with_mcts:
|
|
mcts_args = dict(DEFAULT_ARGS)
|
|
mcts_args.update(
|
|
{
|
|
"num_simulations": args.num_simulations,
|
|
"root_dirichlet_alpha": None,
|
|
"root_exploration_fraction": None,
|
|
}
|
|
)
|
|
root = MCTS(game, model, mcts_args).run(model, current_state, to_play=1)
|
|
action = root.select_action(temperature=0)
|
|
print(f"MCTS best move: {action}")
|
|
|
|
|
|
def ai_action(game, model, state, player, num_simulations):
|
|
current_state = canonical_state(game, state, player)
|
|
mcts_args = dict(DEFAULT_ARGS)
|
|
mcts_args.update(
|
|
{
|
|
"num_simulations": num_simulations,
|
|
"root_dirichlet_alpha": None,
|
|
"root_exploration_fraction": None,
|
|
}
|
|
)
|
|
root = MCTS(game, model, mcts_args).run(model, current_state, to_play=1)
|
|
return root.select_action(temperature=0)
|
|
|
|
|
|
def random_action(game, state):
|
|
legal_actions = [index for index, allowed in enumerate(game.get_valid_moves(state)) if allowed]
|
|
if not legal_actions:
|
|
raise ValueError("No legal actions available.")
|
|
return int(np.random.choice(legal_actions))
|
|
|
|
|
|
def load_player_model(game, checkpoint, device):
|
|
model = build_model(game, device)
|
|
load_checkpoint(model, checkpoint, device)
|
|
return model
|
|
|
|
|
|
def choose_action(game, player_kind, model, state, player, num_simulations):
|
|
if player_kind == "random":
|
|
return random_action(game, state)
|
|
return ai_action(game, model, state, player, num_simulations)
|
|
|
|
|
|
def play_match(game, x_kind, x_model, o_kind, o_model, num_simulations):
|
|
state = game.get_init_board()
|
|
player = 1
|
|
|
|
while True:
|
|
reward = game.get_reward_for_player(state, player)
|
|
if reward is not None:
|
|
if reward == 0:
|
|
return 0
|
|
return player if reward == 1 else -player
|
|
|
|
if player == 1:
|
|
action = choose_action(game, x_kind, x_model, state, player, num_simulations)
|
|
else:
|
|
action = choose_action(game, o_kind, o_model, state, player, num_simulations)
|
|
state, player = game.get_next_state(state, player, action)
|
|
|
|
|
|
def arena_command(args):
|
|
device = get_device(args.device)
|
|
game = UltimateTicTacToe()
|
|
|
|
x_model = None
|
|
o_model = None
|
|
if args.x_player == "checkpoint":
|
|
x_model = load_player_model(game, args.x_checkpoint, device)
|
|
if args.o_player == "checkpoint":
|
|
o_model = load_player_model(game, args.o_checkpoint, device)
|
|
|
|
results = {1: 0, -1: 0, 0: 0}
|
|
for _ in range(args.games):
|
|
winner = play_match(
|
|
game,
|
|
args.x_player,
|
|
x_model,
|
|
args.o_player,
|
|
o_model,
|
|
args.num_simulations,
|
|
)
|
|
results[winner] += 1
|
|
|
|
print(f"Games: {args.games}")
|
|
print(f"X ({args.x_player}) wins: {results[1]}")
|
|
print(f"O ({args.o_player}) wins: {results[-1]}")
|
|
print(f"Draws: {results[0]}")
|
|
|
|
|
|
def play_command(args):
|
|
device = get_device(args.device)
|
|
game = UltimateTicTacToe()
|
|
model = build_model(game, device)
|
|
load_checkpoint(model, args.checkpoint, device)
|
|
|
|
state = game.get_init_board()
|
|
player = 1
|
|
human_player = args.human_player
|
|
|
|
while True:
|
|
print()
|
|
print(format_board(state[0]))
|
|
print(f"Turn: {'X' if player == 1 else 'O'}")
|
|
print(f"Active small board: {state[1]}")
|
|
|
|
reward = game.get_reward_for_player(state, player)
|
|
if reward is not None:
|
|
if reward == 0:
|
|
print("Result: draw")
|
|
else:
|
|
winner = player if reward == 1 else -player
|
|
print(f"Winner: {'X' if winner == 1 else 'O'}")
|
|
return
|
|
|
|
valid_moves = game.get_valid_moves(state)
|
|
legal_actions = [index for index, allowed in enumerate(valid_moves) if allowed]
|
|
print(f"Legal moves: {legal_actions}")
|
|
|
|
if player == human_player:
|
|
while True:
|
|
try:
|
|
action = parse_action(input("Your move (index or 'row col'): "))
|
|
next_state = game.get_next_state(state, player, action, verify_move=True)
|
|
if next_state is False:
|
|
raise ValueError(f"Illegal move: {action}")
|
|
state, player = next_state
|
|
break
|
|
except ValueError as exc:
|
|
print(exc)
|
|
else:
|
|
action = ai_action(game, model, state, player, args.num_simulations)
|
|
print(f"AI move: {action}")
|
|
state, player = game.get_next_state(state, player, action)
|
|
|
|
|
|
def build_parser():
|
|
parser = argparse.ArgumentParser(description="Ultimate Tic-Tac-Toe Runner")
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
|
|
train_parser = subparsers.add_parser("train", help="Train the model with self-play")
|
|
train_parser.add_argument("--device")
|
|
train_parser.add_argument("--checkpoint", default=DEFAULT_ARGS["checkpoint_path"])
|
|
train_parser.add_argument("--resume", action="store_true")
|
|
train_parser.add_argument("--num-simulations", type=int, default=DEFAULT_ARGS["num_simulations"])
|
|
train_parser.add_argument("--num-iters", type=int, default=DEFAULT_ARGS["numIters"])
|
|
train_parser.add_argument("--num-eps", type=int, default=DEFAULT_ARGS["numEps"])
|
|
train_parser.add_argument("--epochs", type=int, default=DEFAULT_ARGS["epochs"])
|
|
train_parser.add_argument("--batch-size", type=int, default=DEFAULT_ARGS["batch_size"])
|
|
train_parser.add_argument("--lr", type=float, default=DEFAULT_ARGS["lr"])
|
|
train_parser.add_argument("--weight-decay", type=float, default=DEFAULT_ARGS["weight_decay"])
|
|
train_parser.add_argument("--replay-buffer-size", type=int, default=DEFAULT_ARGS["replay_buffer_size"])
|
|
train_parser.add_argument("--value-loss-weight", type=float, default=DEFAULT_ARGS["value_loss_weight"])
|
|
train_parser.add_argument("--grad-clip-norm", type=float, default=DEFAULT_ARGS["grad_clip_norm"])
|
|
train_parser.add_argument("--temperature-threshold", type=int, default=DEFAULT_ARGS["temperature_threshold"])
|
|
train_parser.add_argument("--root-dirichlet-alpha", type=float, default=DEFAULT_ARGS["root_dirichlet_alpha"])
|
|
train_parser.add_argument("--root-exploration-fraction", type=float, default=DEFAULT_ARGS["root_exploration_fraction"])
|
|
train_parser.add_argument("--arena-compare-games", type=int, default=DEFAULT_ARGS["arena_compare_games"])
|
|
train_parser.add_argument("--arena-accept-threshold", type=float, default=DEFAULT_ARGS["arena_accept_threshold"])
|
|
train_parser.add_argument("--arena-compare-simulations", type=int, default=DEFAULT_ARGS["arena_compare_simulations"])
|
|
train_parser.set_defaults(func=train_command)
|
|
|
|
eval_parser = subparsers.add_parser("eval", help="Inspect a checkpoint on a position")
|
|
eval_parser.add_argument("--device")
|
|
eval_parser.add_argument("--checkpoint", default=DEFAULT_ARGS["checkpoint_path"])
|
|
eval_parser.add_argument("--moves", default="", help="Comma-separated move sequence")
|
|
eval_parser.add_argument("--top-k", type=int, default=10)
|
|
eval_parser.add_argument("--with-mcts", action="store_true")
|
|
eval_parser.add_argument("--num-simulations", type=int, default=DEFAULT_ARGS["num_simulations"])
|
|
eval_parser.set_defaults(func=eval_command)
|
|
|
|
play_parser = subparsers.add_parser("play", help="Play against the checkpoint")
|
|
play_parser.add_argument("--device")
|
|
play_parser.add_argument("--checkpoint", default=DEFAULT_ARGS["checkpoint_path"])
|
|
play_parser.add_argument("--human-player", type=int, choices=[1, -1], default=1)
|
|
play_parser.add_argument("--num-simulations", type=int, default=DEFAULT_ARGS["num_simulations"])
|
|
play_parser.set_defaults(func=play_command)
|
|
|
|
arena_parser = subparsers.add_parser("arena", help="Run repeated matches between agents")
|
|
arena_parser.add_argument("--device")
|
|
arena_parser.add_argument("--games", type=int, default=20)
|
|
arena_parser.add_argument("--num-simulations", type=int, default=DEFAULT_ARGS["num_simulations"])
|
|
arena_parser.add_argument("--x-player", choices=["checkpoint", "random"], default="checkpoint")
|
|
arena_parser.add_argument("--o-player", choices=["checkpoint", "random"], default="random")
|
|
arena_parser.add_argument("--x-checkpoint", default=DEFAULT_ARGS["checkpoint_path"])
|
|
arena_parser.add_argument("--o-checkpoint", default=DEFAULT_ARGS["checkpoint_path"])
|
|
arena_parser.set_defaults(func=arena_command)
|
|
|
|
return parser
|
|
|
|
|
|
def main():
|
|
parser = build_parser()
|
|
args = parser.parse_args()
|
|
args.func(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|