Connect 4 on the AlphaZero Framework Implementation Plan

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [ ]) syntax for tracking.

Goal: Validate the AlphaZero framework on Connect 4 — train a network that beats a from-scratch minimax-depth-8 opponent in ≥80% of 200 games (alternating colors), reusing the entire Sub-project 1 framework unchanged.

Architecture: Add one new Game subclass (Connect4), one new evaluation opponent (Connect4MinimaxOpponent with α-β pruning), and small generalizations to Trainer.run() (eval-opponent injection) and the CLI (--game flag). The four-layer framework (Game / NN / MCTS / Trainer) is unchanged.

Tech Stack: Same as Sub-project 1 — Python 3.11+, PyTorch 2.2+ (MPS/CUDA portable), NumPy, pytest. No new dependencies.

Spec: docs/superpowers/specs/2026-05-15-connect4-design.md

File Structure

Files this plan creates or modifies:

src/alphazero/games/connect4.py                ← new (~200 lines)
src/alphazero/opponents/__init__.py            ← new (empty package marker)
src/alphazero/opponents/connect4_minimax.py    ← new (~200 lines)
src/alphazero/trainer.py                       ← small change (eval-opponent injection)
src/alphazero/cli.py                           ← small change (--game flag)
configs/connect4.toml                          ← new
tests/test_game_connect4.py                    ← new (~250 lines)
tests/test_connect4_minimax.py                 ← new (~150 lines)
tests/test_e2e_connect4.py                     ← new (~50 lines)

Tasks proceed in dependency order: game first, then opponent (which uses the game's terminal detection), then the framework wiring (Trainer/CLI), then the kill criterion.

Conventions

Files: - Create: src/alphazero/games/connect4.py - Test: tests/test_game_connect4.py

tests/test_game_connect4.py:

import numpy as np
import pytest

from alphazero.games.connect4 import Connect4


@pytest.fixture
def c4() -> Connect4:
    return Connect4()


def test_input_shape_and_action_size(c4: Connect4):
    assert c4.input_shape == (3, 6, 7)
    assert c4.action_size == 7


def test_initial_state_is_empty(c4: Connect4):
    s = c4.initial_state()
    assert s.shape == (6, 7)
    assert s.dtype == np.int8
    assert np.all(s == 0)


def test_current_player_starts_at_plus_one(c4: Connect4):
    assert c4.current_player(c4.initial_state()) == 1


def test_apply_drops_piece_to_bottom_row_of_column(c4: Connect4):
    """First piece in a column should land at row 5 (the bottom)."""
    s = c4.initial_state()
    s = c4.apply(s, action=3)  # player +1 drops in column 3
    assert s[5, 3] == 1
    # All other cells still empty
    assert np.count_nonzero(s) == 1
    # Next player is -1
    assert c4.current_player(s) == -1


def test_apply_stacks_pieces_within_a_column(c4: Connect4):
    s = c4.initial_state()
    s = c4.apply(s, action=3)  # +1 at (5, 3)
    s = c4.apply(s, action=3)  # -1 stacks on top at (4, 3)
    s = c4.apply(s, action=3)  # +1 at (3, 3)
    assert s[5, 3] == 1
    assert s[4, 3] == -1
    assert s[3, 3] == 1
    assert s[2, 3] == 0  # next slot is still empty


def test_apply_does_not_mutate_input(c4: Connect4):
    s = c4.initial_state()
    s_copy = s.copy()
    _ = c4.apply(s, 0)
    np.testing.assert_array_equal(s, s_copy)


def test_apply_raises_on_full_column(c4: Connect4):
    s = c4.initial_state()
    # Fill column 0 (6 pieces)
    for _ in range(6):
        s = c4.apply(s, 0)
    assert s[0, 0] != 0  # top of column 0 is occupied
    with pytest.raises(ValueError, match="full"):
        c4.apply(s, 0)


def test_legal_actions_mask_initial_state_all_seven_legal(c4: Connect4):
    mask = c4.legal_actions_mask(c4.initial_state())
    assert mask.shape == (7,)
    assert mask.dtype == bool
    assert mask.all()


def test_legal_actions_mask_after_filling_one_column(c4: Connect4):
    s = c4.initial_state()
    for _ in range(6):
        s = c4.apply(s, 3)
    mask = c4.legal_actions_mask(s)
    assert mask.sum() == 6
    assert not mask[3]  # column 3 is full

pytest tests/test_game_connect4.py -v

Expected: ImportError on alphazero.games.connect4.

"""Connect 4: 6×7 board with gravity. Pieces drop to the lowest empty cell of the
chosen column.

State representation:
    np.ndarray shape (6, 7), dtype int8, values in {+1, 0, -1}.
    Row 0 is the top of the board; row 5 is the bottom (where pieces land).
    +1 = first player (red); -1 = second player (yellow); 0 = empty.

Action space: integers 0..6 — the column to drop a piece into.
"""
from __future__ import annotations

import numpy as np

from .base import Game

State = np.ndarray

ROWS = 6
COLS = 7


class Connect4(Game):
    @property
    def input_shape(self) -> tuple[int, ...]:
        return (3, ROWS, COLS)

    @property
    def action_size(self) -> int:
        return COLS

    def initial_state(self) -> State:
        return np.zeros((ROWS, COLS), dtype=np.int8)

    def current_player(self, state: State) -> int:
        moves = int(np.count_nonzero(state))
        return 1 if moves % 2 == 0 else -1

    def legal_actions_mask(self, state: State) -> np.ndarray:
        # A column is legal iff its top cell (row 0) is empty.
        return state[0, :] == 0

    def apply(self, state: State, action: int) -> State:
        if state[0, action] != 0:
            raise ValueError(f"Illegal action {action}: column {action} is full")
        new_state = state.copy()
        # Drop: find the lowest empty cell in this column.
        for row in range(ROWS - 1, -1, -1):
            if new_state[row, action] == 0:
                new_state[row, action] = self.current_player(state)
                break
        return new_state

    # Stubs filled in by subsequent tasks
    def terminal_value(self, state: State) -> float | None:
        raise NotImplementedError  # Task 2

    def encode(self, state: State) -> np.ndarray:
        raise NotImplementedError  # Task 3

    def canonical_state(self, state: State) -> State:
        raise NotImplementedError  # Task 3

pytest tests/test_game_connect4.py -v

Expected: 9 passed.

git add src/alphazero/games/connect4.py tests/test_game_connect4.py
git commit -m "$(cat <<'EOF'
Add Connect4 game scaffolding: state, apply (gravity), legal mask

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
EOF
)"

