mujoco-react 10.0.0 → 10.1.0

package/src/hooks/useControlWriter.ts

@@ -0,0 +1,176 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Cooperative actuator/control ownership for policies, IK, teleop, and replay.
+ */
+
+import { useCallback, useEffect, useMemo, useRef } from 'react';
+import type { RefObject } from 'react';
+import { useMujocoContext } from '../core/MujocoSimProvider';
+import type { ControlGroupInfo, ControlGroupSelector, MujocoData, MujocoModel } from '../types';
+
+export interface ControlWriterConflict {
+  actuatorIndex: number;
+  owner: string;
+  requestedBy: string;
+}
+
+export interface ControlWriterOptions {
+  owner: string;
+  selector?: ControlGroupSelector;
+  enabled?: boolean;
+  warnOnConflict?: boolean;
+  allowSameOwner?: boolean;
+  onConflict?: (conflicts: ControlWriterConflict[]) => void;
+}
+
+export interface ControlWriterWriteOptions {
+  force?: boolean;
+}
+
+export interface ControlWriterHandle {
+  owner: string;
+  group: RefObject<ControlGroupInfo | null>;
+  conflicts: RefObject<ControlWriterConflict[]>;
+  canWrite: () => boolean;
+  read: () => Float64Array;
+  write: (values: ArrayLike<number>, options?: ControlWriterWriteOptions) => boolean;
+  release: () => void;
+}
+
+interface ControlWriterClaim {
+  owner: string;
+  token: symbol;
+}
+
+const claimsByModel = new WeakMap<MujocoModel, Map<number, ControlWriterClaim>>();
+
+function getClaims(model: MujocoModel) {
+  let claims = claimsByModel.get(model);
+  if (!claims) {
+    claims = new Map();
+    claimsByModel.set(model, claims);
+  }
+  return claims;
+}
+
+function releaseClaims(model: MujocoModel | null, token: symbol) {
+  if (!model) return;
+  const claims = claimsByModel.get(model);
+  if (!claims) return;
+  for (const [actuatorIndex, claim] of claims) {
+    if (claim.token === token) claims.delete(actuatorIndex);
+  }
+}
+
+export function useControlWriter(options: ControlWriterOptions): ControlWriterHandle {
+  const {
+    owner,
+    selector,
+    enabled = true,
+    warnOnConflict = true,
+    allowSameOwner = true,
+    onConflict,
+  } = options;
+  const { api, mjModelRef, mjDataRef, status } = useMujocoContext();
+  const tokenRef = useRef(Symbol(owner));
+  const claimedModelRef = useRef<MujocoModel | null>(null);
+  const groupRef = useRef<ControlGroupInfo | null>(null);
+  const conflictsRef = useRef<ControlWriterConflict[]>([]);
+  const onConflictRef = useRef(onConflict);
+  onConflictRef.current = onConflict;
+
+  const release = useCallback(() => {
+    releaseClaims(claimedModelRef.current, tokenRef.current);
+    claimedModelRef.current = null;
+    conflictsRef.current = [];
+  }, []);
+
+  useEffect(() => {
+    release();
+    if (!enabled || status !== 'ready') {
+      groupRef.current = null;
+      return;
+    }
+
+    const model = mjModelRef.current;
+    if (!model) {
+      groupRef.current = null;
+      return;
+    }
+
+    const group = selector ? api.resolveControlGroup(selector) : api.getControlMap();
+    groupRef.current = group;
+    if (!group) return;
+
+    const claims = getClaims(model);
+    const conflicts: ControlWriterConflict[] = [];
+    for (const actuatorIndex of group.ctrlAdr) {
+      const existing = claims.get(actuatorIndex);
+      if (
+        existing &&
+        existing.token !== tokenRef.current &&
+        (!allowSameOwner || existing.owner !== owner)
+      ) {
+        conflicts.push({
+          actuatorIndex,
+          owner: existing.owner,
+          requestedBy: owner,
+        });
+      }
+    }
+
+    conflictsRef.current = conflicts;
+    if (conflicts.length > 0) {
+      onConflictRef.current?.(conflicts);
+      if (warnOnConflict) {
+        console.warn(
+          `[mujoco-react] Control writer "${owner}" conflicts with existing writer(s): ${conflicts
+            .map((conflict) => `${conflict.actuatorIndex}:${conflict.owner}`)
+            .join(', ')}`
+        );
+      }
+      return;
+    }
+
+    for (const actuatorIndex of group.ctrlAdr) {
+      claims.set(actuatorIndex, { owner, token: tokenRef.current });
+    }
+    claimedModelRef.current = model;
+
+    return release;
+  }, [allowSameOwner, api, enabled, mjModelRef, owner, release, selector, status, warnOnConflict]);
+
+  const canWrite = useCallback(() => (
+    enabled &&
+    groupRef.current !== null &&
+    conflictsRef.current.length === 0
+  ), [enabled]);
+
+  const read = useCallback(() => {
+    const data = mjDataRef.current;
+    const group = groupRef.current;
+    if (!data || !group) return new Float64Array(0);
+    return group.readCtrl(data);
+  }, [mjDataRef]);
+
+  const write = useCallback((values: ArrayLike<number>, writeOptions: ControlWriterWriteOptions = {}) => {
+    const data: MujocoData | null = mjDataRef.current;
+    const group = groupRef.current;
+    if (!data || !group) return false;
+    if (!writeOptions.force && !canWrite()) return false;
+    group.writeCtrl(data, values);
+    return true;
+  }, [canWrite, mjDataRef]);
+
+  return useMemo(() => ({
+    owner,
+    group: groupRef,
+    conflicts: conflictsRef,
+    canWrite,
+    read,
+    write,
+    release,
+  }), [canWrite, owner, read, release, write]);
+}

