"""Policies for robot sumo. All act on batches of egocentric observations. Interface: ``reset(n, rng)`` once per batch of matches, then ``act(obs)`` with obs of shape (n, OBS_DIM) returning actions (n, 2) = (thrust, turn) in [-1, 1]. Levels ------ L0 RandomPolicy -- random wandering (new random command every 0.5 s) L1 ChargerPolicy -- hand-coded: turn to face the opponent and charge -- DodgerPolicy -- hand-coded "matador": sidestep charges, stay central, hit the opponent from the side L2+ MLPPolicy -- small tanh network; weights found by evolution strategies """ from __future__ import annotations import os import numpy as np from sumo_env import OBS_DIM # Observation indices (see SumoEnv._observe_true) OPP_X, OPP_Y = 0, 1 OPP_HCOS, OPP_HSIN = 4, 5 CEN_X, CEN_Y = 9, 10 OWN_D = 11 def _steer(target_angle: np.ndarray, gain: float = 3.0, cone: float = 0.7, slow: float = 0.25): """Turn toward a body-frame angle; drive fast only when roughly aligned.""" turn = np.clip(gain * target_angle, -1.0, 1.0) thrust = np.where(np.abs(target_angle) < cone, 1.0, slow) return np.stack([thrust, turn], -1) class RandomPolicy: """L0: holds a uniformly random command for `hold` control steps.""" name = "L0-random" def __init__(self, hold: int = 10): self.hold = hold def reset(self, n: int, rng: np.random.Generator | None = None): self.rng = rng if rng is not None else np.random.default_rng() self.k = 0 self.a = self.rng.uniform(-1, 1, (n, 2)) def act(self, obs: np.ndarray) -> np.ndarray: if self.k % self.hold == 0: self.a = self.rng.uniform(-1, 1, (obs.shape[0], 2)) self.k += 1 return self.a class ChargerPolicy: """L1: face the opponent and drive at it full speed.""" name = "L1-charger" def reset(self, n: int, rng=None): pass def act(self, obs: np.ndarray) -> np.ndarray: ang = np.arctan2(obs[:, OPP_Y], obs[:, OPP_X]) return _steer(ang, gain=3.0, cone=0.6, slow=0.2) class DodgerPolicy: """Hand-coded 'matador'. * If I'm near the edge: go back toward the center. * Else if the opponent is pointed at me and close: drive perpendicular to the line between us (choosing the side that keeps me nearer the center). * Else: charge (the opponent is showing its side/back). """ name = "Dodger" def reset(self, n: int, rng=None): pass def act(self, obs: np.ndarray) -> np.ndarray: rel = obs[:, [OPP_X, OPP_Y]] dist = np.linalg.norm(rel, axis=-1) + 1e-9 to_opp = np.arctan2(rel[:, 1], rel[:, 0]) to_center = np.arctan2(obs[:, CEN_Y], obs[:, CEN_X]) # direction the opponent faces, in my frame, vs direction from it to me opp_face = obs[:, [OPP_HCOS, OPP_HSIN]] aiming = np.sum(opp_face * (-rel / dist[:, None]), -1) # 1 = pointed right at me perp = np.stack([-rel[:, 1], rel[:, 0]], -1) # 90 deg left of opponent direction sign = np.sign(np.sum(perp * obs[:, [CEN_X, CEN_Y]], -1) + 1e-9) perp = perp * sign[:, None] dodge = np.arctan2(perp[:, 1], perp[:, 0]) target = to_opp.copy() threatened = (aiming > 0.7) & (dist < 1.0) target = np.where(threatened, dodge, target) target = np.where(obs[:, OWN_D] > 0.8, to_center, target) return _steer(target, gain=3.0, cone=0.8, slow=0.3) # ---------------------------------------------------------------- MLP HIDDEN = (32, 32) LAYER_SIZES = [OBS_DIM, *HIDDEN, 2] def n_params(sizes=LAYER_SIZES) -> int: return sum(a * b + b for a, b in zip(sizes[:-1], sizes[1:])) N_PARAMS = n_params() def init_params(rng: np.random.Generator, sizes=LAYER_SIZES) -> np.ndarray: """Flat parameter vector with fan-in scaled init (small output layer).""" chunks = [] for k, (a, b) in enumerate(zip(sizes[:-1], sizes[1:])): scale = 1.0 / np.sqrt(a) * (0.1 if k == len(sizes) - 2 else 1.0) chunks += [rng.standard_normal(a * b) * scale, np.zeros(b)] return np.concatenate(chunks) def unpack(flat: np.ndarray, sizes=LAYER_SIZES): """(..., D) flat params -> list of (W (..., in, out), b (..., out)).""" layers, i = [], 0 lead = flat.shape[:-1] for a, b in zip(sizes[:-1], sizes[1:]): W = flat[..., i : i + a * b].reshape(*lead, a, b) i += a * b bias = flat[..., i : i + b] i += b layers.append((W, bias)) return layers class MLPPolicy: """Tanh MLP. Either one parameter vector shared by all envs, or one per env. Pass ``params`` of shape (D,) for a single policy or (n, D) to evaluate many different parameter vectors in one vectorized forward pass (this is how all ES population members are evaluated simultaneously). """ def __init__(self, params: np.ndarray, name: str = "MLP"): self.params = np.asarray(params, dtype=float) self.name = name def reset(self, n: int, rng=None): if self.params.ndim == 1: self.layers = unpack(self.params) self.batched = False else: assert self.params.shape[0] == n, "need one parameter row per env" self.layers = unpack(self.params) self.batched = True def act(self, obs: np.ndarray) -> np.ndarray: x = obs for W, b in self.layers: if self.batched: x = np.tanh(np.einsum("ei,eio->eo", x, W) + b) else: x = np.tanh(x @ W + b) return x @staticmethod def load(path: str, name: str | None = None) -> "MLPPolicy": d = np.load(path) if name is None: name = str(d["name"]) if "name" in d.files else path return MLPPolicy(d["params"], name) class MuxPolicy: """Runs different policies on different subsets of the env batch.""" def __init__(self, groups): """groups: list of (policy, env_index_array).""" self.groups = [(p, np.asarray(idx)) for p, idx in groups if len(idx) > 0] self.name = "mux" def reset(self, n: int, rng=None): self.n = n for p, idx in self.groups: p.reset(len(idx), rng) def act(self, obs: np.ndarray) -> np.ndarray: out = np.zeros((obs.shape[0], 2)) for p, idx in self.groups: out[idx] = p.act(obs[idx]) return out class HumanPolicy: """Single-env policy whose action is set from keyboard events (watch.py).""" name = "Human" def __init__(self): self.action = np.zeros(2) def reset(self, n: int, rng=None): self.action[:] = 0 def act(self, obs: np.ndarray) -> np.ndarray: return np.tile(self.action, (obs.shape[0], 1)) # ---------------------------------------------------------------- registry WEIGHTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "weights") LEVELS = ["L0", "L1", "Dodger", "L2", "L3", "L4"] def load_level(level: str, weights_dir: str = WEIGHTS_DIR): """Build a policy from a level name. "L0", "L1", "Dodger" are hand-coded; anything else is loaded from ``weights/.npz`` (e.g. "L2", "L3", "L4", "L3_gen0100"). A path to an .npz file also works. """ key = level.lower() if key in ("l0", "random"): return RandomPolicy() if key in ("l1", "charger"): return ChargerPolicy() if key in ("dodger", "l1b"): return DodgerPolicy() path = level if level.endswith(".npz") else os.path.join(weights_dir, f"{level}.npz") if not os.path.exists(path): raise FileNotFoundError(f"no weights for level {level!r} at {path} (run train.py first)") return MLPPolicy.load(path, name=os.path.basename(path)[:-4]) # ---------------------------------------------------------------- match runner def run_matches(env, pol0, pol1, rng=None, starts=None, record=False): """Play env.n matches to completion (or the time cap). Returns dict with outcome (+1 robot0 win / -1 loss / 0 tie), end_step, final distances from the ring center, and optionally the state trajectory. """ rng = rng if rng is not None else np.random.default_rng() o0, o1 = env.reset(starts) pol0.reset(env.n, rng) pol1.reset(env.n, rng) frames = [env.snapshot()] if record else None final_d = np.zeros((env.n, 2)) captured = np.zeros(env.n, dtype=bool) while True: o0, o1, done, outcome = env.step(pol0.act(o0), pol1.act(o1)) newly = done & ~captured if np.any(newly): final_d[newly] = np.linalg.norm(env.pos[newly], axis=-1) captured |= newly if record: frames.append(env.snapshot()) if done.all(): break return {"outcome": outcome, "end_step": env.end_step.copy(), "final_dist": final_d, "frames": frames}