Download src/sim.js from happynood/quant-control-bench-demo: direct link, hf CLI and curl.
- Browser
- Download file 12.5 kB
-
https://huggingface.co/spaces/happynood/quant-control-bench-demo/resolve/main/src/sim.js
- Command line
-
hf download hf://spaces/happynood/quant-control-bench-demo/src/sim.js
-
curl -L -o sim.js https://huggingface.co/spaces/happynood/quant-control-bench-demo/resolve/main/src/sim.js
12.5 kB
| // MuJoCo WASM wrapper: one independent simulation per precision variant. | |
| // | |
| // Everything that describes the model or the policy interface comes from | |
| // assets/scene/scene.json, which `qcb export-scene` generates from the trained | |
| // environment. Nothing here is hardcoded, because every one of those numbers is a | |
| // silent-failure candidate: a wrong sensor address feeds the policy a different | |
| // robot's velocity and it still walks, just worse. | |
| // MuJoCo exposes some fields as typed-array views into WASM memory and others as | |
| // embind vectors. The two have incompatible `set`: `TypedArray.set(array, offset)` | |
| // copies a sequence, while `vector.set(index, value)` writes one element. Calling | |
| // the wrong one is not a type error — it throws `offset is out of bounds`, or | |
| // worse, silently writes the wrong thing. Views are detected first, explicitly. | |
| const isView = (vec) => ArrayBuffer.isView(vec); | |
| /** Read element `i` from a MuJoCo array. */ | |
| export function at(vec, i) { | |
| return isView(vec) ? vec[i] : vec.get(i); | |
| } | |
| /** Write element `i` of a MuJoCo array. */ | |
| export function put(vec, i, value) { | |
| if (isView(vec)) vec[i] = value; | |
| else vec.set(i, value); | |
| } | |
| export function readBlock(vec, adr, dim) { | |
| const out = new Float64Array(dim); | |
| for (let i = 0; i < dim; i++) out[i] = at(vec, adr + i); | |
| return out; | |
| } | |
| /** | |
| * A deterministic pseudo-random source for observation noise. | |
| * | |
| * The training environment draws its observation noise from JAX's counter-based | |
| * PRNG, which cannot be reproduced here. This is therefore *not* an attempt to | |
| * match Python bit-for-bit — with noise on, no browser rollout can. It exists so | |
| * the demo shows the deployed behaviour (the policy was trained with this noise) | |
| * while staying reproducible across reloads within the browser. The parity test | |
| * runs with `level = 0`, which removes the term entirely and makes the comparison | |
| * meaningful. | |
| */ | |
| export function makeRng(seed) { | |
| let state = seed >>> 0; | |
| return () => { | |
| // xorshift32: cheap, deterministic, adequate for a visual demo. | |
| state ^= state << 13; | |
| state ^= state >>> 17; | |
| state ^= state << 5; | |
| state >>>= 0; | |
| return state / 4294967296; | |
| }; | |
| } | |
| export class Sim { | |
| /** | |
| * @param {object} mujoco loaded WASM module | |
| * @param {object} scene parsed scene.json | |
| * @param {string} label which precision this instance represents | |
| */ | |
| constructor(mujoco, scene, label) { | |
| this.mujoco = mujoco; | |
| this.scene = scene; | |
| this.label = label; | |
| this.iface = scene.policy_interface; | |
| this.model = mujoco.MjModel.loadFromXML(`/working/${scene.top_xml}`); | |
| this.#assertShape(); | |
| this.#applyOverrides(); | |
| this.data = new mujoco.MjData(this.model); | |
| this.nu = scene.nu; | |
| this.lastAction = new Float32Array(this.nu); | |
| this.prevAction = new Float32Array(this.nu); | |
| this.obs = new Float32Array(48); | |
| this.command = new Float32Array([1.0, 0.0, 0.0]); | |
| this.noiseLevel = this.iface.noise.level; | |
| this.rng = makeRng(0x9e3779b9); | |
| // Perturbations, driven by the UI. Model-level ones are applied to the | |
| // compiled model; loop-level ones act between policy and plant. | |
| this.actuatorDelay = 0; | |
| this.delayQueue = []; | |
| this.obsNoiseSigma = 0; | |
| this.reset(); | |
| } | |
| #assertShape() { | |
| const m = this.model; | |
| const expected = [this.scene.nq, this.scene.nv, this.scene.nu]; | |
| const actual = [m.nq, m.nv, m.nu]; | |
| for (let i = 0; i < 3; i++) { | |
| if (expected[i] !== actual[i]) { | |
| throw new Error( | |
| `${this.label}: loaded model has nq/nv/nu ${actual.join('/')}, ` + | |
| `scene.json says ${expected.join('/')} — wrong scene or wrong assets` | |
| ); | |
| } | |
| } | |
| } | |
| /** | |
| * Reapply the edits the training environment makes after parsing the XML. | |
| * | |
| * `Go1Env.__init__` sets the timestep, raises ccd_iterations, and rewrites joint | |
| * damping and every actuator gain and bias from its PD constants. None of it is | |
| * in the XML. Skipping this step produces a robot that walks but is not the one | |
| * that was benchmarked. | |
| */ | |
| #applyOverrides() { | |
| const m = this.model; | |
| const s = this.scene; | |
| m.opt.timestep = s.timestep; | |
| m.opt.ccd_iterations = s.ccd_iterations; | |
| for (let i = 0; i < s.dof_damping.length; i++) put(m.dof_damping, i, s.dof_damping[i]); | |
| // Row widths come from scene.json. Both matrices are (nu, 10) in MuJoCo 3.x; | |
| // assuming 3 for biasprm — which its name and the mjNBIAS constant invite — | |
| // scatters each actuator's bias into another actuator's row. | |
| const gainStride = s.actuator_gainprm_stride; | |
| const biasStride = s.actuator_biasprm_stride; | |
| for (let i = 0; i < s.actuator_gainprm_col0.length; i++) { | |
| put(m.actuator_gainprm, i * gainStride + 0, s.actuator_gainprm_col0[i]); | |
| put(m.actuator_biasprm, i * biasStride + 1, s.actuator_biasprm_col1[i]); | |
| } | |
| } | |
| /** | |
| * Release the WASM-side model and data. | |
| * | |
| * Each Sim holds a compiled 10 MB model in the module's heap. Rebuilding the | |
| * variant list without this leaks one per rebuild, and a visitor toggling | |
| * checkboxes exhausts the heap in a couple of dozen clicks. | |
| */ | |
| dispose() { | |
| for (const handle of [this.data, this.model]) { | |
| if (handle && typeof handle.delete === 'function' && !handle.isDeleted?.()) { | |
| handle.delete(); | |
| } | |
| } | |
| this.data = null; | |
| this.model = null; | |
| } | |
| /** Scale torso mass; 1.0 restores the exported value. */ | |
| setMassScale(scale) { | |
| const bodyId = 1; // first non-world body: the trunk | |
| if (this._baseMass === undefined) this._baseMass = at(this.model.body_mass, bodyId); | |
| put(this.model.body_mass, bodyId, this._baseMass * scale); | |
| } | |
| /** Scale the sliding friction coefficient of every geom. */ | |
| setFrictionScale(scale) { | |
| const n = this.model.ngeom; | |
| if (this._baseFriction === undefined) { | |
| this._baseFriction = new Float64Array(n); | |
| for (let i = 0; i < n; i++) this._baseFriction[i] = at(this.model.geom_friction, i * 3); | |
| } | |
| for (let i = 0; i < n; i++) put(this.model.geom_friction, i * 3, this._baseFriction[i] * scale); | |
| } | |
| setActuatorDelay(steps) { | |
| this.actuatorDelay = Math.max(0, Math.round(steps)); | |
| this.delayQueue = []; | |
| } | |
| setObsNoise(sigma) { | |
| this.obsNoiseSigma = sigma; | |
| } | |
| /** Apply an impulse to the torso's linear velocity, as the eval harness does. */ | |
| push(impulseNs) { | |
| const mass = at(this.model.body_mass, 1); | |
| const dv = impulseNs / Math.max(mass, 1e-9); | |
| // Random direction in the horizontal plane plus a little vertical, unit-normalised. | |
| let v = [this.rng() * 2 - 1, this.rng() * 2 - 1, this.rng() * 2 - 1]; | |
| const norm = Math.hypot(v[0], v[1], v[2]) || 1; | |
| for (let i = 0; i < 3; i++) { | |
| put(this.data.qvel, i, at(this.data.qvel, i) + (v[i] / norm) * dv); | |
| } | |
| } | |
| reset() { | |
| this.mujoco.mj_resetData(this.model, this.data); | |
| // Start from the keyframe the environment uses (`home`), if present. | |
| if (this.model.nkey > 0) { | |
| for (let i = 0; i < this.scene.nq; i++) put(this.data.qpos, i, at(this.model.key_qpos, i)); | |
| } | |
| this.lastAction.fill(0); | |
| this.prevAction.fill(0); | |
| this.delayQueue = []; | |
| this.totalReturn = 0; | |
| this.steps = 0; | |
| this.jitterSum = 0; | |
| this.jitterCount = 0; | |
| this.fallen = false; | |
| this.mujoco.mj_forward(this.model, this.data); | |
| } | |
| /** | |
| * Projected gravity: the world down-vector in the IMU site's frame. | |
| * | |
| * Matches `Go1Env.get_gravity` exactly — | |
| * `data.site_xmat[imu].T @ [0, 0, -1]`, which is the negated third row of that | |
| * row-major 3x3. It is **not** the `upvector` sensor; reading that instead | |
| * gives roughly the negation, in a frame that need not be the IMU's. | |
| * | |
| * These three numbers are the policy's orientation signal. Inverted, the policy | |
| * is told the robot is upside down and does the sensible thing: it stops | |
| * walking and tries to hold still. Measured with the wrong vector, the Go1 | |
| * travelled 0.02 m in five seconds against a 1.0 m/s command. | |
| */ | |
| gravity() { | |
| const base = this.iface.imu_site_id * 9 + 6; | |
| return [ | |
| -at(this.data.site_xmat, base + 0), | |
| -at(this.data.site_xmat, base + 1), | |
| -at(this.data.site_xmat, base + 2), | |
| ]; | |
| } | |
| #noise(scaleName) { | |
| const scale = this.iface.noise.scales[scaleName] ?? 0; | |
| return (2 * this.rng() - 1) * this.noiseLevel * scale; | |
| } | |
| /** Build the 48-dim observation exactly as `_get_obs` does. */ | |
| observe() { | |
| const d = this.data; | |
| const s = this.iface.sensors; | |
| const o = this.obs; | |
| let k = 0; | |
| const linvel = readBlock(d.sensordata, s.local_linvel.adr, 3); | |
| for (let i = 0; i < 3; i++) o[k++] = linvel[i] + this.#noise('linvel'); | |
| const gyro = readBlock(d.sensordata, s.gyro.adr, 3); | |
| for (let i = 0; i < 3; i++) o[k++] = gyro[i] + this.#noise('gyro'); | |
| const g = this.gravity(); | |
| for (let i = 0; i < 3; i++) o[k++] = g[i] + this.#noise('gravity'); | |
| const q0 = this.iface.qpos_joint_start; | |
| for (let i = 0; i < 12; i++) { | |
| o[k++] = at(d.qpos, q0 + i) - this.iface.default_pose[i] + this.#noise('joint_pos'); | |
| } | |
| const v0 = this.iface.qvel_joint_start; | |
| for (let i = 0; i < 12; i++) o[k++] = at(d.qvel, v0 + i) + this.#noise('joint_vel'); | |
| for (let i = 0; i < 12; i++) o[k++] = this.lastAction[i]; | |
| for (let i = 0; i < 3; i++) o[k++] = this.command[i]; | |
| // The perturbation axis adds its own noise on top of the env's, matching | |
| // `perturbed_rollout`, which noises the observation the policy receives. | |
| if (this.obsNoiseSigma > 0) { | |
| for (let i = 0; i < o.length; i++) o[i] += this.obsNoiseSigma * gaussian(this.rng); | |
| } | |
| return o; | |
| } | |
| /** | |
| * Advance one control step: apply `action`, run `n_substeps` physics steps. | |
| * Returns the wall-clock spent inside physics, for the HUD. | |
| */ | |
| step(action) { | |
| let applied = action; | |
| if (this.actuatorDelay > 0) { | |
| this.delayQueue.push(Float32Array.from(action)); | |
| applied = | |
| this.delayQueue.length > this.actuatorDelay | |
| ? this.delayQueue.shift() | |
| : new Float32Array(this.nu); | |
| } | |
| if (this.steps > 0) { | |
| let delta = 0; | |
| for (let i = 0; i < this.nu; i++) delta += (applied[i] - this.prevAction[i]) ** 2; | |
| this.jitterSum += Math.sqrt(delta); | |
| this.jitterCount += 1; | |
| } | |
| this.prevAction = Float32Array.from(applied); | |
| this.lastAction = Float32Array.from(applied); | |
| const pose = this.iface.default_pose; | |
| const scale = this.iface.action_scale; | |
| for (let i = 0; i < this.nu; i++) put(this.data.ctrl, i, pose[i] + applied[i] * scale); | |
| const t0 = performance.now(); | |
| for (let i = 0; i < this.scene.n_substeps; i++) { | |
| this.mujoco.mj_step(this.model, this.data); | |
| } | |
| const physicsMs = performance.now() - t0; | |
| this.steps += 1; | |
| // Uprightness, from the same projected-gravity vector the policy sees: its z | |
| // component is about -1 standing and turns positive once the robot is over. | |
| if (this.gravity()[2] > 0.0) this.fallen = true; | |
| // Command-tracking error: how far the achieved body velocity is from what the | |
| // joystick asked for. Exact, and directly meaningful to a viewer. | |
| // | |
| // The episodic *return* is deliberately not computed here. It is a sum of 16 | |
| // environment-specific terms involving contact state, feet air time, swing | |
| // peaks and actuator forces; a JavaScript reimplementation would produce a | |
| // number that looks like the benchmark's return and is not it, which is worse | |
| // than not showing one. Measured returns live in the results panel, where | |
| // they come from the recorded runs. | |
| const linvel = readBlock(this.data.sensordata, this.iface.sensors.local_linvel.adr, 3); | |
| const gyro = readBlock(this.data.sensordata, this.iface.sensors.gyro.adr, 3); | |
| this.trackingError = Math.hypot( | |
| linvel[0] - this.command[0], | |
| linvel[1] - this.command[1], | |
| gyro[2] - this.command[2] | |
| ); | |
| this.trackingSum = (this.trackingSum ?? 0) + this.trackingError; | |
| return physicsMs; | |
| } | |
| get meanTrackingError() { | |
| return this.steps > 0 ? (this.trackingSum ?? 0) / this.steps : 0; | |
| } | |
| get meanJitter() { | |
| return this.jitterCount > 0 ? this.jitterSum / this.jitterCount : 0; | |
| } | |
| } | |
| /** Box–Muller, for the observation-noise perturbation axis. */ | |
| export function gaussian(rng) { | |
| let u = 0; | |
| while (u === 0) u = rng(); | |
| return Math.sqrt(-2 * Math.log(u)) * Math.cos(2 * Math.PI * rng()); | |
| } | |