package/src/hooks/useNamedObservation.ts

@@ -0,0 +1,42 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Stable React handle for named policy observations.
+ */
+
+import { useMemo, useRef } from 'react';
+import { useMujocoContext } from '../core/MujocoSimProvider';
+import { readNamedObservation } from '../policyObservation';
+import type {
+  NamedObservationOptions,
+  NamedObservationResult,
+} from '../policyObservation';
+
+const EMPTY_NAMED_OBSERVATION: NamedObservationResult = {
+  values: new Float32Array(0),
+  layout: [],
+};
+
+export interface NamedObservationHandle {
+  read: () => NamedObservationResult;
+  readValues: () => Float32Array | Float64Array;
+}
+
+export function useNamedObservation(options: NamedObservationOptions): NamedObservationHandle {
+  const { mjModelRef, mjDataRef } = useMujocoContext();
+  const optionsRef = useRef(options);
+  optionsRef.current = options;
+
+  return useMemo(() => ({
+    read() {
+      const model = mjModelRef.current;
+      const data = mjDataRef.current;
+      if (!model || !data) return EMPTY_NAMED_OBSERVATION;
+      return readNamedObservation(model, data, optionsRef.current);
+    },
+    readValues() {
+      return this.read().values;
+    },
+  }), [mjDataRef, mjModelRef]);
+}

package/src/hooks/usePolicy.ts

@@ -5,9 +5,47 @@
  * usePolicy — policy decimation loop hook (spec 10.1)
  */
 
-import { useRef } from 'react';
+import { useCallback, useMemo, useRef } from 'react';
 import { useBeforePhysicsStep } from '../core/MujocoSimProvider';
-import type { PolicyConfig } from '../types';
+import type { PolicyAPI, PolicyConfig, PolicyInferenceOutput, PolicyVector } from '../types';
+
+type PendingPolicyAction = {
+  action: PolicyVector;
+  observation: PolicyVector;
+};
+
+function isPromiseLike(value: unknown): value is Promise<PolicyInferenceOutput> {
+  return (
+    typeof value === 'object' &&
+    value !== null &&
+    'then' in value &&
+    typeof (value as { then?: unknown }).then === 'function'
+  );
+}
+
+function isPolicyActionChunk(value: PolicyInferenceOutput): value is readonly PolicyVector[] {
+  return (
+    Array.isArray(value) &&
+    value.length > 0 &&
+    (Array.isArray(value[0]) || ArrayBuffer.isView(value[0]))
+  );
+}
+
+function toPolicyActions(output: PolicyInferenceOutput): PolicyVector[] {
+  return isPolicyActionChunk(output) ? [...output] : [output];
+}
+
+function enqueuePolicyActions(
+  queue: PendingPolicyAction[],
+  actions: PolicyVector[],
+  observation: PolicyVector,
+  strategy: PolicyConfig['queueStrategy']
+) {
+  if (strategy === 'replace') {
+    queue.splice(0, queue.length);
+  }
+  queue.push(...actions.map((action) => ({ action, observation })));
+}
 
 /**
  * Framework-agnostic policy execution hook.
@@ -19,15 +57,33 @@ import type { PolicyConfig } from '../types';
  * @param config Policy configuration
  * @returns { step, isRunning } control handles
  */
