mujoco-react 8.3.3 → 8.4.0

package/src/hooks/useIkController.ts

@@ -9,8 +9,8 @@ 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';
+import { createContiguousControlGroup, findSiteByName, resolveControlGroup } from '../core/SceneLoader';
+import type { ControlGroupInfo, IkConfig, IkContextValue, IKSolveFn, MujocoData } from '../types';
 
 // Preallocated temp for syncGizmoToSite
 const _syncMat4 = new THREE.Matrix4();
@@ -40,6 +40,7 @@ export const useIkController = createControllerHook<IkConfig, IkContextValue>(
     const ikCalculatingRef = useRef(false);
     const ikTargetRef = useRef<THREE.Group>(new THREE.Group());
     const siteIdRef = useRef(-1);
+    const controlGroupRef = useRef<ControlGroupInfo | null>(null);
     const genericIkRef = useRef<GenericIK>(new GenericIK(mujocoRef.current));
     const firstIkEnableRef = useRef(true);
     const needsInitialSync = useRef(true);
@@ -54,23 +55,32 @@ export const useIkController = createControllerHook<IkConfig, IkContextValue>(
       duration: 1000,
     });
 
-    // Resolve site ID when model loads or config changes
+    // Resolve site ID and model-aware control group when model loads or config changes.
     useEffect(() => {
       if (!config) {
         siteIdRef.current = -1;
+        controlGroupRef.current = null;
         return;
       }
       const model = mjModelRef.current;
       if (!model || status !== 'ready') {
         siteIdRef.current = -1;
+        controlGroupRef.current = null;
         return;
       }
       siteIdRef.current = findSiteByName(model, config.siteName);
+      controlGroupRef.current = config.numJoints !== undefined
+        ? createContiguousControlGroup(model, config.numJoints)
+        : resolveControlGroup(model, {
+            siteName: config.siteName,
+            joints: config.joints,
+            actuators: config.actuators,
+          });
       const data = mjDataRef.current;
       if (data && ikTargetRef.current) {
         syncGizmoToSite(data, siteIdRef.current, ikTargetRef.current);
       }
-    }, [config?.siteName, status, mjModelRef, mjDataRef, config]);
+    }, [config?.siteName, config?.numJoints, config?.joints, config?.actuators, status, mjModelRef, mjDataRef, config]);
 
     // IK solve function
     const ikSolveFn = useCallback(
@@ -79,9 +89,10 @@ export const useIkController = createControllerHook<IkConfig, IkContextValue>(
         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;
+        const controlGroup = controlGroupRef.current;
+        if (!model || !data || !controlGroup || siteIdRef.current === -1) return null;
         return genericIkRef.current.solve(
-          model, data, siteIdRef.current, config.numJoints,
+          model, data, siteIdRef.current, controlGroup.qposAdr,
           pos, quat, currentQ,
           { damping: config.damping, maxIterations: config.maxIterations },
         );
@@ -126,12 +137,20 @@ export const useIkController = createControllerHook<IkConfig, IkContextValue>(
       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);
+      const controlGroup = controlGroupRef.current;
+      if (!controlGroup) return;
+
+      const currentQ = Array.from(controlGroup.readQpos(data));
+      const solution = config.ikSolveFn
+        ? config.ikSolveFn(target.position, target.quaternion, currentQ, {
+            model,
+            data,
+            siteId: siteIdRef.current,
+            controlGroup,
+          })
+        : ikSolveFnRef.current(target.position, target.quaternion, currentQ);
       if (solution) {
-        for (let i = 0; i < numJoints; i++) data.ctrl[i] = solution[i];
+        controlGroup.writeCtrl(data, solution);
       }
     });
 

package/src/index.ts

@@ -5,6 +5,7 @@
 
 // Core
 export { MujocoProvider, useMujocoWasm } from './core/MujocoProvider';
+export type { MujocoLoader, MujocoLoaderOptions, MujocoProviderProps, MujocoWasmVariant } from './core/MujocoProvider';
 export { MujocoCanvas } from './core/MujocoCanvas';
 export { MujocoPhysics } from './core/MujocoPhysics';
 export type { MujocoPhysicsProps } from './core/MujocoPhysics';
@@ -20,6 +21,10 @@ export {
   findGeomByName,
   findSensorByName,
   findTendonByName,
+  getActuatedJoints,
+  getControlMap,
+  resolveControlGroup,
+  createContiguousControlGroup,
 } from './core/SceneLoader';
 
 // Controller factory
