mujoco-react 8.10.0 → 8.11.0

package/src/hooks/useFrameCapture.ts

@@ -0,0 +1,206 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * useFrameCapture — still-frame capture for canvas-backed MuJoCo/R3F scenes.
+ */
+
+import { useCallback, useState } from 'react';
+import type React from 'react';
+
+export type FrameCaptureStatus = 'idle' | 'capturing' | 'captured' | 'error';
+
+export type FrameCaptureTarget =
+  | HTMLCanvasElement
+  | HTMLElement
+  | null
+  | undefined;
+
+export type FrameCaptureTargetRef =
+  React.RefObject<HTMLCanvasElement | HTMLElement | null>;
+
+export interface FrameCaptureOptions {
+  target?: FrameCaptureTarget | FrameCaptureTargetRef;
+  type?: string;
+  quality?: number;
+  waitForAnimationFrame?: boolean;
+}
+
+export interface FrameCaptureResult {
+  canvas: HTMLCanvasElement;
+  dataUrl: string;
+  type: string;
+}
+
+export interface FrameCaptureBlobResult {
+  canvas: HTMLCanvasElement;
+  blob: Blob;
+  type: string;
+}
+
+export interface FrameCaptureAPI {
+  status: FrameCaptureStatus;
+  error: Error | null;
+  isCapturing: boolean;
+  capture: (options?: FrameCaptureOptions) => Promise<FrameCaptureResult>;
+  captureBlob: (
+    options?: FrameCaptureOptions
+  ) => Promise<FrameCaptureBlobResult>;
+  reset: () => void;
+}
+
+function isTargetRef(
+  target: FrameCaptureOptions['target']
+): target is FrameCaptureTargetRef {
+  return Boolean(target && typeof target === 'object' && 'current' in target);
+}
+
+function resolveCanvasTarget(
+  target: FrameCaptureOptions['target']
+): HTMLCanvasElement {
+  const resolvedTarget = isTargetRef(target) ? target.current : target;
+
+  if (!resolvedTarget) {
+    throw new Error('No frame capture target is available.');
+  }
+
+  if (resolvedTarget instanceof HTMLCanvasElement) {
+    return resolvedTarget;
+  }
+
+  const canvas = resolvedTarget.querySelector('canvas');
+  if (!canvas) {
+    throw new Error('Frame capture target does not contain a canvas.');
+  }
+  return canvas;
+}
+
+function waitForNextAnimationFrame() {
+  return new Promise<void>((resolve) => {
+    requestAnimationFrame(() => resolve());
+  });
+}
+
+/**
+ * Capture the current canvas frame as a data URL.
+ *
+ * For WebGL scenes, create the renderer with `preserveDrawingBuffer: true`
+ * when you need deterministic captures after the frame has presented.
+ */
+export async function captureFrame(
+  options: FrameCaptureOptions
+): Promise<FrameCaptureResult> {
+  const type = options.type ?? 'image/png';
+  const canvas = resolveCanvasTarget(options.target);
+
+  if (options.waitForAnimationFrame ?? true) {
+    await waitForNextAnimationFrame();
+  }
+
+  return {
+    canvas,
+    dataUrl: canvas.toDataURL(type, options.quality),
+    type,
+  };
+}
+
+/**
+ * Capture the current canvas frame as a Blob.
+ */
+export async function captureFrameBlob(
+  options: FrameCaptureOptions
+): Promise<FrameCaptureBlobResult> {
+  const type = options.type ?? 'image/png';
+  const canvas = resolveCanvasTarget(options.target);
+
+  if (options.waitForAnimationFrame ?? true) {
+    await waitForNextAnimationFrame();
+  }
+
+  const blob = await new Promise<Blob>((resolve, reject) => {
+    canvas.toBlob(
+      (nextBlob) => {
+        if (nextBlob) {
+          resolve(nextBlob);
+        } else {
+          reject(new Error('Canvas frame capture did not produce a Blob.'));
+        }
+      },
+      type,
+      options.quality
+    );
+  });
+
+  return { canvas, blob, type };
+}
+
+/**
+ * React state wrapper around `captureFrame` and `captureFrameBlob`.
+ */
+export function useFrameCapture(
+  defaultOptions: FrameCaptureOptions = {}
+): FrameCaptureAPI {
+  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: FrameCaptureOptions = {}) => {
+      setStatus('capturing');
+      setError(null);
+
+      try {
+        const result = await captureFrame({ ...defaultOptions, ...options });
+        setStatus('captured');
+        return result;
+      } catch (nextError) {
+        const error =
+          nextError instanceof Error
+            ? nextError
+            : new Error('Unable to capture the current canvas frame.');
+        setError(error);
+        setStatus('error');
+        throw error;
+      }
+    },
+    [defaultOptions]
+  );
+
+  const captureBlob = useCallback(
+    async (options: FrameCaptureOptions = {}) => {
+      setStatus('capturing');
+      setError(null);
+
+      try {
+        const result = await captureFrameBlob({
+          ...defaultOptions,
+          ...options,
+        });
+        setStatus('captured');
+        return result;
+      } catch (nextError) {
+        const error =
+          nextError instanceof Error
+            ? nextError
+            : new Error('Unable to capture the current canvas frame.');
+        setError(error);
+        setStatus('error');
+        throw error;
+      }
+    },
+    [defaultOptions]
+  );
+
+  return {
+    status,
+    error,
+    isCapturing: status === 'capturing',
+    capture,
+    captureBlob,
+    reset,
+  };
+}

