Q-Learning Cart-Pole

// Cart-pole balanced by tabular Q-learning. The agent starts random,
// fails fast, and slowly converges to a policy that keeps theta in bounds.
// Episode-length history is drawn below the cart.
//
// Pacing: the visible episode advances at a fixed, human-readable rate
// (~67 physics steps/sec, so a 200-step success lasts ~3s). To keep
// learning progressing while the eye watches one episode, we also run a
// burst of "background" training episodes each frame that share the same
// Q-table but are not rendered.

const MC = 1.0, MP = 0.1, L = 1.0, G = 9.81, FMAG = 10;
const X_LIM = 2.4, TH_LIM = 12 * Math.PI / 180;
const MAX_STEPS = 220;
const DT_PHYS = 0.02;

// Visible cart-pole: one physics step every VIS_STEP_DT seconds of wall time.
// 0.015s -> ~67 Hz -> a 200-step episode visible over ~3 seconds.
const VIS_STEP_DT = 0.015;
// Background trainer: how many off-screen training steps per rendered frame.
// Keeps Q-table converging at roughly the old throughput without flicker.
const BG_STEPS_PER_FRAME = 22;
// After a failure, hold the failed pose for this long before respawning.
const FAIL_HOLD_S = 0.6;
// And then fade in the fresh pose over this long.
const FADE_IN_S = 0.35;

// Bin edges: state = [x, xdot, theta, thetadot]
const BX  = [-0.8, 0.8];
const BXD = [-0.5, 0.5];
const BT  = [-6, -1, 0, 1, 6].map(d => d * Math.PI / 180);
const BTD = [-50 * Math.PI / 180, 50 * Math.PI / 180];
const NX = BX.length + 1, NXD = BXD.length + 1, NT = BT.length + 1, NTD = BTD.length + 1;
const NS = NX * NXD * NT * NTD;
const NA = 2;

const ALPHA = 0.5, GAMMA = 0.98;
const EPS0 = 1.0, EPS_END = 0.02, EPS_DECAY = 0.985;

let Q, ep, eps, hist, best, totalSteps;
// Visible episode state.
let vSt, vPrev, vStep, vAcc, vPhase, vPhaseT;
// Background trainer state.
let bgSt, bgStep;

function bin(v, edges) {
  for (let i = 0; i < edges.length; i++) if (v < edges[i]) return i;
  return edges.length;
}
function sidx(s) {
  const i0 = bin(s[0], BX), i1 = bin(s[1], BXD);
  const i2 = bin(s[2], BT), i3 = bin(s[3], BTD);
  return ((i0 * NXD + i1) * NT + i2) * NTD + i3;
}
function deriv(s, F) {
  const [, xd, th, thd] = s;
  const sn = Math.sin(th), cs = Math.cos(th), tot = MC + MP, lc = L / 2;
  const tmp = (F + MP * lc * thd * thd * sn) / tot;
  const thdd = (G * sn - cs * tmp) / (lc * (4 / 3 - MP * cs * cs / tot));
  const xdd = tmp - MP * lc * thdd * cs / tot;
  return [xd, xdd, thd, thdd];
}
function physStep(s, F) {
  const h = DT_PHYS;
  const k1 = deriv(s, F);
  const k2 = deriv(s.map((v, i) => v + h / 2 * k1[i]), F);
  const k3 = deriv(s.map((v, i) => v + h / 2 * k2[i]), F);
  const k4 = deriv(s.map((v, i) => v + h * k3[i]), F);
  return s.map((v, i) => v + h / 6 * (k1[i] + 2 * k2[i] + 2 * k3[i] + k4[i]));
}
function isDone(s) {
  return Math.abs(s[0]) > X_LIM || Math.abs(s[2]) > TH_LIM;
}
function freshState() {
  return [0, 0, (Math.random() - 0.5) * 0.05, (Math.random() - 0.5) * 0.05];
}
function pickAction(si) {
  if (Math.random() < eps) return Math.random() < 0.5 ? 0 : 1;
  const q0 = Q[si * NA], q1 = Q[si * NA + 1];
  return q1 > q0 ? 1 : 0;
}
// Run one Q-learning step from `s`, updating Q in place. Returns
// { next, done, step+1 }.
function qStep(s, step) {
  const si = sidx(s);
  const a = pickAction(si);
  const F = a === 0 ? -FMAG : FMAG;
  const ns = physStep(s, F);
  const d = isDone(ns) || step + 1 >= MAX_STEPS;
  const r = d && step + 1 < MAX_STEPS ? -1 : 0.02;
  const nsi = sidx(ns);
  const target = d ? r : r + GAMMA * Math.max(Q[nsi * NA], Q[nsi * NA + 1]);
  Q[si * NA + a] += ALPHA * (target - Q[si * NA + a]);
  totalSteps++;
  return { next: ns, done: d, step: step + 1 };
}

