Spaces:
Running
Running
Download apps/web/src/sim.ts from osolmaz/microquack: direct link, hf CLI and curl.
- Browser
- Download file 24.5 kB
-
https://huggingface.co/spaces/osolmaz/microquack/resolve/main/apps/web/src/sim.ts
- Command line
-
hf download hf://spaces/osolmaz/microquack/apps/web/src/sim.ts
-
curl -L -o sim.ts https://huggingface.co/spaces/osolmaz/microquack/resolve/main/apps/web/src/sim.ts
24.5 kB
| /** | |
| * 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); | |
| }, | |
| }; | |
| } | |