mujoco-react 0.1.0

package/src/hooks/useJointState.ts

@@ -0,0 +1,64 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * useJointState — per-joint position/velocity access (spec 2.3)
+ */
+
+import { useEffect, useRef } from 'react';
+import { useMujocoSim, useAfterPhysicsStep } from '../core/MujocoSimProvider';
+import { getName } from '../core/SceneLoader';
+import type { JointStateResult } from '../types';
+
+/**
+ * Track a MuJoCo joint's position and velocity by name.
+ * Values are updated every physics frame via refs (no re-renders).
+ *
+ * For hinge/slide joints, position/velocity are scalar (stored as Float64Array of length 1).
+ * For ball joints, position is quat (4), velocity is angular vel (3).
+ * For free joints, position is pos+quat (7), velocity is lin+ang vel (6).
+ */
+export function useJointState(name: string): JointStateResult {
+  const { mjModelRef, mjDataRef, status } = useMujocoSim();
+  const jointIdRef = useRef(-1);
+  const qposAdrRef = useRef(0);
+  const dofAdrRef = useRef(0);
+  const qposDimRef = useRef(1);
+  const dofDimRef = useRef(1);
+  const positionRef = useRef<number | Float64Array>(0);
+  const velocityRef = useRef<number | Float64Array>(0);
+
+  useEffect(() => {
+    const model = mjModelRef.current;
+    if (!model || status !== 'ready') return;
+    for (let i = 0; i < model.njnt; i++) {
+      if (getName(model, model.name_jntadr[i]) === name) {
+        jointIdRef.current = i;
+        qposAdrRef.current = model.jnt_qposadr[i];
+        dofAdrRef.current = model.jnt_dofadr[i];
+        const type = model.jnt_type[i];
+        // Type 0=free (7 qpos, 6 dof), 1=ball (4 qpos, 3 dof), 2=slide (1,1), 3=hinge (1,1)
+        if (type === 0) { qposDimRef.current = 7; dofDimRef.current = 6; }
+        else if (type === 1) { qposDimRef.current = 4; dofDimRef.current = 3; }
+        else { qposDimRef.current = 1; dofDimRef.current = 1; }
+        return;
+      }
+    }
+    jointIdRef.current = -1;
+  }, [name, status, mjModelRef]);
+
+  useAfterPhysicsStep((_model, data) => {
+    if (jointIdRef.current < 0) return;
+    const qa = qposAdrRef.current;
+    const da = dofAdrRef.current;
+    if (qposDimRef.current === 1) {
+      positionRef.current = data.qpos[qa];
+      velocityRef.current = data.qvel[da];
+    } else {
+      positionRef.current = new Float64Array(data.qpos.subarray(qa, qa + qposDimRef.current));
+      velocityRef.current = new Float64Array(data.qvel.subarray(da, da + dofDimRef.current));
+    }
+  });
+
+  return { position: positionRef, velocity: velocityRef };
+}

package/src/hooks/useKeyboardTeleop.ts

