mujoco-react 10.4.0 → 10.6.0

package/src/rendering/cameraFrameCapture.ts

@@ -13,6 +13,29 @@ import type {
   CameraFrameCaptureSource,
   CameraFrameCaptureVector3,
 } from '../types';
+import {
+  pixelsToPolicyImageTensor,
+  type PolicyImageTensorOptions,
+  type PolicyImageTensorResult,
+} from '../policyImageTensors';
+
+/** Options for capturing a camera frame straight into a policy image tensor. */
+export type CameraFrameCaptureTensorOptions = CameraFrameCaptureOptions &
+  Pick<PolicyImageTensorOptions, 'channels' | 'layout' | 'range'>;
+
+export interface CameraFramePixelsResult {
+  /** Raw RGBA pixels, bottom-left origin (reused buffer — consume before the next capture). */
+  pixels: Uint8Array;
+  camera: THREE.Camera;
+  width: number;
+  height: number;
+  source: CameraFrameCaptureSource;
+}
+
+export interface CameraFrameTensorResult extends PolicyImageTensorResult {
+  camera: THREE.Camera;
+  source: CameraFrameCaptureSource;
+}
 
 export interface CameraFrameCaptureSession {
   readonly width: number;
@@ -36,6 +59,14 @@ export interface CameraFrameCaptureSession {
     options?: CameraFrameCaptureOptions
   ): Promise<CameraFrameCaptureResult>;
   captureBlob(options?: CameraFrameCaptureOptions): Promise<CameraFrameCaptureBlobResult>;
+  /**
+   * Render and read raw RGBA pixels without any canvas/PNG round-trip. The
+   * returned buffer is reused between calls — copy or convert it before the
+   * next capture.
+   */
+  capturePixels(options?: CameraFrameCaptureOptions): CameraFramePixelsResult;
+  /** Render straight into a normalized policy image tensor (no canvas/PNG encode). */
+  captureTensor(options?: CameraFrameCaptureTensorOptions): CameraFrameTensorResult;
   dispose(): void;
 }
 
@@ -250,7 +281,7 @@ function applyProjectionMatrix(
   camera.projectionMatrixInverse.copy(camera.projectionMatrix).invert();
 }
 
