Self-play league for Connect Four

git clone https://github.com/unionai/unionai-examples
cd unionai-examples/v2/tutorials/self_play_league
python league.py            # 20 generations, about 25 minutes
python league.py --smoke    # 2 short generations
python app.py               # serve the champion of the latest league run

class Boards:
    """A batch of Connect Four games, stepped together on one device."""

    def __init__(self, n: int, device: str = "cpu"):
        self.n, self.device = n, device
        self.kernels = _K.to(device)
        self.board = torch.zeros(n, ROWS, COLS, device=device)
        self.height = torch.zeros(n, COLS, dtype=torch.long, device=device)
        self.moves = torch.zeros(n, dtype=torch.long, device=device)

    def reset(self, mask: torch.Tensor | None = None) -> None:
        mask = torch.ones(self.n, dtype=torch.bool, device=self.device) if mask is None else mask
        self.board[mask] = 0
        self.height[mask] = 0
        self.moves[mask] = 0

    def legal(self) -> torch.Tensor:
        return self.height < ROWS

    def wins(self, board: torch.Tensor) -> torch.Tensor:
        """True where the +1 player has four in a row."""
        hits = F.conv2d((board == 1).float().unsqueeze(1), self.kernels, padding=3)
        return (hits >= 4).flatten(1).any(1)

    def step(self, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Play ``action`` for the player to move. Returns (mover won, game over)."""
        idx = torch.arange(self.n, device=self.device)
        row = ROWS - 1 - self.height[idx, action]
        self.board[idx, row, action] = 1
        self.height[idx, action] += 1
        self.moves += 1
        won = self.wins(self.board)
        over = won | (self.moves >= ROWS * COLS)
        self.board = -self.board  # the other player moves next
        return won, over

def random_bot(boards: Boards, gen: torch.Generator | None = None) -> torch.Tensor:
    return torch.multinomial(boards.legal().float(), 1, generator=gen).squeeze(1)

CENTER = torch.tensor([1.0, 2.0, 3.0, 4.0, 3.0, 2.0, 1.0])

def tactical_bot(boards: Boards, gen: torch.Generator | None = None) -> torch.Tensor:
    """Win if it can, block if it must, otherwise play a random center-weighted move."""
    win = winning_moves(boards, boards.board)
    block = winning_moves(boards, -boards.board)
    weights = boards.legal().float() * CENTER.to(boards.device)
    weights = torch.where(block.any(1, keepdim=True), block.float(), weights)
    weights = torch.where(win.any(1, keepdim=True), win.float(), weights)
    return torch.multinomial(weights, 1, generator=gen).squeeze(1)

BOTS = {"random": random_bot, "tactical": tactical_bot}

class PolicyValueNet(nn.Module):
    """A small residual conv net: board in, move logits and a value in [-1, 1] out."""

    def __init__(self, channels: int = 64, blocks: int = 4):
        super().__init__()
        self.stem = nn.Sequential(nn.Conv2d(2, channels, 3, padding=1), nn.ReLU())
        self.blocks = nn.ModuleList(
            nn.Sequential(nn.Conv2d(channels, channels, 3, padding=1), nn.ReLU(), nn.Conv2d(channels, channels, 3, padding=1))
            for _ in range(blocks)
        )
        self.policy = nn.Sequential(nn.Conv2d(channels, 2, 1), nn.Flatten(), nn.Linear(2 * ROWS * COLS, COLS))
        self.value = nn.Sequential(nn.Conv2d(channels, 1, 1), nn.Flatten(), nn.Linear(ROWS * COLS, 64), nn.ReLU(),
                                   nn.Linear(64, 1), nn.Tanh())

    def forward(self, board: torch.Tensor, legal: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        x = torch.stack([(board == 1).float(), (board == -1).float()], 1)
        h = self.stem(x)
        for block in self.blocks:
            h = F.relu(h + block(h))
        logits = self.policy(h).masked_fill(~legal, -1e9)
        return logits, self.value(h).squeeze(1)

@torch.no_grad()
def collect(net: PolicyValueNet, opponents: list, games: int, device: str, gen: torch.Generator) -> dict:
    """Play ``games`` games in parallel, split evenly across ``opponents``, and keep the learner's moves."""
    boards = Boards(games, device)
    group = torch.arange(games, device=device) % len(opponents)  # which opponent each game faces
    learner_turn = torch.arange(games, device=device) % 2 == 0  # the learner moves first in half the games
    buf = {k: [] for k in ("board", "legal", "action", "logp", "value", "game")}
    outcome = torch.full((games,), float("nan"), device=device)
    alive = torch.ones(games, dtype=torch.bool, device=device)
    while alive.any():
        legal = boards.legal()
        logits, value = net(boards.board, legal)
        dist = torch.distributions.Categorical(logits=logits)
        action = torch.multinomial(dist.probs, 1, generator=gen).squeeze(1)
        for i, opponent in enumerate(opponents):
            mask = (group == i) & ~learner_turn & alive
            if mask.any():
                action[mask] = opponent(boards, gen)[mask]
        mine = learner_turn & alive
        for k, v in (("board", boards.board), ("legal", legal), ("action", action), ("logp", dist.log_prob(action)),
                     ("value", value), ("game", torch.arange(games, device=device))):
            buf[k].append(v[mine].clone())
        # Finished games keep playing a legal dummy move; their results are already recorded.
        action = torch.where(alive, action, legal.float().argmax(1))
        won, over = boards.step(action)
        newly = over & alive
        outcome[newly & won & learner_turn] = 1.0
        outcome[newly & won & ~learner_turn] = -1.0
        outcome[newly & ~won] = 0.0
        alive &= ~over
        learner_turn = ~learner_turn
    data = {k: torch.cat(v) for k, v in buf.items()}
    data["return"] = outcome[data.pop("game")]
    per_opponent = [outcome[group == i] for i in range(len(opponents))]
    data["results"] = [{"win": float((o == 1).float().mean()), "draw": float((o == 0).float().mean()),
                        "loss": float((o == -1).float().mean())} for o in per_opponent]
    return data

def ppo_update(net: PolicyValueNet, opt: torch.optim.Optimizer, data: dict, epochs: int = 3,
               batch: int = 4096, clip: float = 0.2) -> dict:
    """PPO-clip on Monte Carlo returns, with a value baseline and an entropy bonus."""
    n = data["action"].shape[0]
    adv = data["return"] - data["value"]
    adv = (adv - adv.mean()) / (adv.std() + 1e-8)
    stats = {"policy_loss": 0.0, "value_loss": 0.0, "entropy": 0.0, "updates": 0}
    for _ in range(epochs):
        for idx in torch.randperm(n, device=adv.device).split(batch):
            logits, value = net(data["board"][idx], data["legal"][idx])
            dist = torch.distributions.Categorical(logits=logits)
            ratio = (dist.log_prob(data["action"][idx]) - data["logp"][idx]).exp()
            pg = -torch.min(ratio * adv[idx], ratio.clamp(1 - clip, 1 + clip) * adv[idx]).mean()
            vl = F.mse_loss(value, data["return"][idx])
            ent = dist.entropy().mean()
            loss = pg + 0.5 * vl - 0.01 * ent
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
            opt.step()
            stats["policy_loss"] += float(pg.detach())
            stats["value_loss"] += float(vl.detach())
            stats["entropy"] += float(ent.detach())
            stats["updates"] += 1
    return {k: v / max(stats["updates"], 1) for k, v in stats.items() if k != "updates"}

@train_env.task(report=True)
async def train_generation(
    generation: int, learner: File | None, opponents: list[Player], iterations: int = 25, games: int = 2048,
) -> tuple[File, str]:
    """Train the learner for one generation against the given league opponents and itself, and snapshot it."""
    import torch
    from connect4 import PolicyValueNet, net_bot
    from selfplay import train

    device = "cuda"
    torch.manual_seed(generation)
    net = PolicyValueNet().to(device)
    opt = torch.optim.Adam(net.parameters(), lr=3e-4)
    if learner is not None:
        state = torch.load(await learner.download(), map_location=device, weights_only=True)
        net.load_state_dict(state["model"])
        opt.load_state_dict(state["optimizer"])

    # Games are split evenly across the opponents and the current learner itself.
    names = [p.name for p in opponents] + ["self"]
    players = [await as_bot(p, device) for p in opponents] + [net_bot(net)]

    def live(row: dict) -> None:
        history.append(row)
        flyte.report.replace(reports.training_live(generation, history, names, iterations, games))
        flyte.report.flush()

    history: list[dict] = []
    train(net, opt, players, iterations, games, device, seed=generation, on_iteration=live)

    path = f"{tempfile.mkdtemp()}/gen-{generation}.pt"
    torch.save({"model": net.state_dict(), "optimizer": opt.state_dict(), "generation": generation}, path)
    return await File.from_local(path), json.dumps({"opponents": names, "history": history})

# unionai-reuse provides the bridge that reusable environments need.
image = flyte.Image.from_debian_base(name="connect4-league").with_pip_packages("torch==2.9.1", "numpy<3", "unionai-reuse==0.2.1")
# The game and the trainer are imported inside the tasks, so the code bundle lists them.
CODE = ["connect4.py", "selfplay.py"]

# The learner keeps a warm CUDA context and its optimizer between generations: one reusable
# GPU replica serves every train_generation call of a run.
train_env = flyte.TaskEnvironment(
    name="c4-train",
    image=image,
    resources=flyte.Resources(cpu=4, memory="16Gi", gpu="L4:1"),
    reusable=flyte.ReusePolicy(replicas=1, concurrency=1, idle_ttl=900),
    cache="auto",
    include=CODE,
)

# Matches are CPU work, and a new snapshot plays every member of the league at once. A warm
# pool of CPU replicas takes the fan-out without starting a pod per match.
match_env = flyte.TaskEnvironment(
    name="c4-match",
    image=image,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
    reusable=flyte.ReusePolicy(replicas=(1, 4), concurrency=4, idle_ttl=600, scaledown_ttl=120),
    cache="auto",
    include=CODE,
)

driver_env = flyte.TaskEnvironment(
    name="c4-league-driver",
    image=image,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
    depends_on=[train_env, match_env],
)

@torch.no_grad()
def match(player_a, player_b, games: int, device: str = "cpu", seed: int = 0, opening: int = 4) -> dict:
    """Play ``games`` games between two players, alternating who moves first.

    The first ``opening`` moves of every game are random, so a pair of greedy players does not
    replay the same game every time. Returns A's wins, draws, and losses.
    """
    gen = torch.Generator(device=device).manual_seed(seed)
    boards = Boards(games, device)
    a_turn = torch.arange(games, device=device) % 2 == 0
    alive = torch.ones(games, dtype=torch.bool, device=device)
    result = torch.zeros(games, device=device)
    while alive.any():
        legal = boards.legal()
        if int(boards.moves.max()) < opening:
            action = torch.multinomial(legal.float(), 1, generator=gen).squeeze(1)
        else:
            action = torch.where(a_turn, player_a(boards, gen), player_b(boards, gen))
        action = torch.where(alive & legal.gather(1, action[:, None]).squeeze(1), action, legal.float().argmax(1))
        won, over = boards.step(action)
        newly = over & alive
        result[newly & won] = torch.where(a_turn[newly & won], 1.0, -1.0)
        alive &= ~over
        a_turn = ~a_turn
    return {"win": int((result == 1).sum()), "draw": int((result == 0).sum()), "loss": int((result == -1).sum())}

def fit_elo(results: list[dict], anchor: str = "random", anchor_rating: float = 0.0, iters: int = 2000) -> dict[str, float]:
    """Ratings that best explain every match result at once.

    Each result is ``{"a", "b", "win", "draw", "loss"}`` from A's side, and a draw counts as half
    a win for each player. Every pairing also gets one virtual draw, so a player who wins every
    game gets a large rating instead of an infinite one. Unlike updating Elo one game at a time,
    a joint fit does not depend on the order the matches finished in, so a re-run of the same
    matches gives the same ladder.
    """
    players = sorted({r["a"] for r in results} | {r["b"] for r in results})
    rating = {p: 0.0 for p in players}
    lr = 20.0
    for _ in range(iters):
        grad = {p: 0.0 for p in players}
        for r in results:
            n = r["win"] + r["draw"] + r["loss"] + 1
            expected = 1 / (1 + 10 ** ((rating[r["b"]] - rating[r["a"]]) / 400))
            score = (r["win"] + 0.5 * r["draw"] + 0.5) / n
            grad[r["a"]] += n * (score - expected)
            grad[r["b"]] -= n * (score - expected)
        total = sum(r["win"] + r["draw"] + r["loss"] + 1 for r in results) / max(len(players), 1)
        for p in players:
            rating[p] += lr * grad[p] / max(total, 1)
    shift = anchor_rating - rating.get(anchor, 0.0)
    return {p: round(v + shift, 1) for p, v in rating.items()}

def pick_opponents(members: list[Player], results: list[dict], top: int = 3) -> list[Player]:
    """The bots, plus the strongest past snapshots by Elo.

    Training only against recent snapshots lets the learner forget how to beat older styles of
    play. Training against the strongest members keeps the pressure where the learner is
    weakest, a simple form of the prioritized league play behind AlphaStar.
    """
    ratings = elo.fit_elo(results)
    bots = [m for m in members if m.policy is None]
    snapshots = sorted((m for m in members if m.policy is not None), key=lambda m: -ratings.get(m.name, 0))
    return bots + snapshots[:top]

@driver_env.task(report=True, produces_artifacts=True, retries=2)
async def league(generations: int = 20, iterations: int = 25, games: int = 2048, match_games: int = 400) -> File:
    """Run the self-play league, keep the Elo ladder up to date, and publish the champion."""
    members = [Player("random"), Player("tactical")]
    results = [json.loads(await play_match(members[1], members[0], match_games))]
    ladder: list[dict] = []  # Elo of every member after each generation
    training: list[dict] = []
    learner: File | None = None

    for g in range(1, generations + 1):
        with flyte.group(f"generation-{g}"):
            learner, log = await train_generation(g, learner, pick_opponents(members, results), iterations, games)
            snapshot = Player(f"gen-{g}", learner)
            # The new snapshot plays every existing member at once, on the CPU pool.
            new = await asyncio.gather(*[play_match(snapshot, m, match_games, seed=g) for m in members])
        members.append(snapshot)
        results += [json.loads(r) for r in new]
        training.append({"generation": g, **json.loads(log)})
        ratings = elo.fit_elo(results)
        ladder.append({"generation": g, "ratings": ratings})
        await flyte.report.replace.aio(reports.league_live(ladder, results, training, generations), do_flush=True)

    ratings = ladder[-1]["ratings"]
    champion = max((m for m in members if m.policy is not None), key=lambda m: ratings[m.name])
    game = json.loads(await sample_game(champion, Player("tactical")))
    await flyte.report.replace.aio(reports.league_final(ladder, results, training, champion.name, game), do_flush=True)

    card = await artifacts.Card.create_from.aio(
        content=reports.card(ladder, results, champion.name, game), format="html", card_type="model")
    vs_tactical = next(r for r in results if r["a"] == champion.name and r["b"] == "tactical")
    metadata = artifacts.Metadata.create_model_metadata(
        name="connect4-champion",
        description=f"Connect Four self-play champion {champion.name}: Elo {ratings[champion.name]:.0f}, "
                    f"{vs_tactical['win'] / match_games:.0%} wins against the tactical bot",
        framework="PyTorch",
        model_type="policy-value network",
        architecture="residual CNN, 4 blocks x 64 channels",
        task="Connect Four",
        serial_format="pt",
        card=card,
        attrs={"generation": champion.name, "elo": f"{ratings[champion.name]:.0f}",
               "win_vs_tactical": f"{vs_tactical['win'] / match_games:.3f}"},
    )
    return artifacts.new(champion.policy, metadata)

arena = FastAPIAppEnvironment(
    name="connect4-arena",
    app=app,
    # A CPU-only torch is enough to pick a move, and much smaller than the CUDA build.
    image=flyte.Image.from_debian_base(name="connect4-arena")
    .with_pip_packages("torch==2.9.1", index_url="https://download.pytorch.org/whl/cpu")
    .with_pip_packages("numpy<3", "fastapi", "uvicorn"),
    parameters=[
        # The champion published by the most recent successful league run.
        flyte.app.Parameter(
            name="champion",
            value=flyte.app.RunOutput(type="file", task_name="c4-league-driver.league"),
            mount=CHAMPION_PATH,
            download=True,
        ),
    ],
    # The network is imported when the app starts, so the code bundle lists it.
    include=["connect4.py"],
    resources=flyte.Resources(cpu=1, memory="2Gi"),
    scaling=flyte.app.Scaling(replicas=(0, 1), scaledown_after=600),
    requires_auth=False,
)

@app.post("/move")
async def move(position: Position) -> dict:
    """The champion's move for a position where it is the champion's turn."""
    import torch
    import torch.nn.functional as F

    board = torch.tensor(position.board, dtype=torch.float32).unsqueeze(0)
    legal = board[:, 0, :] == 0
    with torch.no_grad():
        logits, value = state["net"](board, legal)
    probs = F.softmax(logits, 1)[0]
    return {"column": int(probs.argmax()), "probs": [round(float(p), 3) for p in probs], "value": round(float(value), 3)}