import { randomSource } from './star-training';
import { HockeyNetwork } from './hockey-network';
export const HOCKEY_ACTIONS = ['Up', 'Stay', 'Down'];
export const RINK_SIZE = 12;
export const PADDLE_HALF = .13;
export const PUCK_RADIUS = .018;
export type HockeyMode = 'state' | 'pixels';
export type Side = 0 | 1;
export type HockeySettings = { rates: [number, number]; exploration: [number, number] };
export const HOCKEY_DEFAULTS: HockeySettings = { rates: [.01, .01], exploration: [.2, .2] };
export type Rink = { paddles: [number, number]; x: number; y: number; vx: number; vy: number };
type Transition = { state: number[]; action: number; reward: number; next: number[]; done: boolean };
export type MatchStats = { rallies: number; goals: [number, number]; returns: [number, number]; shots: [number, number]; draws: number; recent: [number[], number[]]; lengths: number[] };
const stats = (): MatchStats => ({ rallies: 0, goals: [0, 0], returns: [0, 0], shots: [0, 0], draws: 0, recent: [[], []], lengths: [] });
export const createMatchStats = stats;
export function rinkPixels(rink: Rink, side: Side): number[] {
const frame = Array(RINK_SIZE * RINK_SIZE).fill(0) as number[];
const own = rink.paddles[side], other = rink.paddles[1 - side];
for (let row = 0; row < RINK_SIZE; row++) {
for (const [column, paddle] of [[0, own], [RINK_SIZE - 1, other]]) {
const center = paddle * (RINK_SIZE - 1), half = PADDLE_HALF * (RINK_SIZE - 1);
const coverage = Math.max(0, Math.min(row + .5, center + half) - Math.max(row - .5, center - half));
frame[row * RINK_SIZE + column] = .7 * coverage;
}
}
const x = (side === 0 ? rink.x : 1 - rink.x) * (RINK_SIZE - 1), y = rink.y * (RINK_SIZE - 1);
const left = Math.floor(x), top = Math.floor(y);
for (let dy = 0; dy <= 1; dy++) for (let dx = 0; dx <= 1; dx++) {
const column = left + dx, row = top + dy;
if (column >= 0 && column < RINK_SIZE && row >= 0 && row < RINK_SIZE) {
const brightness = (dx ? x - left : 1 - (x - left)) * (dy ? y - top : 1 - (y - top));
frame[row * RINK_SIZE + column] = Math.max(frame[row * RINK_SIZE + column], brightness);
}
}
return frame;
}
export class HockeyAgent {
network: HockeyNetwork;
target: HockeyNetwork;
memory: Transition[] = [];
cursor = 0;
steps = 0;
loss = 0;
action = 1;
exploring = false;
values = [0, 0, 0];
lastInput: number[];
random: () => number;
constructor(mode: HockeyMode, seed: number, depth = 1, public options = { memoryLimit: 1500, replaySamples: 8, width: 0 }) {
this.random = randomSource(seed);
const inputs = mode === 'pixels' ? 2 * RINK_SIZE ** 2 : 7;
const hidden = options.width || (mode === 'pixels' ? 24 : 16);
this.network = new HockeyNetwork(this.random, inputs, hidden, depth);
this.target = new HockeyNetwork(this.random, inputs, hidden, depth);
this.target.copyFrom(this.network);
this.lastInput = Array(inputs).fill(0);
}
choose(state: number[], epsilon: number) {
this.lastInput = state;
this.values = this.network.forward(state).values;
this.exploring = epsilon > 0 && this.random() < epsilon;
this.action = this.exploring ? Math.floor(this.random() * 3) : this.values.indexOf(Math.max(...this.values));
return this.action;
}
learn(transition: Transition, rate: number) {
this.steps++;
if (this.memory.length < this.options.memoryLimit) this.memory.push(transition);
else { this.memory[this.cursor] = transition; this.cursor = (this.cursor + 1) % this.options.memoryLimit; }
if (this.memory.length >= 32) {
let loss = 0;
for (let i = 0; i < this.options.replaySamples; i++) {
const sample = this.memory[Math.floor(this.random() * this.memory.length)];
const target = sample.reward + (sample.done ? 0 : .95 * Math.max(...this.target.forward(sample.next).values));
loss += this.network.train(sample.state, sample.action, target, rate / 4);
}
this.loss = .95 * this.loss + .05 * loss / this.options.replaySamples;
}
if (this.steps % 120 === 0) this.target.copyFrom(this.network);
}
}
export type HockeyCheckpoint = { rally: number; rates: [number, number]; averageLength: number; goals: [number, number] };
export class HockeyMatch {
agents: [HockeyAgent, HockeyAgent];
rink: Rink = { paddles: [.5, .5], x: .5, y: .5, vx: .028, vy: .01 };
random: () => number;
previous: [number[], number[]];
training = stats();
watching = stats();
history: HockeyCheckpoint[] = [];
lastCompleted: HockeyCheckpoint;
stride = 5;
tick = 0;
terminal = false;
winner: Side | null = null;
hits = 0;
rewards: [number, number] = [0, 0];
serves = 0;
constructor(public mode: HockeyMode = 'state', seed = 42, public depths: [number, number] = [1, 1]) {
this.random = randomSource(seed);
this.agents = [new HockeyAgent(mode, seed + 11, depths[0]), new HockeyAgent(mode, seed + 29, depths[1])];
this.previous = [[], []];
this.newRally();
this.agents.forEach((agent, side) => { agent.lastInput = this.observation(side as Side); });
this.lastCompleted = this.checkpoint();
this.history.push(this.lastCompleted);
}
newRally() {
this.rink = { paddles: [.3 + .4 * this.random(), .3 + .4 * this.random()], x: .5, y: .2 + .6 * this.random(), vx: (this.serves++ % 2 === 0 ? 1 : -1) * .028, vy: (this.random() * 2 - 1) * .018 };
this.tick = 0; this.terminal = false; this.winner = null; this.hits = 0; this.rewards = [0, 0];
this.previous = [rinkPixels(this.rink, 0), rinkPixels(this.rink, 1)];
}
startWatch() { this.watching = stats(); this.newRally(); }
observation(side: Side) {
if (this.mode === 'pixels') return [...rinkPixels(this.rink, side), ...this.previous[side]];
const { x, y, vx, vy, paddles } = this.rink;
return [2 * paddles[side] - 1, 2 * paddles[1 - side] - 1, 2 * (side === 0 ? x : 1 - x) - 1, 2 * y - 1, (side === 0 ? vx : -vx) / .055, vy / .055, y - paddles[side]];
}
advance(actions: [number, number], matchStats: MatchStats) {
const rink = this.rink;
const approaching: Side = rink.vx < 0 ? 0 : 1;
const before = Math.abs(rink.y - rink.paddles[approaching]);
rink.paddles = rink.paddles.map((paddle, side) => Math.max(PADDLE_HALF, Math.min(1 - PADDLE_HALF, paddle + (actions[side] - 1) * .035))) as [number, number];
this.rewards = [0, 0];
// Coaching depends only on movement toward the current puck height, before
// the puck moves. It is supplied by the environment, not as an action label.
this.rewards[approaching] = 2 * (before - Math.abs(rink.y - rink.paddles[approaching]));
const oldX = rink.x;
rink.x += rink.vx; rink.y += rink.vy;
if (rink.y < PUCK_RADIUS) { rink.y = 2 * PUCK_RADIUS - rink.y; rink.vy = Math.abs(rink.vy); }
if (rink.y > 1 - PUCK_RADIUS) { rink.y = 2 * (1 - PUCK_RADIUS) - rink.y; rink.vy = -Math.abs(rink.vy); }
const defender: Side | null = rink.vx < 0 && oldX > .065 && rink.x <= .065 ? 0 : rink.vx > 0 && oldX < .935 && rink.x >= .935 ? 1 : null;
if (defender !== null) {
const returned = Math.abs(rink.y - rink.paddles[defender]) <= PADDLE_HALF + PUCK_RADIUS;
matchStats.shots[defender]++;
matchStats.recent[defender].push(Number(returned));
if (matchStats.recent[defender].length > 100) matchStats.recent[defender].shift();
if (returned) {
matchStats.returns[defender]++; this.hits++; this.rewards[defender] += .25;
const offset = (rink.y - rink.paddles[defender]) / (PADDLE_HALF + PUCK_RADIUS);
rink.vx = (defender === 0 ? 1 : -1) * Math.min(.055, Math.abs(rink.vx) * 1.07);
rink.vy = offset * .05 + (actions[defender] - 1) * .005;
rink.x = defender === 0 ? .065 + (.065 - rink.x) : .935 - (rink.x - .935);
} else {
this.terminal = true; this.winner = (1 - defender) as Side;
this.rewards[this.winner] = 1; this.rewards[defender] = -1;
matchStats.goals[this.winner]++;
}
}
this.tick++;
if (!this.terminal && this.tick >= 400) { this.terminal = true; matchStats.draws++; }
if (this.terminal) {
matchStats.rallies++; matchStats.lengths.push(this.hits);
if (matchStats.lengths.length > 100) matchStats.lengths.shift();
}
}
step(settings: HockeySettings, train = true) {
if (this.terminal) this.newRally();
const states = [this.observation(0), this.observation(1)];
// Both decisions observe the same pre-move state. Neither sees the other's
// action before selecting its own.
const actions = this.agents.map((agent, side) => agent.choose(states[side], train ? settings.exploration[side] : 0)) as [number, number];
const frames = [rinkPixels(this.rink, 0), rinkPixels(this.rink, 1)] as [number[], number[]];
this.advance(actions, train ? this.training : this.watching);
this.previous = frames;
if (train) {
this.agents.forEach((agent, side) => agent.learn({ state: states[side], action: actions[side], reward: this.rewards[side], next: this.observation(side as Side), done: this.terminal }, settings.rates[side]));
if (this.terminal) this.lastCompleted = this.checkpoint();
if (this.terminal && this.training.rallies % this.stride === 0) {
this.history.push(this.lastCompleted);
if (this.history.length > 201) { this.history = this.history.filter((_, i) => i % 2 === 0); this.stride *= 2; }
}
}
}
checkpoint(): HockeyCheckpoint {
const rates = this.training.recent.map(shots => shots.length ? shots.reduce((sum, hit) => sum + hit, 0) / shots.length : NaN) as [number, number];
const lengths = this.training.lengths;
return { rally: this.training.rallies, rates, averageLength: lengths.length ? lengths.reduce((a, b) => a + b, 0) / lengths.length : 0, goals: [...this.training.goals] };
}
}