apps/web/src/sim.ts

/**
 * Real physics for the background crew: the actual microduck MuJoCo model
 * stepped in the browser (official @mujoco/mujoco WASM build) and driven by
 * the walking / skating policies trained in pollen-robotics/microduck
 * (Apache-2.0), executed with onnxruntime-web at 50 Hz — the same
 * obs layout, timestep (0.005 s) and decimation (4) as the robot runtime,
 * ported from pollen-robotics' microduck-simulator.
 *
 * Each crew duck is one MjData over a shared compiled model (legs or
 * rollers variant). A carrot point circling its patrol ellipse turns into
 * twist commands; the policy does the rest — including the stumbles.
 */
import * as THREE from "three";
import { setHinge, type MeshLib, type Robot } from "./robot.js";

const MUJOCO_URL = "https://cdn.jsdelivr.net/npm/@mujoco/mujoco@3.11.0/mujoco.js";
const ORT_URL = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.27.0/dist/ort.min.mjs";
const ORT_WASM_DIR = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.27.0/dist/";

// From the ONNX metadata / mjlab STAND keyframe, identical to the runtime.
const JOINT_NAMES = [
  "left_hip_yaw",
  "left_hip_roll",
  "left_hip_pitch",
  "left_knee",
  "left_ankle",
  "neck_pitch",
  "head_pitch",
  "head_yaw",
  "head_roll",
  "right_hip_yaw",
  "right_hip_roll",
  "right_hip_pitch",
  "right_knee",
  "right_ankle",
];
const DEFAULT_POSE = new Float32Array([
  0, -0.08726646259971647, -0.457924, -0.00494, 0.452984, 0.3490658503988659, 0.3490658503988659, 0,
  0, 0, 0.08726646259971647, 0.457924, 0.00494, -0.452984,
]);
const NUM_JOINTS = 14;
const OBS_SIZE = 61;
const CMD_SIZE = 13;
const TIMESTEP = 0.005;
const DECIMATION = 4;
const CTRL_DT = TIMESTEP * DECIMATION; // 50 Hz
const ACTION_SCALE = 1.0;

// Velocity limits from the runtime (legs) and its roller branch.
const LIMITS = {
  walk: { fwd: 0.25, ang: 1.0 },
  skate: { fwd: 0.5, ang: 0.3 },
};

// ── Minimal typings for the CDN runtimes ─────────────────────────────────
type MjModel = {
  jnt: (name: string) => { qposadr: number; dofadr: number };
  sensor: (name: string) => { adr: number };
};
type MjData = {
  qpos: Float64Array;
  qvel: Float64Array;
  ctrl: Float64Array;
  sensordata: Float64Array;
  body: (id: number) => { xquat: Float64Array };
};
type Mujoco = {
  MjVFS: new () => { addBuffer: (path: string, data: Uint8Array) => void };
  MjModel: { from_xml_string: (xml: string, vfs: unknown) => MjModel };
  MjData: new (model: MjModel) => MjData;
  mj_step: (model: MjModel, data: MjData) => void;
  mj_resetDataKeyframe: (model: MjModel, data: MjData, key: number) => void;
  mj_name2id: (model: MjModel, type: number, name: string) => number;
  mjtObj: { mjOBJ_BODY: { value: number }; mjOBJ_KEY: { value: number } };
};
type OrtTensor = object;
type OrtSession = {
  run: (feeds: Record<string, OrtTensor>) => Promise<{ actions: { data: Float32Array } }>;
};
type Ort = {
  env: { wasm: { wasmPaths: string; numThreads: number } };
  Tensor: new (type: string, data: Float32Array, dims: number[]) => OrtTensor;
  InferenceSession: {
    create: (url: string, opts: { executionProviders: string[] }) => Promise<OrtSession>;
  };
};

/** Walkable rectangle in world units; physics walls fence the crew in. */
export type CrewPen = { minX: number; maxX: number; minZ: number; maxZ: number };

export type CrewSpec = {
  kind: "skate" | "walk";
  robot: Robot;
  /** Patrol ellipse in world units (same numbers the scripted crew used). */
  radius: number;
  speed: number;
  phase: number;
};

