mujoco-react 9.2.0 → 9.3.0

package/src/rendering/cameraFrameSource.ts

@@ -0,0 +1,375 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Helpers for resolving dataset camera streams to mounted MuJoCo resources.
+ */
+
+import type {
+  Bodies,
+  CameraFrameCaptureOptions,
+  CameraFrameCaptureSource,
+  CameraFrameSequenceCamera,
+  CameraFrameSequenceOptions,
+  CameraFrameSequenceResult,
+  Cameras,
+  MujocoSimAPI,
+  Sites,
+} from '../types';
+
+export type MountedCameraFrameCaptureSource = Extract<
+  CameraFrameCaptureSource,
+  | { kind: 'mujoco-camera' }
+  | { kind: 'mujoco-site' }
+  | { kind: 'mujoco-body' }
+>;
+
+export type CameraFrameMountSelector =
+  | { cameraName: Cameras; siteName?: never; bodyName?: never }
+  | { siteName: Sites; cameraName?: never; bodyName?: never }
+  | { bodyName: Bodies; cameraName?: never; siteName?: never };
+
+export interface NamedCameraFrameResource {
+  name: string | null | undefined;
+}
+
+export interface ResolveMountedCameraFrameSourceOptions {
+  cameras?: readonly (Cameras | NamedCameraFrameResource | null | undefined)[];
+  sites?: readonly (Sites | NamedCameraFrameResource | null | undefined)[];
+  bodies?: readonly (Bodies | NamedCameraFrameResource | null | undefined)[];
+  aliases?: Record<
+    string,
+    CameraFrameMountSelector | readonly CameraFrameMountSelector[]
+  >;
+  /**
+   * Accept the first valid alias selector even when the current resource
+   * inventory cannot verify it. This is useful when aliases come from a
+   * previously validated model inventory and the actual provider will validate
+   * again during capture.
+   */
+  allowAliasFallback?: boolean;
+}
+
+export interface ResolvedMountedCameraFrameSource {
+  key: string;
+  selector: CameraFrameMountSelector;
+  source: MountedCameraFrameCaptureSource;
+}
+
+export type MountedCameraFrameSequenceDefaults = Omit<
+  CameraFrameSequenceCamera,
+  'key' | 'cameraName' | 'siteName' | 'bodyName' | 'source'
+>;
+
+export type MountedCameraFrameSequenceCameraOptions = Partial<
+  MountedCameraFrameSequenceDefaults
+>;
+
+export interface CreateMountedCameraFrameSequencePlanOptions
+  extends ResolveMountedCameraFrameSourceOptions {
+  defaults?: MountedCameraFrameSequenceDefaults;
+  cameraOptions?: Record<string, MountedCameraFrameSequenceCameraOptions>;
+  requireAll?: boolean;
+}
+
+export interface MountedCameraFrameSequencePlan {
+  cameraKeys: string[];
+  cameras: CameraFrameSequenceCamera[];
+  resolved: Record<string, ResolvedMountedCameraFrameSource>;
+  missingKeys: string[];
+}
+
+export const MountedCameraFrameSequenceReadinessStatus = {
+  Ready: 'ready',
+  Partial: 'partial',
+  Missing: 'missing',
+} as const;
+
+export type MountedCameraFrameSequenceReadinessStatus =
+  (typeof MountedCameraFrameSequenceReadinessStatus)[keyof typeof MountedCameraFrameSequenceReadinessStatus];
+
+export interface MountedCameraFrameSequenceSourceReadiness {
+  key: string;
+  ready: boolean;
+  selector?: CameraFrameMountSelector;
+  source?: MountedCameraFrameCaptureSource;
+  message: string;
+}
+
+export interface MountedCameraFrameSequenceReadiness {
+  ready: boolean;
+  status: MountedCameraFrameSequenceReadinessStatus;
+  cameraKeys: string[];
+  resolvedKeys: string[];
+  missingKeys: string[];
+  cameras: Record<string, MountedCameraFrameSequenceSourceReadiness>;
+  message: string;
+}
+
+export type MountedCameraFrameSequencePlanOptions = Omit<
+  CreateMountedCameraFrameSequencePlanOptions,
+  'cameras' | 'sites' | 'bodies'
+>;
+
+export interface MountedCameraFrameSequenceRecordOptions
+  extends Omit<CameraFrameSequenceOptions, 'cameras'>,
+    MountedCameraFrameSequencePlanOptions {
+  cameraKeys: readonly string[];
+}
+
+export interface MountedCameraFrameSequenceRecordResult
+  extends CameraFrameSequenceResult {
+  plan: MountedCameraFrameSequencePlan;
+  readiness: MountedCameraFrameSequenceReadiness;
+}
+
+export type MountedCameraFrameSequenceRecorderTarget = Pick<
+  MujocoSimAPI,
+  'getCameras' | 'getSites' | 'getBodies' | 'recordCameraSequence'
+>;
+
+function getResourceName(
+  resource: string | NamedCameraFrameResource | null | undefined
+) {
+  if (!resource) return null;
+  return typeof resource === 'string' ? resource : resource.name ?? null;
+}
+
+function createNameSet(
+  resources:
+    | readonly (string | NamedCameraFrameResource | null | undefined)[]
+    | undefined
+) {
+  return new Set(
+    (resources ?? [])
+      .map((resource) => getResourceName(resource))
+      .filter((name): name is string => Boolean(name))
+  );
+}
+
+function normalizeAliasCandidates(
+  value: CameraFrameMountSelector | readonly CameraFrameMountSelector[] | undefined
+) {
+  if (!value) return [];
+  return Array.isArray(value) ? value : [value];
+}
+
+function countMountedSelectors(selector: CameraFrameMountSelector) {
+  return Number(Boolean(selector.cameraName)) +
+    Number(Boolean(selector.siteName)) +
+    Number(Boolean(selector.bodyName));
+}
+
+export function getMountedCameraFrameCaptureSource(
+  selector: CameraFrameMountSelector
+): MountedCameraFrameCaptureSource | null {
+  if (countMountedSelectors(selector) !== 1) return null;
+  if (selector.cameraName) {
+    return { kind: 'mujoco-camera', cameraName: selector.cameraName };
+  }
+  if (selector.siteName) {
+    return { kind: 'mujoco-site', siteName: selector.siteName };
+  }
+  if (selector.bodyName) {
+    return { kind: 'mujoco-body', bodyName: selector.bodyName };
+  }
+  return null;
+}
+
+export function isMountedCameraFrameCaptureSource(
+  source: CameraFrameCaptureSource
+): source is MountedCameraFrameCaptureSource {
+  return (
+    source.kind === 'mujoco-camera' ||
+    source.kind === 'mujoco-site' ||
+    source.kind === 'mujoco-body'
+  );
+}
+
+export function getCameraFrameCaptureSourceTarget(
+  source: CameraFrameCaptureSource
+) {
+  if (source.kind === 'mujoco-camera') return source.cameraName;
+  if (source.kind === 'mujoco-site') return source.siteName;
+  if (source.kind === 'mujoco-body') return source.bodyName;
+  if (source.kind === 'custom-camera') return 'custom camera';
+  if (source.kind === 'explicit-pose') return 'explicit pose';
+  return 'fallback camera';
+}
+
+function isSelectorMounted(
+  selector: CameraFrameMountSelector,
+  cameraNames: Set<string>,
+  siteNames: Set<string>,
+  bodyNames: Set<string>
+) {
+  if (countMountedSelectors(selector) !== 1) return false;
+  return (
+    (selector.cameraName ? cameraNames.has(selector.cameraName) : false) ||
+    (selector.siteName ? siteNames.has(selector.siteName) : false) ||
+    (selector.bodyName ? bodyNames.has(selector.bodyName) : false)
+  );
+}
+
+export function resolveMountedCameraFrameSource(
+  key: string,
+  options: ResolveMountedCameraFrameSourceOptions
+): ResolvedMountedCameraFrameSource | null {
+  const cameraNames = createNameSet(options.cameras);
+  const siteNames = createNameSet(options.sites);
+  const bodyNames = createNameSet(options.bodies);
+  const directCandidates: CameraFrameMountSelector[] = [
+    { cameraName: key },
+    { siteName: key },
+    { bodyName: key },
+  ];
+  const aliasCandidates = normalizeAliasCandidates(options.aliases?.[key]);
+  const candidates = [...directCandidates, ...aliasCandidates];
+
+  for (const selector of candidates) {
+    if (!isSelectorMounted(selector, cameraNames, siteNames, bodyNames)) {
+      continue;
+    }
+    const source = getMountedCameraFrameCaptureSource(selector);
+    if (!source) continue;
+    return { key, selector, source };
+  }
+
+  if (options.allowAliasFallback) {
+    for (const selector of aliasCandidates) {
+      const source = getMountedCameraFrameCaptureSource(selector);
+      if (!source) continue;
+      return { key, selector, source };
+    }
+  }
+
+  return null;
+}
+
+export function createMountedCameraFrameSequencePlan(
+  cameraKeys: readonly string[],
+  options: CreateMountedCameraFrameSequencePlanOptions
+): MountedCameraFrameSequencePlan {
+  const cameras: CameraFrameSequenceCamera[] = [];
+  const resolved: Record<string, ResolvedMountedCameraFrameSource> = {};
+  const missingKeys: string[] = [];
+
+  for (const key of cameraKeys) {
+    const mountedSource = resolveMountedCameraFrameSource(key, options);
+    if (!mountedSource) {
+      missingKeys.push(key);
+      continue;
+    }
+
+    resolved[key] = mountedSource;
+    cameras.push({
+      key,
+      ...options.defaults,
+      ...options.cameraOptions?.[key],
+      ...mountedSource.selector,
+      source: mountedSource.source,
+    });
+  }
+
+  if (options.requireAll && missingKeys.length > 0) {
+    throw new Error(
+      `Unable to resolve mounted MuJoCo camera source${
+        missingKeys.length === 1 ? '' : 's'
+      } for ${missingKeys.join(', ')}.`
+    );
+  }
+
+  return {
+    cameraKeys: [...cameraKeys],
+    cameras,
+    resolved,
+    missingKeys,
+  };
+}
+
+export function createMountedCameraFrameSequenceReadiness(
+  plan: MountedCameraFrameSequencePlan
+): MountedCameraFrameSequenceReadiness {
+  const cameras: Record<string, MountedCameraFrameSequenceSourceReadiness> = {};
+  const resolvedKeys = plan.cameraKeys.filter((key) => Boolean(plan.resolved[key]));
+
+  for (const key of plan.cameraKeys) {
+    const resolved = plan.resolved[key];
+    cameras[key] = resolved
+      ? {
+          key,
+          ready: true,
+          selector: resolved.selector,
+          source: resolved.source,
+          message: `Camera stream "${key}" resolves to ${resolved.source.kind}:${getCameraFrameCaptureSourceTarget(resolved.source)}.`,
+        }
+      : {
+          key,
+          ready: false,
+          message: `Camera stream "${key}" does not resolve to a mounted MuJoCo camera, site, or body.`,
+        };
+  }
+
+  const missingKeys = [...plan.missingKeys];
+  const ready = missingKeys.length === 0;
+  const status: MountedCameraFrameSequenceReadinessStatus = ready
+    ? MountedCameraFrameSequenceReadinessStatus.Ready
+    : resolvedKeys.length > 0
+      ? MountedCameraFrameSequenceReadinessStatus.Partial
+      : MountedCameraFrameSequenceReadinessStatus.Missing;
+
+  return {
+    ready,
+    status,
+    cameraKeys: [...plan.cameraKeys],
+    resolvedKeys,
+    missingKeys,
+    cameras,
+    message: ready
+      ? `All ${plan.cameraKeys.length} requested camera stream${
+          plan.cameraKeys.length === 1 ? '' : 's'
+        } resolve to mounted MuJoCo sources.`
+      : `Missing mounted MuJoCo source${
+          missingKeys.length === 1 ? '' : 's'
+        } for ${missingKeys.join(', ')}.`,
+  };
+}
+
+export function createMountedCameraFrameSequencePlanFromApi(
+  api: MountedCameraFrameSequenceRecorderTarget,
+  cameraKeys: readonly string[],
+  options: MountedCameraFrameSequencePlanOptions = {}
+): MountedCameraFrameSequencePlan {
+  return createMountedCameraFrameSequencePlan(cameraKeys, {
+    ...options,
+    cameras: api.getCameras(),
+    sites: api.getSites(),
+    bodies: api.getBodies(),
+  });
+}
+
+export async function recordMountedCameraFrameSequence(
+  api: MountedCameraFrameSequenceRecorderTarget,
+  options: MountedCameraFrameSequenceRecordOptions
+): Promise<MountedCameraFrameSequenceRecordResult> {
+  const { cameraKeys, ...restOptions } = options;
+  const requireAll =
+    restOptions.requireAll ?? restOptions.requireMountedSources ?? true;
+  const plan = createMountedCameraFrameSequencePlanFromApi(
+    api,
+    cameraKeys,
+    { ...restOptions, requireAll }
+  );
+  const readiness = createMountedCameraFrameSequenceReadiness(plan);
+  const result = await api.recordCameraSequence({
+    ...restOptions,
+    cameras: plan.cameras,
+    requireMountedSources: restOptions.requireMountedSources ?? true,
+  });
+
+  return {
+    ...result,
+    plan,
+    readiness,
+  };
+}