@@ -0,0 +1,97 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * useKeyboardTeleop — keyboard teleoperation hook (spec 12.1)
+ */
+
+import { useEffect, useRef } from 'react';
+import { useMujocoSim, useBeforePhysicsStep } from '../core/MujocoSimProvider';
+import { findActuatorByName } from '../core/SceneLoader';
+import type { KeyboardTeleopConfig } from '../types';
+
+/**
+ * Map keyboard keys to actuator commands.
+ *
+ * Supports three binding modes:
+ * - `delta`: Add delta to actuator value while key is held
+ * - `toggle`: Toggle between two values on key press
+ * - `set`: Set actuator to a fixed value while key is held
+ */
+export function useKeyboardTeleop(config: KeyboardTeleopConfig) {
+  const { mjModelRef, mjDataRef, status } = useMujocoSim();
+  const pressedRef = useRef(new Set<string>());
+  const toggleStateRef = useRef(new Map<string, boolean>());
+  const enabledRef = useRef(config.enabled ?? true);
+  enabledRef.current = config.enabled ?? true;
+
+  // Resolve actuator IDs
+  const bindingsRef = useRef(config.bindings);
+  bindingsRef.current = config.bindings;
+
+  // Actuator ID cache
+  const actuatorCacheRef = useRef(new Map<string, number>());
+  useEffect(() => {
+    const model = mjModelRef.current;
+    if (!model || status !== 'ready') return;
+    const cache = new Map<string, number>();
+    for (const binding of Object.values(config.bindings)) {
+      if (!cache.has(binding.actuator)) {
+        cache.set(binding.actuator, findActuatorByName(model, binding.actuator));
+      }
+    }
+    actuatorCacheRef.current = cache;
+  }, [config.bindings, status, mjModelRef]);
+
+  // Key event listeners
+  useEffect(() => {
+    const onKeyDown = (e: KeyboardEvent) => {
+      if (!enabledRef.current) return;
+      const key = e.key.toLowerCase();
+      if (bindingsRef.current[key]) {
+        pressedRef.current.add(key);
+        // Handle toggle on keydown
+        const binding = bindingsRef.current[key];
+        if (binding.toggle) {
+          const current = toggleStateRef.current.get(key) ?? false;
+          toggleStateRef.current.set(key, !current);
+        }
+      }
+    };
+    const onKeyUp = (e: KeyboardEvent) => {
+      pressedRef.current.delete(e.key.toLowerCase());
+    };
+    window.addEventListener('keydown', onKeyDown);
+    window.addEventListener('keyup', onKeyUp);
+    return () => {
+      window.removeEventListener('keydown', onKeyDown);
+      window.removeEventListener('keyup', onKeyUp);
+    };
+  }, []);
+
+  // Apply bindings each physics frame
+  useBeforePhysicsStep((_model, data) => {
+    if (!enabledRef.current) return;
+    const bindings = bindingsRef.current;
+    const cache = actuatorCacheRef.current;
+
+    for (const [key, binding] of Object.entries(bindings)) {
+      const actId = cache.get(binding.actuator);
+      if (actId === undefined || actId < 0) continue;
+
+      if (binding.toggle) {
+        // Toggle mode: set value based on toggle state
+        const state = toggleStateRef.current.get(key) ?? false;
+        data.ctrl[actId] = state ? binding.toggle[1] : binding.toggle[0];
+      } else if (pressedRef.current.has(key)) {
+        if (binding.delta !== undefined) {
+          // Delta mode: add delta while held
+          data.ctrl[actId] += binding.delta;
+        } else if (binding.set !== undefined) {
+          // Set mode: set fixed value while held
+          data.ctrl[actId] = binding.set;
+        }
+      }
+    }
+  });
+}

package/src/hooks/usePolicy.ts

@@ -0,0 +1,56 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * usePolicy — policy decimation loop hook (spec 10.1)
+ */
+
+import { useRef } from 'react';
+import { useMujocoSim, useBeforePhysicsStep } from '../core/MujocoSimProvider';
+import type { PolicyConfig } from '../types';
+
+/**
+ * Framework-agnostic policy execution hook.
+ *
+ * Manages a decimation loop: calls `onObservation` to build observations
+ * at the specified frequency, then calls `onAction` to apply the policy output.
+ * The actual inference (ONNX, TF.js, custom) is the consumer's responsibility.
+ *
+ * @param config Policy configuration
+ * @returns { step, isRunning } control handles
+ */
+export function usePolicy(config: PolicyConfig) {
+  const { mjModelRef } = useMujocoSim();
+  const lastActionTimeRef = useRef(0);
+  const lastActionRef = useRef<Float32Array | Float64Array | number[] | null>(null);
+  const isRunningRef = useRef(true);
+  const configRef = useRef(config);
+  configRef.current = config;
+
+  useBeforePhysicsStep((model, data) => {
+    if (!isRunningRef.current) return;
+
+    const cfg = configRef.current;
+    const dt = model.opt?.timestep ?? 0.002;
+    const interval = 1.0 / cfg.frequency;
+
+    // Check if it's time for a new action
+    if (data.time - lastActionTimeRef.current >= interval) {
+      // Build observation
+      const obs = cfg.onObservation(model, data);
+
+      // Apply action (consumer does inference inline or uses cached result)
+      cfg.onAction(obs, model, data);
+
+      lastActionTimeRef.current = data.time;
+      lastActionRef.current = obs;
+    }
+  });
+
+  return {
+    get isRunning() { return isRunningRef.current; },
+    start: () => { isRunningRef.current = true; },
+    stop: () => { isRunningRef.current = false; },
+    get lastObservation() { return lastActionRef.current; },
+  };
+}