Task 2: Connect4 terminal_value (4-in-a-row detection)

Files: - Modify: src/alphazero/games/connect4.py (implement terminal_value and _find_winner) - Modify: tests/test_game_connect4.py (append tests)

def _board(rows: list[list[int]]) -> np.ndarray:
    """Helper: literal board from a 2D list, row 0 at top."""
    return np.array(rows, dtype=np.int8)


def test_terminal_value_none_for_initial(c4: Connect4):
    assert c4.terminal_value(c4.initial_state()) is None


def test_terminal_value_none_when_in_progress(c4: Connect4):
    s = c4.initial_state()
    s = c4.apply(s, 3)
    s = c4.apply(s, 4)
    assert c4.terminal_value(s) is None


def test_terminal_value_horizontal_4_in_a_row(c4: Connect4):
    # +1 has 4 in a row at row 5, columns 0-3. -1 has 3 scattered pieces.
    # Total nonzero = 7 (odd) → current_player = -1 (the loser POV).
    state = _board([
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [-1, -1, -1, 0, 0, 0, 0],
        [1, 1, 1, 1, 0, 0, 0],
    ])
    assert c4.current_player(state) == -1
    assert c4.terminal_value(state) == -1.0


def test_terminal_value_vertical_4_in_a_row(c4: Connect4):
    # +1 has 4 stacked in column 0. -1 has 3 scattered.
    # Nonzero = 7, current_player = -1.
    state = _board([
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0],
        [1, -1, -1, -1, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0],
    ])
    assert c4.current_player(state) == -1
    assert c4.terminal_value(state) == -1.0


def test_terminal_value_diagonal_down_right(c4: Connect4):
    # +1 has 4 on the ↘ diagonal: (2,0), (3,1), (4,2), (5,3).
    state = _board([
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0],
        [-1, 1, 0, 0, 0, 0, 0],
        [-1, -1, 1, 0, 0, 0, 0],
        [-1, -1, -1, 1, 0, 0, 0],
    ])
    # +1 pieces: 4 ; -1 pieces: 6 ; total 10 (even) → current_player = +1
    # Wait, that means it's X's turn next → X (the winner) is current, returns +1.
    # Actually for +1 to have just won, it's now -1's turn → current_player must be -1.
    # Recount: 4 X + 6 O = 10 nonzero, parity even → +1 to move.
    # But that means O just moved. We need X to have just moved → odd parity.
    # Re-construct with +1 winning on the LAST move (odd parity).
    state = _board([
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0],
        [-1, 1, 0, 0, 0, 0, 0],
        [-1, -1, 1, 0, 0, 0, 0],
        [-1, -1, 0, 1, 0, 0, 0],
    ])
    # Count: 4 X + 5 O = 9 nonzero → odd → current_player = -1 (correct, O's turn after X won)
    assert c4.current_player(state) == -1
    assert c4.terminal_value(state) == -1.0


def test_terminal_value_diagonal_up_right(c4: Connect4):
    # +1 has 4 on the ↗ diagonal: (5,0), (4,1), (3,2), (2,3).
    state = _board([
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 1, 0, 0, 0],
        [-1, -1, 1, 0, 0, 0, 0],
        [-1, 1, -1, 0, 0, 0, 0],
        [1, -1, -1, 0, 0, 0, 0],
    ])
    # 4 X + 5 O = 9 nonzero → -1's turn. X just won.
    assert c4.current_player(state) == -1
    assert c4.terminal_value(state) == -1.0


