mujoco-react 0.1.0 → 0.3.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mujoco-react",
-  "version": "0.1.0",
+  "version": "0.3.0",
   "description": "Composable React Three Fiber building blocks for MuJoCo WASM simulations",
   "type": "module",
   "main": "./dist/index.js",

package/src/components/ContactMarkers.tsx

@@ -12,15 +12,16 @@ import { useRef } from 'react';
 import { useFrame } from '@react-three/fiber';
 import * as THREE from 'three';
 import { useMujocoSim } from '../core/MujocoSimProvider';
+import { getContact } from '../types';
 
 const _dummy = new THREE.Object3D();
 
 interface ContactMarkersProps {
   /** Maximum contacts to render. Default: 100. */
   maxContacts?: number;
-  /** Sphere radius. Default: 0.005. */
+  /** Sphere radius. Default: 0.008. */
   radius?: number;
-  /** Color. Default: '#4f46e5'. */
+  /** Color. Default: '#22d3ee'. */
   color?: string;
   /** Show markers. Default: true. */
   visible?: boolean;
@@ -28,8 +29,8 @@ interface ContactMarkersProps {
 
 export function ContactMarkers({
   maxContacts = 100,
-  radius = 0.005,
-  color = '#4f46e5',
+  radius = 0.008,
+  color = '#22d3ee',
   visible = true,
 }: ContactMarkersProps = {}) {
   const { mjDataRef, status } = useMujocoSim();
@@ -47,18 +48,15 @@ export function ContactMarkers({
     const count = Math.min(ncon, maxContacts);
 
     for (let i = 0; i < count; i++) {
-      try {
-        const c = (data.contact as { get(i: number): { pos: Float64Array } | undefined }).get(i);
-        if (!c) break;
-        _dummy.position.set(c.pos[0], c.pos[1], c.pos[2]);
-        _dummy.updateMatrix();
-        mesh.setMatrixAt(i, _dummy.matrix);
-      } catch {
-        // Contact access failed — stop here
+      const c = getContact(data, i);
+      if (!c) {
         mesh.count = i;
         mesh.instanceMatrix.needsUpdate = true;
         return;
       }
+      _dummy.position.set(c.pos[0], c.pos[1], c.pos[2]);
+      _dummy.updateMatrix();
+      mesh.setMatrixAt(i, _dummy.matrix);
     }
 
     mesh.count = count;
@@ -68,14 +66,9 @@ export function ContactMarkers({
   if (status !== 'ready') return null;
 
   return (
-    <instancedMesh ref={meshRef} args={[undefined, undefined, maxContacts]}>
+    <instancedMesh ref={meshRef} args={[undefined, undefined, maxContacts]} frustumCulled={false} renderOrder={999}>
       <sphereGeometry args={[radius, 8, 8]} />
-      <meshStandardMaterial
-        color={color}
-        emissive={color}
-        emissiveIntensity={0.3}
-        roughness={0.5}
-      />
+      <meshBasicMaterial color={color} depthTest={false} />
     </instancedMesh>
   );
 }

package/src/components/Debug.tsx

@@ -10,6 +10,7 @@ import { useFrame, useThree } from '@react-three/fiber';
 import * as THREE from 'three';
 import { useMujocoSim } from '../core/MujocoSimProvider';
 import { getName } from '../core/SceneLoader';
+import { getContact } from '../types';
 import type { DebugProps } from '../types';
 
 const JOINT_COLORS: Record<number, number> = {
@@ -19,6 +20,14 @@ const JOINT_COLORS: Record<number, number> = {
   3: 0xffff00, // hinge - yellow
 };
 
+// Preallocated temps to avoid per-frame GC pressure
+const _v3a = new THREE.Vector3();
+const _v3b = new THREE.Vector3();
+const _quat = new THREE.Quaternion();
+const _contactPos = new THREE.Vector3();
+const _contactNormal = new THREE.Vector3();
+const MAX_CONTACT_ARROWS = 50;
+
 /**
  * Declarative debug visualization component.
  * Renders wireframe geoms, site markers, joint axes, contact forces, COM markers, etc.
@@ -78,27 +87,92 @@ export function Debug({
       }
     }
 
-    // Site markers
+    // Site markers — scale based on site_size if available, else use geom_size of parent body
     if (showSites) {
+      const siteSize = (model as Record<string, unknown>).site_size as Float64Array | undefined;
       for (let i = 0; i < model.nsite; i++) {
-        const geometry = new THREE.OctahedronGeometry(0.01);
-        const mat = new THREE.MeshBasicMaterial({ color: 0xff00ff, transparent: true, opacity: 0.7 });
+        // Determine marker radius: use site_size[3*i] if available, else estimate from parent body's geoms
+        let radius = 0.008;
+        if (siteSize) {
+          radius = Math.max(siteSize[3 * i] * 0.5, 0.004);
+        } else {
+          // Estimate from parent body's geom sizes
+          const bodyId = model.site_bodyid[i];
+          let maxGeomSize = 0;
+          for (let g = 0; g < model.ngeom; g++) {
+            if (model.geom_bodyid[g] === bodyId) {
+              maxGeomSize = Math.max(maxGeomSize, model.geom_size[3 * g]);
+            }
+          }
+          if (maxGeomSize > 0) radius = maxGeomSize * 0.15;
+        }
+
+        const geometry = new THREE.OctahedronGeometry(radius);
+        const mat = new THREE.MeshBasicMaterial({ color: 0xff00ff, depthTest: false });
         const mesh = new THREE.Mesh(geometry, mat);
+        mesh.renderOrder = 999;
+        mesh.frustumCulled = false;
         mesh.userData.siteId = i;
+
+        // Label
+        const canvas = document.createElement('canvas');
+        canvas.width = 256;
+        canvas.height = 64;
+        const ctx = canvas.getContext('2d')!;
+        ctx.fillStyle = '#ff00ff';
+        ctx.font = 'bold 36px monospace';
+        ctx.textAlign = 'center';
+        ctx.fillText(getName(model, model.name_siteadr[i]), 128, 42);
+        const tex = new THREE.CanvasTexture(canvas);
+        const spriteMat = new THREE.SpriteMaterial({ map: tex, depthTest: false, transparent: true });
+        const sprite = new THREE.Sprite(spriteMat);
+        const labelScale = radius * 15;
+        sprite.scale.set(labelScale, labelScale * 0.25, 1);
+        sprite.position.y = radius * 2;
+        sprite.renderOrder = 999;
+        mesh.add(sprite);
+
         sites.push(mesh);
       }
     }
 
-    // Joint axes
+    // Joint axes — scale arrow length based on parent body's geom sizes
     if (showJoints) {
+      // Safely check for jnt_pos/jnt_axis on the WASM model
+      const jntPos = (model as Record<string, unknown>).jnt_pos as Float64Array | undefined;
+      const jntAxis = (model as Record<string, unknown>).jnt_axis as Float64Array | undefined;
+
       for (let i = 0; i < model.njnt; i++) {
         const type = model.jnt_type[i];
         const color = JOINT_COLORS[type] ?? 0xffffff;
+
+        // Scale based on parent body geom size
+        const bodyId = model.jnt_bodyid[i];
+        let maxGeomSize = 0;
+        for (let g = 0; g < model.ngeom; g++) {
+          if (model.geom_bodyid[g] === bodyId) {
+            maxGeomSize = Math.max(maxGeomSize, model.geom_size[3 * g]);
+          }
+        }
+        const arrowLen = Math.max(maxGeomSize * 0.8, 0.05);
+
         const arrow = new THREE.ArrowHelper(
           new THREE.Vector3(0, 0, 1), new THREE.Vector3(),
-          0.05, color, 0.01, 0.005
+          arrowLen, color, arrowLen * 0.25, arrowLen * 0.12
         );
+        // Render on top so arrows show through geometry
+        arrow.renderOrder = 999;
+        arrow.frustumCulled = false;
+        arrow.line.material = new THREE.LineBasicMaterial({ color, depthTest: false });
+        (arrow.cone.material as THREE.MeshBasicMaterial).depthTest = false;
+        arrow.line.renderOrder = 999;
+        arrow.line.frustumCulled = false;
+        arrow.cone.renderOrder = 999;
+        arrow.cone.frustumCulled = false;
         arrow.userData.jointId = i;
+        arrow.userData.bodyId = bodyId;
+        arrow.userData.hasJntPos = !!jntPos;
+        arrow.userData.hasJntAxis = !!jntAxis;
         joints.push(arrow);
       }
     }
@@ -144,6 +218,10 @@ export function Debug({
     const data = mjDataRef.current;
     if (!model || !data || !debugGeometry) return;
 
+    // Safely grab optional arrays once
+    const jntPos = (model as Record<string, unknown>).jnt_pos as Float64Array | undefined;
+    const jntAxis = (model as Record<string, unknown>).jnt_axis as Float64Array | undefined;
+
     // Update geom wireframes
     for (const mesh of debugGeometry.geoms) {
       const bid = mesh.userData.bodyId;
@@ -157,8 +235,9 @@ export function Debug({
       // Apply local geom offset
       const gid = mesh.userData.geomId;
       const gp = model.geom_pos;
-      mesh.position.add(new THREE.Vector3(gp[3 * gid], gp[3 * gid + 1], gp[3 * gid + 2])
-        .applyQuaternion(mesh.quaternion));
+      _v3a.set(gp[3 * gid], gp[3 * gid + 1], gp[3 * gid + 2])
+        .applyQuaternion(mesh.quaternion);
+      mesh.position.add(_v3a);
     }
 
     // Update site markers
@@ -171,6 +250,35 @@ export function Debug({
       );
     }
 
+    // Update joint axes
+    for (const obj of debugGeometry.joints) {
+      const arrow = obj as THREE.ArrowHelper;
+      const jid = arrow.userData.jointId;
+      const bid = arrow.userData.bodyId;
+      const i3 = bid * 3;
+      const i4 = bid * 4;
+
+      _quat.set(
+        data.xquat[i4 + 1], data.xquat[i4 + 2],
+        data.xquat[i4 + 3], data.xquat[i4]
+      );
+
+      // Position: body origin + local joint anchor (if available)
+      arrow.position.set(data.xpos[i3], data.xpos[i3 + 1], data.xpos[i3 + 2]);
+      if (jntPos) {
+        _v3a.set(jntPos[3 * jid], jntPos[3 * jid + 1], jntPos[3 * jid + 2])
+          .applyQuaternion(_quat);
+        arrow.position.add(_v3a);
+      }
+
+      // Orient along joint axis in world frame (if available)
+      if (jntAxis) {
+        _v3a.set(jntAxis[3 * jid], jntAxis[3 * jid + 1], jntAxis[3 * jid + 2])
+          .applyQuaternion(_quat).normalize();
+        arrow.setDirection(_v3a);
+      }
+    }
+
     // Update COM markers
     for (const mesh of debugGeometry.comMarkers) {
       const bid = mesh.userData.bodyId;
@@ -179,41 +287,68 @@ export function Debug({
     }
   });
 
-  // Contact force vectors
+  // Contact force vectors — pre-created pool to avoid per-frame allocation
   const contactGroupRef = useRef<THREE.Group>(null);
-  const contactArrowsRef = useRef<THREE.ArrowHelper[]>([]);
+  const contactPoolRef = useRef<THREE.ArrowHelper[]>([]);
+  const contactPoolInitRef = useRef(false);
 
-  useFrame(() => {
-    if (!showContacts) return;
-    const model = mjModelRef.current;
-    const data = mjDataRef.current;
+  // Initialize arrow pool once
+  useEffect(() => {
     const group = contactGroupRef.current;
-    if (!model || !data || !group) return;
+    if (!group || contactPoolInitRef.current) return;
+    contactPoolInitRef.current = true;
 
-    // Remove old arrows
-    for (const arrow of contactArrowsRef.current) {
-      group.remove(arrow);
-      arrow.dispose();
+    const pool: THREE.ArrowHelper[] = [];
+    for (let i = 0; i < MAX_CONTACT_ARROWS; i++) {
+      const arrow = new THREE.ArrowHelper(
+        new THREE.Vector3(0, 1, 0), new THREE.Vector3(), 0.1, 0xff4444, 0.03, 0.015
+      );
+      arrow.visible = false;
+      group.add(arrow);
+      pool.push(arrow);
     }
-    contactArrowsRef.current = [];
+    contactPoolRef.current = pool;
+
+    return () => {
+      for (const arrow of pool) {
+        group.remove(arrow);
+        arrow.dispose();
+      }
+      contactPoolRef.current = [];
+      contactPoolInitRef.current = false;
+    };
+  }, [showContacts]);
+
+  useFrame(() => {
+    if (!showContacts) return;
+    const data = mjDataRef.current;
+    const pool = contactPoolRef.current;
+    if (!data || pool.length === 0) return;
 
     const ncon = data.ncon;
-    for (let i = 0; i < Math.min(ncon, 50); i++) {
-      try {
-        const c = (data.contact as { get(i: number): { pos: Float64Array; frame: Float64Array; dist: number } }).get(i);
-        const pos = new THREE.Vector3(c.pos[0], c.pos[1], c.pos[2]);
-        const normal = new THREE.Vector3(c.frame[0], c.frame[1], c.frame[2]);
-        const force = Math.abs(c.dist) * 100;
-        const length = Math.min(force * 0.01, 0.1);
-        if (length > 0.001) {
-          const arrow = new THREE.ArrowHelper(normal, pos, length, 0xff4444, length * 0.3, length * 0.15);
-          group.add(arrow);
-          contactArrowsRef.current.push(arrow);
-        }
-      } catch {
-        break;
+    let arrowIdx = 0;
+
+    for (let i = 0; i < Math.min(ncon, MAX_CONTACT_ARROWS); i++) {
+      const c = getContact(data, i);
+      if (!c) break;
+      _contactPos.set(c.pos[0], c.pos[1], c.pos[2]);
+      _contactNormal.set(c.frame[0], c.frame[1], c.frame[2]);
+      const force = Math.abs(c.dist) * 100;
+      const length = Math.min(force * 0.01, 0.1);
+      if (length > 0.001 && arrowIdx < pool.length) {
+        const arrow = pool[arrowIdx];
+        arrow.position.copy(_contactPos);
+        arrow.setDirection(_contactNormal);
+        arrow.setLength(length, length * 0.3, length * 0.15);
+        arrow.visible = true;
+        arrowIdx++;
       }
     }
+
+    // Hide unused arrows
+    for (let i = arrowIdx; i < pool.length; i++) {
+      pool[i].visible = false;
+    }
   });
 
   if (status !== 'ready') return null;

package/src/components/DragInteraction.tsx

@@ -208,7 +208,7 @@ export function DragInteraction({
 
     if (draggingRef.current && bodyIdRef.current > 0) {
       arrow.visible = true;
-      const dir = mouseWorldRef.current.clone().sub(grabWorldRef.current);
+      const dir = _bodyPos.copy(mouseWorldRef.current).sub(grabWorldRef.current);
       const len = dir.length();
       if (len > 0.001) {
         dir.normalize();

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,
+);