package/src/hooks/useSensor.ts

@@ -0,0 +1,83 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * useSensor / useSensors — MuJoCo sensor access hooks (spec 2.1)
+ */
+
+import { useEffect, useRef, useMemo } from 'react';
+import { useMujocoSim, useAfterPhysicsStep } from '../core/MujocoSimProvider';
+import { getName } from '../core/SceneLoader';
+import type { SensorInfo, SensorResult } from '../types';
+
+/**
+ * Access a single MuJoCo sensor by name. Returns a ref-based value
+ * updated every physics frame without causing React re-renders.
+ */
+export function useSensor(name: string): SensorResult {
+  const { mjModelRef, mjDataRef, status } = useMujocoSim();
+  const sensorIdRef = useRef(-1);
+  const sensorAdrRef = useRef(0);
+  const sensorDimRef = useRef(0);
+  const valueRef = useRef<Float64Array>(new Float64Array(0));
+
+  // Resolve sensor ID once model is ready
+  useEffect(() => {
+    const model = mjModelRef.current;
+    if (!model || status !== 'ready') return;
+    for (let i = 0; i < model.nsensor; i++) {
+      if (getName(model, model.name_sensoradr[i]) === name) {
+        sensorIdRef.current = i;
+        sensorAdrRef.current = model.sensor_adr[i];
+        sensorDimRef.current = model.sensor_dim[i];
+        valueRef.current = new Float64Array(model.sensor_dim[i]);
+        return;
+      }
+    }
+    sensorIdRef.current = -1;
+  }, [name, status, mjModelRef]);
+
+  // Update every frame after physics step
+  useAfterPhysicsStep((_model, data) => {
+    if (sensorIdRef.current < 0) return;
+    const adr = sensorAdrRef.current;
+    const dim = sensorDimRef.current;
+    for (let i = 0; i < dim; i++) {
+      valueRef.current[i] = data.sensordata[adr + i];
+    }
+  });
+
+  return { value: valueRef, size: sensorDimRef.current };
+}
+
+/**
+ * Enumerate all sensors in the loaded MuJoCo model.
+ * Returns a stable array recomputed only when the model changes.
+ */
+export function useSensors(): SensorInfo[] {
+  const { mjModelRef, status } = useMujocoSim();
+
+  return useMemo(() => {
+    const model = mjModelRef.current;
+    if (!model || status !== 'ready') return [];
+    const SENSOR_TYPE_NAMES: Record<number, string> = {
+      0: 'touch', 1: 'accelerometer', 2: 'velocimeter', 3: 'gyro',
+      4: 'force', 5: 'torque', 6: 'magnetometer', 7: 'rangefinder',
+      8: 'jointpos', 9: 'jointvel', 10: 'tendonpos', 11: 'tendonvel',
+      12: 'actuatorpos', 13: 'actuatorvel', 14: 'actuatorfrc',
+    };
+    const result: SensorInfo[] = [];
+    for (let i = 0; i < model.nsensor; i++) {
+      const type = model.sensor_type[i];
+      result.push({
+        id: i,
+        name: getName(model, model.name_sensoradr[i]),
+        type,
+        typeName: SENSOR_TYPE_NAMES[type] ?? `unknown(${type})`,
+        dim: model.sensor_dim[i],
+        adr: model.sensor_adr[i],
+      });
+    }
+    return result;
+  }, [mjModelRef, status]);
+}

