mujoco-react 12.0.0 → 12.1.1

package/src/runtime/client.ts

@@ -0,0 +1,196 @@
+import type { MainModule, MjModel } from '@mujoco/mujoco';
+import type { MujocoModule, MujocoModel, MujocoSimAPI, PolicyCameraFrameCaptureOptions } from '../types';
+import { capturePolicyCameraFrames } from '../policyCameraFrames';
+import { displayFields, type DisplayState, type EngineInitialization } from './engine';
+import type { Request, Response } from './protocol';
+import type { RuntimeOptions, RuntimeState, RuntimeSnapshot, RuntimeMetrics, SimulationRuntime, WorkerPort, AdvanceOptions } from './types';
+
+interface Owner { close: () => void; contacts: () => RuntimeState['contacts'] }
+const owners = new WeakMap<MujocoModel, Owner>();
+export const isWorkerOwned = (model: MujocoModel | null) => Boolean(model && owners.has(model));
+export function closeModelRuntime(model: MujocoModel | null) { if (model) owners.get(model)?.close(); }
+export function runtimeContacts(model: MujocoModel) { return owners.get(model)?.contacts(); }
+export function assertMainThreadOwnership(model: MujocoModel | null) {
+  if (isWorkerOwned(model)) throw new Error('Worker runtime owns this simulation. Use runtime operations or dispose it before using viewer physics operations.');
+}
+const abortError = () => new DOMException('Simulation operation cancelled', 'AbortError');
+
+/** Client owns backpressure and cancellation; protocol owns authoritative engine state. */
+export class RuntimeClient implements SimulationRuntime {
+  private pending: { id: number; resolve: (r: Response) => void; reject: (e: Error) => void } | null = null;
+  private nextId = 0;
+  private closed = false;
+  private busy = false;
+  private flight: Promise<unknown> = Promise.resolve();
+  private generation = 0;
+  private latest: RuntimeState = { time: 0, qpos: new Float64Array(), qvel: new Float64Array(), ctrl: new Float64Array(), sensors: new Float64Array(), contacts: [] };
+  private statistics: RuntimeMetrics = { steps: 0, physicsMs: 0, elapsedMs: 0, realTimeFactor: 0, captureMs: 0, busy: false };
+  private readonly modelId = crypto.randomUUID();
+  constructor(
+    private worker: WorkerPort,
+    private apply: (display: DisplayState) => void,
+    private release: (restoreViewer: boolean) => void,
+    private capture?: (options: PolicyCameraFrameCaptureOptions) => ReturnType<typeof capturePolicyCameraFrames>,
+  ) {
+    worker.onmessage = event => {
+      const response = event.data as Response;
+      if (this.closed || response.id !== this.pending?.id) return;
+      if (response.display) {
+        try {
+          this.apply(response.display);
+          const f = response.display.fields;
+          this.latest = { time: response.display.time, qpos: f.qpos, qvel: f.qvel, ctrl: f.ctrl, sensors: f.sensordata,
+            contacts: response.done ? response.display.contacts : this.latest.contacts };
+        } catch (error) { this.close(error instanceof Error ? error : new Error(String(error))); return; }
+      }
+      if (!response.done) this.worker.postMessage({ kind: 'ack', id: response.id });
+      if (response.done) {
+        const pending = this.pending; this.pending = null;
+        if (response.error) pending.reject(new Error(response.error));
+        else pending.resolve(response);
+      }
+    };
+    worker.onerror = event => this.close(new Error(event.message || 'Simulation worker failed'));
+    worker.onmessageerror = () => this.close(new Error('Cannot read simulation worker response'));
+  }
+  get state(): RuntimeState {
+    return { ...this.latest, qpos: this.latest.qpos.slice(), qvel: this.latest.qvel.slice(), ctrl: this.latest.ctrl.slice(), sensors: this.latest.sensors.slice(),
+      contacts: this.latest.contacts.map(c => ({ ...c, pos: [...c.pos] as [number, number, number] })) };
+  }
+  get metrics() { return { ...this.statistics, busy: this.busy }; }
+  private request(request: Omit<Request, 'id'>): Promise<Response> {
+    if (this.closed) return Promise.reject(new Error('Simulation runtime is disposed'));
+    return new Promise((resolve, reject) => {
+      const id = ++this.nextId;
+      this.pending = { id, resolve, reject };
+      try { this.worker.postMessage({ ...request, id }); }
+      catch (error) { this.pending = null; reject(error instanceof Error ? error : new Error(String(error))); }
+    });
+  }
+  private exclusive<T>(operation: () => Promise<T>): Promise<T> {
+    if (this.closed) return Promise.reject(new Error('Simulation runtime is disposed'));
+    if (this.busy) return Promise.reject(new Error('Runtime is busy; await the current operation (work is not queued)'));
+    this.busy = true;
+    const result = operation().finally(() => { this.busy = false; });
+    this.flight = result.catch(() => {});
+    return result;
+  }
+  initialize(initial: EngineInitialization) {
+    return this.exclusive(async () => { await this.request({ kind: 'initialize', initial } as Request); });
+  }
+  advance(options: AdvanceOptions) {
+    return this.exclusive(async () => {
+      if (options.signal?.aborted) throw abortError();
+      const started = performance.now(), simStart = this.latest.time;
+      const abort = () => { this.worker.postMessage({ kind: 'cancel', id: ++this.nextId }); };
+      options.signal?.addEventListener('abort', abort, { once: true });
+      try {
+        const result = await this.request({ kind: 'advance', steps: options.steps,
+          controls: options.controls ? Float64Array.from(options.controls) : undefined } as Request);
+        const elapsedMs = performance.now() - started;
+        this.statistics = { ...this.statistics, steps: result.steps ?? 0, physicsMs: result.physicsMs ?? 0, elapsedMs,
+          realTimeFactor: (this.latest.time - simStart) * 1000 / Math.max(elapsedMs, .001) };
+        if (result.cancelled || options.signal?.aborted) throw abortError();
+        return this.state;
+      } finally { options.signal?.removeEventListener('abort', abort); }
+    });
+  }
+  observe(cameras?: PolicyCameraFrameCaptureOptions) {
+    return this.exclusive(async () => {
+      const generation = this.generation;
+      await this.request({ kind: 'snapshot' });
+      const state = this.state, started = performance.now();
+      const frames = cameras ? await this.capture?.(cameras) : undefined;
+      if (cameras && !frames) throw new Error('Camera observations require an attached canvas');
+      this.statistics.captureMs = performance.now() - started;
+      if (generation !== this.generation || this.closed) throw abortError();
+      return { ...state, ...(frames ? { cameras: frames } : {}) };
+    });
+  }
+  saveState() {
+    return this.exclusive(async () => {
+      const result = await this.request({ kind: 'snapshot' });
+      const state = result.display!.integration!;
+      return { modelId: this.modelId, signature: state.signature, values: state.values.slice() };
+    });
+  }
+  restoreState(snapshot: RuntimeSnapshot) {
+    return this.exclusive(async () => {
+      if (snapshot.modelId !== this.modelId) throw new Error('Snapshot belongs to a different runtime');
+      await this.request({ kind: 'restore', state: { signature: snapshot.signature, values: snapshot.values } } as Request);
+      return this.state;
+    });
+  }
+  async cancel() {
+    if (this.closed) return;
+    this.generation++;
+    this.worker.postMessage({ kind: 'cancel', id: ++this.nextId });
+    await this.flight;
+  }
+  async reset() {
+    await this.cancel();
+    return this.exclusive(async () => { await this.request({ kind: 'reset' }); return this.state; });
+  }
+  async dispose() { await this.cancel(); this.close(); }
+  close(error = new Error('Simulation runtime disposed'), restoreViewer = true) {
+    if (this.closed) return;
+    this.closed = true; this.generation++;
+    this.worker.terminate();
+    this.pending?.reject(error); this.pending = null;
+    this.release(restoreViewer);
+  }
+}
+
+export async function attachSimulationRuntime(
+  module: MujocoModule, api: MujocoSimAPI, waitForDisplay: () => Promise<void>, invalidate: () => void,
+  options: RuntimeOptions = {},
+): Promise<SimulationRuntime> {
+  const model = api.mjModelRef.current, data = api.mjDataRef.current;
+  if (!model || !data) throw new Error('Scene must be ready before creating a runtime');
+  assertMainThreadOwnership(model);
+  api.setPaused(true);
+  const state = api.saveState().integration!;
+  const initialState = { signature: state.signature, values: state.values };
+  const mj = module as unknown as MainModule;
+  const path = `/runtime-${crypto.randomUUID()}.mjb`;
+  let binary: Uint8Array;
+  try { mj.mj_saveModel(model as unknown as MjModel, path, null); binary = mj.FS.readFile(path); }
+  finally { try { mj.FS.unlink(path); } catch { /* save may fail before file creation */ } }
+  const worker = options.createWorker?.() ?? new Worker(new URL('./physics.worker.js', import.meta.url), { type: 'module' });
+  let complete: DisplayState | undefined;
+  const apply = (display: DisplayState) => {
+    if (api.mjModelRef.current !== model) throw new Error('Simulation model changed');
+    if (display.integration) {
+      module.mj_setState(model, data, Array.from(display.integration.values), display.integration.signature);
+      complete = display;
+    }
+    for (const field of displayFields) {
+      const target = data[field];
+      if (!(target instanceof Float64Array || target instanceof Int32Array) || target.length !== display.fields[field].length) throw new Error(`Invalid display field ${field}`);
+      target.set(display.fields[field]);
+    }
+    data.time = display.time;
+    invalidate();
+  };
+  const client = new RuntimeClient(worker, apply, restoreViewer => {
+    // Revert partial display frames to the last acknowledged complete state.
+    try {
+      if (api.mjModelRef.current === model && complete) {
+        apply(complete);
+        // Native contact/constraint structures cannot be mirrored as typed arrays.
+        // Rebuild once when returning ownership, never on progress frames or unmount.
+        if (restoreViewer) {
+          module.mj_forward(model, data);
+          const state = complete.integration!;
+          module.mj_setState(model, data, Array.from(state.values), state.signature);
+        }
+      }
+    } finally { owners.delete(model); }
+  }, async cameras => {
+    await waitForDisplay();
+    return capturePolicyCameraFrames(api, cameras);
+  });
+  owners.set(model, { close: () => client.close(undefined, false), contacts: () => client.state.contacts });
+  try { await client.initialize({ model: binary, state: initialState }); return client; }
+  catch (error) { client.close(); throw error; }
+}