def test_terminal_value_draw_on_full_board_no_winner(c4: Connect4):
    """Construct a full board with no 4-in-a-row. (Easiest: alternate columns.)"""
    state = _board([
        [-1,  1, -1,  1, -1,  1, -1],
        [ 1, -1,  1, -1,  1, -1,  1],
        [-1,  1, -1,  1, -1,  1, -1],
        [ 1, -1,  1, -1,  1, -1,  1],
        [-1,  1, -1,  1, -1,  1, -1],
        [ 1, -1,  1, -1,  1, -1,  1],
    ])
    # Verify no 4-in-a-row exists in this configuration
    # (this striped pattern has no horizontal/vertical/diagonal 4s)
    assert c4.terminal_value(state) == 0.0

pytest tests/test_game_connect4.py -v

Expected: 7 fail with NotImplementedError.

Replace the stub in src/alphazero/games/connect4.py:

    def terminal_value(self, state: State) -> float | None:
        winner = self._find_winner(state)
        if winner is not None:
            return 1.0 if winner == self.current_player(state) else -1.0
        if np.all(state != 0):
            return 0.0  # full board, no winner → draw
        return None

    def _find_winner(self, state: State) -> int | None:
        """Return +1 / -1 if either player has 4 in a row, else None."""
        for player in (1, -1):
            # Horizontal — 4 consecutive in a row
            for r in range(ROWS):
                for c in range(COLS - 3):
                    if np.all(state[r, c:c+4] == player):
                        return player
            # Vertical — 4 consecutive in a column
            for c in range(COLS):
                for r in range(ROWS - 3):
                    if np.all(state[r:r+4, c] == player):
                        return player
            # Diagonal ↘ (top-left to bottom-right)
            for r in range(ROWS - 3):
                for c in range(COLS - 3):
                    if all(state[r+i, c+i] == player for i in range(4)):
                        return player
            # Diagonal ↗ (bottom-left to top-right)
            for r in range(3, ROWS):
                for c in range(COLS - 3):
                    if all(state[r-i, c+i] == player for i in range(4)):
                        return player
        return None

pytest tests/test_game_connect4.py -v

Expected: 16 passed (9 from Task 1 + 7 new).

git add src/alphazero/games/connect4.py tests/test_game_connect4.py
git commit -m "$(cat <<'EOF'
Add Connect4 terminal_value (4-in-a-row in 4 directions, draws)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
EOF
)"

Task 3: Connect4 encode + canonical_state

Files: - Modify: src/alphazero/games/connect4.py - Modify: tests/test_game_connect4.py

def test_canonical_state_is_identity_when_current_player_plus_one(c4: Connect4):
    s = c4.initial_state()
    np.testing.assert_array_equal(c4.canonical_state(s), s)


def test_canonical_state_inverts_when_current_player_minus_one(c4: Connect4):
    s = c4.apply(c4.initial_state(), 3)  # +1 drops in col 3; now -1's turn
    canon = c4.canonical_state(s)
    expected = np.zeros((6, 7), dtype=np.int8)
    expected[5, 3] = -1  # +1's piece becomes -1 in O's canonical view
    np.testing.assert_array_equal(canon, expected)


def test_encode_shape_and_dtype(c4: Connect4):
    enc = c4.encode(c4.canonical_state(c4.initial_state()))
    assert enc.shape == (3, 6, 7)
    assert enc.dtype == np.float32


def test_encode_initial_state_planes(c4: Connect4):
    enc = c4.encode(c4.canonical_state(c4.initial_state()))
    assert enc[0].sum() == 0  # my plane empty
    assert enc[1].sum() == 0  # opp plane empty
    np.testing.assert_array_equal(enc[2], np.ones((6, 7), dtype=np.float32))  # ones plane


def test_encode_treats_canonical_player_as_my_pieces(c4: Connect4):
    """After +1 plays col 3, -1's canonical view: +1's piece appears on the opp plane."""
    s = c4.apply(c4.initial_state(), 3)
    canon = c4.canonical_state(s)
    enc = c4.encode(canon)
    # my (O's pieces): nothing
    assert enc[0].sum() == 0
    # opp (X's piece): at (5, 3)
    assert enc[1, 5, 3] == 1
    assert enc[1].sum() == 1

Replace the stubs in src/alphazero/games/connect4.py:

    def canonical_state(self, state: State) -> State:
        """Rewrite so the current player's pieces are +1, opponent's are -1."""
        return (state * self.current_player(state)).astype(np.int8)

    def encode(self, state: State) -> np.ndarray:
        """Three planes: my pieces (+1 in canonical state), opp pieces (-1), ones.

        Caller is expected to pass a canonical state.
        """
        my = (state == 1).astype(np.float32)
        opp = (state == -1).astype(np.float32)
        ones = np.ones((ROWS, COLS), dtype=np.float32)
        return np.stack([my, opp, ones], axis=0)

pytest tests/test_game_connect4.py -v

Expected: 21 passed.

git add src/alphazero/games/connect4.py tests/test_game_connect4.py
git commit -m "$(cat <<'EOF'
Add Connect4 encode + canonical_state (perspective transform)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
EOF
)"