package/src/hooks/useSitePosition.ts

@@ -0,0 +1,62 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { useEffect, useRef } from 'react';
+import { useFrame } from '@react-three/fiber';
+import * as THREE from 'three';
+import { useMujocoSim } from '../core/MujocoSimProvider';
+import { findSiteByName } from '../core/SceneLoader';
+import type { SitePositionResult } from '../types';
+
+// Preallocated temp for rotation matrix extraction
+const _mat4 = new THREE.Matrix4();
+
+/**
+ * Returns reactive refs for a MuJoCo site's world position and orientation.
+ * Refs are updated every frame without triggering React re-renders.
+ */
+export function useSitePosition(siteName: string): SitePositionResult {
+  const { mjModelRef, mjDataRef, status } = useMujocoSim();
+  const siteIdRef = useRef(-1);
+  const positionRef = useRef(new THREE.Vector3());
+  const quaternionRef = useRef(new THREE.Quaternion());
+
+  // Resolve site ID when model is ready
+  useEffect(() => {
+    const model = mjModelRef.current;
+    if (!model || status !== 'ready') {
+      siteIdRef.current = -1;
+      return;
+    }
+    siteIdRef.current = findSiteByName(model, siteName);
+  }, [siteName, status, mjModelRef]);
+
+  // Update refs every frame
+  useFrame(() => {
+    const data = mjDataRef.current;
+    const sid = siteIdRef.current;
+    if (!data || sid < 0) return;
+
+    const i3 = sid * 3;
+    const i9 = sid * 9;
+
+    positionRef.current.set(
+      data.site_xpos[i3],
+      data.site_xpos[i3 + 1],
+      data.site_xpos[i3 + 2]
+    );
+
+    const m = data.site_xmat;
+    _mat4.set(
+      m[i9],     m[i9 + 1], m[i9 + 2], 0,
+      m[i9 + 3], m[i9 + 4], m[i9 + 5], 0,
+      m[i9 + 6], m[i9 + 7], m[i9 + 8], 0,
+      0,          0,          0,          1,
+    );
+    quaternionRef.current.setFromRotationMatrix(_mat4);
+  });
+
+  return { position: positionRef, quaternion: quaternionRef };
+}

package/src/hooks/useTrajectoryPlayer.ts