package/src/runtime/engine.ts

@@ -0,0 +1,93 @@
+import type { MainModule, MjModel, MjData } from '@mujoco/mujoco';
+import type { ContactInfo } from '../types';
+
+export const displayFields = ['qpos', 'qvel', 'ctrl', 'act', 'plastic', 'sensordata', 'xpos', 'xquat', 'xmat',
+  'geom_xpos', 'geom_xmat', 'site_xpos', 'site_xmat', 'cam_xpos', 'cam_xmat',
+  'flexvert_xpos', 'ten_length', 'wrap_xpos', 'cvel', 'cfrc_ext', 'qfrc_applied', 'xfrc_applied', 'qfrc_bias', 'ten_wrapadr', 'ten_wrapnum'] as const;
+export interface EngineState { signature: number; values: Float64Array }
+export interface EngineInitialization { model: Uint8Array; state: EngineState }
+export interface DisplayState {
+  time: number;
+  fields: Record<typeof displayFields[number], Float64Array>;
+  contacts: ContactInfo[];
+  integration?: EngineState;
+}
+export class SimulationEngine {
+  readonly model: MjModel;
+  readonly data: MjData;
+  constructor(private mj: MainModule, private initial: EngineInitialization) {
+    mj.FS.writeFile('/runtime.mjb', initial.model);
+    const vfs = new mj.MjVFS();
+    try { this.model = mj.MjModel.from_binary_path('/runtime.mjb', vfs); }
+    finally { vfs.delete(); mj.FS.unlink('/runtime.mjb'); }
+    this.data = new mj.MjData(this.model);
+    try { this.reset(); } catch (error) { this.dispose(); throw error; }
+  }
+  reset() {
+    this.mj.mj_resetData(this.model, this.data);
+    this.restore(this.initial.state);
+  }
+  restore(state: EngineState) {
+    if (state.signature !== this.mj.mjtState.mjSTATE_INTEGRATION.value ||
+      state.values.length !== this.mj.mj_stateSize(this.model, state.signature) ||
+      !state.values.every(Number.isFinite)) throw new Error('Invalid runtime snapshot');
+    this.mj.mj_setState(this.model, this.data, Array.from(state.values), state.signature);
+    this.mj.mj_forward(this.model, this.data);
+  }
+  setControls(controls?: Float64Array) {
+    if (controls) {
+      if (controls.length !== this.data.ctrl.length || !controls.every(Number.isFinite)) throw new Error('Invalid runtime controls');
+      this.data.ctrl.set(controls);
+    }
+  }
+  step() {
+    const expected = this.data.time + this.model.opt.timestep;
+    this.mj.mj_step(this.model, this.data);
+    if (Math.abs(this.data.time - expected) > 1e-7 || !this.data.qpos.every(Number.isFinite)) {
+      throw new Error('Physics became unstable or reset unexpectedly');
+    }
+  }
+  save(): EngineState {
+    const signature = this.mj.mjtState.mjSTATE_INTEGRATION.value;
+    const buffer = new this.mj.DoubleBuffer(this.mj.mj_stateSize(this.model, signature));
+    try {
+      this.mj.mj_getState(this.model, this.data, buffer, signature);
+      return { signature, values: new Float64Array(buffer.GetView()) };
+    } finally { buffer.delete(); }
+  }
+  snapshot(complete = true): DisplayState {
+    const integration = complete ? this.save() : undefined;
+    // mj_step integrates qpos after computing derived fields. Refresh completed
+    // observations at their reported time, preserving solver warmstart exactly.
+    if (integration) this.mj.mj_forward(this.model, this.data);
+    try { return {
+      time: this.data.time,
+      fields: Object.fromEntries(displayFields.map(key => [key, new Float64Array(this.data[key])])) as DisplayState['fields'],
+      contacts: complete ? this.contacts() : [],
+      ...(integration ? { integration } : {}),
+    }; } finally {
+      if (integration) this.mj.mj_setState(this.model, this.data, Array.from(integration.values), integration.signature);
+    }
+  }
+  private contacts(): ContactInfo[] {
+    const result: ContactInfo[] = [], vector = this.data.contact;
+    const names = this.model.names, addresses = this.model.name_geomadr;
+    const name = (id: number) => {
+      if (id < 0) return 'flex';
+      const start = addresses[id]; let end = start;
+      while (end < names.length && names[end]) end++;
+      return new TextDecoder().decode(names.subarray(start, end));
+    };
+    try {
+      for (let i = 0; i < this.data.ncon; i++) {
+        const c = vector.get(i);
+        if (!c) continue;
+        try { result.push({ geom1: c.geom1, geom2: c.geom2, geom1Name: name(c.geom1), geom2Name: name(c.geom2),
+          pos: Array.from(c.pos) as [number, number, number], depth: c.dist }); }
+        finally { c.delete(); }
+      }
+    } finally { vector.delete(); }
+    return result;
+  }
+  dispose() { this.data.delete(); this.model.delete(); }
+}

