mujoco-react 6.0.1 → 7.0.1

package/src/hooks/useIkController.ts

@@ -0,0 +1,242 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { useCallback, useEffect, useMemo, useRef } from 'react';
+import { useFrame } from '@react-three/fiber';
+import * as THREE from 'three';
+import { createControllerHook } from '../core/createController';
+import { useMujocoContext, useBeforePhysicsStep } from '../core/MujocoSimProvider';
+import { GenericIK } from '../core/GenericIK';
+import { findSiteByName } from '../core/SceneLoader';
+import type { IkConfig, IkContextValue, 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);
+}
+
+export const useIkController = createControllerHook<IkConfig, IkContextValue>(
+  { name: 'useIkController', defaultConfig: { damping: 0.01, maxIterations: 50 } },
+  function useIkControllerImpl(config) {
+    const { mjModelRef, mjDataRef, mujocoRef, resetCallbacks, status } =
+      useMujocoContext();
+
+    // All IK state lives here
+    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(() => {
+      if (!config) {
+        siteIdRef.current = -1;
+        return;
+      }
+      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, config]);
+
+    // IK solve function
+    const ikSolveFn = useCallback(
+      (pos: THREE.Vector3, quat: THREE.Quaternion, currentQ: number[]): number[] | null => {
+        if (!config) return 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, mjModelRef, mjDataRef],
+    );
+    const ikSolveFnRef = useRef<IKSolveFn>(ikSolveFn);
+    ikSolveFnRef.current = ikSolveFn;
+
+    // Gizmo animation + one-time initial sync
+    useFrame(() => {
+      if (!config) return;
+
+      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 (!config || !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
+    useEffect(() => {
+      if (!config) return;
+      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, config]);
+
+    // --- 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],
+    );
+
+    if (!config) return null;
+
+    return contextValue;
+  },
+);

package/src/hooks/useJointState.ts

@@ -6,7 +6,7 @@
  */
 
 import { useEffect, useRef } from 'react';
-import { useMujoco, useAfterPhysicsStep } from '../core/MujocoSimProvider';
+import { useMujocoContext, useAfterPhysicsStep } from '../core/MujocoSimProvider';
 import { getName } from '../core/SceneLoader';
 import type { JointStateResult } from '../types';
 
@@ -19,7 +19,7 @@ import type { JointStateResult } from '../types';
  * For free joints, position is pos+quat (7), velocity is lin+ang vel (6).
  */
 export function useJointState(name: string): JointStateResult {
-  const { mjModelRef, mjDataRef, status } = useMujoco();
+  const { mjModelRef, mjDataRef, status } = useMujocoContext();
   const jointIdRef = useRef(-1);
   const qposAdrRef = useRef(0);
   const dofAdrRef = useRef(0);

package/src/hooks/useKeyboardTeleop.ts

@@ -6,7 +6,7 @@
  */
 
 import { useEffect, useRef } from 'react';
-import { useMujoco, useBeforePhysicsStep } from '../core/MujocoSimProvider';
+import { useMujocoContext, useBeforePhysicsStep } from '../core/MujocoSimProvider';
 import { findActuatorByName } from '../core/SceneLoader';
 import type { KeyboardTeleopConfig } from '../types';
 
@@ -19,7 +19,7 @@ import type { KeyboardTeleopConfig } from '../types';
  * - `set`: Set actuator to a fixed value while key is held
  */
 export function useKeyboardTeleop(config: KeyboardTeleopConfig) {
-  const { mjModelRef, mjDataRef, status } = useMujoco();
+  const { mjModelRef, mjDataRef, status } = useMujocoContext();
   const pressedRef = useRef(new Set<string>());
   const toggleStateRef = useRef(new Map<string, boolean>());
   const enabledRef = useRef(config.enabled ?? true);

package/src/hooks/usePolicy.ts

@@ -6,7 +6,7 @@
  */
 
 import { useRef } from 'react';
-import { useMujoco, useBeforePhysicsStep } from '../core/MujocoSimProvider';
+import { useMujocoContext, useBeforePhysicsStep } from '../core/MujocoSimProvider';
 import type { PolicyConfig } from '../types';
 
 /**
@@ -20,7 +20,7 @@ import type { PolicyConfig } from '../types';
  * @returns { step, isRunning } control handles
  */
 export function usePolicy(config: PolicyConfig) {
-  const { mjModelRef } = useMujoco();
+  const { mjModelRef } = useMujocoContext();
   const lastActionTimeRef = useRef(0);
   const lastActionRef = useRef<Float32Array | Float64Array | number[] | null>(null);
   const isRunningRef = useRef(true);

package/src/hooks/useSceneLights.ts

@@ -10,10 +10,10 @@
 import { useEffect, useRef } from 'react';
 import * as THREE from 'three';
 import { useThree } from '@react-three/fiber';
-import { useMujoco } from '../core/MujocoSimProvider';
+import { useMujocoContext } from '../core/MujocoSimProvider';
 
 export function useSceneLights(intensity = 1.0) {
-  const { mjModelRef, status } = useMujoco();
+  const { mjModelRef, status } = useMujocoContext();
   const { scene } = useThree();
   const lightsRef = useRef<THREE.Light[]>([]);
   const targetsRef = useRef<THREE.Object3D[]>([]);

package/src/hooks/useSensor.ts

@@ -6,7 +6,7 @@
  */
 
 import { useEffect, useRef, useMemo } from 'react';
-import { useMujoco, useAfterPhysicsStep } from '../core/MujocoSimProvider';
+import { useMujocoContext, useAfterPhysicsStep } from '../core/MujocoSimProvider';
 import { getName } from '../core/SceneLoader';
 import type { SensorInfo, SensorResult } from '../types';
 
@@ -15,7 +15,7 @@ import type { SensorInfo, SensorResult } from '../types';
  * updated every physics frame without causing React re-renders.
  */
 export function useSensor(name: string): SensorResult {
-  const { mjModelRef, mjDataRef, status } = useMujoco();
+  const { mjModelRef, mjDataRef, status } = useMujocoContext();
   const sensorIdRef = useRef(-1);
   const sensorAdrRef = useRef(0);
   const sensorDimRef = useRef(0);
@@ -55,7 +55,7 @@ export function useSensor(name: string): SensorResult {
  * Returns a stable array recomputed only when the model changes.
  */
 export function useSensors(): SensorInfo[] {
-  const { mjModelRef, status } = useMujoco();
+  const { mjModelRef, status } = useMujocoContext();
 
   return useMemo(() => {
     const model = mjModelRef.current;

package/src/hooks/useSitePosition.ts

@@ -6,7 +6,7 @@
 import { useEffect, useRef } from 'react';
 import { useFrame } from '@react-three/fiber';
 import * as THREE from 'three';
-import { useMujoco } from '../core/MujocoSimProvider';
+import { useMujocoContext } from '../core/MujocoSimProvider';
 import { findSiteByName } from '../core/SceneLoader';
 import type { SitePositionResult } from '../types';
 
@@ -18,7 +18,7 @@ const _mat4 = new THREE.Matrix4();
  * Refs are updated every frame without triggering React re-renders.
  */
 export function useSitePosition(siteName: string): SitePositionResult {
-  const { mjModelRef, mjDataRef, status } = useMujoco();
+  const { mjModelRef, mjDataRef, status } = useMujocoContext();
   const siteIdRef = useRef(-1);
   const positionRef = useRef(new THREE.Vector3());
   const quaternionRef = useRef(new THREE.Quaternion());

package/src/hooks/useTrajectoryPlayer.ts

@@ -7,7 +7,7 @@
 
 import { useCallback, useRef } from 'react';
 import { useFrame } from '@react-three/fiber';
-import { useMujoco } from '../core/MujocoSimProvider';
+import { useMujocoContext } from '../core/MujocoSimProvider';
 
 interface TrajectoryPlayerOptions {
   fps?: number;
@@ -24,7 +24,7 @@ export function useTrajectoryPlayer(
   trajectory: number[][],
   options: TrajectoryPlayerOptions = {},
 ) {
-  const { mjModelRef, mjDataRef, mujocoRef, pausedRef } = useMujoco();
+  const { mjModelRef, mjDataRef, mujocoRef, pausedRef } = useMujocoContext();
   const fps = options.fps ?? 30;
   const loop = options.loop ?? false;
 

package/src/hooks/useTrajectoryRecorder.ts

@@ -6,7 +6,7 @@
  */
 
 import { useCallback, useRef } from 'react';
-import { useMujoco, useAfterPhysicsStep } from '../core/MujocoSimProvider';
+import { useMujocoContext, useAfterPhysicsStep } from '../core/MujocoSimProvider';
 import type { TrajectoryFrame } from '../types';
 
 interface RecorderOptions {
@@ -17,7 +17,7 @@ interface RecorderOptions {
  * Record simulation trajectories for analysis, replay, or training data.
  */
 export function useTrajectoryRecorder(options: RecorderOptions = {}) {
-  const { mjModelRef } = useMujoco();
+  const { mjModelRef } = useMujocoContext();
   const recordingRef = useRef(false);
   const framesRef = useRef<TrajectoryFrame[]>([]);
   const fields = options.fields ?? ['qpos'];

package/src/index.ts

@@ -26,12 +26,11 @@ export {
 export { createController } from './core/createController';
 export type { ControllerOptions, ControllerComponent } from './core/createController';
 
-// IK controller plugin
-export { IkController } from './components/IkController';
-export { useIk } from './core/IkContext';
-export type { IkContextValue } from './core/IkContext';
+// IK controller hook
+export { useIkController } from './hooks/useIkController';
 
 // Components
+export { Body } from './components/Body';
 export { IkGizmo } from './components/IkGizmo';
 export { ContactMarkers } from './components/ContactMarkers';
 export { DragInteraction } from './components/DragInteraction';
@@ -74,6 +73,7 @@ export type {
   PhysicsConfig,
   // IK
   IkConfig,
+  IkContextValue,
   IKSolveFn,
   // Callbacks
   PhysicsStepCallback,
@@ -101,6 +101,7 @@ export type {
   // Policy
   PolicyConfig,
   // Component props
+  BodyProps,
   IkGizmoProps,
   DragInteractionProps,
   DebugProps,

package/src/types.ts

@@ -3,6 +3,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import type React from 'react';
+import type { ReactNode } from 'react';
 import type { CanvasProps } from '@react-three/fiber';
 import * as THREE from 'three';
 
@@ -309,6 +311,18 @@ export interface IkConfig {
   maxIterations?: number;
 }
 
+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 interface SceneMarker {
   id: number;
   position: THREE.Vector3;
@@ -482,6 +496,7 @@ export interface DebugProps {
 // ---- Component Props ----
 
 export interface IkGizmoProps {
+  controller: IkContextValue;
   siteName?: string;
   scale?: number;
   onDrag?: (position: THREE.Vector3, quaternion: THREE.Quaternion) => void;
@@ -517,6 +532,21 @@ export interface ContactListenerProps {
   onContactExit?: (info: ContactInfo) => void;
 }
 
+export interface BodyProps {
+  name: string;
+  type: 'box' | 'sphere' | 'cylinder';
+  size: [number, number, number];
+  position?: [number, number, number];
+  rgba?: [number, number, number, number];
+  mass?: number;
+  freejoint?: boolean;
+  friction?: string;
+  solref?: string;
+  solimp?: string;
+  condim?: number;
+  children?: ReactNode;
+}
+
 // ---- Public API (spec: full surface) ----
 
 export interface MujocoSimAPI {