mujoco-react 0.2.0 → 1.0.0

package/src/components/IkController.tsx

@@ -0,0 +1,262 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * IkController — composable IK controller plugin.
+ * Extracts all IK logic from MujocoSimProvider into an opt-in component.
+ */
+
+import { useCallback, useEffect, useMemo, useRef } from 'react';
+import { useFrame } from '@react-three/fiber';
+import * as THREE from 'three';
+import { createController } from '../core/createController';
+import { IkContext, type IkContextValue } from '../core/IkContext';
+import { useMujocoSim, useBeforePhysicsStep } from '../core/MujocoSimProvider';
+import { GenericIK } from '../core/GenericIK';
+import { findSiteByName } from '../core/SceneLoader';
+import type { IkConfig, IKSolveFn, MujocoData } from '../types';
+
+// Preallocated temp for syncGizmoToSite
+const _syncMat4 = new THREE.Matrix4();
+
+function syncGizmoToSite(data: MujocoData, siteId: number, target: THREE.Group) {
+  if (siteId === -1) return;
+  const sitePos = data.site_xpos.subarray(siteId * 3, siteId * 3 + 3);
+  const siteMat = data.site_xmat.subarray(siteId * 9, siteId * 9 + 9);
+  target.position.set(sitePos[0], sitePos[1], sitePos[2]);
+  _syncMat4.set(
+    siteMat[0], siteMat[1], siteMat[2], 0,
+    siteMat[3], siteMat[4], siteMat[5], 0,
+    siteMat[6], siteMat[7], siteMat[8], 0,
+    0, 0, 0, 1,
+  );
+  target.quaternion.setFromRotationMatrix(_syncMat4);
+}
+
+function IkControllerImpl({
+  config,
+  children,
+}: {
+  config: IkConfig;
+  children?: React.ReactNode;
+}) {
+  const { mjModelRef, mjDataRef, mujocoRef, configRef, resetCallbacks, status } =
+    useMujocoSim();
+
+  // All IK state lives here, NOT in the provider
+  const ikEnabledRef = useRef(false);
+  const ikCalculatingRef = useRef(false);
+  const ikTargetRef = useRef<THREE.Group>(new THREE.Group());
+  const siteIdRef = useRef(-1);
+  const genericIkRef = useRef<GenericIK>(new GenericIK(mujocoRef.current));
+  const firstIkEnableRef = useRef(true);
+
+  const needsInitialSync = useRef(true);
+
+  const gizmoAnimRef = useRef({
+    active: false,
+    startPos: new THREE.Vector3(),
+    endPos: new THREE.Vector3(),
+    startRot: new THREE.Quaternion(),
+    endRot: new THREE.Quaternion(),
+    startTime: 0,
+    duration: 1000,
+  });
+
+  // Resolve site ID when model loads or config changes
+  useEffect(() => {
+    const model = mjModelRef.current;
+    if (!model || status !== 'ready') {
+      siteIdRef.current = -1;
+      return;
+    }
+    siteIdRef.current = findSiteByName(model, config.siteName);
+    const data = mjDataRef.current;
+    if (data && ikTargetRef.current) {
+      syncGizmoToSite(data, siteIdRef.current, ikTargetRef.current);
+    }
+  }, [config.siteName, status, mjModelRef, mjDataRef]);
+
+  // IK solve function — use custom solver if provided, otherwise built-in GenericIK
+  const ikSolveFn = useCallback(
+    (pos: THREE.Vector3, quat: THREE.Quaternion, currentQ: number[]): number[] | null => {
+      if (config.ikSolveFn) return config.ikSolveFn(pos, quat, currentQ);
+      const model = mjModelRef.current;
+      const data = mjDataRef.current;
+      if (!model || !data || siteIdRef.current === -1) return null;
+      return genericIkRef.current.solve(
+        model,
+        data,
+        siteIdRef.current,
+        config.numJoints,
+        pos,
+        quat,
+        currentQ,
+        {
+          damping: config.damping,
+          maxIterations: config.maxIterations,
+        },
+      );
+    },
+    [config.ikSolveFn, config.numJoints, config.damping, config.maxIterations, mjModelRef, mjDataRef],
+  );
+  const ikSolveFnRef = useRef<IKSolveFn>(ikSolveFn);
+  ikSolveFnRef.current = ikSolveFn;
+
+  // Gizmo animation + one-time initial sync in useFrame
+  useFrame(() => {
+    // Ensure the gizmo is positioned at the site after the first physics step
+    if (needsInitialSync.current && siteIdRef.current !== -1) {
+      const data = mjDataRef.current;
+      if (data && ikTargetRef.current) {
+        syncGizmoToSite(data, siteIdRef.current, ikTargetRef.current);
+        needsInitialSync.current = false;
+      }
+    }
+
+    const ga = gizmoAnimRef.current;
+    const target = ikTargetRef.current;
+    if (!ga.active || !target) return;
+
+    const now = performance.now();
+    const elapsed = now - ga.startTime;
+    const t = Math.min(elapsed / ga.duration, 1.0);
+    const ease = 1 - Math.pow(1 - t, 3);
+    target.position.lerpVectors(ga.startPos, ga.endPos, ease);
+    target.quaternion.slerpQuaternions(ga.startRot, ga.endRot, ease);
+    if (t >= 1.0) ga.active = false;
+  });
+
+  // IK solve in physics loop
+  useBeforePhysicsStep((model, data) => {
+    if (!ikEnabledRef.current) {
+      ikCalculatingRef.current = false;
+      return;
+    }
+    const target = ikTargetRef.current;
+    if (!target) return;
+
+    ikCalculatingRef.current = true;
+    const numJoints = config.numJoints;
+    const currentQ: number[] = [];
+    for (let i = 0; i < numJoints; i++) currentQ.push(data.qpos[i]);
+    const solution = ikSolveFnRef.current(target.position, target.quaternion, currentQ);
+    if (solution) {
+      for (let i = 0; i < numJoints; i++) data.ctrl[i] = solution[i];
+    }
+  });
+
+  // Reset callback — sync gizmo and reset IK state
+  useEffect(() => {
+    const cb = () => {
+      const data = mjDataRef.current;
+      if (data && ikTargetRef.current) {
+        syncGizmoToSite(data, siteIdRef.current, ikTargetRef.current);
+      }
+      gizmoAnimRef.current.active = false;
+      firstIkEnableRef.current = true;
+      ikEnabledRef.current = false;
+      needsInitialSync.current = true;
+    };
+    resetCallbacks.current.add(cb);
+    return () => {
+      resetCallbacks.current.delete(cb);
+    };
+  }, [resetCallbacks, mjDataRef]);
+
+  // --- API methods ---
+
+  const setIkEnabled = useCallback(
+    (enabled: boolean) => {
+      ikEnabledRef.current = enabled;
+      const data = mjDataRef.current;
+      if (enabled && data && !gizmoAnimRef.current.active && ikTargetRef.current) {
+        syncGizmoToSite(data, siteIdRef.current, ikTargetRef.current);
+        firstIkEnableRef.current = false;
+      }
+    },
+    [mjDataRef],
+  );
+
+  const syncTargetToSiteApi = useCallback(() => {
+    const data = mjDataRef.current;
+    const target = ikTargetRef.current;
+    if (data && target) syncGizmoToSite(data, siteIdRef.current, target);
+  }, [mjDataRef]);
+
+  const solveIK = useCallback(
+    (pos: THREE.Vector3, quat: THREE.Quaternion, currentQ: number[]): number[] | null => {
+      return ikSolveFnRef.current(pos, quat, currentQ);
+    },
+    [],
+  );
+
+  const moveTarget = useCallback(
+    (pos: THREE.Vector3, duration = 0) => {
+      if (!ikEnabledRef.current) setIkEnabled(true);
+      const target = ikTargetRef.current;
+      if (!target) return;
+
+      const targetPos = pos.clone();
+      const targetRot = new THREE.Quaternion().setFromEuler(
+        new THREE.Euler(Math.PI, 0, 0),
+      );
+
+      if (duration > 0) {
+        const ga = gizmoAnimRef.current;
+        ga.active = true;
+        ga.startPos.copy(target.position);
+        ga.endPos.copy(targetPos);
+        ga.startRot.copy(target.quaternion);
+        ga.endRot.copy(targetRot);
+        ga.startTime = performance.now();
+        ga.duration = duration;
+      } else {
+        gizmoAnimRef.current.active = false;
+        target.position.copy(targetPos);
+        target.quaternion.copy(targetRot);
+      }
+    },
+    [setIkEnabled],
+  );
+
+  const getGizmoStats = useCallback(
+    (): { pos: THREE.Vector3; rot: THREE.Euler } | null => {
+      const target = ikTargetRef.current;
+      if (!ikCalculatingRef.current || !target) return null;
+      return {
+        pos: target.position.clone(),
+        rot: new THREE.Euler().setFromQuaternion(target.quaternion),
+      };
+    },
+    [],
+  );
+
+  const contextValue = useMemo<IkContextValue>(
+    () => ({
+      ikEnabledRef,
+      ikCalculatingRef,
+      ikTargetRef,
+      siteIdRef,
+      setIkEnabled,
+      moveTarget,
+      syncTargetToSite: syncTargetToSiteApi,
+      solveIK,
+      getGizmoStats,
+    }),
+    [setIkEnabled, moveTarget, syncTargetToSiteApi, solveIK, getGizmoStats],
+  );
+
+  return <IkContext.Provider value={contextValue}>{children}</IkContext.Provider>;
+}
+
+export const IkController = createController<IkConfig>(
+  {
+    name: 'IkController',
+    defaultConfig: {
+      damping: 0.01,
+      maxIterations: 50,
+    },
+  },
+  IkControllerImpl,
+);

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