package/src/runtime/protocol.ts

@@ -0,0 +1,59 @@
+import type { MainModule } from '@mujoco/mujoco';
+import { SimulationEngine, type DisplayState, type EngineInitialization, type EngineState } from './engine';
+export type Request = { id: number } & (
+  { kind: 'initialize'; initial: EngineInitialization } |
+  { kind: 'advance'; steps: number; controls?: Float64Array } |
+  { kind: 'restore'; state: EngineState } |
+  { kind: 'reset' | 'snapshot' | 'cancel' | 'ack' }
+);
+export interface Response {
+  id: number; done: boolean; display?: DisplayState; error?: string;
+  cancelled?: boolean; steps?: number; physicsMs?: number;
+}
+/** Same protocol is exercised with a real engine in worker-thread tests. */
+export function createRuntimeHandler(load: () => Promise<MainModule>, send: (response: Response) => void) {
+  let engine: SimulationEngine | undefined;
+  let generation = 0;
+  let busy = false;
+  let displayPending = false;
+  return async (request: Request) => {
+    if (request.kind === 'ack') { displayPending = false; return; }
+    if (request.kind === 'cancel') { generation++; return; }
+    if (busy) { send({ id: request.id, done: true, error: 'Runtime is busy; await the current operation' }); return; }
+    busy = true;
+    const current = generation;
+    let steps = 0, physicsMs = 0;
+    try {
+      if (request.kind === 'initialize') {
+        if (engine) throw new Error('Runtime is already initialized');
+        engine = new SimulationEngine(await load(), request.initial);
+      } else if (!engine) throw new Error('Runtime is not initialized');
+      else if (request.kind === 'reset') engine.reset();
+      else if (request.kind === 'restore') engine.restore(request.state);
+      else if (request.kind === 'advance') {
+        if (!Number.isSafeInteger(request.steps) || request.steps <= 0) throw new Error('Steps must be a positive safe integer');
+        engine.setControls(request.controls);
+        let lastDisplay = performance.now();
+        while (steps < request.steps && generation === current) {
+          const sliceStart = performance.now();
+          do {
+            const start = performance.now();
+            engine.step();
+            physicsMs += performance.now() - start;
+            steps++;
+          } while (steps < request.steps && performance.now() - sliceStart < 8);
+          if (!displayPending && performance.now() - lastDisplay >= 32) {
+            displayPending = true;
+            send({ id: request.id, done: false, display: engine.snapshot(false) });
+            lastDisplay = performance.now();
+          }
+          // A task yield is required: a microtask alone starves cancel messages.
+          if (steps < request.steps) await new Promise(resolve => setTimeout(resolve, 0));
+        }
+      }
+      send({ id: request.id, done: true, display: engine!.snapshot(), cancelled: generation !== current, steps, physicsMs });
+    } catch (error) {
+      send({ id: request.id, done: true, error: error instanceof Error ? error.message : String(error) });
+    } finally { busy = false; displayPending = false; }
+  };
+}