Task 4: Connect4 symmetries (horizontal mirror)

Files: - Modify: src/alphazero/games/connect4.py - Modify: tests/test_game_connect4.py

def test_symmetries_yields_two_tuples(c4: Connect4):
    """Connect 4 has 2 symmetries: identity + horizontal (left-right) mirror."""
    s = c4.apply(c4.initial_state(), 0)  # asymmetric: piece in column 0
    enc = c4.encode(c4.canonical_state(s))
    policy = np.array([0.5, 0.1, 0.1, 0.1, 0.1, 0.05, 0.05], dtype=np.float32)
    syms = c4.symmetries(enc, policy)
    assert len(syms) == 2


def test_symmetries_first_is_identity(c4: Connect4):
    enc = c4.encode(c4.canonical_state(c4.initial_state()))
    policy = np.full(7, 1 / 7, dtype=np.float32)
    syms = c4.symmetries(enc, policy)
    np.testing.assert_array_equal(syms[0][0], enc)
    np.testing.assert_array_equal(syms[0][1], policy)


def test_symmetries_mirror_flips_columns(c4: Connect4):
    """Mirror reflects across the vertical axis: column 0 ↔ column 6."""
    s = c4.apply(c4.initial_state(), 0)  # +1's piece at (5, 0)
    enc = c4.encode(c4.canonical_state(s))
    policy = np.zeros(7, dtype=np.float32)
    policy[0] = 1.0  # all mass on leftmost column
    syms = c4.symmetries(enc, policy)
    mirror_enc, mirror_policy = syms[1]
    # In the mirrored encoding, the piece is at the rightmost column
    # (column 6 instead of 0). Plane 1 = opp = X's pieces (from O's POV).
    assert mirror_enc[1, 5, 6] == 1
    assert mirror_enc[1, 5, 0] == 0
    # Mirror policy now puts mass on column 6
    np.testing.assert_array_equal(mirror_policy, [0, 0, 0, 0, 0, 0, 1])


def test_symmetries_preserve_policy_sum(c4: Connect4):
    enc = c4.encode(c4.canonical_state(c4.initial_state()))
    policy = np.full(7, 1 / 7, dtype=np.float32)
    for _, p in c4.symmetries(enc, policy):
        np.testing.assert_allclose(p.sum(), 1.0, rtol=1e-6)

pytest tests/test_game_connect4.py -v

    def symmetries(
        self, encoded: np.ndarray, policy: np.ndarray
    ) -> list[tuple[np.ndarray, np.ndarray]]:
        """Connect 4 has 2 D-group elements: identity + horizontal mirror.

        (Unlike TTT's 8, Connect 4 has no vertical symmetry because gravity
        is asymmetric — pieces only fall downward.)

        Mirror flips columns left↔right in both the board planes (axis=2 of
        the encoded tensor) and the policy (a 1D array of column indices).
        """
        identity = (encoded.copy(), policy.copy())
        mirrored = (
            np.flip(encoded, axis=2).copy(),
            np.flip(policy, axis=0).copy(),
        )
        return [identity, mirrored]

pytest tests/test_game_connect4.py -v

Expected: 25 passed.

git add src/alphazero/games/connect4.py tests/test_game_connect4.py
git commit -m "$(cat <<'EOF'
Add Connect4 symmetries (identity + horizontal mirror)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
EOF
)"

Task 5: Connect4MinimaxOpponent (α-β pruning with heuristic eval)

Files: - Create: src/alphazero/opponents/__init__.py (empty marker) - Create: src/alphazero/opponents/connect4_minimax.py - Test: tests/test_connect4_minimax.py

mkdir -p src/alphazero/opponents
touch src/alphazero/opponents/__init__.py

tests/test_connect4_minimax.py:

import numpy as np
import pytest

from alphazero.games.connect4 import Connect4
from alphazero.opponents.connect4_minimax import Connect4MinimaxOpponent


@pytest.fixture
def c4() -> Connect4:
    return Connect4()


def _board(rows: list[list[int]]) -> np.ndarray:
    return np.array(rows, dtype=np.int8)


def test_only_plays_legal_moves(c4: Connect4):
    """Minimax must never propose an illegal action."""
    opp = Connect4MinimaxOpponent(depth=4)
    s = c4.initial_state()
    # Fill column 0 entirely
    for _ in range(6):
        s = c4.apply(s, 0)
    action = opp(c4, s)
    legal = c4.legal_actions_mask(s)
    assert legal[action], f"Minimax returned illegal action {action}"


def test_takes_immediate_winning_move(c4: Connect4):
    """+1 can win at column 3 by completing the bottom row 4-in-a-row.

    Setup: +1 has pieces at (5, 0..2), and -1 has 3 pieces stacked elsewhere.
    Total: 3+3 = 6 nonzero → even parity → current_player = +1.
    """
    state = _board([
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [-1, 0, 0, 0, 0, 0, 0],
        [-1, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 0, -1, 0, 0],
    ])
    assert c4.current_player(state) == 1
    opp = Connect4MinimaxOpponent(depth=4)
    action = opp(c4, state)
    assert action == 3, f"Should play winning move (col 3), got {action}"


