Add learning methods (GA, ES, CEM, PPO)

719b42f verified 29 days ago

6.96 kB

"""Trois façons d'apprendre sans gradient, en ne regardant que le score :
- algorithme génétique (Evolution) : population, sélection, croisement, mutation
- stratégies d'évolution OpenAI-ES (Pole-Zero) : un seul cerveau, on estime dans quelle
  direction bouger ses poids en essayant des petites secousses au hasard
- entropie croisée CEM (Fly) : une moyenne et un écart-type par poids, on recentre sur l'élite
"""
from __future__ import annotations

import numpy as np

from ..nets import MLP
from .common import Logger, eval_flats, line, validate


INIT_OUT = 0.1          # dernière couche petite (comme d'habitude en RL) : au départ le bot tire presque droit


def _start(bot, rng, pop):
    """Population de départ : des cerveaux tirés au hasard (ou repris d'une sauvegarde)."""
    sizes = bot.policy.sizes
    a, b = sizes[-2], sizes[-1]
    flats = []
    for _ in range(pop):
        f = MLP(sizes, rng=rng).flat.copy()
        f[-(a * b + b):] *= INIT_OUT
        flats.append(f)
    if bot.trained_iters > 0:
        flats[0] = bot.policy.flat.copy()
    return np.array(flats)


def _finish(bot, flat, it, v, best, log, save):
    """À chaque contrôle : sauvegarde le record s'il est battu, et toujours l'état courant (_last)
    pour pouvoir reprendre avec --resume, même si l'entraînement est arrêté avec Ctrl+C."""
    bot.policy.set_flat(flat)
    bot.trained_iters = it + log.offset
    if v["touche"] > best:
        best = v["touche"]
        save(bot, "best")
    save(bot, "last")
    return best