@@ -82,6 +87,11 @@ export type {
   // Model introspection
   BodyInfo,
   JointInfo,
+  ActuatedJointInfo,
+  ControlJointInfo,
+  ControlGroupInfo,
+  ControlGroupSelector,
+  ResourceSelector,
   GeomInfo,
   SiteInfo,
   ActuatorInfo,

package/src/types.ts

@@ -344,11 +344,27 @@ export interface SceneConfig {
 
 // ---- IK Controller Config ----
 
+export type ResourceSelector<TInfo, TName extends string = string> =
+  | TName
+  | readonly TName[]
+  | RegExp
+  | ((info: TInfo) => boolean);
+
 export interface IkConfig {
   /** MuJoCo site name for IK target. */
   siteName: Sites;
-  /** Number of joints to solve for. */
-  numJoints: number;
+  /**
+   * Explicit joints for IK. When omitted, the controller infers scalar hinge/slide
+   * joints by walking from the site body to the model root.
+   */
+  joints?: ResourceSelector<JointInfo, Joints>;
+  /** Explicit actuators for IK control output. */
+  actuators?: ResourceSelector<ActuatorInfo, Actuators>;
+  /**
+   * Number of joints to solve for, assuming legacy contiguous qpos/ctrl layout
+   * starting at index 0. Prefer inferred IK or `joints`/`actuators`.
+   */
+  numJoints?: number;
   /** Custom IK solver. When omitted, uses built-in Damped Least-Squares solver. */
   ikSolveFn?: IKSolveFn;
   /** DLS damping. Default: 0.01. */
@@ -390,9 +406,17 @@ export interface PhysicsConfig {
 export type IKSolveFn = (
   pos: THREE.Vector3,
   quat: THREE.Quaternion,
-  currentQ: number[]
+  currentQ: number[],
+  context?: IKSolveContext
 ) => number[] | null;
 
+export interface IKSolveContext {
+  model: MujocoModel;
+  data: MujocoData;
+  siteId: number;
+  controlGroup: ControlGroupInfo;
+}
+
 // ---- Callbacks ----
 
 export type PhysicsStepCallback = (
@@ -453,6 +477,48 @@ export interface ActuatorInfo {
   range: [number, number];
 }
 
+export interface ActuatedJointInfo extends JointInfo {
+  actuatorId: number;
+  actuatorName: string;
+  ctrlAdr: number;
+  ctrlRange: [number, number];
+}
+
+export interface ControlJointInfo extends JointInfo {
+  actuatorId: number | null;
+  actuatorName: string | null;
+  ctrlAdr: number | null;
+  ctrlRange: [number, number] | null;
+}
+
+export interface ControlGroupSelector {
+  /** Infer a kinematic chain from a MuJoCo site. */
+  siteName?: Sites;
+  /** Infer a kinematic chain from a body. */
+  bodyName?: Bodies;
+  /** Select joints by name, names, regex, or predicate. */
+  joints?: ResourceSelector<JointInfo, Joints>;
+  /** Select actuators by name, names, regex, or predicate. */
+  actuators?: ResourceSelector<ActuatorInfo, Actuators>;
+}
+
+export interface ControlGroupInfo {
+  /** Joints in solve/control order. */
+  joints: ControlJointInfo[];
+  /** Actuators in control output order. */
+  actuators: ActuatorInfo[];
+  /** qpos addresses for scalar hinge/slide joints. */
+  qposAdr: number[];
+  /** dof addresses for scalar hinge/slide joints. */
+  dofAdr: number[];
+  /** ctrl addresses matching writable actuators. */
+  ctrlAdr: number[];
+  readQpos(data: MujocoData): Float64Array;
+  readCtrl(data: MujocoData): Float64Array;
+  writeQpos(data: MujocoData, values: ArrayLike<number>): void;
+  writeCtrl(data: MujocoData, values: ArrayLike<number>): void;
+}
+
 export interface SensorInfo {
   id: number;
   name: string;
@@ -629,6 +695,9 @@ export interface MujocoSimAPI {
   // Actuator / control (spec 3.1)
   setCtrl(nameOrValues: Actuators | Record<Actuators, number>, value?: number): void;
   getCtrl(): Float64Array;
+  getControlMap(): ControlGroupInfo;
+  getActuatedJoints(): ActuatedJointInfo[];
+  resolveControlGroup(selector: ControlGroupSelector): ControlGroupInfo | null;
 
   // Force application (spec 8.1)
   applyForce(bodyName: Bodies, force: THREE.Vector3, point?: THREE.Vector3): void;