export type CrewSim = {
  /** Copy sim state into the visual rigs; call once per rendered frame. */
  sync: () => void;
  /** Return serializable physics coordinates for browser diagnostics. */
  debug: () => object;
  stop: () => void;
};

/** MuJoCo wants binary STL in its VFS; rebuild from the loaded geometry. */
function geometryToBinaryStl(geometry: THREE.BufferGeometry): ArrayBuffer {
  const pos = geometry.getAttribute("position");
  if (pos.count === 0) return new ArrayBuffer(84);
  const idx = geometry.index;
  const triCount = (idx ? idx.count : pos.count) / 3;
  const buf = new ArrayBuffer(84 + triCount * 50);
  const view = new DataView(buf);
  view.setUint32(80, triCount, true);
  let off = 84;
  const vx = (i: number): [number, number, number] => {
    const j = idx ? idx.getX(i) : i;
    return [pos.getX(j), pos.getY(j), pos.getZ(j)];
  };
  for (let t = 0; t < triCount; t++) {
    const a = vx(t * 3);
    const b = vx(t * 3 + 1);
    const c = vx(t * 3 + 2);
    const nx = (b[1] - a[1]) * (c[2] - a[2]) - (b[2] - a[2]) * (c[1] - a[1]);
    const ny = (b[2] - a[2]) * (c[0] - a[0]) - (b[0] - a[0]) * (c[2] - a[2]);
    const nz = (b[0] - a[0]) * (c[1] - a[1]) - (b[1] - a[1]) * (c[0] - a[0]);
    const len = Math.hypot(nx, ny, nz) || 1;
    const tri = [nx / len, ny / len, nz / len, ...a, ...b, ...c];
    for (let k = 0; k < 12; k++) view.setFloat32(off + k * 4, tri[k] ?? 0, true);
    view.setUint16(off + 48, 0, true);
    off += 50;
  }
  return buf;
}

/**
 * Prepare the MJCF for physics: drop visual geoms (every body carries an
 * explicit inertial), keep only meshes that collision geoms use, add the
 * floor plane, timestep, and the STAND keyframe.
 */
type Box = { name: string; pos: string; size: string };

function cellDeckHeight(
  deckHeight: (x: number, z: number) => number,
  x: number,
  z: number,
  cell: number,
): number {
  let height = Infinity;
  for (const [offsetX, offsetZ] of [
    [0, 0],
    [-0.45, -0.45],
    [0.45, -0.45],
    [-0.45, 0.45],
    [0.45, 0.45],
  ] as const) {
    height = Math.min(height, deckHeight(x + (offsetX * cell) / 2, z + (offsetZ * cell) / 2));
  }
  return height;
}

/**
 * Bake the visual deck contour into the physics: sample the deck height on
 * a grid over the pen and emit a static box per raised cell. Feet then
 * climb the real lips and plates instead of clipping through them.
 */
function terrainBoxes(
  pen: CrewPen,
  scale: number,
  deckHeight: (x: number, z: number) => number,
): { boxes: Box[]; groundAt: (sx: number, sy: number) => number; baseY: number } {
  const cell = 0.45; // world units
  const boxes: Box[] = [];
  const cols = Math.ceil((pen.maxX - pen.minX) / cell);
  const rows = Math.ceil((pen.maxZ - pen.minZ) / cell);
  const worldHeights = new Float32Array(cols * rows);
  let baseY = Infinity;
  for (let i = 0; i < cols; i++) {
    for (let j = 0; j < rows; j++) {
      const wx = pen.minX + (i + 0.5) * cell;
      const wz = pen.minZ + (j + 0.5) * cell;
      // MIN of center + corners: a box may never stand taller than the
      // lowest deck inside its cell, or ducks walk on invisible plateaus
      // and appear to hover at height steps.
      const height = cellDeckHeight(deckHeight, wx, wz, cell);
      worldHeights[i * rows + j] = height;
      baseY = Math.min(baseY, height);
    }
  }
  if (!Number.isFinite(baseY)) baseY = 0;

  const heights = new Float32Array(cols * rows);
  for (let i = 0; i < cols; i++) {
    for (let j = 0; j < rows; j++) {
      const wx = pen.minX + (i + 0.5) * cell;
      const wz = pen.minZ + (j + 0.5) * cell;
      const rel = Math.min(0.55, Math.max(0, (worldHeights[i * rows + j] ?? baseY) - baseY));
      heights[i * rows + j] = rel / scale;
      if (rel < 0.02) continue;
      const hz = rel / (2 * scale);
      const half = cell / (2 * scale);
      boxes.push({
        name: `deck_${String(i)}_${String(j)}`,
        pos: `${String(wx / scale)} ${String(-wz / scale)} ${String(hz)}`,
        size: `${String(half)} ${String(half)} ${String(hz)}`,
      });
    }
  }
  const groundAt = (sx: number, sy: number): number => {
    const i = Math.round((sx * scale - pen.minX) / cell - 0.5);
    const j = Math.round((-sy * scale - pen.minZ) / cell - 0.5);
    if (i < 0 || j < 0 || i >= cols || j >= rows) return 0;
    return heights[i * rows + j] ?? 0;
  };
  return { boxes, groundAt, baseY };
}

