Blame for graphwar/bots.py

"""Les 8 bots de l'arène.

Deux bots suivent une règle écrite à la main (Dot, Locus) et servent de référence.
Les six autres partent d'un cerveau aux poids tirés au hasard et apprennent seuls,
chacun avec sa propre méthode et sa propre famille de fonctions.
"""
from __future__ import annotations

import os
from dataclasses import dataclass, field

import numpy as np

from . import config as C
from . import sensors, vocab as V
from .nets import MLP
from .world import World

CKPT_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "checkpoints")


@dataclass
class Shot:
    g: object
    z: object
    text: str
    ztext: str
    target: int
    params: dict = field(default_factory=dict)


def line_blocked(world: World, shooter: int, target: int) -> float:
    """Part du segment tireur → cible qui traverse un obstacle (sert aux choix de cible simples)."""
    s, t = world.soldiers[shooter], world.soldiers[target]
    k = np.linspace(0, 1, 60)
    return float(np.mean([not world.is_free(s.x + (t.x - s.x) * a, s.y + (t.y - s.y) * a) for a in k]))


def clearest_target(world: World, shooter: int) -> int:
    s = world.soldiers[shooter]
    enemies = world.alive(1 - s.team)
    return min(enemies, key=lambda j: (round(line_blocked(world, shooter, j), 2),
                                       np.hypot(world.soldiers[j].x - s.x, world.soldiers[j].y - s.y)))


class Bot:
    key = "?"
    name = "?"
    color = (0, 0, 0)
    method = ""
    blurb = ""
    vocab: V.Vocab = V.Line()
    learns = False

    def choose_target(self, world: World, shooter: int) -> int:
        return clearest_target(world, shooter)

    def act(self, world: World, shooter: int, target: int | None = None) -> Shot: ...


# ============================================================ règles fixes

class DotBot(Bot):
    key, name, color = "dot", "Dot", (90, 90, 100)
    method = "règle fixe"
    blurb = "Tire tout droit sur la cible la plus dégagée. Aucun apprentissage : c'est l'étalon."

    def act(self, world, shooter, target=None):
        target = self.choose_target(world, shooter) if target is None else target
        s, t = world.soldiers[shooter], world.soldiers[target]
        ut = max((t.x - s.x) * s.dir, 1e-3)
        p = {"a": (t.y - s.y) / ut}
        return Shot(self.vocab.g(p), None, self.vocab.text(p), "", target, p)