def test_blocks_immediate_opponent_threat(c4: Connect4):
    """-1 has 3 in a row at the bottom (cols 0-2). +1 to move MUST block at col 3.

    Setup: -1 has bottom-row 0..2; +1 has 3 random pieces. Total 6, even → +1's turn.
    """
    state = _board([
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0],
        [-1, -1, -1, 0, 1, 0, 0],
    ])
    assert c4.current_player(state) == 1
    opp = Connect4MinimaxOpponent(depth=4)
    action = opp(c4, state)
    assert action == 3, f"Must block at col 3, got {action}"


def test_prefers_center_on_empty_board(c4: Connect4):
    """Minimax-with-center-bias should pick column 3 on the first move."""
    opp = Connect4MinimaxOpponent(depth=4)
    action = opp(c4, c4.initial_state())
    assert action == 3, f"Expected center (col 3) opening, got {action}"


def test_depth_8_beats_depth_2(c4: Connect4):
    """Deeper search should produce stronger play."""
    from alphazero.arena import play_match
    strong = Connect4MinimaxOpponent(depth=8)
    weak = Connect4MinimaxOpponent(depth=2)
    # 10 games is enough for a strong-vs-weak signal in Connect 4.
    np.random.seed(0)
    result = play_match(c4, strong, weak, num_games=10)
    # win_rate counts draws as 0.5; strong should be clearly ahead.
    assert result.win_rate >= 0.80, (
        f"depth-8 only achieved win_rate {result.win_rate:.2%} vs depth-2"
    )

pytest tests/test_connect4_minimax.py -v

src/alphazero/opponents/connect4_minimax.py:

"""Connect 4 α-β minimax opponent with heuristic position evaluation.

Used as the eval ground truth for Sub-project 2 (we don't have a perfect
solver for Connect 4 from scratch — depth-8 minimax is a strong amateur).

Heuristic eval at non-terminal leaves:
    - Score sliding 4-cell windows in all 4 directions
        4-of-mine        → +1000 (effectively a forced win)
        3-of-mine + 1 empty → +5
        2-of-mine + 2 empties → +2
        symmetric negatives for opponent
        opp 3-in-a-row gets -50 (block aggressively even before center bias)
    - +3 per piece in the center column (mine); -3 per opponent center piece
"""
from __future__ import annotations

import numpy as np

from ..games.base import Game, State
from ..games.connect4 import COLS, ROWS, Connect4

CENTER_COL = COLS // 2  # column 3 (the center of a 7-col board)
WIN_SCORE = 1000.0      # heuristic value of having 4-in-a-row in a leaf window
BLOCK_SCORE = -50.0     # extra emphasis for blocking opp's 3-in-a-row


class Connect4MinimaxOpponent:
    """α-β minimax to a fixed depth with a center-bias position heuristic."""

    def __init__(self, depth: int = 8):
        self.depth = depth
        self._game = Connect4()

    def __call__(self, game: Game, state: State) -> int:
        """Return the best action for current_player at this state."""
        legal = game.legal_actions_mask(state)
        legal_actions = np.where(legal)[0]
        if len(legal_actions) == 0:
            raise ValueError("No legal actions at this state")

        best_action = int(legal_actions[0])
        best_value = -float("inf")
        alpha = -float("inf")
        beta = float("inf")
        for action in legal_actions:
            next_state = game.apply(state, int(action))
            # Negamax: value from child's POV is negated to get parent's POV.
            value = -self._minimax(next_state, self.depth - 1, -beta, -alpha)
            if value > best_value:
                best_value = value
                best_action = int(action)
            alpha = max(alpha, value)
        return best_action

    def _minimax(self, state: State, depth: int, alpha: float, beta: float) -> float:
        """Return the negamax value from current_player(state)'s POV."""
        terminal = self._game.terminal_value(state)
        if terminal is not None:
            return float(terminal) * WIN_SCORE if terminal != 0 else 0.0
        if depth == 0:
            return self._evaluate(state)

        legal = self._game.legal_actions_mask(state)
        best_value = -float("inf")
        for action in np.where(legal)[0]:
            next_state = self._game.apply(state, int(action))
            value = -self._minimax(next_state, depth - 1, -beta, -alpha)
            best_value = max(best_value, value)
            alpha = max(alpha, value)
            if alpha >= beta:
                break  # β-cutoff
        return best_value

    def _evaluate(self, state: State) -> float:
        """Heuristic score from current_player's POV."""
        player = self._game.current_player(state)
        score = 0.0
        for window in _all_windows(state):
            score += _window_score(window, player)
        # Center column bonus
        score += 3 * int((state[:, CENTER_COL] == player).sum())
        score -= 3 * int((state[:, CENTER_COL] == -player).sum())
        return score