/**
 * Static boxes fencing the walkable aisle, so the policies physically bump
 * into the boundary instead of wandering up the set dressing.
 */
function penWalls(pen: {
  minX: number;
  maxX: number;
  minY: number;
  maxY: number;
}): { name: string; pos: string; size: string }[] {
  const cx = (pen.minX + pen.maxX) / 2;
  const cy = (pen.minY + pen.maxY) / 2;
  const hx = (pen.maxX - pen.minX) / 2 + 0.1;
  const hy = (pen.maxY - pen.minY) / 2 + 0.1;
  return [
    {
      name: "pen_px",
      pos: `${String(pen.maxX + 0.05)} ${String(cy)} 0.15`,
      size: `0.05 ${String(hy)} 0.15`,
    },
    {
      name: "pen_nx",
      pos: `${String(pen.minX - 0.05)} ${String(cy)} 0.15`,
      size: `0.05 ${String(hy)} 0.15`,
    },
    {
      name: "pen_py",
      pos: `${String(cx)} ${String(pen.maxY + 0.05)} 0.15`,
      size: `${String(hx)} 0.05 0.15`,
    },
    {
      name: "pen_ny",
      pos: `${String(cx)} ${String(pen.minY - 0.05)} 0.15`,
      size: `${String(hx)} 0.05 0.15`,
    },
  ];
}

/** Drop visual geoms and any meshes only they referenced. */
function pruneToCollisionMeshes(doc: Document): void {
  for (const g of [...doc.querySelectorAll('geom[class="visual"]')]) g.remove();
  const used = new Set([...doc.querySelectorAll("geom[mesh]")].map((g) => g.getAttribute("mesh")));
  for (const m of [...doc.querySelectorAll("asset > mesh")]) {
    const name = m.getAttribute("name") ?? m.getAttribute("file")?.replace(/\.stl$/i, "");
    if (!used.has(name ?? "")) m.remove();
  }
}

function buildPhysicsXml(
  src: string,
  pen: { minX: number; maxX: number; minY: number; maxY: number },
  terrain: Box[],
): { xml: string; meshFiles: string[] } {
  const doc = new DOMParser().parseFromString(src, "text/xml");
  pruneToCollisionMeshes(doc);
  const el = (tag: string, attrs: Record<string, string>): Element => {
    const e = doc.createElement(tag);
    for (const [k, v] of Object.entries(attrs)) e.setAttribute(k, v);
    return e;
  };
  const world = doc.querySelector("worldbody");
  if (!world) throw new Error("MJCF has no worldbody");
  doc.documentElement.appendChild(el("option", { timestep: String(TIMESTEP) }));
  world.appendChild(el("geom", { name: "floor", type: "plane", size: "0 0 0.05", pos: "0 0 0" }));
  for (const w of [...penWalls(pen), ...terrain]) {
    world.appendChild(el("geom", { name: w.name, type: "box", pos: w.pos, size: w.size }));
  }
  const poseByName = new Map(JOINT_NAMES.map((n, i) => [n, DEFAULT_POSE[i] ?? 0]));
  const qposJoints = [...doc.querySelectorAll("body > joint")]
    .map((j) => poseByName.get(j.getAttribute("name") ?? "") ?? 0)
    .join(" ");
  const kf = doc.createElement("keyframe");
  kf.appendChild(
    el("key", {
      name: "STAND",
      qpos: `0 0 0.12 1 0 0 0 ${qposJoints}`,
      ctrl: Array.from(DEFAULT_POSE).join(" "),
    }),
  );
  doc.documentElement.appendChild(kf);
  const meshFiles = [...doc.querySelectorAll("asset > mesh")]
    .map((m) => m.getAttribute("file"))
    .filter((f): f is string => f !== null);
  return { xml: new XMLSerializer().serializeToString(doc), meshFiles };
}