-function createCaptureCamera(
+export function createCaptureCamera(
   options: CameraFrameCaptureOptions,
   fallbackCamera: THREE.Camera,
   width: number,
@@ -290,7 +321,7 @@ function getCaptureDimensions(
   return { width, height };
 }
 
-function prepareCaptureCamera(
+export function prepareCaptureCamera(
   camera: THREE.Camera,
   options: CameraFrameCaptureOptions,
   fallbackCamera: THREE.Camera,
@@ -646,7 +677,10 @@ export function createCameraFrameCaptureSession(
     return captureOptions;
   }
 
-  function renderPreparedCapture(captureOptions: CameraFrameCaptureOptions) {
+  function renderCaptureToTarget(
+    captureOptions: CameraFrameCaptureOptions,
+    readback: () => void
+  ) {
     const previousState = saveRendererState(sessionRenderer);
     const previousSceneState = applyCaptureVisualOverrides(
       sessionRenderer,
@@ -676,6 +710,16 @@ export function createCameraFrameCaptureSession(
       }
       sessionRenderer.clear();
       sessionRenderer.render(scene, camera);
+      readback();
+    } finally {
+      restoreObjectVisibility(hidden);
+      if (previousSceneState) restoreSceneVisualState(scene, previousSceneState);
+      restoreRendererState(sessionRenderer, previousState);
+    }
+  }
+
+  function renderPreparedCapture(captureOptions: CameraFrameCaptureOptions) {
+    renderCaptureToTarget(captureOptions, () => {
       readRenderTargetToCanvas(
         sessionRenderer,
         target,
@@ -688,24 +732,50 @@ export function createCameraFrameCaptureSession(
         sessionRenderer.outputColorSpace,
         captureOptions.flipX ?? false
       );
-      return {
-        canvas,
-        camera,
-        width,
-        height,
-        source: getCameraFrameCaptureSource(captureOptions),
-      };
-    } finally {
-      restoreObjectVisibility(hidden);
-      if (previousSceneState) restoreSceneVisualState(scene, previousSceneState);
-      restoreRendererState(sessionRenderer, previousState);
-    }
+    });
+    return {
+      canvas,
+      camera,
+      width,
+      height,
+      source: getCameraFrameCaptureSource(captureOptions),
+    };
   }
 
   function capture(nextOptions: CameraFrameCaptureOptions = {}) {
     return renderPreparedCapture(resolveCaptureOptions(nextOptions));
   }
 
+  function capturePixels(nextOptions: CameraFrameCaptureOptions = {}): CameraFramePixelsResult {
+    const captureOptions = resolveCaptureOptions(nextOptions);
+    renderCaptureToTarget(captureOptions, () => {
+      sessionRenderer.readRenderTargetPixels(target, 0, 0, width, height, pixels);
+    });
+    return {
+      pixels,
+      camera,
+      width,
+      height,
+      source: getCameraFrameCaptureSource(captureOptions),
+    };
+  }
+
+  function captureTensor(
+    nextOptions: CameraFrameCaptureTensorOptions = {}
+  ): CameraFrameTensorResult {
+    const result = capturePixels(nextOptions);
+    const tensor = pixelsToPolicyImageTensor(pixels, {
+      width,
+      height,
+      channels: nextOptions.channels,
+      layout: nextOptions.layout,
+      range: nextOptions.range,
+      sourceOrigin: 'bottom-left',
+      flipX: nextOptions.flipX,
+    });
+    return { ...tensor, camera, source: result.source };
+  }
+
   async function captureAsync(nextOptions: CameraFrameCaptureOptions = {}) {
     const captureOptions = resolveCaptureOptions(nextOptions);
     runCapturePreRenderHooks(scene);
@@ -779,6 +849,8 @@ export function createCameraFrameCaptureSession(
     height,
     capture,
     captureAsync,
+    capturePixels,
+    captureTensor,
     captureDataUrl(nextOptions = {}) {
       const type = nextOptions.type ?? options.type ?? 'image/png';
       const result = capture(nextOptions);
@@ -889,3 +961,28 @@ export async function captureCameraFrameBlob(
     session.dispose();
   }
 }
+
+/**
+ * One-shot camera frame capture straight into a policy image tensor, skipping
+ * the canvas/PNG round-trip. For repeated captures (live inference, recording),
+ * create a session once with {@link createCameraFrameCaptureSession} and call
+ * `session.captureTensor()` so the render target and buffers are reused.
+ */
+export function captureCameraFrameTensor(
+  renderer: THREE.WebGLRenderer,
+  scene: THREE.Scene,
+  fallbackCamera: THREE.Camera,
+  options: CameraFrameCaptureTensorOptions = {}
+): CameraFrameTensorResult {
+  const session = createCameraFrameCaptureSession(
+    renderer,
+    scene,
+    fallbackCamera,
+    options
+  );
+  try {
+    return session.captureTensor(options);
+  } finally {
+    session.dispose();
+  }
+}

package/src/types.ts

@@ -7,6 +7,11 @@ import type React from 'react';
 import type { ReactNode } from 'react';
 import type { CanvasProps, ThreeElements } from '@react-three/fiber';
 import * as THREE from 'three';
+import type {
+  CameraFrameCaptureSession,
+  CameraFrameCaptureTensorOptions,
+  CameraFrameTensorResult,
+} from './rendering/cameraFrameCapture';
 
 // ---- Register (type-safe named resources) ----
 
@@ -459,6 +464,10 @@ export interface SceneObject {
   solref?: string;
   solimp?: string;
   condim?: number;
+  /** MuJoCo geom contact type bitmask. Defaults to 1 for generated objects. */
+  contype?: number;
+  /** MuJoCo geom contact affinity bitmask. Defaults to 1 for generated objects. */
+  conaffinity?: number;
   /** MuJoCo geom group. Group 3 is conventionally used for collision-only helper geoms. */
   group?: number;
 }
@@ -527,6 +536,12 @@ export interface IkConfig {
    * starting at index 0. Prefer inferred IK or `joints`/`actuators`.
    */
   numJoints?: number;
+  /**
+   * Optional solve-space joint limits in the same order as the resolved joints.
+   * Use this when MJCF limits are intentionally broad or when a setup/calibration
+   * tool should stay within a narrower envelope.
+   */
+  jointLimits?: ReadonlyArray<readonly [number, number] | null | undefined>;
   /** Custom IK solver. When omitted, uses built-in Damped Least-Squares solver. */
   ikSolveFn?: IKSolveFn;
   /** DLS damping. Default: 0.01. */
@@ -549,7 +564,7 @@ export interface IkContextValue {
   ikTargetRef: React.RefObject<THREE.Group>;
   siteIdRef: React.RefObject<number>;
   setIkEnabled: (enabled: boolean) => void;
-  moveTarget: (pos: THREE.Vector3, duration?: number) => void;
+  moveTarget: (pos: IkTargetPosition, duration?: number) => void;
   syncTargetToSite: () => void;
   solveIK: (input: IkSolveInput) => number[] | null;
   getGizmoStats: () => { pos: THREE.Vector3; rot: THREE.Euler } | null;
@@ -577,9 +592,28 @@ export type IKSolveFn = (
   input: IkSolveInput
 ) => number[] | null;
 
+export type IkTargetPosition =
+  | THREE.Vector3
+  | readonly [number, number, number]
+  | {
+      readonly x: number;
+      readonly y: number;
+      readonly z: number;
+    };
+
+export type IkTargetQuaternion =
+  | THREE.Quaternion
+  | readonly [number, number, number, number]
+  | {
+      readonly x: number;
+      readonly y: number;
+      readonly z: number;
+      readonly w: number;
+    };
+
 export interface IkSolveInput {
-  position: THREE.Vector3;
-  quaternion: THREE.Quaternion;
+  position: IkTargetPosition;
+  quaternion: IkTargetQuaternion;
   currentQ: number[];
   context?: IKSolveContext;
 }
@@ -1458,6 +1492,19 @@ export interface MujocoSimAPI {
   captureFrameBlob(options?: MujocoFrameCaptureOptions): Promise<FrameCaptureBlobResult>;
   captureCameraFrame(options?: CameraFrameCaptureOptions): Promise<CameraFrameCaptureResult>;
   captureCameraFrameBlob(options?: CameraFrameCaptureOptions): Promise<CameraFrameCaptureBlobResult>;
+  /** Capture a camera frame straight into a policy image tensor (no canvas/PNG encode). */
+  captureCameraFrameTensor(options?: CameraFrameCaptureTensorOptions): CameraFrameTensorResult;
+  /**
+   * Create a reusable offscreen capture session bound to this scene. Reuse it
+   * for live inference/recording so the render target and buffers persist
+   * across frames; call `session.captureTensor()` / `capturePixels()` each step.
+   */
+  createCameraFrameCaptureSession(options?: CameraFrameCaptureOptions): CameraFrameCaptureSession;
+  /**
+   * Resolve a named MuJoCo camera/site/body into concrete capture options with
+   * the current world pose. Useful for re-aiming a persistent session each step.
+   */
+  resolveCameraCaptureOptions(options?: CameraFrameCaptureOptions): CameraFrameCaptureOptions;
   recordCameraSequence(options: CameraFrameSequenceOptions): Promise<CameraFrameSequenceResult>;
   project2DTo3D(
     x: number,