package/src/runtime/types.ts

@@ -0,0 +1,51 @@
+import type { ContactInfo, PolicyCameraFrameCaptureOptions, PolicyCameraFrameCaptureResult } from '../types';
+
+export interface RuntimeMetrics {
+  steps: number;
+  physicsMs: number;
+  elapsedMs: number;
+  realTimeFactor: number;
+  captureMs: number;
+  busy: boolean;
+}
+export interface RuntimeState {
+  time: number;
+  qpos: Float64Array;
+  qvel: Float64Array;
+  ctrl: Float64Array;
+  sensors: Float64Array;
+  contacts: ContactInfo[];
+}
+export interface RuntimeSnapshot {
+  readonly modelId: string;
+  readonly signature: number;
+  readonly values: Float64Array;
+}
+export interface AdvanceOptions {
+  steps: number;
+  controls?: ArrayLike<number>;
+  signal?: AbortSignal;
+}
+/** One operation at a time. Await advance/observe rather than queueing work. */
+export interface SimulationRuntime {
+  readonly state: RuntimeState;
+  readonly metrics: Readonly<RuntimeMetrics>;
+  advance(options: AdvanceOptions): Promise<RuntimeState>;
+  observe(cameras?: PolicyCameraFrameCaptureOptions): Promise<RuntimeState & { cameras?: PolicyCameraFrameCaptureResult }>;
+  saveState(): Promise<RuntimeSnapshot>;
+  restoreState(snapshot: RuntimeSnapshot): Promise<RuntimeState>;
+  cancel(): Promise<void>;
+  reset(): Promise<RuntimeState>;
+  dispose(): Promise<void>;
+}
+export interface WorkerPort {
+  postMessage(message: unknown): void;
+  terminate(): void;
+  onmessage: ((event: MessageEvent) => void) | null;
+  onerror: ((event: ErrorEvent) => void) | null;
+  onmessageerror: ((event: MessageEvent) => void) | null;
+}
+export interface RuntimeOptions {
+  /** Override worker construction for custom bundlers or runtime tests. */
+  createWorker?: () => WorkerPort;
+}