type Variant = {
  model: MjModel;
  scale: number;
  terrain: ReturnType<typeof terrainBoxes>;
  penSim: { minX: number; maxX: number; minY: number; maxY: number };
  qposAdr: number[];
  dofAdr: number[];
  gyroAdr: number;
  trunkId: number;
  standKeyId: number;
};

type SimDuck = {
  spec: CrewSpec;
  variant: Variant;
  session: OrtSession;
  data: MjData;
  lastAction: Float32Array;
  obs: Float32Array;
  /** Wheel/passive hinges present in the rig but outside the 14 actuated. */
  extraAdr: { name: string; adr: number }[];
};

const _q = new THREE.Quaternion();
const _g = new THREE.Vector3();

/** Indexed read with the noUncheckedIndexedAccess fallback in one place. */
function at(arr: ArrayLike<number>, i: number): number {
  return arr[i] ?? 0;
}

/** World -z in the trunk frame (projected gravity), z component & vector. */
function projectGravity(data: MjData, trunkId: number): THREE.Vector3 {
  const xq = data.body(trunkId).xquat; // [w, x, y, z]
  _q.set(at(xq, 1), at(xq, 2), at(xq, 3), xq[0] ?? 1).conjugate();
  return _g.set(0, 0, -1).applyQuaternion(_q);
}

/** World patrol ellipse (matches the scripted crew) in sim meters. */
function carrot(
  spec: CrewSpec,
  simTime: number,
  scale: number,
  pen: CrewPen,
): { x: number; y: number } {
  const a = simTime * spec.speed + spec.phase + 0.4; // lead the duck a little
  const margin = 0.5;
  const wx = Math.min(
    pen.maxX - margin,
    Math.max(pen.minX + margin, Math.cos(a) * spec.radius * 0.78 - 2.2),
  );
  const wz = Math.min(
    pen.maxZ - margin,
    Math.max(pen.minZ + margin, Math.sin(a) * spec.radius - 8.5),
  );
  return { x: wx / scale, y: -wz / scale };
}

function buildObs(duck: SimDuck): Float32Array {
  const { data, variant, obs, lastAction } = duck;
  let i = 0;
  for (let a = 0; a < 3; a++) obs[i++] = at(data.sensordata, variant.gyroAdr + a);
  const g = projectGravity(data, variant.trunkId);
  obs[i++] = g.x;
  obs[i++] = g.y;
  obs[i++] = g.z;
  for (let j = 0; j < NUM_JOINTS; j++) {
    obs[i++] = at(data.qpos, at(variant.qposAdr, j)) - at(DEFAULT_POSE, j);
  }
  for (let j = 0; j < NUM_JOINTS; j++) obs[i++] = at(data.qvel, at(variant.dofAdr, j));
  for (let j = 0; j < NUM_JOINTS; j++) obs[i++] = at(lastAction, j);
  // command slots start zeroed; [vx, vy, wz] filled by the carrot follower
  for (let c = 0; c < CMD_SIZE; c++) obs[i + c] = 0;
  return obs;
}