function bgTrainStep() {
  const out = qStep(bgSt, bgStep);
  bgSt = out.next; bgStep = out.step;
  if (out.done) {
    hist.push(bgStep);
    if (hist.length > 600) hist.shift();
    if (bgStep > best) best = bgStep;
    eps = Math.max(EPS_END, eps * EPS_DECAY);
    ep++;
    bgSt = freshState();
    bgStep = 0;
  }
}

function visStep() {
  // Phase machine for the rendered episode.
  if (vPhase === "run") {
    vPrev = vSt.slice();
    const out = qStep(vSt, vStep);
    vSt = out.next; vStep = out.step;
    if (out.done) { vPhase = "fail"; vPhaseT = 0; }
  } else if (vPhase === "fail") {
    // Hold pose; the wall-clock timer is advanced by tick().
  } else if (vPhase === "respawn") {
    // Faded in; no physics, just timer.
  }
}

function init() {
  Q = new Float32Array(NS * NA);
  ep = 0; eps = EPS0; hist = []; best = 0; totalSteps = 0;
  vSt = freshState(); vPrev = vSt.slice(); vStep = 0; vAcc = 0;
  vPhase = "run"; vPhaseT = 0;
  bgSt = freshState(); bgStep = 0;
}

function drawCart(ctx, W, panelH, displaySt, fail, alpha) {
  const cy = panelH * 0.7;
  const pxM = Math.min(W, panelH) * 0.18;
  ctx.strokeStyle = "rgba(120,140,180,0.55)"; ctx.lineWidth = 2;
  ctx.beginPath(); ctx.moveTo(0, cy); ctx.lineTo(W, cy); ctx.stroke();
  const limPx = X_LIM * pxM;
  ctx.strokeStyle = "rgba(255,90,90,0.6)"; ctx.lineWidth = 1.5;
  ctx.beginPath();
  ctx.moveTo(W / 2 - limPx, cy - 18); ctx.lineTo(W / 2 - limPx, cy + 10);
  ctx.moveTo(W / 2 + limPx, cy - 18); ctx.lineTo(W / 2 + limPx, cy + 10);
  ctx.stroke();

  const prevAlpha = ctx.globalAlpha;
  ctx.globalAlpha = prevAlpha * alpha;
  const cx = W / 2 + displaySt[0] * pxM;
  const px = cx + Math.sin(displaySt[2]) * L * pxM;
  const py = cy - 12 - Math.cos(displaySt[2]) * L * pxM;
  ctx.strokeStyle = fail ? "rgba(255,120,120,0.95)" : "rgba(255,210,120,0.95)";
  ctx.lineWidth = 5; ctx.lineCap = "round";
  ctx.beginPath(); ctx.moveTo(cx, cy - 12); ctx.lineTo(px, py); ctx.stroke();
  ctx.fillStyle = fail ? "#ff8080" : "#ffcb5a";
  ctx.beginPath(); ctx.arc(px, py, 8, 0, Math.PI * 2); ctx.fill();
  ctx.fillStyle = fail ? "#a05060" : "#5a8fd6";
  ctx.fillRect(cx - 26, cy - 22, 52, 22);
  ctx.strokeStyle = "rgba(255,255,255,0.5)"; ctx.lineWidth = 1.5;
  ctx.strokeRect(cx - 25.5, cy - 21.5, 51, 21);
  ctx.fillStyle = "#222";
  ctx.beginPath(); ctx.arc(cx - 16, cy, 5, 0, Math.PI * 2); ctx.fill();
  ctx.beginPath(); ctx.arc(cx + 16, cy, 5, 0, Math.PI * 2); ctx.fill();
  ctx.globalAlpha = prevAlpha;
}
function drawChart(ctx, W, H, yTop) {
  const padL = 36, padR = 12, padT = 14, padB = 22;
  const x0 = padL, y0 = yTop + padT, x1 = W - padR, y1 = H - padB;
  ctx.fillStyle = "rgba(0,0,0,0.45)"; ctx.fillRect(x0 - 6, y0 - 6, x1 - x0 + 12, y1 - y0 + 12);
  ctx.strokeStyle = "rgba(120,140,180,0.55)"; ctx.lineWidth = 1;
  ctx.beginPath();
  ctx.moveTo(x0, y0); ctx.lineTo(x0, y1); ctx.lineTo(x1, y1); ctx.stroke();
  const yMax = MAX_STEPS;
  ctx.strokeStyle = "rgba(120,255,180,0.25)";
  ctx.setLineDash([4, 4]);
  const y200 = y1 - (200 / yMax) * (y1 - y0);
  ctx.beginPath(); ctx.moveTo(x0, y200); ctx.lineTo(x1, y200); ctx.stroke();
  ctx.setLineDash([]);
  ctx.fillStyle = "rgba(220,230,255,0.55)";
  ctx.font = "10px ui-monospace, monospace";
  ctx.textAlign = "right";
  ctx.fillText("200", x0 - 4, y200 + 3);
  ctx.fillText("0", x0 - 4, y1 + 3);
  ctx.textAlign = "left";
  ctx.fillText("episode length", x0 + 4, y0 - 2);
  if (hist.length > 1) {
    const n = hist.length;
    ctx.lineWidth = 1.4;
    ctx.beginPath();
    for (let i = 0; i < n; i++) {
      const xp = x0 + (i / (n - 1 || 1)) * (x1 - x0);
      const yp = y1 - Math.min(1, hist[i] / yMax) * (y1 - y0);
      if (i === 0) ctx.moveTo(xp, yp); else ctx.lineTo(xp, yp);
    }
    ctx.strokeStyle = "rgba(120,200,255,0.55)";
    ctx.stroke();
    // moving average overlay
    const win = Math.max(3, Math.floor(n / 30));
    let sum = 0; const ma = [];
    for (let i = 0; i < n; i++) {
      sum += hist[i];
      if (i >= win) sum -= hist[i - win];
      ma.push(sum / Math.min(i + 1, win));
    }
    ctx.beginPath();
    for (let i = 0; i < n; i++) {
      const xp = x0 + (i / (n - 1 || 1)) * (x1 - x0);
      const yp = y1 - Math.min(1, ma[i] / yMax) * (y1 - y0);
      if (i === 0) ctx.moveTo(xp, yp); else ctx.lineTo(xp, yp);
    }
    ctx.strokeStyle = "rgba(255,210,120,0.95)";
    ctx.lineWidth = 1.8;
    ctx.stroke();
  }
  ctx.fillStyle = "rgba(220,230,255,0.55)";
  ctx.font = "10px ui-monospace, monospace";
  ctx.fillText(`${hist.length} episodes shown`, x0 + 4, y1 + 14);
}
function drawHUD(ctx, W) {
  const p = 10;
  ctx.fillStyle = "rgba(0,0,0,0.55)"; ctx.fillRect(p, p, 178, 96);
  ctx.fillStyle = "#fff"; ctx.font = "12px ui-monospace, monospace";
  ctx.fillText(`episode  ${ep}`, p + 8, p + 18);
  ctx.fillText(`step     ${vStep}/${MAX_STEPS}`, p + 8, p + 34);
  ctx.fillText(`epsilon  ${eps.toFixed(3)}`, p + 8, p + 50);
  ctx.fillText(`best     ${best}`, p + 8, p + 66);
  ctx.fillText(`frames   ${totalSteps}`, p + 8, p + 82);
}