package/src/hooks/usePolicy.ts

@@ -6,7 +6,7 @@
  */
 
 import { useRef } from 'react';
-import { useMujocoContext, useBeforePhysicsStep } from '../core/MujocoSimProvider';
+import { useBeforePhysicsStep } from '../core/MujocoSimProvider';
 import type { PolicyConfig } from '../types';
 
 /**
@@ -20,12 +20,13 @@ import type { PolicyConfig } from '../types';
  * @returns { step, isRunning } control handles
  */
 export function usePolicy(config: PolicyConfig) {
-  const { mjModelRef } = useMujocoContext();
   const lastActionTimeRef = useRef(0);
+  const lastObservationRef = useRef<ReturnType<PolicyConfig['onObservation']> | null>(null);
   const lastActionRef = useRef<Float32Array | Float64Array | number[] | null>(null);
-  const isRunningRef = useRef(true);
+  const isRunningRef = useRef(config.enabled ?? true);
   const configRef = useRef(config);
   configRef.current = config;
+  isRunningRef.current = config.enabled ?? isRunningRef.current;
 
   useBeforePhysicsStep((model, data) => {
     if (!isRunningRef.current) return;
@@ -37,13 +38,15 @@ export function usePolicy(config: PolicyConfig) {
     // Check if it's time for a new action
     if (data.time - lastActionTimeRef.current >= interval) {
       // Build observation
-      const obs = cfg.onObservation(model, data);
+      const observation = cfg.onObservation({ model, data });
+      const action = cfg.infer ? cfg.infer({ observation, model, data }) : observation;
 
-      // Apply action (consumer does inference inline or uses cached result)
-      cfg.onAction(obs, model, data);
+      // Apply action. If `infer` is omitted, this preserves the legacy inline-controller path.
+      cfg.onAction({ action, observation, model, data });
 
       lastActionTimeRef.current = data.time;
-      lastActionRef.current = obs;
+      lastObservationRef.current = observation;
+      lastActionRef.current = action;
     }
   });
 
@@ -51,6 +54,7 @@ export function usePolicy(config: PolicyConfig) {
     get isRunning() { return isRunningRef.current; },
     start: () => { isRunningRef.current = true; },
     stop: () => { isRunningRef.current = false; },
-    get lastObservation() { return lastActionRef.current; },
+    get lastObservation() { return lastObservationRef.current; },
+    get lastAction() { return lastActionRef.current; },
   };
 }

package/src/index.ts

@@ -44,11 +44,15 @@ export { SceneLights } from './components/SceneLights';
 export {
   ScenarioLighting,
   SplatEnvironment,
+  VisualScenarioEffects,
+  createPairedSplatEnvironment,
   createSparkSplatViewerUrl,
   createSplatEnvironmentUserData,
   getScenarioBackground,
   getScenarioCameraPosition,
   useSplatEnvironment,
+  useVisualScenarioEffects,
+  withSplatEnvironment,
 } from './components/VisualScenario';
 export { Debug } from './components/Debug';
 export { TendonRenderer } from './components/TendonRenderer';
@@ -73,6 +77,20 @@ export { useTrajectoryPlayer } from './hooks/useTrajectoryPlayer';
 export { useTrajectoryRecorder } from './hooks/useTrajectoryRecorder';
 export { useGamepad } from './hooks/useGamepad';
 export { useVideoRecorder } from './hooks/useVideoRecorder';
+export {
+  captureFrame,
+  captureFrameBlob,
+  useFrameCapture,
+} from './hooks/useFrameCapture';
+export type {
+  FrameCaptureAPI,
+  FrameCaptureBlobResult,
+  FrameCaptureOptions,
+  FrameCaptureResult,
+  FrameCaptureStatus,
+  FrameCaptureTarget,
+  FrameCaptureTargetRef,
+} from './hooks/useFrameCapture';
 export { useCtrlNoise } from './hooks/useCtrlNoise';
 export { useBodyMeshes } from './hooks/useBodyMeshes';
 export { useSelectionHighlight } from './hooks/useSelectionHighlight';
@@ -124,6 +142,10 @@ export type {
   KeyboardTeleopConfig,
   // Policy
   PolicyConfig,
+  PolicyVector,
+  PolicyObservationInput,
+  PolicyInferenceInput,
+  PolicyActionInput,
   // Observations
   ObservationConfig,
   ObservationHandle,
@@ -147,9 +169,12 @@ export type {
   PairedSplatEnvironmentConfig,
   SplatEnvironmentMetadataInput,
   SplatEnvironmentMetadata,
+  SplatSceneInput,
   VisualScenarioConfig,
   ScenarioLightingProps,
+  ScenarioMaterialConfig,
   SplatEnvironmentProps,
+  VisualScenarioEffectsProps,
   TrajectoryPlayerProps,
   ContactListenerProps,
   // API

package/src/spark.tsx

@@ -4,7 +4,13 @@
  */
 
 import { useThree } from '@react-three/fiber';
-import { useEffect, useRef, useState } from 'react';
+import {
+  useCallback,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+} from 'react';
 import * as THREE from 'three';
 import {
   SplatEnvironment,
@@ -17,9 +23,22 @@ import type {
 type SparkModule = typeof import('@sparkjsdev/spark');
 type SparkRendererInstance = InstanceType<SparkModule['SparkRenderer']>;
 type SparkSplatMeshInstance = InstanceType<SparkModule['SplatMesh']>;
+type SparkDisposable = {
+  dispose?: () => unknown;
+};
 
 export type SparkSplatStatus = 'idle' | 'loading' | 'ready' | 'error';
 
+export interface SparkSplatLifecycle {
+  status: SparkSplatStatus;
+  error: Error | null;
+  isLoading: boolean;
+  isReady: boolean;
+  isError: boolean;
+  props: Pick<SparkSplatEnvironmentProps, 'onStatusChange' | 'onError'>;
+  reset: () => void;
+}
+
 export interface SparkSplatEnvironmentProps extends SplatEnvironmentProps {
   /** Enable Spark LoD handling for large splat assets. Default: true. */
   lod?: boolean | 'quality';
@@ -34,6 +53,75 @@ export interface SparkSplatEnvironmentProps extends SplatEnvironmentProps {
   onError?: (error: Error) => void;
 }
 
+/**
+ * Tracks Spark 3DGS loading state for UI that wraps `SparkSplatEnvironment`.
+ *
+ * Use the returned `props` with `<SparkSplatEnvironment {...lifecycle.props} />`
+ * to avoid repeating status/error state in app code.
+ */
+export function useSparkSplatLifecycle({
+  enabled = true,
+  initialStatus,
+  onError,
+  onStatusChange,
+}: {
+  enabled?: boolean;
+  initialStatus?: SparkSplatStatus;
+  onError?: (error: Error) => void;
+  onStatusChange?: (status: SparkSplatStatus) => void;
+} = {}): SparkSplatLifecycle {
+  const [status, setStatus] = useState<SparkSplatStatus>(
+    initialStatus ?? (enabled ? 'loading' : 'idle')
+  );
+  const [error, setError] = useState<Error | null>(null);
+
+  useEffect(() => {
+    setStatus(enabled ? initialStatus ?? 'loading' : 'idle');
+    setError(null);
+  }, [enabled, initialStatus]);
+
+  const handleStatusChange = useCallback(
+    (nextStatus: SparkSplatStatus) => {
+      setStatus(nextStatus);
+      if (nextStatus !== 'error') {
+        setError(null);
+      }
+      onStatusChange?.(nextStatus);
+    },
+    [onStatusChange]
+  );
+
+  const handleError = useCallback(
+    (nextError: Error) => {
+      setError(nextError);
+      setStatus('error');
+      onError?.(nextError);
+    },
+    [onError]
+  );
+
+  const reset = useCallback(() => {
+    setStatus(enabled ? initialStatus ?? 'loading' : 'idle');
+    setError(null);
+  }, [enabled, initialStatus]);
+
+  return useMemo(
+    () => ({
+      status,
+      error,
+      isLoading: status === 'loading',
+      isReady: status === 'ready',
+      isError: status === 'error',
+      props: {
+        onStatusChange: handleStatusChange,
+        onError: handleError,
+      },
+      reset,
+    }),
+    [error, handleError, handleStatusChange, reset, status]
+  );
+}
+
 /**
  * Optional SparkJS-backed Gaussian splat renderer for React Three Fiber scenes.
  *
@@ -42,6 +130,8 @@ export interface SparkSplatEnvironmentProps extends SplatEnvironmentProps {
  */
 export function SparkSplatEnvironment({
   environment,
+  scenario,
+  renderer = 'spark',
   src,
   format,
   collisionProxy,
@@ -66,6 +156,8 @@ export function SparkSplatEnvironment({
   const { gl, invalidate } = useThree();
   const metadata = useSplatEnvironment({
     environment,
+    scenario,
+    renderer,
     src,
     format,
     collisionProxy: collisionProxyMetadata,
@@ -99,11 +191,20 @@ export function SparkSplatEnvironment({
     }
 
     async function loadSplat() {
-      if (!metadata.src || metadata.format !== 'spz') {
+      if (!metadata.src) {
         setLifecycleStatus('idle');
         return;
       }
 
+      if (metadata.format !== 'spz') {
+        const unsupportedFormatError = new Error(
+          `SparkSplatEnvironment only supports .spz assets; received "${metadata.format}".`
+        );
+        setLifecycleStatus('error');
+        onErrorRef.current?.(unsupportedFormatError);
+        return;
+      }
+
       setLifecycleStatus('loading');
 
       try {
@@ -166,13 +267,13 @@ export function SparkSplatEnvironment({
 
       if (meshRef.current) {
         groupRef.current?.remove(meshRef.current);
-        meshRef.current.dispose?.();
+        safelyDisposeSparkResource(meshRef.current);
         meshRef.current = null;
       }
 
       if (sparkRef.current) {
         groupRef.current?.remove(sparkRef.current);
-        sparkRef.current.dispose?.();
+        safelyDisposeSparkResource(sparkRef.current);
         sparkRef.current = null;
       }
     };
@@ -189,6 +290,8 @@ export function SparkSplatEnvironment({
     <SplatEnvironment
       {...groupProps}
       environment={environment}
+      scenario={scenario}
+      renderer={renderer}
       src={metadata.src}
       format={metadata.format}
       collisionProxyMetadata={metadata.collisionProxy}
@@ -200,3 +303,34 @@ export function SparkSplatEnvironment({
     </SplatEnvironment>
   );
 }
+
+function safelyDisposeSparkResource(resource: SparkDisposable) {
+  try {
+    const result = resource.dispose?.();
+    if (isPromiseLike(result)) {
+      void Promise.resolve(result).catch(handleSparkDisposeError);
+    }
+  } catch (error) {
+    handleSparkDisposeError(error);
+  }
+}
+
+function isPromiseLike(value: unknown): value is PromiseLike<unknown> {
+  return (
+    typeof value === 'object' &&
+    value !== null &&
+    'then' in value &&
+    typeof (value as { then?: unknown }).then === 'function'
+  );
+}
+
+function handleSparkDisposeError(error: unknown) {
+  if (
+    error instanceof Error &&
+    error.message.toLowerCase().includes('worker terminate')
+  ) {
+    return;
+  }
+
+  console.warn('[mujoco-react] Spark resource disposal failed.', error);
+}

package/src/types.ts

@@ -429,6 +429,8 @@ export interface SceneObject {
   solref?: string;
   solimp?: string;
   condim?: number;
+  /** MuJoCo geom group. Group 3 is conventionally used for collision-only helper geoms. */
+  group?: number;
 }
 
 export interface XmlPatch {
@@ -443,6 +445,8 @@ export type LocalMujocoFile = File;
 export interface LoadFromFilesOptions {
   /** Entry MJCF/URDF file. Inferred from scene.xml, model.xml, robot.xml, or the first XML/URDF file when omitted. */
   sceneFile?: string;
+  /** Additional MJCF environment XML files merged into the entry scene before MuJoCo compilation. */
+  environmentFiles?: string[];
   homeJoints?: number[];
   xmlPatches?: XmlPatch[];
   sceneObjects?: SceneObject[];
@@ -456,6 +460,14 @@ export interface SceneConfig {
   sceneFile: string;
   /** Browser-selected files for local MJCF/URDF loading. Preserves webkitRelativePath when available. */
   files?: readonly LocalMujocoFile[];
+  /**
+   * Additional MJCF environment XML files merged into the entry scene before compilation.
+   *
+   * Use this for static collision/physics layers such as a Gaussian-splat
+   * environment's proxy `scene.xml`; render the splat itself as a separate
+   * visual layer.
+   */
+  environmentFiles?: string[];
   sceneObjects?: SceneObject[];
   homeJoints?: number[];
   xmlPatches?: XmlPatch[];
@@ -709,10 +721,28 @@ export interface KeyboardTeleopConfig {
 
 // ---- Policy (spec 10.1) ----
 
+export type PolicyVector = Float32Array | Float64Array | number[];
+
+export interface PolicyObservationInput {
+  model: MujocoModel;
+  data: MujocoData;
+}
+
+export interface PolicyInferenceInput extends PolicyObservationInput {
+  observation: PolicyVector;
+}
+
+export interface PolicyActionInput extends PolicyInferenceInput {
+  action: PolicyVector;
+}
+
 export interface PolicyConfig {
   frequency: number;
-  onObservation: (model: MujocoModel, data: MujocoData) => Float32Array | Float64Array | number[];
-  onAction: (action: Float32Array | Float64Array | number[], model: MujocoModel, data: MujocoData) => void;
+  enabled?: boolean;
+  onObservation: (input: PolicyObservationInput) => PolicyVector;
+  /** Run policy inference. Omit to pass observations directly to `onAction` for custom inline controllers. */
+  infer?: (input: PolicyInferenceInput) => PolicyVector;
+  onAction: (input: PolicyActionInput) => void;
 }
 
 // ---- Observation Builder ----
@@ -805,6 +835,13 @@ export interface ScenarioCameraConfig {
   blur?: number;
 }
 
+export interface ScenarioMaterialConfig {
+  randomizeObjectColors?: boolean;
+  randomizeTableMaterial?: boolean;
+  roughness?: number;
+  metalness?: number;
+}
+
 export interface SplatAssetConfig {
   src: string;
   /** Common browser-friendly splat format. Renderer-specific loaders may accept more. */
@@ -819,7 +856,7 @@ export interface SplatScenarioConfig {
   format?: SplatFormat;
   src?: string;
   requiresCollisionProxy?: boolean;
-  collisionProxy?: SplatCollisionProxyConfig;
+  collisionProxy?: SplatCollisionProxyConfig | null;
 }
 
 export interface SplatCollisionProxyConfig {
@@ -845,6 +882,8 @@ export interface PairedSplatEnvironmentConfig {
 
 export interface SplatEnvironmentMetadataInput {
   environment?: PairedSplatEnvironmentConfig;
+  scenario?: VisualScenarioConfig;
+  renderer?: SplatRendererKind;
   src?: string;
   format?: SplatFormat;
   collisionProxy?: SplatCollisionProxyConfig;
@@ -857,6 +896,12 @@ export interface SplatEnvironmentMetadata {
   userData: Record<string, unknown>;
 }
 
+export type SplatSceneInput =
+  | PairedSplatEnvironmentConfig
+  | VisualScenarioConfig
+  | undefined
+  | null;
+
 export interface VisualScenarioConfig {
   id?: string;
   label?: string;
@@ -864,7 +909,8 @@ export interface VisualScenarioConfig {
   lighting?: ScenarioLightingPreset;
   environment?: string;
   camera?: ScenarioCameraConfig;
-  splat?: SplatScenarioConfig;
+  materials?: ScenarioMaterialConfig;
+  splat?: SplatScenarioConfig | null;
 }
 
 export interface ScenarioLightingProps {
@@ -875,6 +921,8 @@ export interface ScenarioLightingProps {
 
 export interface SplatEnvironmentProps extends Omit<ThreeElements['group'], 'ref'> {
   environment?: PairedSplatEnvironmentConfig;
+  scenario?: VisualScenarioConfig;
+  renderer?: SplatRendererKind;
   src?: string;
   format?: SplatFormat;
   collisionProxy?: ReactNode;
@@ -882,6 +930,19 @@ export interface SplatEnvironmentProps extends Omit<ThreeElements['group'], 'ref
   showPlaceholder?: boolean;
 }
 
+export interface VisualScenarioEffectsProps {
+  scenario?: VisualScenarioConfig;
+  enabled?: boolean;
+  applyBackground?: boolean;
+  applyFog?: boolean;
+  applyRenderer?: boolean;
+  applyMaterials?: boolean;
+  background?: THREE.ColorRepresentation;
+  fogNear?: number;
+  fogFar?: number;
+  materialFilter?: (object: THREE.Object3D, material: THREE.Material) => boolean;
+}
+
 export type TrajectoryInput = TrajectoryFrame[] | number[][];
 
 export interface TrajectoryPlayerProps {
@@ -920,6 +981,8 @@ export interface BodyProps {
   solref?: string;
   solimp?: string;
   condim?: number;
+  /** MuJoCo geom group. Group 3 is conventionally used for collision-only helper geoms. */
+  group?: number;
   children?: ReactNode;
 }
 
@@ -1017,6 +1080,8 @@ export interface MujocoSimAPI {
 
 export type MujocoCanvasProps = Omit<CanvasProps, 'onError'> & {
   config: SceneConfig;
+  /** R3F content rendered while the MuJoCo WASM module is still loading. */
+  loadingFallback?: ReactNode;
   onReady?: (api: MujocoSimAPI) => void;
   onError?: (error: Error) => void;
   onStep?: (time: number) => void;