function steerToCarrot(duck: SimDuck, simTime: number, scale: number, pen: CrewPen): void {
  const { data, obs, spec } = duck;
  const target = carrot(spec, simTime, scale, pen);
  const qpos = data.qpos;
  const dx = target.x - at(qpos, 0);
  const dy = target.y - at(qpos, 1);
  const w = qpos[3] ?? 1;
  const zq = at(qpos, 6);
  const yaw = Math.atan2(
    2 * (w * zq + at(qpos, 4) * at(qpos, 5)),
    1 - 2 * (at(qpos, 5) * at(qpos, 5) + zq * zq),
  );
  let err = Math.atan2(dy, dx) - yaw;
  while (err > Math.PI) err -= 2 * Math.PI;
  while (err < -Math.PI) err += 2 * Math.PI;
  const lim = LIMITS[spec.kind];
  const dist = Math.hypot(dx, dy);
  const base = OBS_SIZE - CMD_SIZE;
  obs[base] = lim.fwd * Math.max(0, Math.cos(err)) * Math.min(1, dist / 0.12);
  obs[base + 1] = 0;
  obs[base + 2] = Math.max(-lim.ang, Math.min(lim.ang, 1.6 * err));
}

/** Fallen or exploded? Same thresholds as the reference simulator. */
function poseIsDead(duck: SimDuck, ground: number): boolean {
  const z = at(duck.data.qpos, 2) - ground;
  const gz = projectGravity(duck.data, duck.variant.trunkId).z;
  return !Number.isFinite(z) || !Number.isFinite(gz) || gz > -0.5 || z < 0.02;
}

function respawn(duck: SimDuck, mujoco: Mujoco, simTime: number, pen: CrewPen): void {
  const { variant } = duck;
  mujoco.mj_resetDataKeyframe(variant.model, duck.data, variant.standKeyId);
  const target = carrot(duck.spec, simTime, variant.scale, pen);
  duck.data.qpos[0] = target.x;
  duck.data.qpos[1] = target.y;
  duck.data.qpos[2] = 0.12 + variant.terrain.groundAt(target.x, target.y);
  duck.lastAction.fill(0);
}

/** Copy one duck's sim state into its visual rig. */
function syncDuck(duck: SimDuck): void {
  const { robot } = duck.spec;
  const qpos = duck.data.qpos;
  const trunk = robot.root.getObjectByName("trunk_base");
  if (!trunk) return;
  trunk.position.set(at(qpos, 0), at(qpos, 1), at(qpos, 2));
  trunk.quaternion.set(at(qpos, 4), at(qpos, 5), at(qpos, 6), qpos[3] ?? 1);
  for (let j = 0; j < NUM_JOINTS; j++) {
    const joint = robot.joints.get(JOINT_NAMES[j] ?? "");
    if (joint) setHinge(joint, at(qpos, at(duck.variant.qposAdr, j)));
  }
  for (const extra of duck.extraAdr) {
    const joint = robot.joints.get(extra.name);
    if (joint) setHinge(joint, at(qpos, extra.adr));
  }
}

/** One 50 Hz control step for one duck: observe, infer, actuate, step. */
async function stepDuck(
  duck: SimDuck,
  mujoco: Mujoco,
  ort: Ort,
  simTime: number,
  pen: CrewPen,
): Promise<void> {
  buildObs(duck);
  steerToCarrot(duck, simTime, duck.variant.scale, pen);
  const feeds = { obs: new ort.Tensor("float32", duck.obs, [1, OBS_SIZE]) };
  const out = await duck.session.run(feeds);
  duck.lastAction.set(out.actions.data);
  for (let j = 0; j < NUM_JOINTS; j++) {
    duck.data.ctrl[j] = (DEFAULT_POSE[j] ?? 0) + (out.actions.data[j] ?? 0) * ACTION_SCALE;
  }
  for (let s = 0; s < DECIMATION; s++) mujoco.mj_step(duck.variant.model, duck.data);
  const ground = duck.variant.terrain.groundAt(at(duck.data.qpos, 0), at(duck.data.qpos, 1));
  if (poseIsDead(duck, ground)) respawn(duck, mujoco, simTime, pen);
}