class LocusBot(Bot):
    """Cherche un chemin libre colonne par colonne (programmation dynamique),
    puis l'écrit comme une ligne brisée avec des valeurs absolues."""
    key, name, color = "locus", "Locus", (40, 120, 120)
    method = "planificateur"
    blurb = "Calcule le lieu des points libres et trace une ligne brisée qui l'emprunte. Aucun apprentissage."
    vocab = V.PiecewiseAbs()
    DU, DY, MAXJ, MARGIN = 1.0, 0.5, 6, 1.25

    def _plan(self, world, shooter, target):
        s, t = world.soldiers[shooter], world.soldiers[target]
        d = s.dir
        ut = (t.x - s.x) * d
        ncol = int(ut // self.DU)
        if ncol < 2:
            return None
        ys = np.arange(-C.BOARD_H + 0.75, C.BOARD_H - 0.74, self.DY)
        us = np.arange(ncol + 1) * self.DU
        us[-1] = ut
        X = s.x + d * us[:, None]
        Y = np.broadcast_to(ys, (len(us), len(ys)))
        o = world.obstacles
        dd = np.hypot(X[..., None] - o[:, 0], Y[..., None] - o[:, 1]) - o[:, 2]
        blocked = (dd < self.MARGIN).any(-1)
        if len(world.craters):
            c = world.craters
            in_hole = (np.hypot(X[..., None] - c[:, 0], Y[..., None] - c[:, 1]) < c[:, 2] - 0.5).any(-1)
            blocked &= ~in_hole
        for i, so in enumerate(world.soldiers):            # ne pas traverser les autres soldats
            if i in (shooter, target) or not so.alive:
                continue
            blocked |= np.hypot(X - so.x, Y - so.y) < 1.8
        j0 = int(np.argmin(np.abs(ys - s.y)))
        j1 = int(np.argmin(np.abs(ys - t.y)))
        blocked[0, j0] = False
        blocked[-1, j1] = False
        INF = 1e18
        cost = np.full(len(ys), INF); cost[j0] = 0
        back = np.zeros((len(us), len(ys)), dtype=np.int64)
        jj = np.arange(len(ys))
        for c_ in range(1, len(us)):
            best = np.full(len(ys), INF); arg = np.zeros(len(ys), dtype=np.int64)
            for dj in range(-self.MAXJ, self.MAXJ + 1):
                src = jj - dj
                ok = (src >= 0) & (src < len(ys))
                cand = np.full(len(ys), INF)
                cand[ok] = cost[src[ok]] + (dj * self.DY) ** 2 + 0.02
                upd = cand < best
                best[upd], arg[upd] = cand[upd], src[upd]
            best[blocked[c_]] = INF
            cost, back[c_] = best, arg
        if cost[j1] >= INF:
            return None
        path = [j1]
        for c_ in range(len(us) - 1, 0, -1):
            path.append(back[c_, path[-1]])
        path = path[::-1]
        pts = np.column_stack([us, ys[path]])
        pts[0, 1], pts[-1, 1] = s.y, t.y
        return _simplify(pts, 0.35)

    def act(self, world, shooter, target=None):
        s = world.soldiers[shooter]
        order = [target] if target is not None else sorted(
            world.alive(1 - s.team), key=lambda j: abs(world.soldiers[j].x - s.x))
        for tgt in order:
            pts = self._plan(world, shooter, tgt)
            if pts is not None:
                slopes = np.diff(pts[:, 1]) / np.maximum(np.diff(pts[:, 0]), 1e-6)
                cs = (slopes[1:] - slopes[:-1]) / 2
                p = {"A": float(slopes[0] + cs.sum()), "c": cs.tolist(), "u": pts[1:-1, 0].tolist()}
                return Shot(self.vocab.g(p), None, self.vocab.text(p), "", tgt, p)
        return DotBot().act(world, shooter, order[0])


def _simplify(pts, eps):
    """Ramer-Douglas-Peucker : garde seulement les coins utiles de la ligne brisée."""
    if len(pts) < 3:
        return pts
    a, b = pts[0], pts[-1]
    ab = b - a
    n = np.hypot(*ab) or 1.0
    d = np.abs(ab[0] * (pts[:, 1] - a[1]) - ab[1] * (pts[:, 0] - a[0])) / n
    k = int(np.argmax(d))
    if d[k] > eps:
        left, right = _simplify(pts[:k + 1], eps), _simplify(pts[k:], eps)
        return np.vstack([left[:-1], right])
    return np.vstack([a, b])


# ============================================================ bots qui apprennent

class LearnedBot(Bot):
    learns = True
    channels = ("couloir",)
    hidden = (32,)
    learner = "ga"
    ally_penalty = 1.0

    def __init__(self, path: str | None = None, seed: int = 0):
        n_in = sensors.obs_size(self.channels)
        out_act = "linear" if self.learner == "ppo" else "tanh"
        self.policy = MLP([n_in, *self.hidden, self.vocab.n_out], out_act, rng=np.random.default_rng(seed))
        self.value = None
        self.trained_iters = 0
        path = path or os.path.join(CKPT_DIR, f"{self.key}.npz")
        if os.path.exists(path):
            self.load(path)

    # -- cerveau
    def brain_out(self, obs):
        o = self.policy(obs)
        return np.tanh(o) if self.policy.out_act == "linear" else o

    def obs(self, world, shooter, target):
        return sensors.normalize(sensors.observe(world, shooter, target, self.channels), self.channels)

    def choose_target(self, world, shooter):
        enemies = world.alive(1 - world.soldiers[shooter].team)
        if self.value is not None and len(enemies) > 1:
            scores = [float(self.value(self.obs(world, shooter, j))[0]) for j in enemies]
            return enemies[int(np.argmax(scores))]
        return clearest_target(world, shooter)

    def act(self, world, shooter, target=None):
        target = self.choose_target(world, shooter) if target is None else target
        out = self.brain_out(self.obs(world, shooter, target))
        return self.shot_from(out, V.ctx_of(world, shooter, target), target)

    def shot_from(self, out, ctx, target):
        p = self.vocab.decode(out, ctx)
        return Shot(self.vocab.g(p), self.vocab.z(p), self.vocab.text(p), self.vocab.ztext(p), target, p)

    # -- sauvegarde
    def save(self, path=None, **extra):
        path = path or os.path.join(CKPT_DIR, f"{self.key}.npz")
        os.makedirs(os.path.dirname(path), exist_ok=True)
        d = self.policy.to_dict("pi_")
        if self.value is not None:
            d.update(self.value.to_dict("v_"))
        d["iters"] = np.array(self.trained_iters)
        for k, v in extra.items():
            d[k] = np.asarray(v)
        np.savez(path, **d)

    def load(self, path):
        d = np.load(path, allow_pickle=False)
        self.policy = MLP.from_dict(d, "pi_")
        self.value = MLP.from_dict(d, "v_") if "v_flat" in d else None
        self.trained_iters = int(d["iters"]) if "iters" in d else 0


class EvolutionBot(LearnedBot):
    key, name, color = "evolution", "Evolution", (200, 60, 50)
    method = "algorithme génétique"
    blurb = ("Une population de 48 cerveaux. À chaque génération les meilleurs se reproduisent, "
             "leurs enfants sont mutés, les nuls disparaissent. Les champions gardent la moyenne de tous leurs scores.")
    vocab, channels, hidden, learner = V.Gauss(), ("couloir",), (8,), "ga"


class PoleZeroBot(LearnedBot):
    key, name, color = "polezero", "Pole-Zero", (120, 70, 170)
    method = "stratégies d'évolution (ES)"
    blurb = ("Un seul cerveau qu'on secoue dans 48 directions au hasard ; on le pousse vers "
             "les secousses qui ont le mieux visé. Écrit des fractions à pôles complexes.")
    vocab, channels, hidden, learner = V.PoleZero(), ("couloir", "allies"), (16,), "es"


class FlyBot(LearnedBot):
    key, name, color = "fly", "Fly", (30, 140, 200)
    method = "entropie croisée (CEM)"
    blurb = ("Tire 64 cerveaux autour d'une moyenne, garde les 16 meilleurs et recentre la moyenne "
             "dessus. Écrit des ondes : il vole.")
    vocab, channels, hidden, learner = V.Wave(), ("couloir",), (8,), "cem"


class IrrationalBot(LearnedBot):
    key, name, color = "irrational", "Irrational", (220, 140, 30)
    method = "apprentissage par renforcement (PPO)"
    blurb = ("Apprend par essais et erreurs avec PPO, comme un agent de jeu vidéo. "
             "Écrit des racines et des valeurs absolues.")
    vocab, channels, hidden, learner = V.Irrational(), ("couloir", "allies"), (64, 64), "ppo"


class PeaceBot(LearnedBot):
    key, name, color = "peace", "Peace", (70, 160, 80)
    method = "PPO + punition des tirs fratricides"
    blurb = ("Même méthode que Irrational, mais tuer un coéquipier lui coûte trois fois une victoire. "
             "Il voit ses alliés et préfère rater que trahir.")
    vocab, channels, hidden, learner, ally_penalty = V.Steps(), ("couloir", "allies", "enemies"), (64, 64), "ppo", 3.0


class ThreeDBot(LearnedBot):
    key, name, color = "3d", "3D", (230, 60, 140)
    method = "PPO en 3 dimensions"
    blurb = ("Ajoute z(x) à son tir : il peut sauter par-dessus les obstacles bas. "
             "Il voit la hauteur de chaque obstacle.")
    vocab, channels, hidden, learner = V.Lob3D(), ("couloir", "hauteur", "allies"), (64, 64), "ppo"


ALL = [DotBot, LocusBot, EvolutionBot, PoleZeroBot, FlyBot, IrrationalBot, PeaceBot, ThreeDBot]
BY_KEY = {b.key: b for b in ALL}
LEARNERS = [b.key for b in ALL if b.learns]


def make(key: str, **kw) -> Bot:
    cls = BY_KEY[key.lower()]
    return cls(**kw) if cls.learns else cls()