mujoco-react 9.1.0 → 9.3.0

package/src/rendering/cameraFrameCapture.ts

@@ -0,0 +1,353 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Offscreen camera-frame capture for R3F/MuJoCo scenes.
+ */
+
+import * as THREE from 'three';
+import type {
+  CameraFrameCaptureBlobResult,
+  CameraFrameCaptureOptions,
+  CameraFrameCaptureResult,
+  CameraFrameCaptureSource,
+  CameraFrameCaptureVector3,
+} from '../types';
+
+export interface CameraFrameCaptureSession {
+  readonly width: number;
+  readonly height: number;
+  capture(options?: CameraFrameCaptureOptions): {
+    canvas: HTMLCanvasElement;
+    camera: THREE.Camera;
+    width: number;
+    height: number;
+    source: CameraFrameCaptureSource;
+  };
+  captureDataUrl(options?: CameraFrameCaptureOptions): CameraFrameCaptureResult;
+  captureBlob(options?: CameraFrameCaptureOptions): Promise<CameraFrameCaptureBlobResult>;
+  dispose(): void;
+}
+
+function toVector3(
+  value: CameraFrameCaptureVector3 | undefined,
+  fallback: THREE.Vector3
+): THREE.Vector3 {
+  if (!value) return fallback.clone();
+  return value instanceof THREE.Vector3
+    ? value.clone()
+    : new THREE.Vector3(value[0], value[1], value[2]);
+}
+
+function applyCameraPose(
+  camera: THREE.Camera,
+  options: CameraFrameCaptureOptions,
+  fallbackCamera: THREE.Camera
+) {
+  camera.position.copy(toVector3(options.position, fallbackCamera.position));
+  camera.up.copy(toVector3(options.up, fallbackCamera.up));
+
+  if (options.quaternion) {
+    if (options.quaternion instanceof THREE.Quaternion) {
+      camera.quaternion.copy(options.quaternion);
+    } else {
+      camera.quaternion.set(
+        options.quaternion[0],
+        options.quaternion[1],
+        options.quaternion[2],
+        options.quaternion[3]
+      );
+    }
+  } else if (options.lookAt) {
+    camera.lookAt(toVector3(options.lookAt, new THREE.Vector3()));
+  } else {
+    camera.quaternion.copy(fallbackCamera.quaternion);
+  }
+
+  camera.updateMatrixWorld();
+}
+
+function createCaptureCamera(
+  options: CameraFrameCaptureOptions,
+  fallbackCamera: THREE.Camera,
+  width: number,
+  height: number
+): THREE.Camera {
+  const camera = options.camera
+    ? options.camera.clone()
+    : fallbackCamera instanceof THREE.PerspectiveCamera
+      ? fallbackCamera.clone()
+      : new THREE.PerspectiveCamera(45, width / height, 0.01, 100);
+
+  if (camera instanceof THREE.PerspectiveCamera) {
+    camera.aspect = width / height;
+    camera.fov = options.fov ?? camera.fov;
+    camera.near = options.near ?? camera.near;
+    camera.far = options.far ?? camera.far;
+    camera.updateProjectionMatrix();
+  }
+
+  applyCameraPose(camera, options, fallbackCamera);
+  return camera;
+}
+
+function getCaptureDimensions(
+  renderer: THREE.WebGLRenderer,
+  options: CameraFrameCaptureOptions
+) {
+  const width = Math.max(
+    1,
+    Math.floor(options.width ?? renderer.domElement.width)
+  );
+  const height = Math.max(
+    1,
+    Math.floor(options.height ?? renderer.domElement.height)
+  );
+  return { width, height };
+}
+
+function prepareCaptureCamera(
+  camera: THREE.Camera,
+  options: CameraFrameCaptureOptions,
+  fallbackCamera: THREE.Camera,
+  width: number,
+  height: number
+) {
+  if (options.camera) {
+    camera.copy(options.camera);
+  }
+
+  if (camera instanceof THREE.PerspectiveCamera) {
+    camera.aspect = width / height;
+    camera.fov = options.fov ?? camera.fov;
+    camera.near = options.near ?? camera.near;
+    camera.far = options.far ?? camera.far;
+    camera.updateProjectionMatrix();
+  }
+
+  applyCameraPose(camera, options, fallbackCamera);
+}
+
+function readRenderTargetToCanvas(
+  renderer: THREE.WebGLRenderer,
+  target: THREE.WebGLRenderTarget,
+  canvas: HTMLCanvasElement,
+  context: CanvasRenderingContext2D,
+  pixels: Uint8Array,
+  imageData: ImageData,
+  width: number,
+  height: number
+) {
+  renderer.readRenderTargetPixels(target, 0, 0, width, height, pixels);
+
+  const rowBytes = width * 4;
+  for (let y = 0; y < height; y += 1) {
+    const sourceStart = (height - y - 1) * rowBytes;
+    const targetStart = y * rowBytes;
+    imageData.data.set(
+      pixels.subarray(sourceStart, sourceStart + rowBytes),
+      targetStart
+    );
+  }
+  context.putImageData(imageData, 0, 0);
+  return canvas;
+}
+
+function getCameraFrameCaptureSource(
+  options: CameraFrameCaptureOptions
+): CameraFrameCaptureSource {
+  if (options.source) return options.source;
+  if (options.cameraName) {
+    return { kind: 'mujoco-camera', cameraName: options.cameraName };
+  }
+  if (options.siteName) {
+    return { kind: 'mujoco-site', siteName: options.siteName };
+  }
+  if (options.bodyName) {
+    return { kind: 'mujoco-body', bodyName: options.bodyName };
+  }
+  if (options.camera) return { kind: 'custom-camera' };
+  if (options.position || options.lookAt || options.quaternion) {
+    return { kind: 'explicit-pose' };
+  }
+  return { kind: 'fallback-camera' };
+}
+
+export function createCameraFrameCaptureSession(
+  renderer: THREE.WebGLRenderer,
+  scene: THREE.Scene,
+  fallbackCamera: THREE.Camera,
+  options: CameraFrameCaptureOptions = {}
+): CameraFrameCaptureSession {
+  const { width, height } = getCaptureDimensions(renderer, options);
+  const camera = createCaptureCamera(options, fallbackCamera, width, height);
+  const target = new THREE.WebGLRenderTarget(width, height, {
+    format: THREE.RGBAFormat,
+    type: THREE.UnsignedByteType,
+  });
+  const canvas = document.createElement('canvas');
+  canvas.width = width;
+  canvas.height = height;
+  const context = canvas.getContext('2d');
+  if (!context) {
+    target.dispose();
+    throw new Error('Unable to create a 2D canvas for camera frame capture.');
+  }
+  const drawContext = context;
+
+  const pixels = new Uint8Array(width * height * 4);
+  const imageData = drawContext.createImageData(width, height);
+
+  function capture(nextOptions: CameraFrameCaptureOptions = {}) {
+    const captureOptions = { ...options, ...nextOptions };
+    const nextDimensions = getCaptureDimensions(renderer, captureOptions);
+    if (
+      nextDimensions.width !== width ||
+      nextDimensions.height !== height
+    ) {
+      throw new Error(
+        'Camera frame capture sessions require stable width and height.'
+      );
+    }
+
+    prepareCaptureCamera(
+      camera,
+      captureOptions,
+      fallbackCamera,
+      width,
+      height
+    );
+
+    const previousTarget = renderer.getRenderTarget();
+    const previousXrEnabled = renderer.xr.enabled;
+
+    scene.updateMatrixWorld(true);
+    try {
+      renderer.xr.enabled = false;
+      renderer.setRenderTarget(target);
+      renderer.clear();
+      renderer.render(scene, camera);
+      readRenderTargetToCanvas(
+        renderer,
+        target,
+        canvas,
+        drawContext,
+        pixels,
+        imageData,
+        width,
+        height
+      );
+      return {
+        canvas,
+        camera,
+        width,
+        height,
+        source: getCameraFrameCaptureSource(captureOptions),
+      };
+    } finally {
+      renderer.setRenderTarget(previousTarget);
+      renderer.xr.enabled = previousXrEnabled;
+    }
+  }
+
+  return {
+    width,
+    height,
+    capture,
+    captureDataUrl(nextOptions = {}) {
+      const type = nextOptions.type ?? options.type ?? 'image/png';
+      const result = capture(nextOptions);
+      return {
+        ...result,
+        dataUrl: result.canvas.toDataURL(
+          type,
+          nextOptions.quality ?? options.quality
+        ),
+        type,
+      };
+    },
+    async captureBlob(nextOptions = {}) {
+      const type = nextOptions.type ?? options.type ?? 'image/png';
+      const result = capture(nextOptions);
+      const blob = await new Promise<Blob>((resolve, reject) => {
+        result.canvas.toBlob(
+          (nextBlob) => {
+            if (nextBlob) resolve(nextBlob);
+            else reject(new Error('Camera frame capture did not produce a Blob.'));
+          },
+          type,
+          nextOptions.quality ?? options.quality
+        );
+      });
+      return { ...result, blob, type };
+    },
+    dispose() {
+      target.dispose();
+    },
+  };
+}
+
+export function renderCameraFrameToCanvas(
+  renderer: THREE.WebGLRenderer,
+  scene: THREE.Scene,
+  fallbackCamera: THREE.Camera,
+  options: CameraFrameCaptureOptions = {}
+) {
+  const session = createCameraFrameCaptureSession(
+    renderer,
+    scene,
+    fallbackCamera,
+    options
+  );
+  try {
+    return session.capture();
+  } finally {
+    session.dispose();
+  }
+}
+
+export async function captureCameraFrame(
+  renderer: THREE.WebGLRenderer,
+  scene: THREE.Scene,
+  fallbackCamera: THREE.Camera,
+  options: CameraFrameCaptureOptions = {}
+): Promise<CameraFrameCaptureResult> {
+  const type = options.type ?? 'image/png';
+  const result = renderCameraFrameToCanvas(
+    renderer,
+    scene,
+    fallbackCamera,
+    options
+  );
+  return {
+    ...result,
+    dataUrl: result.canvas.toDataURL(type, options.quality),
+    type,
+  };
+}
+
+export async function captureCameraFrameBlob(
+  renderer: THREE.WebGLRenderer,
+  scene: THREE.Scene,
+  fallbackCamera: THREE.Camera,
+  options: CameraFrameCaptureOptions = {}
+): Promise<CameraFrameCaptureBlobResult> {
+  const type = options.type ?? 'image/png';
+  const result = renderCameraFrameToCanvas(
+    renderer,
+    scene,
+    fallbackCamera,
+    options
+  );
+  const blob = await new Promise<Blob>((resolve, reject) => {
+    result.canvas.toBlob(
+      (nextBlob) => {
+        if (nextBlob) resolve(nextBlob);
+        else reject(new Error('Camera frame capture did not produce a Blob.'));
+      },
+      type,
+      options.quality
+    );
+  });
+  return { ...result, blob, type };
+}

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,
+  };
+}