-export function usePolicy(config: PolicyConfig) {
+export function usePolicy(config: PolicyConfig): PolicyAPI {
   const lastActionTimeRef = useRef(0);
   const lastObservationRef = useRef<ReturnType<PolicyConfig['onObservation']> | null>(null);
   const lastActionRef = useRef<Float32Array | Float64Array | number[] | null>(null);
+  const actionQueueRef = useRef<PendingPolicyAction[]>([]);
+  const inFlightRef = useRef(false);
+  const lastErrorRef = useRef<unknown>(null);
+  const epochRef = useRef(0);
   const isRunningRef = useRef(config.enabled ?? true);
   const configRef = useRef(config);
   configRef.current = config;
   isRunningRef.current = config.enabled ?? isRunningRef.current;
 
+  const clearQueue = useCallback(() => {
+    epochRef.current += 1;
+    actionQueueRef.current.splice(0, actionQueueRef.current.length);
+    inFlightRef.current = false;
+    lastErrorRef.current = null;
+  }, []);
+
+  const reset = useCallback(() => {
+    clearQueue();
+    lastActionTimeRef.current = 0;
+    lastObservationRef.current = null;
+    lastActionRef.current = null;
+  }, [clearQueue]);
+
   useBeforePhysicsStep(({ model, data }) => {
     if (!isRunningRef.current) return;
 
@@ -37,24 +93,91 @@ export function usePolicy(config: PolicyConfig) {
 
     // Check if it's time for a new action
     if (data.time - lastActionTimeRef.current >= interval) {
+      const queuedAction = actionQueueRef.current.shift();
+      if (queuedAction) {
+        cfg.onAction({
+          action: queuedAction.action,
+          observation: queuedAction.observation,
+          model,
+          data,
+        });
+        lastActionTimeRef.current = data.time;
+        lastActionRef.current = queuedAction.action;
+      }
+
+      const prefetchThreshold = cfg.prefetchThreshold ?? 0;
+      const shouldInfer = !inFlightRef.current && (!queuedAction || actionQueueRef.current.length <= prefetchThreshold);
+      if (!shouldInfer) return;
+
       // Build observation
       const observation = cfg.onObservation({ model, data });
-      const action = cfg.infer ? cfg.infer({ observation, model, data }) : observation;
+      const result = cfg.infer ? cfg.infer({ observation, model, data }) : observation;
 
-      // Apply action. If `infer` is omitted, this preserves the legacy inline-controller path.
-      cfg.onAction({ action, observation, model, data });
+      if (isPromiseLike(result)) {
+        const epoch = epochRef.current;
+        inFlightRef.current = true;
+        result
+          .then((output) => {
+            if (epoch !== epochRef.current) return;
+            enqueuePolicyActions(
+              actionQueueRef.current,
+              toPolicyActions(output),
+              observation,
+              cfg.queueStrategy ?? 'append'
+            );
+            lastErrorRef.current = null;
+          })
+          .catch((error: unknown) => {
+            if (epoch !== epochRef.current) return;
+            lastErrorRef.current = error;
+            cfg.onError?.(error);
+          })
+          .finally(() => {
+            if (epoch !== epochRef.current) return;
+            inFlightRef.current = false;
+          });
+      } else {
+        const actions = toPolicyActions(result);
+        if (queuedAction) {
+          enqueuePolicyActions(
+            actionQueueRef.current,
+            actions,
+            observation,
+            cfg.queueStrategy ?? 'append'
+          );
+        } else {
+          const [action, ...queuedActions] = actions;
+          if (!action) return;
+          enqueuePolicyActions(
+            actionQueueRef.current,
+            queuedActions,
+            observation,
+            cfg.queueStrategy ?? 'append'
+          );
+          // Apply action. If `infer` is omitted, this preserves the legacy inline-controller path.
+          cfg.onAction({ action, observation, model, data });
+          lastActionRef.current = action;
+        }
+      }
 
       lastActionTimeRef.current = data.time;
       lastObservationRef.current = observation;
-      lastActionRef.current = action;
     }
   });
 
-  return {
+  return useMemo(() => ({
     get isRunning() { return isRunningRef.current; },
     start: () => { isRunningRef.current = true; },
-    stop: () => { isRunningRef.current = false; },
+    stop: () => {
+      isRunningRef.current = false;
+      if (configRef.current.clearQueueOnStop) reset();
+    },
+    clearQueue,
+    reset,
+    get inFlight() { return inFlightRef.current; },
+    get queuedActions() { return actionQueueRef.current.length; },
     get lastObservation() { return lastObservationRef.current; },
     get lastAction() { return lastActionRef.current; },
-  };
+    get lastError() { return lastErrorRef.current; },
+  }), [clearQueue, reset]);
 }

package/src/hooks/usePolicyCameraFrames.ts

@@ -0,0 +1,162 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * React wrapper for capturing policy image payloads from Three/MuJoCo cameras.
+ */
+
+import { useCallback, useState } from 'react';
+import { useMujoco } from '../core/MujocoSimProvider';
+import {
+  capturePolicyCameraFrames,
+  capturePolicyCameraFramesFromMountedStreams,
+} from '../policyCameraFrames';
+import type {
+  CreatePolicyCameraFrameCapturePlanOptions,
+  PolicyCameraFrameCapturePlan,
+} from '../policyCameraFrames';
+import type {
+  FrameCaptureStatus,
+  PolicyCameraFrameCaptureAPI,
+  PolicyCameraFrameCaptureOptions,
+  PolicyCameraFrameCaptureResult,
+} from '../types';
+
+export type MountedPolicyCameraFrameCaptureOptions = Omit<
+  CreatePolicyCameraFrameCapturePlanOptions,
+  'cameras' | 'sites' | 'bodies'
+>;
+
+export interface MountedPolicyCameraFrameCaptureAPI {
+  status: FrameCaptureStatus;
+  error: Error | null;
+  isCapturing: boolean;
+  capture: (
+    options?: Partial<MountedPolicyCameraFrameCaptureOptions>
+  ) => Promise<PolicyCameraFrameCaptureResult & { plan: PolicyCameraFrameCapturePlan }>;
+  reset: () => void;
+}
+
+function mergePolicyCameraFrameCaptureOptions(
+  defaultOptions: MountedPolicyCameraFrameCaptureOptions,
+  options: Partial<MountedPolicyCameraFrameCaptureOptions>
+): MountedPolicyCameraFrameCaptureOptions {
+  return {
+    ...defaultOptions,
+    ...options,
+    cameraKeys: options.cameraKeys ?? defaultOptions.cameraKeys,
+    aliases: {
+      ...defaultOptions.aliases,
+      ...options.aliases,
+    },
+    defaults: {
+      ...defaultOptions.defaults,
+      ...options.defaults,
+    },
+    streamOptions: {
+      ...defaultOptions.streamOptions,
+      ...options.streamOptions,
+    },
+  };
+}
+
+export function usePolicyCameraFrames(
+  defaultOptions: PolicyCameraFrameCaptureOptions
+): PolicyCameraFrameCaptureAPI {
+  const mujoco = useMujoco();
+  const [status, setStatus] = useState<FrameCaptureStatus>('idle');
+  const [error, setError] = useState<Error | null>(null);
+
+  const reset = useCallback(() => {
+    setStatus('idle');
+    setError(null);
+  }, []);
+
+  const capture = useCallback(
+    async (options: Partial<PolicyCameraFrameCaptureOptions> = {}) => {
+      if (!mujoco.api) {
+        throw new Error('MuJoCo scene is not ready for policy camera capture.');
+      }
+
+      setStatus('capturing');
+      setError(null);
+
+      try {
+        const result = await capturePolicyCameraFrames(mujoco.api, {
+          ...defaultOptions,
+          ...options,
+          streams: options.streams ?? defaultOptions.streams,
+        });
+        setStatus('captured');
+        return result;
+      } catch (nextError) {
+        const error =
+          nextError instanceof Error
+            ? nextError
+            : new Error('Unable to capture policy camera frames.');
+        setError(error);
+        setStatus('error');
+        throw error;
+      }
+    },
+    [defaultOptions, mujoco.api]
+  );
+
+  return {
+    status,
+    error,
+    isCapturing: status === 'capturing',
+    capture,
+    reset,
+  };
+}
+
+export function usePolicyCameraFramesFromMountedStreams(
+  defaultOptions: MountedPolicyCameraFrameCaptureOptions
+): MountedPolicyCameraFrameCaptureAPI {
+  const mujoco = useMujoco();
+  const [status, setStatus] = useState<FrameCaptureStatus>('idle');
+  const [error, setError] = useState<Error | null>(null);
+
+  const reset = useCallback(() => {
+    setStatus('idle');
+    setError(null);
+  }, []);
+
+  const capture = useCallback(
+    async (options: Partial<MountedPolicyCameraFrameCaptureOptions> = {}) => {
+      if (!mujoco.api) {
+        throw new Error('MuJoCo scene is not ready for mounted policy camera capture.');
+      }
+
+      setStatus('capturing');
+      setError(null);
+
+      try {
+        const result = await capturePolicyCameraFramesFromMountedStreams(
+          mujoco.api,
+          mergePolicyCameraFrameCaptureOptions(defaultOptions, options)
+        );
+        setStatus('captured');
+        return result;
+      } catch (nextError) {
+        const error =
+          nextError instanceof Error
+            ? nextError
+            : new Error('Unable to capture mounted policy camera frames.');
+        setError(error);
+        setStatus('error');
+        throw error;
+      }
+    },
+    [defaultOptions, mujoco.api]
+  );
+
+  return {
+    status,
+    error,
+    isCapturing: status === 'capturing',
+    capture,
+    reset,
+  };
+}

package/src/hooks/usePose.ts

@@ -0,0 +1,119 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Ref-based world pose hooks for named MuJoCo bodies, geoms, and sites.
+ */
+
+import { useEffect, useRef } from 'react';
+import type { RefObject } from 'react';
+import * as THREE from 'three';
+import { useAfterPhysicsStep, useMujocoContext } from '../core/MujocoSimProvider';
+import { findBodyByName, findGeomByName, findSiteByName } from '../core/SceneLoader';
+import type { Bodies, Geoms, Sites } from '../types';
+
+export type PoseResourceKind = 'body' | 'geom' | 'site';
+
+export interface PoseReadout {
+  id: RefObject<number>;
+  found: RefObject<boolean>;
+  position: RefObject<THREE.Vector3>;
+  quaternion: RefObject<THREE.Quaternion>;
+}
+
+const _matrix = new THREE.Matrix4();
+
+function quaternionFromMatrixArray(
+  target: THREE.Quaternion,
+  values: ArrayLike<number>,
+  offset: number
+) {
+  _matrix.set(
+    values[offset], values[offset + 1], values[offset + 2], 0,
+    values[offset + 3], values[offset + 4], values[offset + 5], 0,
+    values[offset + 6], values[offset + 7], values[offset + 8], 0,
+    0, 0, 0, 1
+  );
+  target.setFromRotationMatrix(_matrix);
+}
+
+function quaternionFromMujocoQuat(
+  target: THREE.Quaternion,
+  values: ArrayLike<number>,
+  offset: number
+) {
+  target.set(
+    values[offset + 1] ?? 0,
+    values[offset + 2] ?? 0,
+    values[offset + 3] ?? 0,
+    values[offset] ?? 1
+  );
+}
+
+function useNamedPose(kind: PoseResourceKind, name: string): PoseReadout {
+  const { mjModelRef, status } = useMujocoContext();
+  const idRef = useRef(-1);
+  const foundRef = useRef(false);
+  const positionRef = useRef(new THREE.Vector3());
+  const quaternionRef = useRef(new THREE.Quaternion());
+
+  useEffect(() => {
+    const model = mjModelRef.current;
+    if (!model || status !== 'ready') {
+      idRef.current = -1;
+      foundRef.current = false;
+      return;
+    }
+
+    if (kind === 'body') idRef.current = findBodyByName(model, name);
+    else if (kind === 'geom') idRef.current = findGeomByName(model, name);
+    else idRef.current = findSiteByName(model, name);
+    foundRef.current = idRef.current >= 0;
+  }, [kind, name, status, mjModelRef]);
+
+  useAfterPhysicsStep(({ data }) => {
+    const id = idRef.current;
+    if (id < 0) return;
+
+    if (kind === 'body') {
+      const p = id * 3;
+      positionRef.current.set(data.xpos[p], data.xpos[p + 1], data.xpos[p + 2]);
+      if (data.xmat) {
+        quaternionFromMatrixArray(quaternionRef.current, data.xmat, id * 9);
+      } else {
+        quaternionFromMujocoQuat(quaternionRef.current, data.xquat, id * 4);
+      }
+      return;
+    }
+
+    if (kind === 'geom') {
+      const p = id * 3;
+      positionRef.current.set(data.geom_xpos[p], data.geom_xpos[p + 1], data.geom_xpos[p + 2]);
+      quaternionFromMatrixArray(quaternionRef.current, data.geom_xmat, id * 9);
+      return;
+    }
+
+    const p = id * 3;
+    positionRef.current.set(data.site_xpos[p], data.site_xpos[p + 1], data.site_xpos[p + 2]);
+    quaternionFromMatrixArray(quaternionRef.current, data.site_xmat, id * 9);
+  });
+
+  return {
+    id: idRef,
+    found: foundRef,
+    position: positionRef,
+    quaternion: quaternionRef,
+  };
+}
+
+export function useBodyPose(name: Bodies): PoseReadout {
+  return useNamedPose('body', name);
+}
+
+export function useGeomPose(name: Geoms): PoseReadout {
+  return useNamedPose('geom', name);
+}
+
+export function useSitePose(name: Sites): PoseReadout {
+  return useNamedPose('site', name);
+}