def _all_windows(state: State) -> list[np.ndarray]:
    """Generate all length-4 windows in all 4 directions."""
    windows: list[np.ndarray] = []
    # Horizontal
    for r in range(ROWS):
        for c in range(COLS - 3):
            windows.append(state[r, c:c + 4])
    # Vertical
    for c in range(COLS):
        for r in range(ROWS - 3):
            windows.append(state[r:r + 4, c])
    # Diagonal ↘
    for r in range(ROWS - 3):
        for c in range(COLS - 3):
            windows.append(np.array([state[r + i, c + i] for i in range(4)]))
    # Diagonal ↗
    for r in range(3, ROWS):
        for c in range(COLS - 3):
            windows.append(np.array([state[r - i, c + i] for i in range(4)]))
    return windows


def _window_score(window: np.ndarray, player: int) -> float:
    """Score a single 4-cell window from `player`'s POV."""
    my = int((window == player).sum())
    opp = int((window == -player).sum())
    if my == 4:
        return WIN_SCORE
    if opp == 4:
        return -WIN_SCORE
    if opp > 0 and my > 0:
        return 0.0  # mixed window — no threat either way
    if my == 3:
        return 5.0
    if my == 2:
        return 2.0
    if opp == 3:
        return BLOCK_SCORE
    if opp == 2:
        return -2.0
    return 0.0

pytest tests/test_connect4_minimax.py -v

Expected: 5 passed. (Note: the depth_8_beats_depth_2 test is the slowest — ~5-10 seconds.)

git add src/alphazero/opponents/__init__.py src/alphazero/opponents/connect4_minimax.py tests/test_connect4_minimax.py
git commit -m "$(cat <<'EOF'
Add Connect4MinimaxOpponent (α-β depth-8 + center-bias heuristic)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
EOF
)"

Task 6: Generalize Trainer eval hook (_eval_vs_opponent)

Files: - Modify: src/alphazero/trainer.py - Modify: tests/test_trainer.py

def test_eval_vs_opponent_uses_provided_agent(tiny_config):
    """Trainer._eval_vs_opponent should play best_net vs an arbitrary agent."""
    from alphazero.games.tictactoe import TicTacToe
    trainer = Trainer(TicTacToe(), tiny_config)

    # Sentinel opponent: always plays the first legal action.
    def sentinel_agent(game, state):
        import numpy as np
        legal = game.legal_actions_mask(state)
        return int(np.where(legal)[0][0])

    result = trainer._eval_vs_opponent(sentinel_agent, num_games=4)
    assert set(result.keys()) == {"wins", "draws", "losses"}
    assert result["wins"] + result["draws"] + result["losses"] == 4

pytest tests/test_trainer.py::test_eval_vs_opponent_uses_provided_agent -v

Replace the existing _eval_vs_solver in src/alphazero/trainer.py with:

    def _eval_vs_opponent(self, opponent, num_games: int) -> dict:
        """Play current best_net vs an arbitrary opponent agent.

        Generic alternative to `_eval_vs_solver` — used for games (like
        Connect 4) where the eval baseline is a minimax opponent, not a
        perfect-play solver.

        Args:
            opponent: callable `(game, state) -> int` returning chosen action.
            num_games: how many games to play (alternating colors).

        Returns:
            dict with keys "wins", "draws", "losses" from best_net's POV.
        """
        from .arena import play_match
        net_agent = self._make_argmax_mcts_agent(self.best_net)
        result = play_match(self.game, net_agent, opponent, num_games=num_games)
        return {
            "wins": result.wins_a,
            "draws": result.draws,
            "losses": result.losses_a,
        }

    def _eval_vs_solver(self) -> dict:
        """Play current best_net vs the perfect TTT solver. Thin wrapper over
        _eval_vs_opponent for backward-compatibility with the Sub-project 1
        E2E test."""
        from .solvers.tictactoe_solver import solve_tictactoe_action

        def solver_agent(_game, state):
            return solve_tictactoe_action(state)

        return self._eval_vs_opponent(solver_agent, self.config.eval_games)

In src/alphazero/trainer.py, change the run signature and the eval block. Replace:

    def run(self, verbose: bool = True) -> None:

with:

    def run(self, verbose: bool = True, eval_opponent=None) -> None:
        """Train for num_iterations.

        Args:
            verbose: print per-iteration progress.
            eval_opponent: optional callable `(game, state) -> int`. If given,
                eval-vs-solver hook is replaced by eval-vs-this-opponent.
                Used by Connect 4 (Sub-project 2) to inject the minimax baseline.
        """

And replace the existing eval block:

            if self.iteration % self.config.eval_interval == 0:
                if verbose:
                    print(f"  eval vs solver ({self.config.eval_games} games)...", flush=True)
                result = self._eval_vs_solver()
                if verbose:
                    print(f"  eval: wins={result['wins']} draws={result['draws']} "
                          f"losses={result['losses']}", flush=True)