def train_ga(bot, pool, iters, rng, say=print, save=None, resume=False, val_every=5, best0=-1.0,
             pop=48, n_seeds=128, elite=6, mut_sigma=0.02, cross_p=0.2):
    """Algorithme génétique.
    Les tirs sont très bruités (une carte facile ou difficile change tout) : pour ne pas garder
    un cerveau juste parce qu'il a eu de la chance, chaque élite garde la moyenne de tous ses
    scores depuis sa naissance. Plus il survit, plus son score est fiable.
    """
    log = Logger(bot.key, resume)
    P = _start(bot, rng, pop)
    h_sum, h_n = np.zeros(pop), np.zeros(pop)
    best, champion = best0, P[0]
    for it in range(1, iters + 1):
        seeds = rng.integers(0, 2 ** 31 - 1, n_seeds)
        F, _ = eval_flats(pool, bot.key, P, seeds)
        h_sum += F.mean(axis=1)
        h_n += 1
        fit = h_sum / h_n
        order = np.argsort(-fit)
        champion = P[order[0]].copy()

        # ---- génération suivante : les élites survivent (avec leur historique), les autres sont remplacés
        keep = order[:elite]
        children, cs, cn = [P[i].copy() for i in keep], list(h_sum[keep]), list(h_n[keep])
        parents = order[: pop // 4]                                          # le meilleur quart se reproduit
        while len(children) < pop:
            a = P[rng.choice(parents)]
            if rng.random() < cross_p:                                       # croisement : moitié des poids de chacun
                b = P[rng.choice(parents)]
                child = np.where(rng.random(a.size) < 0.5, a, b)
            else:
                child = a.copy()
            child = child + rng.normal(0, mut_sigma, child.size).astype(np.float32)
            children.append(child); cs.append(0.0); cn.append(0.0)
        P, h_sum, h_n = np.array(children, dtype=np.float32), np.array(cs), np.array(cn)

        if it % val_every == 0 or it == iters:
            v = validate(pool, bot.key, champion)
            log.log(it, float(fit[order[0]]), v)
            best = _finish(bot, champion, it, v, best, log, save)
            say(line(bot.key, it + log.offset, iters + log.offset, float(fit[order[0]]), v, best,
                     f"(champion vieux de {int(h_n[0])} gén.)"))
    bot.policy.set_flat(champion); bot.trained_iters = iters + log.offset; save(bot, "last")
    return best


class _Adam:
    def __init__(self, n, lr):
        self.m, self.v, self.t, self.lr = np.zeros(n), np.zeros(n), 0, lr

    def step(self, g):
        self.t += 1
        self.m = 0.9 * self.m + 0.1 * g
        self.v = 0.999 * self.v + 0.001 * g * g
        mh, vh = self.m / (1 - 0.9 ** self.t), self.v / (1 - 0.999 ** self.t)
        return self.lr * mh / (np.sqrt(vh) + 1e-8)


def _ranks(x):
    r = np.empty(len(x)); r[np.argsort(x)] = np.arange(len(x))
    return r / (len(x) - 1) - 0.5


def train_es(bot, pool, iters, rng, say=print, save=None, resume=False, val_every=5, best0=-1.0,
             pairs=24, n_seeds=64, sigma=0.02, lr=0.01, decay=0.005):
    log = Logger(bot.key, resume)
    theta = _start(bot, rng, 1)[0].astype(np.float64)
    opt, best = _Adam(theta.size, lr), best0
    for it in range(1, iters + 1):
        seeds = rng.integers(0, 2 ** 31 - 1, n_seeds)
        eps = rng.standard_normal((pairs, theta.size))
        cand = np.vstack([theta + sigma * eps, theta - sigma * eps, theta[None]])   # secousses + / -, + le centre
        F, _ = eval_flats(pool, bot.key, cand.astype(np.float32), seeds)
        fit = F.mean(axis=1)
        r = _ranks(fit[:-1])
        grad = (r[:pairs] - r[pairs:]) @ eps / (2 * pairs * sigma)
        theta += opt.step(grad - decay * theta)
        if it % val_every == 0 or it == iters:
            v = validate(pool, bot.key, theta.astype(np.float32))
            log.log(it, float(fit[-1]), v)
            best = _finish(bot, theta.astype(np.float32), it, v, best, log, save)
            say(line(bot.key, it + log.offset, iters + log.offset, float(fit[-1]), v, best))
    bot.policy.set_flat(theta.astype(np.float32)); bot.trained_iters = iters + log.offset; save(bot, "last")
    return best


def train_cem(bot, pool, iters, rng, say=print, save=None, resume=False, val_every=5, best0=-1.0,
              pop=64, n_seeds=96, elite_frac=0.25, sigma0=0.04, extra_noise=0.005):
    log = Logger(bot.key, resume)
    mu = _start(bot, rng, 1)[0].astype(np.float64)
    std = np.full(mu.size, sigma0)
    n_el, best = max(2, int(pop * elite_frac)), best0
    for it in range(1, iters + 1):
        seeds = rng.integers(0, 2 ** 31 - 1, n_seeds)
        cand = mu + std * rng.standard_normal((pop, mu.size))
        cand[0] = mu
        F, _ = eval_flats(pool, bot.key, cand.astype(np.float32), seeds)
        fit = F.mean(axis=1)
        el = cand[np.argsort(-fit)[:n_el]]
        noise = extra_noise * max(0.0, 1 - it / (0.8 * iters))
        mu = 0.5 * el.mean(axis=0) + 0.5 * mu
        std = np.sqrt(0.5 * el.var(axis=0) + 0.5 * std ** 2) + noise
        if it % val_every == 0 or it == iters:
            v = validate(pool, bot.key, mu.astype(np.float32))
            log.log(it, float(fit[0]), v)
            best = _finish(bot, mu.astype(np.float32), it, v, best, log, save)
            say(line(bot.key, it + log.offset, iters + log.offset, float(fit[0]), v, best, f"(σ moyen {std.mean():.3f})"))
    bot.policy.set_flat(mu.astype(np.float32)); bot.trained_iters = iters + log.offset; save(bot, "last")
    return best

	"""Trois façons d'apprendre sans gradient, en ne regardant que le score :

	- algorithme génétique (Evolution) : population, sélection, croisement, mutation
	- stratégies d'évolution OpenAI-ES (Pole-Zero) : un seul cerveau, on estime dans quelle
	direction bouger ses poids en essayant des petites secousses au hasard
	- entropie croisée CEM (Fly) : une moyenne et un écart-type par poids, on recentre sur l'élite
	"""
	from __future__ import annotations

	import numpy as np

	from ..nets import MLP
	from .common import Logger, eval_flats, line, validate


	INIT_OUT = 0.1 # dernière couche petite (comme d'habitude en RL) : au départ le bot tire presque droit


	def _start(bot, rng, pop):
	"""Population de départ : des cerveaux tirés au hasard (ou repris d'une sauvegarde)."""
	sizes = bot.policy.sizes
	a, b = sizes[-2], sizes[-1]
	flats = []
	for _ in range(pop):
	f = MLP(sizes, rng=rng).flat.copy()
	f[-(a * b + b):] *= INIT_OUT
	flats.append(f)
	if bot.trained_iters > 0:
	flats[0] = bot.policy.flat.copy()
	return np.array(flats)


	def _finish(bot, flat, it, v, best, log, save):
	"""À chaque contrôle : sauvegarde le record s'il est battu, et toujours l'état courant (_last)
	pour pouvoir reprendre avec --resume, même si l'entraînement est arrêté avec Ctrl+C."""
	bot.policy.set_flat(flat)
	bot.trained_iters = it + log.offset
	if v["touche"] > best:
	best = v["touche"]
	save(bot, "best")
	save(bot, "last")
	return best


	def train_ga(bot, pool, iters, rng, say=print, save=None, resume=False, val_every=5, best0=-1.0,
	pop=48, n_seeds=128, elite=6, mut_sigma=0.02, cross_p=0.2):
	"""Algorithme génétique.

	Les tirs sont très bruités (une carte facile ou difficile change tout) : pour ne pas garder
	un cerveau juste parce qu'il a eu de la chance, chaque élite garde la moyenne de tous ses
	scores depuis sa naissance. Plus il survit, plus son score est fiable.
	"""
	log = Logger(bot.key, resume)
	P = _start(bot, rng, pop)
	h_sum, h_n = np.zeros(pop), np.zeros(pop)
	best, champion = best0, P[0]
	for it in range(1, iters + 1):
	seeds = rng.integers(0, 2 ** 31 - 1, n_seeds)
	F, _ = eval_flats(pool, bot.key, P, seeds)
	h_sum += F.mean(axis=1)
	h_n += 1
	fit = h_sum / h_n
	order = np.argsort(-fit)
	champion = P[order[0]].copy()

	# ---- génération suivante : les élites survivent (avec leur historique), les autres sont remplacés
	keep = order[:elite]
	children, cs, cn = [P[i].copy() for i in keep], list(h_sum[keep]), list(h_n[keep])
	parents = order[: pop // 4] # le meilleur quart se reproduit
	while len(children) < pop:
	a = P[rng.choice(parents)]
	if rng.random() < cross_p: # croisement : moitié des poids de chacun
	b = P[rng.choice(parents)]
	child = np.where(rng.random(a.size) < 0.5, a, b)
	else:
	child = a.copy()
	child = child + rng.normal(0, mut_sigma, child.size).astype(np.float32)
	children.append(child); cs.append(0.0); cn.append(0.0)
	P, h_sum, h_n = np.array(children, dtype=np.float32), np.array(cs), np.array(cn)

	if it % val_every == 0 or it == iters:
	v = validate(pool, bot.key, champion)
	log.log(it, float(fit[order[0]]), v)
	best = _finish(bot, champion, it, v, best, log, save)
	say(line(bot.key, it + log.offset, iters + log.offset, float(fit[order[0]]), v, best,
	f"(champion vieux de {int(h_n[0])} gén.)"))
	bot.policy.set_flat(champion); bot.trained_iters = iters + log.offset; save(bot, "last")
	return best


	class _Adam:
	def __init__(self, n, lr):
	self.m, self.v, self.t, self.lr = np.zeros(n), np.zeros(n), 0, lr

	def step(self, g):
	self.t += 1
	self.m = 0.9 * self.m + 0.1 * g
	self.v = 0.999 * self.v + 0.001 * g * g
	mh, vh = self.m / (1 - 0.9 self.t), self.v / (1 - 0.999 self.t)
	return self.lr * mh / (np.sqrt(vh) + 1e-8)


	def _ranks(x):
	r = np.empty(len(x)); r[np.argsort(x)] = np.arange(len(x))
	return r / (len(x) - 1) - 0.5


	def train_es(bot, pool, iters, rng, say=print, save=None, resume=False, val_every=5, best0=-1.0,
	pairs=24, n_seeds=64, sigma=0.02, lr=0.01, decay=0.005):
	log = Logger(bot.key, resume)
	theta = _start(bot, rng, 1)[0].astype(np.float64)
	opt, best = _Adam(theta.size, lr), best0
	for it in range(1, iters + 1):
	seeds = rng.integers(0, 2 ** 31 - 1, n_seeds)
	eps = rng.standard_normal((pairs, theta.size))
	cand = np.vstack([theta + sigma * eps, theta - sigma * eps, theta[None]]) # secousses + / -, + le centre
	F, _ = eval_flats(pool, bot.key, cand.astype(np.float32), seeds)
	fit = F.mean(axis=1)
	r = _ranks(fit[:-1])
	grad = (r[:pairs] - r[pairs:]) @ eps / (2 * pairs * sigma)
	theta += opt.step(grad - decay * theta)
	if it % val_every == 0 or it == iters:
	v = validate(pool, bot.key, theta.astype(np.float32))
	log.log(it, float(fit[-1]), v)
	best = _finish(bot, theta.astype(np.float32), it, v, best, log, save)
	say(line(bot.key, it + log.offset, iters + log.offset, float(fit[-1]), v, best))
	bot.policy.set_flat(theta.astype(np.float32)); bot.trained_iters = iters + log.offset; save(bot, "last")
	return best


	def train_cem(bot, pool, iters, rng, say=print, save=None, resume=False, val_every=5, best0=-1.0,
	pop=64, n_seeds=96, elite_frac=0.25, sigma0=0.04, extra_noise=0.005):
	log = Logger(bot.key, resume)
	mu = _start(bot, rng, 1)[0].astype(np.float64)
	std = np.full(mu.size, sigma0)
	n_el, best = max(2, int(pop * elite_frac)), best0
	for it in range(1, iters + 1):
	seeds = rng.integers(0, 2 ** 31 - 1, n_seeds)
	cand = mu + std * rng.standard_normal((pop, mu.size))
	cand[0] = mu
	F, _ = eval_flats(pool, bot.key, cand.astype(np.float32), seeds)
	fit = F.mean(axis=1)
	el = cand[np.argsort(-fit)[:n_el]]
	noise = extra_noise * max(0.0, 1 - it / (0.8 * iters))
	mu = 0.5 * el.mean(axis=0) + 0.5 * mu
	std = np.sqrt(0.5 * el.var(axis=0) + 0.5 * std ** 2) + noise
	if it % val_every == 0 or it == iters:
	v = validate(pool, bot.key, mu.astype(np.float32))
	log.log(it, float(fit[0]), v)
	best = _finish(bot, mu.astype(np.float32), it, v, best, log, save)
	say(line(bot.key, it + log.offset, iters + log.offset, float(fit[0]), v, best, f"(σ moyen {std.mean():.3f})"))
	bot.policy.set_flat(mu.astype(np.float32)); bot.trained_iters = iters + log.offset; save(bot, "last")
	return best