mujoco-react 12.1.1 → 12.2.0

package/src/rendering/cameraFrameCapture.ts

@@ -6,6 +6,7 @@
  */
 
 import * as THREE from 'three';
+import { createMetricDepthReader } from './metricDepth';
 import type {
   CameraFrameCaptureBlobResult,
   CameraFrameCaptureOptions,
@@ -32,6 +33,26 @@ export interface CameraFramePixelsResult {
   source: CameraFrameCaptureSource;
 }
 
+/** An owned, aligned RGB-D observation from one render. Scene units must be meters. */
+export interface CameraRGBDResult {
+  /** Interleaved sRGB bytes (RGB), row-major, top-left origin. */
+  rgb: Uint8Array;
+  /** Positive camera-axis distance (-view Z) in meters, not Euclidean range.
+   * Background / no depth-writing surface is +Infinity. Same pixel order as rgb. */
+  depth: Float32Array;
+  width: number;
+  height: number;
+  source: CameraFrameCaptureSource;
+  origin: 'top-left';
+  depthUnits: 'meters';
+  depthConvention: 'camera-z';
+  /** Horizontal mirroring applied to both arrays. Undo before unprojecting. */
+  flipX: boolean;
+  /** Owned column-major matrices. Camera coordinates: +X right, +Y up, -Z forward. */
+  projectionMatrix: number[];
+  cameraWorldMatrix: number[];
+}
+
 export interface CameraFrameTensorResult extends PolicyImageTensorResult {
   camera: THREE.Camera;
   source: CameraFrameCaptureSource;
@@ -67,6 +88,8 @@ export interface CameraFrameCaptureSession {
   capturePixels(options?: CameraFrameCaptureOptions): CameraFramePixelsResult;
   /** Render straight into a normalized policy image tensor (no canvas/PNG encode). */
   captureTensor(options?: CameraFrameCaptureTensorOptions): CameraFrameTensorResult;
+  /** Synchronous RGB + metric depth from one render. Returned arrays are owned. */
+  captureRGBD(options?: CameraFrameCaptureOptions): CameraRGBDResult;
   dispose(): void;
 }
 
@@ -78,6 +101,12 @@ export const CAPTURE_EXCLUDE_KEY =
   'mujoco.capture.exclude';
 
 export type CameraFrameCaptureRenderInput = {
+  renderMode: 'snapshot' | 'preview';
+  /** Apply capture visibility/visual overrides only for a synchronous render.
+   * Async adapters must wrap each render after awaiting preparation, and await
+   * readback outside this scope so the interactive viewport stays untouched.
+   */
+  withCaptureState: (render: () => void) => void;
   renderer: THREE.WebGLRenderer;
   scene: THREE.Scene;
   camera: THREE.Camera;
@@ -300,6 +329,11 @@ export function createCaptureCamera(
     camera.far = options.far ?? camera.far;
     camera.updateProjectionMatrix();
   }
+  if (camera instanceof THREE.OrthographicCamera) {
+    camera.near = options.near ?? camera.near;
+    camera.far = options.far ?? camera.far;
+    camera.updateProjectionMatrix();
+  }
   applyProjectionMatrix(camera, options.projectionMatrix);
 
   applyCameraPose(camera, options, fallbackCamera);
@@ -339,6 +373,11 @@ export function prepareCaptureCamera(
     camera.far = options.far ?? camera.far;
     camera.updateProjectionMatrix();
   }
+  if (camera instanceof THREE.OrthographicCamera) {
+    camera.near = options.near ?? camera.near;
+    camera.far = options.far ?? camera.far;
+    camera.updateProjectionMatrix();
+  }
   applyProjectionMatrix(camera, options.projectionMatrix);
 
   applyCameraPose(camera, options, fallbackCamera);
@@ -647,12 +686,14 @@ export function createCameraFrameCaptureSession(
 
   const pixels = new Uint8Array(width * height * 4);
   const imageData = drawContext.createImageData(width, height);
+  let depthReader: ReturnType<typeof createMetricDepthReader> | undefined;
   let inFlight = false, disposed = false;
   function assertAvailable() {
     if (disposed) throw new Error('Camera capture session is disposed');
     if (inFlight) throw new Error('Camera capture session is busy; await the current capture');
   }
   function releaseResources() {
+    depthReader?.dispose();
     target.dispose();
     if (isolatedRenderer && !isolatedRenderer.cached) isolatedRenderer.renderer.dispose();
   }
@@ -688,29 +729,17 @@ export function createCameraFrameCaptureSession(
     return captureOptions;
   }
 
-  function renderCaptureToTarget(
-    captureOptions: CameraFrameCaptureOptions,
-    readback: () => void
-  ) {
+  function withCaptureState(captureOptions: CameraFrameCaptureOptions, render: () => void) {
     const previousState = saveRendererState(sessionRenderer);
-    const previousSceneState = applyCaptureVisualOverrides(
-      sessionRenderer,
-      scene,
-      captureOptions
-    );
+    const previousSceneState = applyCaptureVisualOverrides(sessionRenderer, scene, captureOptions);
     const hidden = [
       ...hideExcludedCaptureObjects(scene),
       ...hideCaptureGeomGroups(scene, captureOptions),
     ];
-
-    runCapturePreRenderHooks(scene);
-    scene.updateMatrixWorld(true);
     try {
+      runCapturePreRenderHooks(scene);
+      scene.updateMatrixWorld(true);
       sessionRenderer.xr.enabled = false;
-      sessionRenderer.setRenderTarget(target);
-      sessionRenderer.setViewport(0, 0, width, height);
-      sessionRenderer.setScissor(0, 0, width, height);
-      sessionRenderer.setScissorTest(false);
       if (captureOptions.background !== undefined) {
         sessionRenderer.setClearColor(
           new THREE.Color(captureOptions.background),
@@ -719,9 +748,7 @@ export function createCameraFrameCaptureSession(
       } else if (captureOptions.backgroundAlpha !== undefined) {
         sessionRenderer.setClearColor(previousState.clearColor, captureOptions.backgroundAlpha);
       }
-      sessionRenderer.clear();
-      sessionRenderer.render(scene, camera);
-      readback();
+      render();
     } finally {
       restoreObjectVisibility(hidden);
       if (previousSceneState) restoreSceneVisualState(scene, previousSceneState);
@@ -729,6 +756,20 @@ export function createCameraFrameCaptureSession(
     }
   }
 
+  function renderCaptureToTarget(
+    captureOptions: CameraFrameCaptureOptions,
+    readback: () => void
+  ) {
+    withCaptureState(captureOptions, () => {
+      sessionRenderer.setRenderTarget(target);
+      // The render target installs its pixel viewport; setters would apply DPR twice.
+      sessionRenderer.setScissorTest(false);
+      sessionRenderer.clear();
+      sessionRenderer.render(scene, camera);
+      readback();
+    });
+  }
+
   function renderPreparedCapture(captureOptions: CameraFrameCaptureOptions) {
     renderCaptureToTarget(captureOptions, () => {
       readRenderTargetToCanvas(
@@ -773,6 +814,51 @@ export function createCameraFrameCaptureSession(
     };
   }
 
+  function captureRGBD(nextOptions: CameraFrameCaptureOptions = {}): CameraRGBDResult {
+    assertAvailable();
+    const captureOptions = resolveCaptureOptions(nextOptions);
+    if (getCaptureRenderer(scene)) {
+      throw new Error('Metric RGB-D capture requires the standard Three renderer; custom capture renderers do not provide aligned depth.');
+    }
+    if (sessionRenderer.capabilities.logarithmicDepthBuffer || sessionRenderer.capabilities.reversedDepthBuffer) {
+      throw new Error('Metric RGB-D capture does not support logarithmic or reversed depth buffers.');
+    }
+    if (!camera.projectionMatrix.elements.every(Number.isFinite) ||
+        !Number.isFinite(camera.projectionMatrix.determinant()) || camera.projectionMatrix.determinant() === 0) {
+      throw new Error('Metric RGB-D capture requires a finite, invertible projection matrix.');
+    }
+    if (!depthReader) {
+      // Recreate the framebuffer if this session has already rendered RGB-only frames.
+      target.dispose();
+      target.depthTexture = new THREE.DepthTexture(width, height, THREE.UnsignedIntType);
+      depthReader = createMetricDepthReader(width, height, target.depthTexture);
+    }
+    const rgb = new Uint8Array(width * height * 3);
+    let depth!: Float32Array;
+    renderCaptureToTarget(captureOptions, () => {
+      sessionRenderer.readRenderTargetPixels(target, 0, 0, width, height, pixels);
+      for (let y = 0; y < height; y++) {
+        for (let x = 0; x < width; x++) {
+          const sx = captureOptions.flipX ? width - x - 1 : x;
+          const src = ((height - y - 1) * width + sx) * 4;
+          const dst = (y * width + x) * 3;
+          for (let c = 0; c < 3; c++) {
+            rgb[dst + c] = linearByteToSrgbByte(pixels[src + c]);
+          }
+        }
+      }
+      depth = depthReader!.read(sessionRenderer, camera, captureOptions.flipX ?? false);
+    });
+    return {
+      rgb, depth, width, height,
+      source: getCameraFrameCaptureSource(captureOptions),
+      origin: 'top-left', depthUnits: 'meters', depthConvention: 'camera-z',
+      flipX: captureOptions.flipX ?? false,
+      projectionMatrix: camera.projectionMatrix.toArray(),
+      cameraWorldMatrix: camera.matrixWorld.toArray(),
+    };
+  }
+
   function captureTensor(
     nextOptions: CameraFrameCaptureTensorOptions = {}
   ): CameraFrameTensorResult {
@@ -802,63 +888,35 @@ export function createCameraFrameCaptureSession(
     scene.updateMatrixWorld(true);
     const captureRenderer = getCaptureRenderer(scene);
     if (captureRenderer) {
-      const previousState = saveRendererState(sessionRenderer);
-      const previousSceneState = applyCaptureVisualOverrides(
-        sessionRenderer,
-        scene,
-        captureOptions
-      );
-      const hidden = [
-        ...hideExcludedCaptureObjects(scene),
-        ...hideCaptureGeomGroups(scene, captureOptions),
-      ];
-      try {
-        sessionRenderer.xr.enabled = false;
-        if (captureOptions.background !== undefined) {
-          sessionRenderer.setClearColor(
-            new THREE.Color(captureOptions.background),
-            captureOptions.backgroundAlpha ?? previousState.clearAlpha
-          );
-        } else if (captureOptions.backgroundAlpha !== undefined) {
-          sessionRenderer.setClearColor(previousState.clearColor, captureOptions.backgroundAlpha);
-        }
-        const captureResult = await captureRenderer({
+      let pending: ReturnType<CameraFrameCaptureRender>;
+      const renderWithCaptureState = (render: () => void) => withCaptureState(captureOptions, render);
+      // Preserve synchronous adapters while restoring shared state before any
+      // asynchronous preparation/readback yields to the viewer's next frame.
+      renderWithCaptureState(() => {
+        pending = captureRenderer({
           renderer: sessionRenderer,
           scene,
           camera,
           target,
           width,
           height,
+          withCaptureState: renderWithCaptureState,
+          renderMode: captureOptions.renderMode ?? 'snapshot',
         });
-        if (captureResult) {
-          const captureWidth = captureResult.width ?? width;
-          const captureHeight = captureResult.height ?? height;
-          if (captureWidth !== width || captureHeight !== height) {
-            throw new Error(
-              'Camera frame capture renderer returned unexpected dimensions.'
-            );
-          }
-          readPixelsToCanvas(
-            captureResult.pixels,
-            drawContext,
-            imageData,
-            width,
-            height,
-            captureResult.flipY ?? true,
-            captureResult.flipX ?? captureOptions.flipX ?? false
-          );
-          return {
-            canvas,
-            camera,
-            width,
-            height,
-            source: getCameraFrameCaptureSource(captureOptions),
-          };
+      });
+      const captureResult = await pending;
+      if (captureResult) {
+        const captureWidth = captureResult.width ?? width;
+        const captureHeight = captureResult.height ?? height;
+        if (captureWidth !== width || captureHeight !== height) {
+          throw new Error('Camera frame capture renderer returned unexpected dimensions.');
         }
-      } finally {
-        restoreObjectVisibility(hidden);
-        if (previousSceneState) restoreSceneVisualState(scene, previousSceneState);
-        restoreRendererState(sessionRenderer, previousState);
+        readPixelsToCanvas(
+          captureResult.pixels, drawContext, imageData, width, height,
+          captureResult.flipY ?? true,
+          captureResult.flipX ?? captureOptions.flipX ?? false
+        );
+        return { canvas, camera, width, height, source: getCameraFrameCaptureSource(captureOptions) };
       }
     }
     let readback: Promise<unknown> | undefined;
@@ -878,6 +936,7 @@ export function createCameraFrameCaptureSession(
     captureAsync,
     capturePixels,
     captureTensor,
+    captureRGBD,
     captureDataUrl(nextOptions = {}) {
       const type = nextOptions.type ?? options.type ?? 'image/png';
       const result = capture(nextOptions);
@@ -1012,3 +1071,17 @@ export function captureCameraFrameTensor(
     session.dispose();
   }
 }
+/** One-shot aligned RGB-D observation. Use a session for repeated observations. */
+export function captureCameraRGBD(
+  renderer: THREE.WebGLRenderer,
+  scene: THREE.Scene,
+  fallbackCamera: THREE.Camera,
+  options: CameraFrameCaptureOptions = {}
+): CameraRGBDResult {
+  const session = createCameraFrameCaptureSession(renderer, scene, fallbackCamera, options);
+  try {
+    return session.captureRGBD();
+  } finally {
+    session.dispose();
+  }
+}

package/src/rendering/metricDepth.ts

@@ -0,0 +1,56 @@
+/** @license SPDX-License-Identifier: Apache-2.0 */
+import * as THREE from 'three';
+
+/** Internal GPU readback adapter. Read the real depth attachment, preserving
+ * alpha tests, displacement, skinning, and depthWrite from the RGB render. */
+export function createMetricDepthReader(width: number, height: number, texture: THREE.DepthTexture) {
+  const target = new THREE.WebGLRenderTarget(width, height, { depthBuffer: false });
+  const material = new THREE.ShaderMaterial({
+    uniforms: { depthTexture: { value: texture } },
+    vertexShader: `varying vec2 vUv;
+      void main() { vUv = uv; gl_Position = vec4(position.xy, 0.0, 1.0); }`,
+    fragmentShader: `#include <packing>
+      uniform sampler2D depthTexture;
+      varying vec2 vUv;
+      void main() { gl_FragColor = packDepthToRGBA(texture2D(depthTexture, vUv).r); }`,
+    depthTest: false, depthWrite: false, blending: THREE.NoBlending,
+    toneMapped: false,
+  });
+  const geometry = new THREE.PlaneGeometry(2, 2);
+  const scene = new THREE.Scene();
+  scene.add(new THREE.Mesh(geometry, material));
+  const screenCamera = new THREE.Camera();
+  const packed = new Uint8Array(width * height * 4);
+  return {
+    read(renderer: THREE.WebGLRenderer, camera: THREE.Camera, flipX: boolean) {
+      renderer.setRenderTarget(target);
+      renderer.clear();
+      renderer.render(scene, screenCamera);
+      renderer.readRenderTargetPixels(target, 0, 0, width, height, packed);
+      const result = new Float32Array(width * height);
+      const inverse = camera.projectionMatrixInverse.elements;
+      for (let y = 0; y < height; y++) {
+        for (let x = 0; x < width; x++) {
+          const sx = flipX ? width - x - 1 : x;
+          const sy = height - y - 1;
+          const offset = (sy * width + sx) * 4;
+          const r = packed[offset], g = packed[offset + 1], b = packed[offset + 2], a = packed[offset + 3];
+          if (r === 255 && g === 255 && b === 255 && a === 255) {
+            result[y * width + x] = Infinity;
+            continue;
+          }
+          // Mirrors Three's unpackRGBAToDepth (packing.glsl.js).
+          const d = r / 256 + g / 65536 + b / 16777216 + a / (255 * 16777216);
+          const nx = (sx + 0.5) / width * 2 - 1;
+          const ny = (sy + 0.5) / height * 2 - 1;
+          const nz = d * 2 - 1;
+          const z = inverse[2] * nx + inverse[6] * ny + inverse[10] * nz + inverse[14];
+          const w = inverse[3] * nx + inverse[7] * ny + inverse[11] * nz + inverse[15];
+          result[y * width + x] = -z / w;
+        }
+      }
+      return result;
+    },
+    dispose() { target.dispose(); material.dispose(); geometry.dispose(); },
+  };
+}

package/src/spark.tsx

@@ -373,7 +373,7 @@ export function SparkSplatEnvironment({
 }: SparkSplatEnvironmentProps) {
   const groupRef = useRef<THREE.Group>(null);
   const sparkRef = useRef<SparkRendererInstance | null>(null);
-  const captureSparkRef = useRef<SparkCaptureRenderer | null>(null);
+  const captureSparkRef = useRef(new Map<'snapshot' | 'preview', SparkCaptureRenderer>());
   const meshRef = useRef<SparkSplatMeshInstance | null>(null);
   const hiddenMeshesRef = useRef<THREE.Mesh[]>([]);
   const onStatusChangeRef = useRef(onStatusChange);
@@ -434,8 +434,10 @@ export function SparkSplatEnvironment({
     }
 
     function disposeCaptureSpark() {
-      safelyDisposeSparkResource(captureSparkRef.current?.renderer);
-      captureSparkRef.current = null;
+      for (const capture of captureSparkRef.current.values()) {
+        safelyDisposeSparkResource(capture.renderer);
+      }
+      captureSparkRef.current.clear();
     }
 
     async function loadSplat() {
@@ -465,25 +467,32 @@ export function SparkSplatEnvironment({
           ...resolvedRenderTuning,
         });
         spark.userData[CAMERA_FRAME_CAPTURE_RENDER_USER_DATA_KEY] = async ({
+          withCaptureState,
+          renderMode,
           scene,
           camera,
           width,
           height,
         }: CameraFrameCaptureRenderInput): Promise<CameraFrameCaptureRenderResult> => {
-          let captureSpark = captureSparkRef.current;
+          let captureSpark = captureSparkRef.current.get(renderMode);
           if (
             !captureSpark ||
             captureSpark.width !== width ||
             captureSpark.height !== height
           ) {
-            disposeCaptureSpark();
+            safelyDisposeSparkResource(captureSpark?.renderer);
             captureSpark = {
               renderer: new sparkModule.SparkRenderer({
                 renderer: gl,
-                autoUpdate: false,
-                lodSplatScale: resolvedCaptureTuning.lodSplatScale,
-                lodRenderScale: resolvedCaptureTuning.lodRenderScale,
-                minSortIntervalMs: resolvedCaptureTuning.minSortIntervalMs,
+                // Preview uses Spark's change detection and background sorting;
+                // forced update() regenerates even an unchanged static scene.
+                autoUpdate: renderMode === 'preview',
+                preUpdate: false,
+                ...(renderMode === 'preview' ? resolvedRenderTuning : {
+                  lodSplatScale: resolvedCaptureTuning.lodSplatScale,
+                  lodRenderScale: resolvedCaptureTuning.lodRenderScale,
+                  minSortIntervalMs: resolvedCaptureTuning.minSortIntervalMs,
+                }),
                 target: {
                   width,
                   height,
@@ -492,7 +501,7 @@ export function SparkSplatEnvironment({
               width,
               height,
             };
-            captureSparkRef.current = captureSpark;
+            captureSparkRef.current.set(renderMode, captureSpark);
           }
 
           const captureRenderer = captureSpark.renderer;
@@ -503,11 +512,11 @@ export function SparkSplatEnvironment({
           let pixels: Uint8Array | undefined;
           for (let attempt = 0; attempt < maxWarmupFrames; attempt += 1) {
             captureRenderer.renderSize.set(width, height);
-            await captureRenderer.update({ scene, camera });
-            pixels = await captureRenderer.renderReadTarget({
-              scene,
-              camera,
-            });
+            if (renderMode === 'snapshot') {
+              await captureRenderer.update({ scene, camera });
+            }
+            withCaptureState(() => captureRenderer.renderTarget({ scene, camera }));
+            pixels = await captureRenderer.readTarget();
             if (
               captureRenderer.activeSplats > 0 &&
               !isMostlyBlankCapture(pixels, resolvedCaptureTuning)

package/src/types.ts

@@ -13,6 +13,7 @@ import type {
   CameraFrameCaptureSession,
   CameraFrameCaptureTensorOptions,
   CameraFrameTensorResult,
+  CameraRGBDResult,
 } from './rendering/cameraFrameCapture';
 
 // ---- Register (type-safe named resources) ----
@@ -1548,6 +1549,8 @@ export interface MujocoSimAPI {
   captureFrameBlob(options?: MujocoFrameCaptureOptions): Promise<FrameCaptureBlobResult>;
   captureCameraFrame(options?: CameraFrameCaptureOptions): Promise<CameraFrameCaptureResult>;
   captureCameraFrameBlob(options?: CameraFrameCaptureOptions): Promise<CameraFrameCaptureBlobResult>;
+  /** Aligned sRGB + camera-axis depth in meters from one render. */
+  captureCameraRGBD(options?: CameraFrameCaptureOptions): CameraRGBDResult;
   /** Capture a camera frame straight into a policy image tensor (no canvas/PNG encode). */
   captureCameraFrameTensor(options?: CameraFrameCaptureTensorOptions): CameraFrameTensorResult;
   /**
@@ -1677,6 +1680,11 @@ export interface CameraFrameRenderIsolationOptions {
 }
 
 export interface CameraFrameCaptureOptions {
+  /** Snapshot (default) waits for splat preparation. Preview reuses the latest
+   * sorted splats while updating in the background for responsive live motion.
+   * useCameraStream defaults to preview; policy/dataset captures stay snapshot.
+   */
+  renderMode?: 'snapshot' | 'preview';
   /** Existing Three camera to clone before applying pose overrides. */
   camera?: THREE.Camera;
   /** Named MuJoCo `<camera>` to render from when available in the loaded model. */
@@ -1791,6 +1799,7 @@ export interface CameraFrameCaptureAPI {
   captureBlob: (
     options?: CameraFrameCaptureOptions
   ) => Promise<CameraFrameCaptureBlobResult>;
+  captureRGBD: (options?: CameraFrameCaptureOptions) => Promise<CameraRGBDResult>;
   reset: () => void;
 }