with:

            if self.iteration % self.config.eval_interval == 0:
                if eval_opponent is not None:
                    if verbose:
                        print(f"  eval vs opponent ({self.config.eval_games} games)...",
                              flush=True)
                    result = self._eval_vs_opponent(eval_opponent, self.config.eval_games)
                else:
                    if verbose:
                        print(f"  eval vs solver ({self.config.eval_games} games)...",
                              flush=True)
                    result = self._eval_vs_solver()
                if verbose:
                    print(f"  eval: wins={result['wins']} draws={result['draws']} "
                          f"losses={result['losses']}", flush=True)

pytest tests/test_trainer.py -v

Expected: 10 passed (9 existing + 1 new).

git add src/alphazero/trainer.py tests/test_trainer.py
git commit -m "$(cat <<'EOF'
Generalize Trainer eval hook to accept arbitrary opponent agent

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
EOF
)"

Task 7: CLI --game flag

Files: - Modify: src/alphazero/cli.py

Edit src/alphazero/cli.py. After the imports, add:

from .games.connect4 import Connect4

GAMES = {
    "tictactoe": TicTacToe,
    "connect4": Connect4,
}

Then in each subcommand function (_cmd_train, _cmd_eval, _cmd_play), replace the hardcoded game = TicTacToe() with:

game_cls = GAMES[args.game]
game = game_cls()

And in main(), add the --game flag to each subparser:

    p_train = subs.add_parser("train", help="Run training loop")
    p_train.add_argument("--game", choices=list(GAMES), default="tictactoe")
    p_train.add_argument("--config", type=Path, default=Path("configs/tictactoe.toml"))
    p_train.set_defaults(func=_cmd_train)

    p_eval = subs.add_parser("eval", help="Evaluate a checkpoint vs perfect solver")
    p_eval.add_argument("--game", choices=list(GAMES), default="tictactoe")
    p_eval.add_argument("--checkpoint", type=Path, required=True)
    p_eval.add_argument("--config", type=Path, default=None,
                        help="Override architecture config (needed for older checkpoints "
                             "that didn't save their config inline).")
    p_eval.add_argument("--num-games", type=int, default=200)
    p_eval.set_defaults(func=_cmd_eval)

    p_play = subs.add_parser("play", help="Play interactively against best_net")
    p_play.add_argument("--game", choices=list(GAMES), default="tictactoe")
    p_play.add_argument("--checkpoint", type=Path, required=True)
    p_play.add_argument("--config", type=Path, default=None,
                        help="Override architecture config (needed for older checkpoints "
                             "that didn't save their config inline).")
    p_play.add_argument("--num-simulations", type=int, default=None,
                        help="Override MCTS simulations per move at play time.")
    p_play.add_argument("--as", dest="as_player", choices=["x", "o", "X", "O"], default="x",
                        help="Play as X (moves first) or O (moves second). Default: x.")
    p_play.set_defaults(func=_cmd_play)

In _cmd_eval, also generalize the eval logic. Currently it hardcodes the TTT solver as opponent. Change it to dispatch by --game:

def _cmd_eval(args: argparse.Namespace) -> int:
    config = _config_for_checkpoint(args.checkpoint, args.config)
    game_cls = GAMES[args.game]
    game = game_cls()
    trainer = Trainer(game, config)

    ckpt = torch.load(args.checkpoint, map_location=trainer.device, weights_only=False)
    trainer.best_net.load_state_dict(ckpt["best_net"])
    trainer.best_net.eval()

    agent = trainer._make_argmax_mcts_agent(trainer.best_net)

    if args.game == "tictactoe":
        def opponent(_g, state):
            return solve_tictactoe_action(state)
        label = "perfect solver"
    elif args.game == "connect4":
        from .opponents.connect4_minimax import Connect4MinimaxOpponent
        opponent = Connect4MinimaxOpponent(depth=8)
        label = "minimax-depth-8"
    else:
        raise ValueError(f"No eval opponent defined for game={args.game}")

    result = play_match(game, agent, opponent, num_games=args.num_games)
    print(
        f"vs {label} ({args.num_games} games): "
        f"wins={result.wins_a} draws={result.draws} losses={result.losses_a}"
    )
    return 0 if result.losses_a == 0 else 1

python3 -m alphazero --help
python3 -m alphazero train --help
python3 -m alphazero eval --help
python3 -m alphazero play --help

Expected: all four exit 0 and show the --game flag.

pytest -m "not slow" -v 2>&1 | tail -5

Expected: all tests pass (no test added in this task; we're verifying nothing regressed).

git add src/alphazero/cli.py
git commit -m "$(cat <<'EOF'
Add --game {tictactoe,connect4} flag to CLI; route to right opponent

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
EOF
)"

Task 8: Connect 4 training config

Files: - Create: configs/connect4.toml

# Connect 4 training config — Sub-project 2.
# Defaults to AlphaZero paper conventions; tuned down for M4 Pro CPU/MPS.

# Network — bigger than TTT (richer position features needed for 6×7 board)
n_blocks = 6
n_channels = 64

# MCTS
num_simulations = 100
c_puct = 2.5
dirichlet_alpha = 0.5
dirichlet_weight = 0.25

# Self-play
games_per_iteration = 100
temperature_threshold = 15  # Connect 4 games are ~20-30 plies; cover the opening