@@ -4,6 +4,7 @@
  */
 
 import { useFrame } from '@react-three/fiber';
+import type { ThreeElements } from '@react-three/fiber';
 import { useEffect, useMemo, useRef } from 'react';
 import * as THREE from 'three';
 import { GeomBuilder } from '../rendering/GeomBuilder';
@@ -13,21 +14,22 @@ import { useMujocoSim } from '../core/MujocoSimProvider';
 
 /**
  * SceneRenderer — creates and syncs MuJoCo body meshes every frame.
- * Replaces RenderSystem.initScene() + RenderSystem.update() body loop.
+ * Accepts standard R3F group props (position, rotation, scale, visible, etc.).
  */
-export function SceneRenderer() {
+export function SceneRenderer(props: Omit<ThreeElements['group'], 'ref'>) {
   const { mjModelRef, mjDataRef, mujocoRef, onSelectionRef, status } = useMujocoSim();
   const groupRef = useRef<THREE.Group>(null);
   const bodyRefs = useRef<(THREE.Group | null)[]>([]);
   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;
@@ -84,18 +86,21 @@ export function SceneRenderer() {
 
   return (
     <group
+      {...props}
       ref={groupRef}
       onDoubleClick={(e) => {
+        if (typeof props.onDoubleClick === 'function') props.onDoubleClick(e);
         e.stopPropagation();
         let obj: THREE.Object3D | null = e.object;
         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;
 }