function tick({ dt, ctx, width, height }) {
  // Clamp dt so a stalled tab can't burn through episodes invisibly.
  const wdt = Math.min(dt || 0.016, 0.1);

  // Background training: keep Q converging while the eye watches the visible cart.
  for (let i = 0; i < BG_STEPS_PER_FRAME; i++) bgTrainStep();

  // Visible cart: advance phase machine on wall clock.
  if (vPhase === "run") {
    vAcc += wdt;
    // Cap the number of substeps per frame so big dt doesn't fast-forward.
    let safety = 6;
    while (vAcc >= VIS_STEP_DT && safety-- > 0 && vPhase === "run") {
      vAcc -= VIS_STEP_DT;
      visStep();
    }
    if (vPhase !== "run") vAcc = 0;
  } else if (vPhase === "fail") {
    vPhaseT += wdt;
    if (vPhaseT >= FAIL_HOLD_S) {
      vPhase = "respawn"; vPhaseT = 0;
      vPrev = freshState(); vSt = vPrev.slice(); vStep = 0;
    }
  } else if (vPhase === "respawn") {
    vPhaseT += wdt;
    if (vPhaseT >= FADE_IN_S) {
      vPhase = "run"; vPhaseT = 0; vAcc = 0;
    }
  }

  // Render.
  ctx.fillStyle = "rgba(10,12,22,1)"; ctx.fillRect(0, 0, width, height);
  const grd = ctx.createLinearGradient(0, 0, 0, height);
  grd.addColorStop(0, "rgba(40,30,70,0.35)"); grd.addColorStop(1, "rgba(8,10,18,0)");
  ctx.fillStyle = grd; ctx.fillRect(0, 0, width, height);
  const panelH = Math.round(height * 0.55);

  // Interpolated display state for buttery motion between physics ticks.
  let displaySt, fail, alpha;
  if (vPhase === "run") {
    const t = Math.min(1, vAcc / VIS_STEP_DT);
    displaySt = [
      vPrev[0] + (vSt[0] - vPrev[0]) * t,
      vPrev[1] + (vSt[1] - vPrev[1]) * t,
      vPrev[2] + (vSt[2] - vPrev[2]) * t,
      vPrev[3] + (vSt[3] - vPrev[3]) * t,
    ];
    fail = false;
    alpha = 1;
  } else if (vPhase === "fail") {
    displaySt = vSt;
    fail = true;
    // Brief flash: bright at t=0, fading to ~0.55 by hold's end.
    const k = vPhaseT / FAIL_HOLD_S;
    alpha = 1 - 0.45 * k;
  } else {
    // respawn: fade in fresh pose
    displaySt = vSt;
    fail = false;
    alpha = vPhaseT / FADE_IN_S;
  }

  drawCart(ctx, width, panelH, displaySt, fail, alpha);
  drawChart(ctx, width, height, panelH);
  drawHUD(ctx, width);
}