@@ -0,0 +1,105 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * useTrajectoryPlayer — trajectory playback/scrubbing (spec 13.2)
+ */
+
+import { useCallback, useRef } from 'react';
+import { useFrame } from '@react-three/fiber';
+import { useMujocoSim } from '../core/MujocoSimProvider';
+
+interface TrajectoryPlayerOptions {
+  fps?: number;
+  loop?: boolean;
+}
+
+/**
+ * Play back a sequence of qpos frames, overriding simulation state.
+ *
+ * When playing, the simulation is effectively paused and qpos is set
+ * from the trajectory each render frame at the specified FPS.
+ */
+export function useTrajectoryPlayer(
+  trajectory: number[][],
+  options: TrajectoryPlayerOptions = {},
+) {
+  const { mjModelRef, mjDataRef, mujocoRef, pausedRef } = useMujocoSim();
+  const fps = options.fps ?? 30;
+  const loop = options.loop ?? false;
+
+  const playingRef = useRef(false);
+  const frameRef = useRef(0);
+  const lastFrameTimeRef = useRef(0);
+
+  const play = useCallback(() => {
+    playingRef.current = true;
+    pausedRef.current = true; // Pause sim during playback
+    lastFrameTimeRef.current = performance.now();
+  }, [pausedRef]);
+
+  const pause = useCallback(() => {
+    playingRef.current = false;
+  }, []);
+
+  const seek = useCallback((frameIdx: number) => {
+    frameRef.current = Math.max(0, Math.min(frameIdx, trajectory.length - 1));
+    const model = mjModelRef.current;
+    const data = mjDataRef.current;
+    if (!model || !data || !trajectory[frameRef.current]) return;
+    const qpos = trajectory[frameRef.current];
+    for (let i = 0; i < Math.min(qpos.length, model.nq); i++) {
+      data.qpos[i] = qpos[i];
+    }
+    mujocoRef.current.mj_forward(model, data);
+  }, [trajectory, mjModelRef, mjDataRef, mujocoRef]);
+
+  const reset = useCallback(() => {
+    frameRef.current = 0;
+    playingRef.current = false;
+    pausedRef.current = false;
+  }, [pausedRef]);
+
+  useFrame(() => {
+    if (!playingRef.current || trajectory.length === 0) return;
+
+    const now = performance.now();
+    const elapsed = now - lastFrameTimeRef.current;
+    const frameInterval = 1000 / fps;
+
+    if (elapsed < frameInterval) return;
+    lastFrameTimeRef.current = now;
+
+    const model = mjModelRef.current;
+    const data = mjDataRef.current;
+    if (!model || !data) return;
+
+    const qpos = trajectory[frameRef.current];
+    if (!qpos) return;
+
+    for (let i = 0; i < Math.min(qpos.length, model.nq); i++) {
+      data.qpos[i] = qpos[i];
+    }
+    mujocoRef.current.mj_forward(model, data);
+
+    frameRef.current++;
+    if (frameRef.current >= trajectory.length) {
+      if (loop) {
+        frameRef.current = 0;
+      } else {
+        playingRef.current = false;
+        pausedRef.current = false;
+      }
+    }
+  });
+
+  return {
+    play,
+    pause,
+    seek,
+    reset,
+    get frame() { return frameRef.current; },
+    get playing() { return playingRef.current; },
+    get totalFrames() { return trajectory.length; },
+  };
+}

package/src/hooks/useTrajectoryRecorder.ts

@@ -0,0 +1,97 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * useTrajectoryRecorder — trajectory recording hook (spec 13.1)
+ */
+
+import { useCallback, useRef } from 'react';
+import { useMujocoSim, useAfterPhysicsStep } from '../core/MujocoSimProvider';
+import type { TrajectoryFrame } from '../types';
+
+interface RecorderOptions {
+  fields?: ('qpos' | 'qvel' | 'ctrl' | 'sensordata')[];
+}
+
+/**
+ * Record simulation trajectories for analysis, replay, or training data.
+ */
+export function useTrajectoryRecorder(options: RecorderOptions = {}) {
+  const { mjModelRef } = useMujocoSim();
+  const recordingRef = useRef(false);
+  const framesRef = useRef<TrajectoryFrame[]>([]);
+  const fields = options.fields ?? ['qpos'];
+
+  useAfterPhysicsStep((_model, data) => {
+    if (!recordingRef.current) return;
+
+    const frame: TrajectoryFrame = {
+      time: data.time,
+      qpos: new Float64Array(data.qpos),
+    };
+
+    if (fields.includes('qvel')) frame.qvel = new Float64Array(data.qvel);
+    if (fields.includes('ctrl')) frame.ctrl = new Float64Array(data.ctrl);
+    if (fields.includes('sensordata') && data.sensordata) {
+      frame.sensordata = new Float64Array(data.sensordata);
+    }
+
+    framesRef.current.push(frame);
+  });
+
+  const start = useCallback(() => {
+    framesRef.current = [];
+    recordingRef.current = true;
+  }, []);
+
+  const stop = useCallback(() => {
+    recordingRef.current = false;
+    return framesRef.current;
+  }, []);
+
+  const downloadJSON = useCallback(() => {
+    const frames = framesRef.current;
+    const data = frames.map(f => ({
+      time: f.time,
+      qpos: Array.from(f.qpos),
+      ...(f.qvel ? { qvel: Array.from(f.qvel) } : {}),
+      ...(f.ctrl ? { ctrl: Array.from(f.ctrl) } : {}),
+      ...(f.sensordata ? { sensordata: Array.from(f.sensordata) } : {}),
+    }));
+    const blob = new Blob([JSON.stringify(data, null, 2)], { type: 'application/json' });
+    const url = URL.createObjectURL(blob);
+    const a = document.createElement('a');
+    a.href = url;
+    a.download = 'trajectory.json';
+    a.click();
+    URL.revokeObjectURL(url);
+  }, []);
+
+  const downloadCSV = useCallback(() => {
+    const frames = framesRef.current;
+    if (frames.length === 0) return;
+    const nq = frames[0].qpos.length;
+    const headers = ['time', ...Array.from({ length: nq }, (_, i) => `qpos_${i}`)];
+    const rows = frames.map(f =>
+      [f.time, ...Array.from(f.qpos)].join(',')
+    );
+    const csv = [headers.join(','), ...rows].join('\n');
+    const blob = new Blob([csv], { type: 'text/csv' });
+    const url = URL.createObjectURL(blob);
+    const a = document.createElement('a');
+    a.href = url;
+    a.download = 'trajectory.csv';
+    a.click();
+    URL.revokeObjectURL(url);
+  }, []);
+
+  return {
+    start,
+    stop,
+    downloadJSON,
+    downloadCSV,
+    get recording() { return recordingRef.current; },
+    get frameCount() { return framesRef.current.length; },
+    get frames() { return framesRef.current; },
+  };
+}