# Training
training_steps_per_iteration = 1000
batch_size = 64
learning_rate = 1e-3
weight_decay = 1e-4

# Replay buffer
replay_buffer_capacity = 200_000
min_buffer_size = 20_000

# Eval (vs minimax-depth-8)
eval_interval = 5
eval_games = 200

# Schedule
num_iterations = 50
seed = 42
device = "auto"
checkpoint_dir = "checkpoints"
log_dir = "runs"

python3 -c "from alphazero.config import load_config; print(load_config('configs/connect4.toml'))"

Expected: prints a TrainingConfig(...) with the values above.

git add configs/connect4.toml
git commit -m "$(cat <<'EOF'
Add Connect 4 training config (configs/connect4.toml)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
EOF
)"

Task 9: E2E kill-criterion test (the success bar)

Files: - Create: tests/test_e2e_connect4.py

tests/test_e2e_connect4.py:

"""End-to-end test: train an AlphaZero agent on Connect 4 and verify it beats
the minimax-depth-8 opponent in ≥80% of 200 games.

This is the kill criterion for Sub-project 2. Marked slow — a full training
run takes ~3-6 hours on M4 Pro MPS / ~12-20 hours on CPU.

Run with:
    pytest -m slow tests/test_e2e_connect4.py -v -s
"""
from __future__ import annotations

from pathlib import Path

import pytest

from alphazero.arena import play_match
from alphazero.config import load_config
from alphazero.games.connect4 import Connect4
from alphazero.opponents.connect4_minimax import Connect4MinimaxOpponent
from alphazero.trainer import Trainer

REPO_ROOT = Path(__file__).resolve().parent.parent


@pytest.mark.slow
def test_e2e_connect4_beats_minimax_depth_8(tmp_path, monkeypatch):
    """Train, then play 200 games vs the minimax-depth-8 opponent.
    Pass criterion: win_rate >= 0.80 (alternating colors)."""
    monkeypatch.chdir(tmp_path)

    config = load_config(REPO_ROOT / "configs/connect4.toml")
    opponent = Connect4MinimaxOpponent(depth=8)

    trainer = Trainer(Connect4(), config)
    trainer.run(eval_opponent=opponent)

    agent = trainer._make_argmax_mcts_agent(trainer.best_net)
    result = play_match(Connect4(), agent, opponent, num_games=200)

    print(
        f"\nFinal: wins={result.wins_a} draws={result.draws} "
        f"losses={result.losses_a} win_rate={result.win_rate:.2%}"
    )
    assert result.win_rate >= 0.80, (
        f"Agent win rate {result.win_rate:.2%} below 0.80 target. "
        "Increase num_iterations or num_simulations in configs/connect4.toml."
    )

pytest --collect-only -m "not slow" 2>&1 | grep test_e2e_connect4 || echo "(excluded ✓)"
pytest --collect-only -m slow tests/test_e2e_connect4.py 2>&1 | tail -3

Expected: first command prints "(excluded ✓)" (test not in default collection); second collects 1 test.

git add tests/test_e2e_connect4.py
git commit -m "$(cat <<'EOF'
Add E2E kill-criterion test (win_rate ≥ 0.80 vs minimax-depth-8)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
EOF
)"

Task 10: Final unit-test sweep + push

pytest -m "not slow" -v 2>&1 | tail -10

Expected: all tests pass. Note the total count — should be ~100+ (84 from Sub-project 1 + ~25 new game tests + 5 new minimax tests + 1 new trainer test = ~115).

ruff check src tests 2>&1 | tail -10

Note any issues. If there are real errors, fix them before continuing. Ignore minor style suggestions.

git push 2>&1 | tail -3

This is the formal Sub-project 2 success bar. It takes 3-6 hours on MPS.

pytest -m slow tests/test_e2e_connect4.py -v -s

Pass criterion: win_rate ≥ 0.80 against minimax-depth-8 over 200 games.

If it passes: Sub-project 2 is officially done. Sub-project 3 (chess) can begin.

If it doesn't reach 80%, candidate causes to investigate (in order of likelihood): 1. Undertraining — bump num_iterations from 50 to 80 or 100. 2. Too-shallow MCTS — bump num_simulations from 100 to 200. 3. Wrong network sizing — try 8 blocks × 96 channels. 4. Heuristic too weak — debug Connect4MinimaxOpponent._evaluate by playing it against a known stronger heuristic.

Final state

After Task 9, the repo has:

Crucial property: the framework code from Sub-project 1 (Game ABC, AlphaZeroNet, MCTS, SelfPlayWorker, ReplayBuffer, Arena) remains completely unchanged. The only modifications are: (1) Trainer's run() signature gains an eval_opponent parameter; (2) _eval_vs_solver becomes a thin wrapper around the new _eval_vs_opponent; (3) CLI dispatches on --game. These are extensions, not breaking changes.

Sub-project 2 is done when pytest -m slow tests/test_e2e_connect4.py reaches win_rate >= 0.80. Then Sub-project 3 (chess) can begin — and again, only the Game subclass changes.