async function loadRuntimes(): Promise<{ mujoco: Mujoco; ort: Ort }> {
  /* eslint-disable @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-call */
  const [mujocoModule, ortModule] = await Promise.all([
    import(/* @vite-ignore */ MUJOCO_URL),
    import(/* @vite-ignore */ ORT_URL),
  ]);
  const mujoco = (await mujocoModule.default()) as Mujoco;
  const ort = ortModule as unknown as Ort;
  /* eslint-enable @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-call */
  ort.env.wasm.wasmPaths = ORT_WASM_DIR;
  ort.env.wasm.numThreads = 1; // static hosting: no COOP/COEP headers
  return { mujoco, ort };
}

/**
 * Boot physics for the crew. Throws on any failure — the caller keeps the
 * scripted crew as fallback.
 */
export async function startCrewSim(
  specs: CrewSpec[],
  meshLib: MeshLib,
  assetBase: string,
  pen: CrewPen,
  initialPoses: () => { x: number; z: number; heading: number }[],
  deckHeight: (x: number, z: number) => number,
): Promise<CrewSim> {
  console.info("[sim] loading runtimes");
  const { mujoco, ort } = await loadRuntimes();
  console.info("[sim] runtimes ready");

  const fetchText = async (url: string): Promise<string> => {
    const res = await fetch(url);
    if (!res.ok) throw new Error(`fetch ${url}: ${String(res.status)}`);
    return res.text();
  };
  const [legsXmlSrc, rollersXmlSrc] = await Promise.all([
    fetchText(`${assetBase}/sim/robot_allcollisions.xml`),
    fetchText(`${assetBase}/sim/robot_allcollisions_rollers.xml`),
  ]);
  // Sim meters -> world units. Walkers and skaters normalize to different
  // scales, so each compiled model needs matching terrain and pen units.
  const makeSpace = (kind: CrewSpec["kind"]) => {
    const scale = specs.find((spec) => spec.kind === kind)?.robot.root.scale.x ?? 1;
    const penSim = {
      minX: pen.minX / scale,
      maxX: pen.maxX / scale,
      minY: -pen.maxZ / scale,
      maxY: -pen.minZ / scale,
    };
    return { scale, penSim, terrain: terrainBoxes(pen, scale, deckHeight) };
  };
  const spaces = { walk: makeSpace("walk"), skate: makeSpace("skate") };
  const legs = buildPhysicsXml(legsXmlSrc, spaces.walk.penSim, spaces.walk.terrain.boxes);
  const rollers = buildPhysicsXml(rollersXmlSrc, spaces.skate.penSim, spaces.skate.terrain.boxes);

  // Collision meshes: legs/body parts come from the already-loaded GLB
  // library (rebuilt as binary STL in memory); roller-only parts are
  // fetched from the rollers asset directory.
  const vfs = new mujoco.MjVFS();
  const added = new Set<string>();
  const addMeshes = async (files: string[]): Promise<void> => {
    await Promise.all(
      files.map(async (f) => {
        if (added.has(f)) return;
        added.add(f);
        const key = f.toLowerCase().replace(/[^a-z0-9]/g, "");
        const geometry = meshLib.get(key);
        const buf = geometry
          ? geometryToBinaryStl(geometry)
          : await (await fetch(`${assetBase}/rollers/${f}`)).arrayBuffer();
        vfs.addBuffer(`assets/${f}`, new Uint8Array(buf));
      }),
    );
  };
  await addMeshes(legs.meshFiles);
  await addMeshes(rollers.meshFiles);

  const resolve = (model: MjModel, space: (typeof spaces)[CrewSpec["kind"]]): Variant => ({
    model,
    ...space,
    qposAdr: JOINT_NAMES.map((n) => model.jnt(n).qposadr),
    dofAdr: JOINT_NAMES.map((n) => model.jnt(n).dofadr),
    gyroAdr: model.sensor("imu_ang_vel").adr,
    trunkId: mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY.value, "trunk_base"),
    standKeyId: mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY.value, "STAND"),
  });
  console.info("[sim] meshes in vfs, compiling models");
  const variants = {
    walk: resolve(mujoco.MjModel.from_xml_string(legs.xml, vfs), spaces.walk),
    skate: resolve(mujoco.MjModel.from_xml_string(rollers.xml, vfs), spaces.skate),
  };
  console.info("[sim] models compiled, loading policies");

  const sessionOpts = { executionProviders: ["wasm"] };
  const [walkSession, skateSession] = await Promise.all([
    ort.InferenceSession.create(`${assetBase}/sim/BEST_alpha_walking.onnx`, sessionOpts),
    ort.InferenceSession.create(`${assetBase}/sim/BEST_roller.onnx`, sessionOpts),
  ]);
  const sessions = { walk: walkSession, skate: skateSession };
  console.info("[sim] policies ready");

  let simTime = 0;
  const startPoses = initialPoses();
  const ducks: SimDuck[] = specs.map((spec, index) => {
    const variant = variants[spec.kind];
    const data = new mujoco.MjData(variant.model);
    const duck: SimDuck = {
      spec,
      variant,
      session: sessions[spec.kind],
      data,
      lastAction: new Float32Array(NUM_JOINTS),
      obs: new Float32Array(OBS_SIZE),
      extraAdr: [...spec.robot.joints.keys()]
        .filter((n) => !JOINT_NAMES.includes(n))
        .map((n) => ({ name: n, adr: variant.model.jnt(n).qposadr })),
    };
    respawn(duck, mujoco, 0, pen);
    // take over exactly where the scripted crew stands right now
    const pose = startPoses[index];
    if (pose) {
      const { penSim, scale } = variant;
      const margin = 0.3;
      duck.data.qpos[0] = Math.min(
        penSim.maxX - margin,
        Math.max(penSim.minX + margin, pose.x / scale),
      );
      duck.data.qpos[1] = Math.min(
        penSim.maxY - margin,
        Math.max(penSim.minY + margin, -pose.z / scale),
      );
      duck.data.qpos[3] = Math.cos(pose.heading / 2);
      duck.data.qpos[4] = 0;
      duck.data.qpos[5] = 0;
      duck.data.qpos[6] = Math.sin(pose.heading / 2);
      duck.data.qpos[2] =
        0.12 + variant.terrain.groundAt(at(duck.data.qpos, 0), at(duck.data.qpos, 1));
    }
    return duck;
  });

  // Rigs go under sim control with a fully clean root transform so the zUp
  // group maps MuJoCo coordinates straight into the world (x -> x, y -> -z,
  // z -> y). Scripted skaters leave a root roll behind during async boot;
  // keeping it would tilt the complete physics frame above or below the deck.
  for (const duck of ducks) {
    duck.spec.robot.root.rotation.set(0, 0, 0);
    duck.spec.robot.root.position.set(0, duck.variant.terrain.baseY, 0);
  }

  // ── 50 Hz control loop, independent of the render loop ───────────────
  let busy = false;
  let stopped = false;
  const timer = setInterval(() => {
    if (busy || stopped) return;
    busy = true;
    void (async () => {
      try {
        simTime += CTRL_DT;
        // sequential: ducks share ONNX sessions, and session.run is not
        // safely reentrant. Five tiny MLPs stay well under the 20 ms budget.
        for (const duck of ducks) {
          await stepDuck(duck, mujoco, ort, simTime, pen);
        }
      } catch (err) {
        console.error("[sim] control tick failed", err);
        stopped = true;
        clearInterval(timer);
      } finally {
        busy = false;
      }
    })();
  }, CTRL_DT * 1000);

  const sync = (): void => {
    for (const duck of ducks) syncDuck(duck);
  };

  return {
    sync,
    debug: () => ({
      spaces: {
        walk: { scale: spaces.walk.scale, baseY: spaces.walk.terrain.baseY },
        skate: { scale: spaces.skate.scale, baseY: spaces.skate.terrain.baseY },
      },
      ducks: ducks.map((duck) => {
        const { scale, terrain } = duck.variant;
        return {
          kind: duck.spec.kind,
          position: {
            x: at(duck.data.qpos, 0) * scale,
            y: terrain.baseY + at(duck.data.qpos, 2) * scale,
            z: -at(duck.data.qpos, 1) * scale,
          },
          groundY:
            terrain.baseY + terrain.groundAt(at(duck.data.qpos, 0), at(duck.data.qpos, 1)) * scale,
        };
      }),
    }),
    stop: () => {
      stopped = true;
      clearInterval(timer);
    },
  };
}