PID Cart-Pole Balancer

// Inverted pendulum on a cart, stabilised by a PID controller on theta.
// Drag the cart to disturb. Keys 1/2/3 toggle P/I/D terms; 4 cycles presets.

const PRESETS = [
  { name: "off",        Kp: 0,   Ki: 0,  Kd: 0  },
  { name: "balanced",   Kp: 60,  Ki: 8,  Kd: 14 },
  { name: "P-only",     Kp: 35,  Ki: 0,  Kd: 0  },
  { name: "aggressive", Kp: 140, Ki: 30, Kd: 26 },
];
const MC = 1.0, MP = 0.15, L = 1.2, G = 9.81, MU = 0.6, FMAX = 80;

let st, pxM, gy, preset, useP, useI, useD, integ, lastErr;
let trail, drag, dragOff, Ffilt, tAcc, flash;

function init({ width, height }) {
  layout(width, height);
  preset = 1; useP = true; useI = true; useD = true;
  reset(); trail = []; flash = null;
}
function layout(w, h) { pxM = Math.min(w, h) * 0.22; gy = h * 0.62; }
function reset() {
  st = [0, 0, (Math.random() - 0.5) * 0.12, 0];
  integ = 0; lastErr = 0; Ffilt = 0; tAcc = 0;
}
function gains() {
  const p = PRESETS[preset];
  return { Kp: useP ? p.Kp : 0, Ki: useI ? p.Ki : 0, Kd: useD ? p.Kd : 0, name: p.name };
}
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 - MU * xd + 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 rk4(s, F, h) {
  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 pid(dt) {
  const { Kp, Ki, Kd } = gains();
  const err = -st[2];
  integ = Math.max(-1.5, Math.min(1.5, integ + err * dt));
  const dErr = dt > 1e-6 ? (err - lastErr) / dt : 0;
  lastErr = err;
  const center = -st[0] * 1.2 - st[1] * 1.5;
  let F = Kp * err + Ki * integ + Kd * dErr + center;
  return Math.max(-FMAX, Math.min(FMAX, F));
}
const w2s = (x, W) => W / 2 + x * pxM;
const s2w = (px, W) => (px - W / 2) / pxM;
function keys(input) {
  if (input.justPressed("1")) { useP = !useP; flash = { k: "P", t: 1 }; }
  if (input.justPressed("2")) { useI = !useI; integ = 0; flash = { k: "I", t: 1 }; }
  if (input.justPressed("3")) { useD = !useD; flash = { k: "D", t: 1 }; }
  if (input.justPressed("4")) {
    preset = (preset + 1) % PRESETS.length; integ = 0;
    flash = { k: PRESETS[preset].name, t: 1.2 };
  }
}
function drawTrack(ctx, W) {
  ctx.strokeStyle = "rgba(120,140,180,0.55)"; ctx.lineWidth = 2;
  ctx.beginPath(); ctx.moveTo(0, gy); ctx.lineTo(W, gy); ctx.stroke();
  ctx.strokeStyle = "rgba(80,100,140,0.35)"; ctx.lineWidth = 1;
  for (let gx = 0; gx <= W; gx += pxM * 0.5) {
    ctx.beginPath(); ctx.moveTo(gx, gy); ctx.lineTo(gx, gy + 6); ctx.stroke();
  }
  ctx.strokeStyle = "rgba(255,255,255,0.45)";
  ctx.beginPath(); ctx.moveTo(W / 2, gy - 4); ctx.lineTo(W / 2, gy + 8); ctx.stroke();
}
function drawCartPole(ctx, W) {
  const cx = w2s(st[0], W), cy = gy - 14;
  const cw = 56, ch = 26;
  const px = cx + Math.sin(st[2]) * L * pxM;
  const py = cy - Math.cos(st[2]) * L * pxM;
  ctx.strokeStyle = "rgba(255,210,120,0.95)"; ctx.lineWidth = 5; ctx.lineCap = "round";
  ctx.beginPath(); ctx.moveTo(cx, cy); ctx.lineTo(px, py); ctx.stroke();
  ctx.fillStyle = "#ffcb5a";
  ctx.beginPath(); ctx.arc(px, py, 9, 0, Math.PI * 2); ctx.fill();
  ctx.strokeStyle = "rgba(0,0,0,0.4)"; ctx.lineWidth = 1.5; ctx.stroke();
  ctx.fillStyle = drag ? "#7ab8ff" : "#5a8fd6";
  ctx.fillRect(cx - cw / 2, cy - ch / 2, cw, ch);
  ctx.strokeStyle = "rgba(255,255,255,0.55)"; ctx.lineWidth = 1.5;
  ctx.strokeRect(cx - cw / 2 + 0.5, cy - ch / 2 + 0.5, cw - 1, ch - 1);
  ctx.fillStyle = "#222";
  ctx.beginPath(); ctx.arc(cx - cw / 2 + 10, cy + ch / 2, 6, 0, Math.PI * 2); ctx.fill();
  ctx.beginPath(); ctx.arc(cx + cw / 2 - 10, cy + ch / 2, 6, 0, Math.PI * 2); ctx.fill();
  ctx.fillStyle = "#fff";
  ctx.beginPath(); ctx.arc(cx, cy, 3, 0, Math.PI * 2); ctx.fill();
  return { cx, cy };
}
function drawForce(ctx, cx, cy, F) {
  if (Math.abs(F) < 0.5) return;
  const mag = Math.min(80, Math.abs(F)) * 0.9, dir = F > 0 ? 1 : -1;
  const x0 = cx + dir * 32, x1 = x0 + dir * mag;
  ctx.strokeStyle = "rgba(120,255,180,0.85)"; ctx.lineWidth = 3;
  ctx.beginPath(); ctx.moveTo(x0, cy); ctx.lineTo(x1, cy); ctx.stroke();
  ctx.fillStyle = "rgba(120,255,180,0.85)";
  ctx.beginPath();
  ctx.moveTo(x1, cy); ctx.lineTo(x1 - dir * 8, cy - 5); ctx.lineTo(x1 - dir * 8, cy + 5);
  ctx.closePath(); ctx.fill();
}
function drawTrail(ctx, W) {
  if (trail.length < 2) return;
  ctx.lineWidth = 1.4;
  for (let i = 1; i < trail.length; i++) {
    const a = trail[i - 1], b = trail[i];
    const age = i / trail.length;
    const f = Math.max(-1, Math.min(1, b.f / 40));
    const hue = f >= 0 ? 150 + f * 60 : 200 + f * 60;
    ctx.strokeStyle = `hsla(${hue},90%,65%,${0.15 + age * 0.55})`;
    ctx.beginPath();
    ctx.moveTo(w2s(a.x, W), gy + 18); ctx.lineTo(w2s(b.x, W), gy + 18); ctx.stroke();
  }
}
function badge(ctx, x, y, label, on, color) {
  ctx.fillStyle = on ? color : "rgba(40,46,60,0.85)";
  ctx.fillRect(x, y, 26, 22);
  ctx.strokeStyle = on ? "rgba(255,255,255,0.85)" : "rgba(255,255,255,0.25)";
  ctx.lineWidth = 1; ctx.strokeRect(x + 0.5, y + 0.5, 25, 21);
  ctx.fillStyle = on ? "#0b0d14" : "rgba(220,230,255,0.7)";
  ctx.font = "bold 13px ui-monospace, monospace";
  ctx.textAlign = "center"; ctx.textBaseline = "middle";
  ctx.fillText(label, x + 13, y + 12);
  ctx.textAlign = "left"; ctx.textBaseline = "alphabetic";
}
function drawHUD(ctx, W, H, F, g) {
  const p = 12;
  ctx.fillStyle = "rgba(0,0,0,0.55)"; ctx.fillRect(p, p, 188, 104);
  ctx.fillStyle = "#fff"; ctx.font = "12px ui-monospace, monospace";
  ctx.fillText(`theta = ${(st[2] * 180 / Math.PI).toFixed(1)} deg`, p + 8, p + 18);
  ctx.fillText(`F     = ${F.toFixed(1)} N`, p + 8, p + 34);
  ctx.fillText(`x     = ${st[0].toFixed(2)} m`, p + 8, p + 50);
  ctx.fillText(`preset: ${g.name}`, p + 8, p + 70);
  badge(ctx, p + 8, p + 78, "P", useP, "#7ee787");
  badge(ctx, p + 38, p + 78, "I", useI, "#ffd866");
  badge(ctx, p + 68, p + 78, "D", useD, "#79c0ff");
  const rx = W - 152 - p;
  ctx.fillStyle = "rgba(0,0,0,0.55)"; ctx.fillRect(rx, p, 152, 64);
  ctx.fillStyle = "#fff";
  ctx.fillText(`Kp = ${g.Kp.toFixed(0)}`, rx + 8, p + 18);
  ctx.fillText(`Ki = ${g.Ki.toFixed(0)}`, rx + 8, p + 34);
  ctx.fillText(`Kd = ${g.Kd.toFixed(0)}`, rx + 8, p + 50);
  ctx.fillStyle = "rgba(220,230,255,0.6)"; ctx.font = "11px ui-monospace, monospace";
  ctx.fillText("drag cart  |  1 P  2 I  3 D  4 preset", p, H - 10);
}
function tick({ ctx, dt, width, height, input }) {
  if (gy !== height * 0.62) layout(width, height);
  keys(input);
  const cx = w2s(st[0], width), cy = gy - 14;
  if (input.mouseDown) {
    if (!drag) {
      const dx = input.mouseX - cx, dy = input.mouseY - cy;
      if (dx * dx + dy * dy < 1600) { drag = true; dragOff = input.mouseX - cx; }
    }
    if (drag) {
      const lim = (width / 2 - 40) / pxM;
      st[0] = Math.max(-lim, Math.min(lim, s2w(input.mouseX - dragOff, width)));
      st[1] = 0;
    }
  } else { drag = false; }
  for (const c of input.consumeClicks()) {
    const dx = c.x - cx, dy = c.y - cy;
    if (dx * dx + dy * dy > 3600) st[3] += (c.x < cx ? 1 : -1) * 2.5;
  }
  const steps = 6, h = Math.min(dt, 0.033) / steps;
  let F = 0;
  for (let i = 0; i < steps; i++) {
    F = drag ? 0 : pid(h);
    st = rk4(st, F, h);
    const lim = (width / 2 - 40) / pxM;
    if (st[0] > lim) { st[0] = lim; st[1] = -Math.abs(st[1]) * 0.3; }
    if (st[0] < -lim) { st[0] = -lim; st[1] = Math.abs(st[1]) * 0.3; }
  }
  Ffilt = Ffilt * 0.7 + F * 0.3;
  tAcc += dt;
  if (tAcc > 1 / 30) {
    trail.push({ x: st[0], f: Ffilt });
    if (trail.length > 220) trail.shift();
    tAcc = 0;
  }
  if (Math.abs(st[2]) > Math.PI * 0.55 && !drag) {
    st[3] *= 0.985;
    if (Math.abs(st[3]) < 0.05 && Math.abs(st[1]) < 0.05) reset();
  }
  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);
  drawTrack(ctx, width);
  drawTrail(ctx, width);
  const cart = drawCartPole(ctx, width);
  drawForce(ctx, cart.cx, cart.cy, Ffilt);
  drawHUD(ctx, width, height, Ffilt, gains());
  if (flash) {
    flash.t -= dt;
    if (flash.t <= 0) flash = null;
    else {
      ctx.fillStyle = `rgba(255,210,120,${Math.min(1, flash.t)})`;
      ctx.font = "bold 22px ui-monospace, monospace";
      ctx.textAlign = "center";
      ctx.fillText(flash.k, width / 2, 34);
      ctx.textAlign = "left";
    }
  }
}