mujoco-react 0.2.0 → 0.3.0

package/src/components/IkGizmo.tsx

@@ -8,6 +8,7 @@ import { useFrame, useThree } from '@react-three/fiber';
 import { useEffect, useRef } from 'react';
 import * as THREE from 'three';
 import { useMujocoSim } from '../core/MujocoSimProvider';
+import { useIk } from '../core/IkContext';
 import { findSiteByName } from '../core/SceneLoader';
 import type { IkGizmoProps } from '../types';
 
@@ -20,25 +21,18 @@ const _scale = new THREE.Vector3(1, 1, 1);
 /**
  * IkGizmo — drei PivotControls that tracks a MuJoCo site.
  *
+ * Must be rendered inside an `<IkController>`.
+ *
  * Props:
- * - `siteName` — MuJoCo site to track. Defaults to `SceneConfig.tcpSiteName`.
+ * - `siteName` — MuJoCo site to track. Defaults to the IkController's configured site.
  * - `scale` — Gizmo handle scale. Default: 0.18.
  * - `onDrag` — Custom drag callback `(pos, quat) => void`.
- *   When omitted, dragging enables IK and writes to the provider's IK target.
+ *   When omitted, dragging enables IK and writes to the IK target.
  *   When provided, the consumer handles what happens during drag.
- *
- * Multiple gizmos can be rendered — each tracks its own site.
- * Zero gizmos is fine — programmatic IK control works via the provider API.
- *
- * Uses a tiny invisible mesh as child instead of axesHelper — PivotControls
- * computes an anchor offset from children's bounding box, and axesHelper's
- * (0→0.15) bounds would shift the handles away from the TCP origin.
  */
 export function IkGizmo({ siteName, scale = 0.18, onDrag }: IkGizmoProps) {
-  const {
-    ikTargetRef, mjModelRef, mjDataRef, siteIdRef,
-    api, ikEnabledRef, status,
-  } = useMujocoSim();
+  const { mjModelRef, mjDataRef, status } = useMujocoSim();
+  const { ikTargetRef, siteIdRef, ikEnabledRef, setIkEnabled } = useIk();
 
   const wrapperRef = useRef<THREE.Group>(null);
   const pivotRef = useRef<THREE.Group>(null);
@@ -46,25 +40,23 @@ export function IkGizmo({ siteName, scale = 0.18, onDrag }: IkGizmoProps) {
   const localSiteIdRef = useRef(-1);
   const { controls } = useThree();
 
-  // Resolve the site ID from siteName (or fall back to provider's tcpSiteName)
+  // Resolve the site ID from siteName (only when an explicit siteName override is given)
   useEffect(() => {
     const model = mjModelRef.current;
-    if (!model || status !== 'ready') {
+    if (!model || status !== 'ready' || !siteName) {
       localSiteIdRef.current = -1;
       return;
     }
-    if (siteName) {
-      localSiteIdRef.current = findSiteByName(model, siteName);
-    } else {
-      // Default: use the provider's siteIdRef (from SceneConfig.tcpSiteName)
-      localSiteIdRef.current = siteIdRef.current;
-    }
-  }, [siteName, status, mjModelRef, siteIdRef]);
+    localSiteIdRef.current = findSiteByName(model, siteName);
+  }, [siteName, status, mjModelRef]);
 
   // Every frame: sync the visual wrapper to the tracked site (when not dragging)
   useFrame(() => {
     const data = mjDataRef.current;
-    const sid = localSiteIdRef.current;
+    // Read IkController's siteIdRef directly in useFrame — avoids useEffect timing
+    // issues (React runs child effects before parent effects, so reading siteIdRef
+    // in a useEffect would see -1 before IkController resolves it).
+    const sid = siteName ? localSiteIdRef.current : siteIdRef.current;
     if (!data || sid < 0 || !wrapperRef.current) return;
 
     if (!draggingRef.current) {
@@ -107,7 +99,7 @@ export function IkGizmo({ siteName, scale = 0.18, onDrag }: IkGizmoProps) {
           draggingRef.current = true;
           if (!onDrag) {
             // Default: enable IK so the robot follows
-            if (!ikEnabledRef.current) api.setIkEnabled(true);
+            if (!ikEnabledRef.current) setIkEnabled(true);
           }
           if (controls) (controls as unknown as { enabled: boolean }).enabled = false;
         }}
@@ -126,7 +118,7 @@ export function IkGizmo({ siteName, scale = 0.18, onDrag }: IkGizmoProps) {
             // Custom: consumer handles the drag
             onDrag(_pos.clone(), _quat.clone());
           } else {
-            // Default: write to provider's IK target
+            // Default: write to IK target
             const target = ikTargetRef.current;
             if (target) {
               target.position.copy(_pos);

package/src/components/SceneLights.tsx

@@ -12,120 +12,10 @@
  * Note: light_directional does NOT exist in WASM — use light_type instead.
  */
 
-import { useEffect, useRef } from 'react';
-import * as THREE from 'three';
-import { useThree } from '@react-three/fiber';
-import { useMujocoSim } from '../core/MujocoSimProvider';
+import { useSceneLights } from '../hooks/useSceneLights';
 import type { SceneLightsProps } from '../types';
 
 export function SceneLights({ intensity = 1.0 }: SceneLightsProps) {
-  const { mjModelRef, status } = useMujocoSim();
-  const { scene } = useThree();
-  const lightsRef = useRef<THREE.Light[]>([]);
-  const targetsRef = useRef<THREE.Object3D[]>([]);
-
-  useEffect(() => {
-    const model = mjModelRef.current;
-    if (!model || status !== 'ready') return;
-
-    // Clean up previous lights
-    for (const light of lightsRef.current) {
-      scene.remove(light);
-      light.dispose();
-    }
-    for (const t of targetsRef.current) scene.remove(t);
-    lightsRef.current = [];
-    targetsRef.current = [];
-
-    const nlight = model.nlight ?? 0;
-    if (nlight === 0) return;
-
-    for (let i = 0; i < nlight; i++) {
-      // Check if light is active
-      const active = model.light_active ? model.light_active[i] : 1;
-      if (!active) continue;
-
-      // light_type: 0 = directional, 1 = spot (no light_directional in WASM)
-      const lightType = model.light_type ? model.light_type[i] : 0;
-      const isDirectional = lightType === 0;
-      const castShadow = model.light_castshadow ? model.light_castshadow[i] !== 0 : false;
-
-      // Read intensity from model if available, otherwise use prop
-      const mjIntensity = model.light_intensity ? model.light_intensity[i] : 1.0;
-      const finalIntensity = intensity * mjIntensity;
-
-      // Read diffuse color
-      const dr = model.light_diffuse ? model.light_diffuse[3 * i] : 1;
-      const dg = model.light_diffuse ? model.light_diffuse[3 * i + 1] : 1;
-      const db = model.light_diffuse ? model.light_diffuse[3 * i + 2] : 1;
-      const color = new THREE.Color(dr, dg, db);
-
-      // Read position and direction
-      const px = model.light_pos[3 * i];
-      const py = model.light_pos[3 * i + 1];
-      const pz = model.light_pos[3 * i + 2];
-      const dx = model.light_dir[3 * i];
-      const dy = model.light_dir[3 * i + 1];
-      const dz = model.light_dir[3 * i + 2];
-
-      if (isDirectional) {
-        const light = new THREE.DirectionalLight(color, finalIntensity);
-        light.position.set(px, py, pz);
-        light.target.position.set(px + dx, py + dy, pz + dz);
-        light.castShadow = castShadow;
-        if (castShadow) {
-          light.shadow.mapSize.width = 1024;
-          light.shadow.mapSize.height = 1024;
-          light.shadow.camera.near = 0.1;
-          light.shadow.camera.far = 50;
-          const d = 5;
-          light.shadow.camera.left = -d;
-          light.shadow.camera.right = d;
-          light.shadow.camera.top = d;
-          light.shadow.camera.bottom = -d;
-        }
-        scene.add(light);
-        scene.add(light.target);
-        lightsRef.current.push(light);
-        targetsRef.current.push(light.target);
-      } else {
-        // Spot light
-        const cutoff = model.light_cutoff ? model.light_cutoff[i] : 45;
-        const exponent = model.light_exponent ? model.light_exponent[i] : 10;
-        const angle = (cutoff * Math.PI) / 180;
-        const light = new THREE.SpotLight(color, finalIntensity, 0, angle, exponent / 128);
-        light.position.set(px, py, pz);
-        light.target.position.set(px + dx, py + dy, pz + dz);
-        light.castShadow = castShadow;
-
-        if (model.light_attenuation) {
-          const att1 = model.light_attenuation[3 * i + 1]; // linear
-          const att2 = model.light_attenuation[3 * i + 2]; // quadratic
-          light.decay = att2 > 0 ? 2 : (att1 > 0 ? 1 : 0);
-          light.distance = att1 > 0 ? 1 / att1 : 0;
-        }
-
-        if (castShadow) {
-          light.shadow.mapSize.width = 512;
-          light.shadow.mapSize.height = 512;
-        }
-        scene.add(light);
-        scene.add(light.target);
-        lightsRef.current.push(light);
-        targetsRef.current.push(light.target);
-      }
-    }
-
-    return () => {
-      for (const light of lightsRef.current) {
-        scene.remove(light);
-        light.dispose();
-      }
-      for (const t of targetsRef.current) scene.remove(t);
-      lightsRef.current = [];
-      targetsRef.current = [];
-    };
-  }, [status, mjModelRef, scene, intensity]);
-
+  useSceneLights(intensity);
   return null;
 }

package/src/components/SceneRenderer.tsx

@@ -22,12 +22,13 @@ export function SceneRenderer() {
   const prevModelRef = useRef<MujocoModel | null>(null);
 
   const geomBuilder = useMemo(() => {
+    if (status !== 'ready') return null;
     return new GeomBuilder(mujocoRef.current);
-  }, [mujocoRef.current]);
+  }, [status, mujocoRef]);
 
   // Build body groups when model loads
   useEffect(() => {
-    if (status !== 'ready') return;
+    if (status !== 'ready' || !geomBuilder) return;
     const model = mjModelRef.current;
     const group = groupRef.current;
     if (!model || !group) return;
@@ -91,11 +92,12 @@ export function SceneRenderer() {
         while (obj && obj.userData.bodyID === undefined && obj.parent) {
           obj = obj.parent;
         }
-        if (obj && obj.userData.bodyID !== undefined && obj.userData.bodyID > 0) {
+        const bodyID = obj?.userData.bodyID;
+        if (typeof bodyID === 'number' && bodyID > 0) {
           const model = mjModelRef.current;
-          if (model && onSelectionRef.current) {
-            const name = getName(model, model.name_bodyadr[obj.userData.bodyID]);
-            onSelectionRef.current(obj.userData.bodyID, name);
+          if (model && bodyID < model.nbody && onSelectionRef.current) {
+            const name = getName(model, model.name_bodyadr[bodyID]);
+            onSelectionRef.current(bodyID, name);
           }
         }
       }}

package/src/components/SelectionHighlight.tsx

@@ -5,9 +5,7 @@
  * SelectionHighlight — highlight a selected body with emissive color (spec 6.5)
  */
 
-import { useEffect, useRef } from 'react';
-import { useThree } from '@react-three/fiber';
-import * as THREE from 'three';
+import { useSelectionHighlight } from '../hooks/useSelectionHighlight';
 import type { SelectionHighlightProps } from '../types';
 
 /**
@@ -19,51 +17,6 @@ export function SelectionHighlight({
   color = '#ff4444',
   emissiveIntensity = 0.3,
 }: SelectionHighlightProps) {
-  const { scene } = useThree();
-  const prevMeshesRef = useRef<{ mesh: THREE.Mesh; originalEmissive: THREE.Color; originalIntensity: number }[]>([]);
-
-  useEffect(() => {
-    // Restore previous highlights
-    for (const entry of prevMeshesRef.current) {
-      const mat = entry.mesh.material as THREE.MeshStandardMaterial;
-      if (mat.emissive) {
-        mat.emissive.copy(entry.originalEmissive);
-        mat.emissiveIntensity = entry.originalIntensity;
-      }
-    }
-    prevMeshesRef.current = [];
-
-    if (bodyId === null || bodyId < 0) return;
-
-    // Find all meshes belonging to this body
-    const highlightColor = new THREE.Color(color);
-    scene.traverse((obj) => {
-      if (obj.userData.bodyID === bodyId && (obj as THREE.Mesh).isMesh) {
-        const mesh = obj as THREE.Mesh;
-        const mat = mesh.material as THREE.MeshStandardMaterial;
-        if (mat.emissive) {
-          prevMeshesRef.current.push({
-            mesh,
-            originalEmissive: mat.emissive.clone(),
-            originalIntensity: mat.emissiveIntensity ?? 0,
-          });
-          mat.emissive.copy(highlightColor);
-          mat.emissiveIntensity = emissiveIntensity;
-        }
-      }
-    });
-
-    return () => {
-      for (const entry of prevMeshesRef.current) {
-        const mat = entry.mesh.material as THREE.MeshStandardMaterial;
-        if (mat.emissive) {
-          mat.emissive.copy(entry.originalEmissive);
-          mat.emissiveIntensity = entry.originalIntensity;
-        }
-      }
-      prevMeshesRef.current = [];
-    };
-  }, [bodyId, color, emissiveIntensity, scene]);
-
+  useSelectionHighlight(bodyId, { color, emissiveIntensity });
   return null;
 }

package/src/components/TendonRenderer.tsx

@@ -11,7 +11,7 @@
  * Tendons use a default color and width.
  */
 
-import { useRef } from 'react';
+import { useEffect, useRef } from 'react';
 import { useFrame } from '@react-three/fiber';
 import * as THREE from 'three';
 import { useMujocoSim } from '../core/MujocoSimProvider';
@@ -19,12 +19,18 @@ import { useMujocoSim } from '../core/MujocoSimProvider';
 const DEFAULT_TENDON_COLOR = new THREE.Color(0.3, 0.3, 0.8);
 const DEFAULT_TENDON_WIDTH = 0.002;
 
+// Preallocated temp vector to avoid per-frame allocations
+const _tmpVec = new THREE.Vector3();
+
 export function TendonRenderer() {
   const { mjModelRef, mjDataRef, status } = useMujocoSim();
   const groupRef = useRef<THREE.Group>(null);
   const meshesRef = useRef<THREE.Mesh[]>([]);
+  const curvesRef = useRef<THREE.CatmullRomCurve3[]>([]);
+  const materialRef = useRef<THREE.MeshStandardMaterial | null>(null);
 
-  useFrame(() => {
+  // Build tendon meshes once when model loads
+  useEffect(() => {
     const model = mjModelRef.current;
     const data = mjDataRef.current;
     const group = groupRef.current;
@@ -33,49 +39,107 @@ export function TendonRenderer() {
     const ntendon = model.ntendon ?? 0;
     if (ntendon === 0) return;
 
-    // Clean up old meshes
-    for (const mesh of meshesRef.current) {
-      group.remove(mesh);
-      mesh.geometry.dispose();
-      (mesh.material as THREE.Material).dispose();
+    // Shared material for all tendons
+    const material = new THREE.MeshStandardMaterial({
+      color: DEFAULT_TENDON_COLOR,
+      roughness: 0.6,
+      metalness: 0.1,
+    });
+    materialRef.current = material;
+
+    const meshes: THREE.Mesh[] = [];
+    const curves: THREE.CatmullRomCurve3[] = [];
+
+    for (let t = 0; t < ntendon; t++) {
+      const wrapNum = model.ten_wrapnum[t];
+      if (wrapNum < 2) {
+        meshes.push(null!);
+        curves.push(null!);
+        continue;
+      }
+
+      // Initial dummy points — will be overwritten in useFrame
+      const points = Array.from({ length: wrapNum }, () => new THREE.Vector3());
+      const curve = new THREE.CatmullRomCurve3(points, false);
+      const segments = Math.max(wrapNum * 2, 4);
+      const geometry = new THREE.TubeGeometry(curve, segments, DEFAULT_TENDON_WIDTH, 6, false);
+      const mesh = new THREE.Mesh(geometry, material);
+      mesh.frustumCulled = false;
+      group.add(mesh);
+      meshes.push(mesh);
+      curves.push(curve);
     }
-    meshesRef.current = [];
+
+    meshesRef.current = meshes;
+    curvesRef.current = curves;
+
+    return () => {
+      for (const mesh of meshes) {
+        if (!mesh) continue;
+        group.remove(mesh);
+        mesh.geometry.dispose();
+      }
+      material.dispose();
+      meshesRef.current = [];
+      curvesRef.current = [];
+      materialRef.current = null;
+    };
+  }, [status, mjModelRef, mjDataRef]);
+
+  // Update curve control points and rebuild geometry each frame
+  useFrame(() => {
+    const model = mjModelRef.current;
+    const data = mjDataRef.current;
+    if (!model || !data) return;
+
+    const ntendon = model.ntendon ?? 0;
+    const meshes = meshesRef.current;
+    const curves = curvesRef.current;
 
     for (let t = 0; t < ntendon; t++) {
+      const mesh = meshes[t];
+      const curve = curves[t];
+      if (!mesh || !curve) continue;
+
       const wrapAdr = model.ten_wrapadr[t];
       const wrapNum = model.ten_wrapnum[t];
-      if (wrapNum < 2) continue;
 
-      // Get wrap path points from data
-      const points: THREE.Vector3[] = [];
+      // Update existing control points in-place
+      let validCount = 0;
       for (let w = 0; w < wrapNum; w++) {
         const idx = (wrapAdr + w) * 3;
         if (data.wrap_xpos && idx + 2 < data.wrap_xpos.length) {
           const x = data.wrap_xpos[idx];
           const y = data.wrap_xpos[idx + 1];
           const z = data.wrap_xpos[idx + 2];
-          // Skip zero points (uninitialized wrap points)
           if (x !== 0 || y !== 0 || z !== 0) {
-            points.push(new THREE.Vector3(x, y, z));
+            if (validCount < curve.points.length) {
+              curve.points[validCount].set(x, y, z);
+            }
+            validCount++;
           }
         }
       }
 
-      if (points.length < 2) continue;
+      if (validCount < 2) {
+        mesh.visible = false;
+        continue;
+      }
 
-      const curve = new THREE.CatmullRomCurve3(points, false);
-      const geometry = new THREE.TubeGeometry(
-        curve, Math.max(points.length * 2, 4), DEFAULT_TENDON_WIDTH, 6, false
-      );
+      // Trim or pad points array to match valid count
+      if (curve.points.length !== validCount) {
+        curve.points.length = validCount;
+        while (curve.points.length < validCount) {
+          curve.points.push(new THREE.Vector3());
+        }
+      }
 
-      const material = new THREE.MeshStandardMaterial({
-        color: DEFAULT_TENDON_COLOR,
-        roughness: 0.6,
-        metalness: 0.1,
-      });
-      const mesh = new THREE.Mesh(geometry, material);
-      group.add(mesh);
-      meshesRef.current.push(mesh);
+      // Rebuild geometry from updated curve
+      mesh.geometry.dispose();
+      mesh.geometry = new THREE.TubeGeometry(
+        curve, Math.max(validCount * 2, 4), DEFAULT_TENDON_WIDTH, 6, false
+      );
+      mesh.visible = true;
     }
   });
 

package/src/components/TrajectoryPlayer.tsx

@@ -5,7 +5,8 @@
  * TrajectoryPlayer — component form of trajectory playback (spec 13.2)
  */
 
-import { useEffect } from 'react';
+import { useEffect, useRef } from 'react';
+import { useFrame } from '@react-three/fiber';
 import { useTrajectoryPlayer } from '../hooks/useTrajectoryPlayer';
 import type { TrajectoryPlayerProps } from '../types';
 
@@ -21,6 +22,9 @@ export function TrajectoryPlayer({
   onFrame,
 }: TrajectoryPlayerProps) {
   const player = useTrajectoryPlayer(trajectory, { fps, loop });
+  const onFrameRef = useRef(onFrame);
+  onFrameRef.current = onFrame;
+  const lastReportedFrameRef = useRef(-1);
 
   useEffect(() => {
     if (playing) {
@@ -28,17 +32,17 @@ export function TrajectoryPlayer({
     } else {
       player.pause();
     }
-  }, [playing]);
+  }, [playing, player]);
 
-  useEffect(() => {
-    if (onFrame) {
-      // Poll frame changes (lightweight, no extra useFrame needed)
-      const interval = setInterval(() => {
-        if (player.playing) onFrame(player.frame);
-      }, 1000 / fps);
-      return () => clearInterval(interval);
+  // Use useFrame instead of setInterval to sync with the render loop
+  useFrame(() => {
+    if (!onFrameRef.current) return;
+    const currentFrame = player.frame;
+    if (currentFrame !== lastReportedFrameRef.current && player.playing) {
+      lastReportedFrameRef.current = currentFrame;
+      onFrameRef.current(currentFrame);
     }
-  }, [onFrame, fps]);
+  });
 
   return null;
 }

package/src/core/IkContext.tsx

@@ -0,0 +1,40 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * IkContext — React context for the IK controller plugin.
+ */
+
+import { createContext, useContext } from 'react';
+import * as THREE from 'three';
+
+export interface IkContextValue {
+  ikEnabledRef: React.RefObject<boolean>;
+  ikCalculatingRef: React.RefObject<boolean>;
+  ikTargetRef: React.RefObject<THREE.Group>;
+  siteIdRef: React.RefObject<number>;
+  setIkEnabled(enabled: boolean): void;
+  moveTarget(pos: THREE.Vector3, duration?: number): void;
+  syncTargetToSite(): void;
+  solveIK(pos: THREE.Vector3, quat: THREE.Quaternion, currentQ: number[]): number[] | null;
+  getGizmoStats(): { pos: THREE.Vector3; rot: THREE.Euler } | null;
+}
+
+export const IkContext = createContext<IkContextValue | null>(null);
+
+/**
+ * Access the IK controller context.
+ *
+ * - `useIk()` — throws if no `<IkController>` ancestor (use inside `<IkController>`)
+ * - `useIk({ optional: true })` — returns `null` if no ancestor (use in components
+ *   that optionally interact with IK, e.g. keyboard controllers that disable IK)
+ */
+export function useIk(): IkContextValue;
+export function useIk(options: { optional: true }): IkContextValue | null;
+export function useIk(options?: { optional?: boolean }): IkContextValue | null {
+  const ctx = useContext(IkContext);
+  if (!ctx && !options?.optional) {
+    throw new Error('useIk() must be used inside an <IkController>');
+  }
+  return ctx;
+}

package/src/core/MujocoProvider.tsx

@@ -22,6 +22,8 @@ export function useMujoco(): MujocoContextValue {
 
 interface MujocoProviderProps {
   wasmUrl?: string;
+  /** Timeout in ms for WASM module load. Default: 30000. */
+  timeout?: number;
   children: React.ReactNode;
   onError?: (error: Error) => void;
 }
@@ -30,7 +32,7 @@ interface MujocoProviderProps {
  * MujocoProvider — WASM / module lifecycle.
  * Loads the MuJoCo WASM module on mount and provides it to children via context.
  */
-export function MujocoProvider({ wasmUrl, children, onError }: MujocoProviderProps) {
+export function MujocoProvider({ wasmUrl, timeout = 30000, children, onError }: MujocoProviderProps) {
   const [status, setStatus] = useState<'loading' | 'ready' | 'error'>('loading');
   const [error, setError] = useState<string | null>(null);
   const moduleRef = useRef<MujocoModule | null>(null);
@@ -39,7 +41,7 @@ export function MujocoProvider({ wasmUrl, children, onError }: MujocoProviderPro
   useEffect(() => {
     isMounted.current = true;
 
-    loadMujoco({
+    const wasmPromise = loadMujoco({
       ...(wasmUrl ? { locateFile: (path: string) => path.endsWith('.wasm') ? wasmUrl : path } : {}),
       printErr: (text: string) => {
         if (text.includes('Aborted') && isMounted.current) {
@@ -47,7 +49,13 @@ export function MujocoProvider({ wasmUrl, children, onError }: MujocoProviderPro
           setStatus('error');
         }
       },
-    })
+    });
+
+    const timeoutPromise = new Promise<never>((_, reject) =>
+      setTimeout(() => reject(new Error(`WASM module load timed out after ${timeout}ms`)), timeout)
+    );
+
+    Promise.race([wasmPromise, timeoutPromise])
       .then((inst: unknown) => {
         if (isMounted.current) {
           moduleRef.current = inst as MujocoModule;
@@ -66,7 +74,7 @@ export function MujocoProvider({ wasmUrl, children, onError }: MujocoProviderPro
     return () => {
       isMounted.current = false;
     };
-  }, [wasmUrl]);
+  }, [wasmUrl, timeout]);
 
   return (
     <MujocoContext.Provider