package/src/spark.tsx

@@ -15,9 +15,14 @@ import * as THREE from 'three';
 import {
   SplatEnvironment,
   useSplatEnvironment,
+  useSplatSceneConfig,
 } from './components/VisualScenario';
 import type {
+  PairedSplatEnvironmentConfig,
+  SceneConfig,
   SplatEnvironmentProps,
+  SplatEnvironmentReadiness,
+  VisualScenarioConfig,
 } from './types';
 
 type SparkModule = typeof import('@sparkjsdev/spark');
@@ -26,9 +31,22 @@ type SparkSplatMeshInstance = InstanceType<SparkModule['SplatMesh']>;
 type SparkDisposable = {
   dispose?: () => unknown;
 };
+type SparkWorkerMessage = {
+  reject?: (error: unknown) => void;
+};
+type SparkWorkerLike = {
+  messages?: Record<string, SparkWorkerMessage>;
+};
+type SparkResourceWithWorkers = SparkDisposable & {
+  worker?: SparkWorkerLike;
+  sortWorker?: SparkWorkerLike;
+  lodWorker?: SparkWorkerLike;
+};
 
 export type SparkSplatStatus = 'idle' | 'loading' | 'ready' | 'error';
 
+let sparkDisposeRejectionHandlerRegistered = false;
+
 export interface SparkSplatLifecycle {
   status: SparkSplatStatus;
   error: Error | null;
@@ -39,6 +57,18 @@ export interface SparkSplatLifecycle {
   reset: () => void;
 }
 
+export interface SparkSplatEnvironmentState {
+  environment: PairedSplatEnvironmentConfig | undefined;
+  sceneConfig: SceneConfig;
+  readiness: SplatEnvironmentReadiness;
+  lifecycle: SparkSplatLifecycle;
+  props: Pick<
+    SparkSplatEnvironmentProps,
+    'environment' | 'scenario' | 'src' | 'format' | 'onStatusChange' | 'onError'
+  >;
+  enabled: boolean;
+}
+
 export interface SparkSplatEnvironmentProps extends SplatEnvironmentProps {
   /** Enable Spark LoD handling for large splat assets. Default: true. */
   lod?: boolean | 'quality';
@@ -53,6 +83,67 @@ export interface SparkSplatEnvironmentProps extends SplatEnvironmentProps {
   onError?: (error: Error) => void;
 }
 
+/**
+ * Resolve a visual scenario's paired splat environment, compose its MJCF
+ * collision proxy into the MuJoCo scene config, and expose Spark lifecycle
+ * props for `<SparkSplatEnvironment />`.
+ */
+export function useSparkSplatEnvironment({
+  sceneConfig,
+  scenario,
+  environment,
+  enabled = true,
+  renderer = 'spark',
+  onError,
+  onStatusChange,
+}: {
+  sceneConfig: SceneConfig;
+  scenario?: VisualScenarioConfig;
+  environment?: PairedSplatEnvironmentConfig;
+  enabled?: boolean;
+  renderer?: 'spark';
+  onError?: (error: Error) => void;
+  onStatusChange?: (status: SparkSplatStatus) => void;
+}): SparkSplatEnvironmentState {
+  const splatScene = useSplatSceneConfig({
+    sceneConfig,
+    scenario,
+    environment,
+    enabled,
+    renderer,
+  });
+  const metadata = useSplatEnvironment({
+    scenario,
+    environment: splatScene.environment,
+    renderer,
+  });
+  const renderEnabled = enabled && Boolean(metadata.src);
+  const readiness = enabled ? metadata.readiness : splatScene.readiness;
+  const lifecycle = useSparkSplatLifecycle({
+    enabled: renderEnabled,
+    onError,
+    onStatusChange,
+  });
+
+  return useMemo(
+    () => ({
+      environment: splatScene.environment,
+      sceneConfig: splatScene.sceneConfig,
+      readiness,
+      lifecycle,
+      props: {
+        environment: splatScene.environment,
+        scenario: enabled ? scenario : undefined,
+        src: enabled ? metadata.src : undefined,
+        format: metadata.format,
+        ...lifecycle.props,
+      },
+      enabled: renderEnabled,
+    }),
+    [enabled, lifecycle, metadata, readiness, renderEnabled, scenario, splatScene]
+  );
+}
+
 /**
  * Tracks Spark 3DGS loading state for UI that wraps `SparkSplatEnvironment`.
  *
@@ -177,6 +268,7 @@ export function SparkSplatEnvironment({
 
   useEffect(() => {
     let disposed = false;
+    ensureSparkDisposeRejectionHandler();
 
     function setLifecycleStatus(nextStatus: SparkSplatStatus) {
       setStatus(nextStatus);
@@ -306,6 +398,7 @@ export function SparkSplatEnvironment({
 
 function safelyDisposeSparkResource(resource: SparkDisposable) {
   try {
+    silenceSparkWorkerTerminateRejections(resource);
     const result = resource.dispose?.();
     if (isPromiseLike(result)) {
       void Promise.resolve(result).catch(handleSparkDisposeError);
@@ -324,6 +417,33 @@ function isPromiseLike(value: unknown): value is PromiseLike<unknown> {
   );
 }
 
+function silenceSparkWorkerTerminateRejections(resource: SparkDisposable) {
+  const workers = getSparkWorkers(resource);
+  for (const worker of workers) {
+    if (!worker.messages) continue;
+
+    for (const message of Object.values(worker.messages)) {
+      const reject = message.reject;
+      if (!reject) continue;
+
+      message.reject = (error: unknown) => {
+        if (!isSparkWorkerTerminateError(error)) {
+          reject(error);
+        }
+      };
+    }
+  }
+}
+
+function getSparkWorkers(resource: SparkDisposable): SparkWorkerLike[] {
+  const sparkResource = resource as SparkResourceWithWorkers;
+  return [
+    sparkResource.worker,
+    sparkResource.sortWorker,
+    sparkResource.lodWorker,
+  ].filter((worker): worker is SparkWorkerLike => Boolean(worker));
+}
+
 function handleSparkDisposeError(error: unknown) {
   if (
     error instanceof Error &&
@@ -334,3 +454,27 @@ function handleSparkDisposeError(error: unknown) {
 
   console.warn('[mujoco-react] Spark resource disposal failed.', error);
 }
+
+function ensureSparkDisposeRejectionHandler() {
+  if (
+    sparkDisposeRejectionHandlerRegistered ||
+    typeof window === 'undefined' ||
+    typeof window.addEventListener !== 'function'
+  ) {
+    return;
+  }
+
+  sparkDisposeRejectionHandlerRegistered = true;
+  window.addEventListener('unhandledrejection', (event) => {
+    if (isSparkWorkerTerminateError(event.reason)) {
+      event.preventDefault();
+    }
+  });
+}
+
+function isSparkWorkerTerminateError(reason: unknown) {
+  return (
+    reason instanceof Error &&
+    reason.message.toLowerCase().includes('worker terminate')
+  );
+}