package/src/hooks/useVideoRecorder.ts

@@ -0,0 +1,82 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * useVideoRecorder — canvas video recording hook (spec 13.3)
+ */
+
+import { useCallback, useRef } from 'react';
+import { useThree } from '@react-three/fiber';
+
+interface VideoRecorderOptions {
+  fps?: number;
+  mimeType?: string;
+}
+
+/**
+ * Record the R3F canvas to a video file using MediaRecorder.
+ */
+export function useVideoRecorder(options: VideoRecorderOptions = {}) {
+  const { gl } = useThree();
+  const recorderRef = useRef<MediaRecorder | null>(null);
+  const chunksRef = useRef<Blob[]>([]);
+  const recordingRef = useRef(false);
+
+  const start = useCallback(() => {
+    const canvas = gl.domElement;
+    const fps = options.fps ?? 30;
+    const mimeType = options.mimeType ?? 'video/webm';
+
+    const stream = canvas.captureStream(fps);
+    const recorder = new MediaRecorder(stream, {
+      mimeType: MediaRecorder.isTypeSupported(mimeType) ? mimeType : 'video/webm',
+    });
+
+    chunksRef.current = [];
+    recorder.ondataavailable = (e) => {
+      if (e.data.size > 0) chunksRef.current.push(e.data);
+    };
+
+    recorder.start();
+    recorderRef.current = recorder;
+    recordingRef.current = true;
+  }, [gl, options.fps, options.mimeType]);
+
+  const stop = useCallback((): Promise<Blob> => {
+    return new Promise((resolve) => {
+      const recorder = recorderRef.current;
+      if (!recorder || recorder.state === 'inactive') {
+        resolve(new Blob([]));
+        return;
+      }
+
+      recorder.onstop = () => {
+        const blob = new Blob(chunksRef.current, { type: recorder.mimeType });
+        chunksRef.current = [];
+        recordingRef.current = false;
+        recorderRef.current = null;
+        resolve(blob);
+      };
+
+      recorder.stop();
+    });
+  }, []);
+
+  const download = useCallback(async (filename = 'recording.webm') => {
+    const blob = await stop();
+    if (blob.size === 0) return;
+    const url = URL.createObjectURL(blob);
+    const a = document.createElement('a');
+    a.href = url;
+    a.download = filename;
+    a.click();
+    URL.revokeObjectURL(url);
+  }, [stop]);
+
+  return {
+    start,
+    stop,
+    download,
+    get recording() { return recordingRef.current; },
+  };
+}