mujoco-react 8.9.2 → 8.11.0

package/src/spark.tsx

@@ -0,0 +1,336 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { useThree } from '@react-three/fiber';
+import {
+  useCallback,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+} from 'react';
+import * as THREE from 'three';
+import {
+  SplatEnvironment,
+  useSplatEnvironment,
+} from './components/VisualScenario';
+import type {
+  SplatEnvironmentProps,
+} from './types';
+
+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';
+  /**
+   * Hide meshes whose names include floor, ground, or plane while the splat is
+   * active. This mirrors the common hybrid-rendering setup where MJCF keeps
+   * collision geometry but the splat owns the visual environment.
+   */
+  hideGroundMeshes?: boolean;
+  onStatusChange?: (status: SparkSplatStatus) => void;
+  onLoad?: (mesh: SparkSplatMeshInstance) => void;
+  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.
+ *
+ * Import from `mujoco-react/spark` and install `@sparkjsdev/spark` in the app
+ * that uses it. The core `mujoco-react` entrypoint does not depend on Spark.
+ */
+export function SparkSplatEnvironment({
+  environment,
+  scenario,
+  renderer = 'spark',
+  src,
+  format,
+  collisionProxy,
+  collisionProxyMetadata,
+  showPlaceholder,
+  children,
+  lod = true,
+  hideGroundMeshes = false,
+  onStatusChange,
+  onLoad,
+  onError,
+  ...groupProps
+}: SparkSplatEnvironmentProps) {
+  const groupRef = useRef<THREE.Group>(null);
+  const sparkRef = useRef<SparkRendererInstance | null>(null);
+  const meshRef = useRef<SparkSplatMeshInstance | null>(null);
+  const hiddenMeshesRef = useRef<THREE.Mesh[]>([]);
+  const onStatusChangeRef = useRef(onStatusChange);
+  const onLoadRef = useRef(onLoad);
+  const onErrorRef = useRef(onError);
+  const [status, setStatus] = useState<SparkSplatStatus>('idle');
+  const { gl, invalidate } = useThree();
+  const metadata = useSplatEnvironment({
+    environment,
+    scenario,
+    renderer,
+    src,
+    format,
+    collisionProxy: collisionProxyMetadata,
+  });
+
+  useEffect(() => {
+    onStatusChangeRef.current = onStatusChange;
+  }, [onStatusChange]);
+
+  useEffect(() => {
+    onLoadRef.current = onLoad;
+  }, [onLoad]);
+
+  useEffect(() => {
+    onErrorRef.current = onError;
+  }, [onError]);
+
+  useEffect(() => {
+    let disposed = false;
+
+    function setLifecycleStatus(nextStatus: SparkSplatStatus) {
+      setStatus(nextStatus);
+      onStatusChangeRef.current?.(nextStatus);
+    }
+
+    function restoreHiddenMeshes() {
+      for (const mesh of hiddenMeshesRef.current) {
+        mesh.visible = true;
+      }
+      hiddenMeshesRef.current = [];
+    }
+
+    async function loadSplat() {
+      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 {
+        const sparkModule = await import('@sparkjsdev/spark');
+        if (disposed || !groupRef.current) return;
+
+        const spark = new sparkModule.SparkRenderer({
+          renderer: gl,
+          onDirty: invalidate,
+        });
+        const mesh = new sparkModule.SplatMesh({
+          url: metadata.src,
+          lod,
+        });
+        mesh.name = 'GaussianSplatMesh';
+
+        groupRef.current.add(spark);
+        groupRef.current.add(mesh);
+        sparkRef.current = spark;
+        meshRef.current = mesh;
+
+        if (hideGroundMeshes && groupRef.current.parent) {
+          groupRef.current.parent.traverse((object) => {
+            if (
+              !(object instanceof THREE.Mesh) ||
+              object === (mesh as unknown as THREE.Object3D)
+            ) {
+              return;
+            }
+            const name = object.name.toLowerCase();
+            if (
+              name.includes('floor') ||
+              name.includes('ground') ||
+              name.includes('plane')
+            ) {
+              object.visible = false;
+              hiddenMeshesRef.current.push(object);
+            }
+          });
+        }
+
+        await mesh.initialized;
+        if (disposed) return;
+        setLifecycleStatus('ready');
+        onLoadRef.current?.(mesh);
+        invalidate();
+      } catch (error) {
+        const normalizedError =
+          error instanceof Error ? error : new Error(String(error));
+        setLifecycleStatus('error');
+        onErrorRef.current?.(normalizedError);
+      }
+    }
+
+    void loadSplat();
+
+    return () => {
+      disposed = true;
+      restoreHiddenMeshes();
+
+      if (meshRef.current) {
+        groupRef.current?.remove(meshRef.current);
+        safelyDisposeSparkResource(meshRef.current);
+        meshRef.current = null;
+      }
+
+      if (sparkRef.current) {
+        groupRef.current?.remove(sparkRef.current);
+        safelyDisposeSparkResource(sparkRef.current);
+        sparkRef.current = null;
+      }
+    };
+  }, [
+    gl,
+    hideGroundMeshes,
+    invalidate,
+    lod,
+    metadata.format,
+    metadata.src,
+  ]);
+
+  return (
+    <SplatEnvironment
+      {...groupProps}
+      environment={environment}
+      scenario={scenario}
+      renderer={renderer}
+      src={metadata.src}
+      format={metadata.format}
+      collisionProxyMetadata={metadata.collisionProxy}
+      collisionProxy={collisionProxy}
+      showPlaceholder={showPlaceholder ?? status !== 'ready'}
+    >
+      <group ref={groupRef} />
+      {children}
+    </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

@@ -5,7 +5,7 @@
 
 import type React from 'react';
 import type { ReactNode } from 'react';
-import type { CanvasProps } from '@react-three/fiber';
+import type { CanvasProps, ThreeElements } from '@react-three/fiber';
 import * as THREE from 'three';
 
 // ---- Register (type-safe named resources) ----
@@ -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 ----
@@ -791,6 +821,128 @@ export interface SceneLightsProps {
   intensity?: number;
 }
 
+// ---- Visual scenarios / 3DGS composition ----
+
+export type ScenarioLightingPreset = 'studio' | 'warehouse' | 'low-light' | 'splat';
+export type SplatFormat = 'spz' | 'ply' | 'splat';
+export type SplatRendererKind = 'spark' | 'custom';
+export type SplatCollisionPrimitive = 'plane' | 'box' | 'sphere' | 'capsule' | 'mesh';
+
+export interface ScenarioCameraConfig {
+  jitter?: number;
+  exposure?: number;
+  noise?: number;
+  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. */
+  format?: SplatFormat;
+  /** Optional renderer hint. The library does not import renderer-specific code. */
+  renderer?: SplatRendererKind;
+}
+
+export interface SplatScenarioConfig {
+  enabled: boolean;
+  /** Common browser-friendly splat format. Renderer-specific loaders may accept more. */
+  format?: SplatFormat;
+  src?: string;
+  requiresCollisionProxy?: boolean;
+  collisionProxy?: SplatCollisionProxyConfig | null;
+}
+
+export interface SplatCollisionProxyConfig {
+  /** MJCF/XML file or artifact path that provides physics collision for the visual splat. */
+  xmlPath?: string;
+  /** Human-readable status for authoring and validation flows. */
+  status?: 'missing' | 'planned' | 'generated' | 'validated';
+  /** Primitive proxy shapes expected in the MJCF collision proxy. */
+  primitives?: SplatCollisionPrimitive[];
+  /** Optional notes that should travel with scene variants and rollout metadata. */
+  notes?: string[];
+}
+
+export interface PairedSplatEnvironmentConfig {
+  id: string;
+  label: string;
+  description?: string;
+  /** Visual-only Gaussian splat asset. */
+  splat: SplatAssetConfig;
+  /** MJCF/XML contact geometry paired with the visual splat. */
+  collisionProxy: SplatCollisionProxyConfig & { xmlPath: string };
+}
+
+export interface SplatEnvironmentMetadataInput {
+  environment?: PairedSplatEnvironmentConfig;
+  scenario?: VisualScenarioConfig;
+  renderer?: SplatRendererKind;
+  src?: string;
+  format?: SplatFormat;
+  collisionProxy?: SplatCollisionProxyConfig;
+}
+
+export interface SplatEnvironmentMetadata {
+  src?: string;
+  format: SplatFormat;
+  collisionProxy?: SplatCollisionProxyConfig;
+  userData: Record<string, unknown>;
+}
+
+export type SplatSceneInput =
+  | PairedSplatEnvironmentConfig
+  | VisualScenarioConfig
+  | undefined
+  | null;
+
+export interface VisualScenarioConfig {
+  id?: string;
+  label?: string;
+  seed?: number;
+  lighting?: ScenarioLightingPreset;
+  environment?: string;
+  camera?: ScenarioCameraConfig;
+  materials?: ScenarioMaterialConfig;
+  splat?: SplatScenarioConfig | null;
+}
+
+export interface ScenarioLightingProps {
+  preset?: ScenarioLightingPreset;
+  intensity?: number;
+  castShadow?: boolean;
+}
+
+export interface SplatEnvironmentProps extends Omit<ThreeElements['group'], 'ref'> {
+  environment?: PairedSplatEnvironmentConfig;
+  scenario?: VisualScenarioConfig;
+  renderer?: SplatRendererKind;
+  src?: string;
+  format?: SplatFormat;
+  collisionProxy?: ReactNode;
+  collisionProxyMetadata?: SplatCollisionProxyConfig;
+  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 {
@@ -829,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;
 }
 
@@ -926,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;

package/src/vite.ts

@@ -86,6 +86,14 @@ export function mujocoReact(options: MujocoReactPluginOptions) {
   return {
     name: 'mujoco-react',
     enforce: 'pre' as const,
+    config(userConfig: { build?: { chunkSizeWarningLimit?: number } }) {
+      // three + drei + MuJoCo WASM glue are inherently large; the large-chunk
+      // warning is expected, not a failure. Raise the limit so consumers don't
+      // see it. Vite merges plugin config on top of user config, so only set a
+      // default when the consumer hasn't specified their own limit.
+      if (userConfig.build?.chunkSizeWarningLimit !== undefined) return;
+      return { build: { chunkSizeWarningLimit: 4000 } };
+    },
     configResolved(config: ViteConfig) {
       root = config.root;
       generatedRegister = path.resolve(root, generatedRegister);