package/src/types.ts

@@ -3,6 +3,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import type { SimulationRuntime, RuntimeOptions } from './runtime/types';
+import type { PhysicsPerformance } from './core/stepBudget';
 import type React from 'react';
 import type { ReactNode } from 'react';
 import type { CanvasProps, ThreeElements } from '@react-three/fiber';
@@ -618,6 +620,8 @@ export interface PhysicsConfig {
   timestep?: number;
   /** Override the MJCF integrator. Omit to preserve model settings. */
   integrator?: MujocoIntegrator;
+  /** Main-thread physics time budget per frame. Default 8 ms; at most 100 steps. */
+  physicsBudgetMs?: number;
   substeps?: number;
   paused?: boolean;
   speed?: number;
@@ -956,6 +960,8 @@ export interface PolicyConfig {
    * - replace is useful for receding-horizon policies where a fresh chunk should supersede stale queued actions.
    */
   queueStrategy?: 'append' | 'replace';
+  /** Maximum buffered actions; overflow fails explicitly. Default: 256. */
+  maxQueuedActions?: number;
   /**
    * Clear queued actions and ignore in-flight async results when `stop()` is called.
    * Defaults to false so callers can choose pause/resume behavior explicitly.
@@ -1463,6 +1469,9 @@ export interface BodyProps {
 // ---- Public API (spec: full surface) ----
 
 export interface MujocoSimAPI {
+  /** Dedicated worker owner for exact stepping and policy observations. */
+  createRuntime(options?: RuntimeOptions): Promise<SimulationRuntime>;
+  getPerformance(): PhysicsPerformance;
   // State
   readonly status: 'loading' | 'ready' | 'error';
   readonly config: SceneConfig;
@@ -1887,6 +1896,8 @@ export type MujocoCanvasProps = Omit<CanvasProps, 'onError'> & {
   timestep?: number;
   /** Override the MJCF integrator. Omit to preserve model settings. */
   integrator?: MujocoIntegrator;
+  /** Main-thread physics time budget per frame. Default 8 ms; at most 100 steps. */
+  physicsBudgetMs?: number;
   substeps?: number;
   paused?: boolean;
   speed?: number;

package/src/vite.ts

@@ -92,8 +92,13 @@ export function mujocoReact(options: MujocoReactPluginOptions) {
       // warning is expected, not a failure. Raise the limit so consumers don't
       // see it. Vite merges plugin config on top of user config, so only set a
       // default when the consumer hasn't specified their own limit.
-      if (userConfig.build?.chunkSizeWarningLimit !== undefined) return;
-      return { build: { chunkSizeWarningLimit: 4000 } };
+      return {
+        // Keep the worker URL relative to its real package file in dev. Vite's
+        // dependency prebundler otherwise relocates it into .vite/deps.
+        optimizeDeps: { exclude: ['mujoco-react'] },
+        ...(userConfig.build?.chunkSizeWarningLimit === undefined
+          ? { build: { chunkSizeWarningLimit: 4000 } } : {}),
+      };
     },
     configResolved(config: ViteConfig) {
       root = config.root;