mujoco-react 10.4.0 → 10.5.0

package/src/onnx.ts

@@ -0,0 +1,126 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Optional ONNX Runtime Web helpers for browser policy demos.
+ *
+ * This entry point is exported as `mujoco-react/onnx` so the main package does
+ * not import or bundle `onnxruntime-web`.
+ */
+
+import type * as ort from 'onnxruntime-web';
+import type { PolicyActionChunk } from './types';
+
+export type OnnxPolicyDtype = 'float32' | 'float64' | 'int32' | 'int64' | 'bool' | string;
+
+export interface OnnxPolicyTensorSpec {
+  name: string;
+  shape: number[];
+  dtype: OnnxPolicyDtype;
+}
+
+export interface OnnxPolicyImageSpec {
+  width: number;
+  height: number;
+  channels?: number;
+  layout?: 'CHW' | 'HWC' | string;
+  range?: readonly [number, number];
+}
+
+export interface OnnxPolicyManifest {
+  model: string;
+  variants?: Record<string, string>;
+  fps?: number;
+  joints?: string[];
+  cameras?: string[];
+  image?: OnnxPolicyImageSpec;
+  chunk_size?: number;
+  n_action_steps?: number;
+  inputs: OnnxPolicyTensorSpec[];
+  output: OnnxPolicyTensorSpec & {
+    units?: string;
+  };
+  [key: string]: unknown;
+}
+
+export interface LoadOnnxPolicyManifestResult<TManifest extends OnnxPolicyManifest = OnnxPolicyManifest> {
+  manifest: TManifest;
+  manifestUrl: URL;
+  modelUrl: URL;
+}
+
+export interface CreateOnnxPolicySessionOptions<TManifest extends OnnxPolicyManifest = OnnxPolicyManifest> {
+  manifestUrl: string | URL;
+  variant?: string;
+  runtime: typeof ort;
+  sessionOptions?: ort.InferenceSession.SessionOptions;
+  fetcher?: typeof fetch;
+  readManifest?: (response: Response) => Promise<TManifest>;
+}
+
+export interface OnnxPolicySession<TManifest extends OnnxPolicyManifest = OnnxPolicyManifest>
+  extends LoadOnnxPolicyManifestResult<TManifest> {
+  session: ort.InferenceSession;
+}
+
+function asUrl(value: string | URL, base = globalThis.location?.href) {
+  return value instanceof URL ? value : new URL(value, base);
+}
+
+function resolveModelPath(manifest: OnnxPolicyManifest, variant: string | undefined) {
+  if (variant && manifest.variants?.[variant]) return manifest.variants[variant];
+  return manifest.model;
+}
+
+export async function loadOnnxPolicyManifest<TManifest extends OnnxPolicyManifest = OnnxPolicyManifest>(
+  manifestUrlInput: string | URL,
+  options: Pick<CreateOnnxPolicySessionOptions<TManifest>, 'variant' | 'fetcher' | 'readManifest'> = {}
+): Promise<LoadOnnxPolicyManifestResult<TManifest>> {
+  const fetcher = options.fetcher ?? fetch;
+  const manifestUrl = asUrl(manifestUrlInput);
+  const response = await fetcher(manifestUrl);
+  if (!response.ok) {
+    throw new Error(`Unable to load ONNX policy manifest from ${manifestUrl.href} (${response.status}).`);
+  }
+  const manifest = options.readManifest
+    ? await options.readManifest(response)
+    : await response.json() as TManifest;
+  const modelPath = resolveModelPath(manifest, options.variant);
+  const modelUrl = asUrl(modelPath, manifestUrl.href);
+  return { manifest, manifestUrl, modelUrl };
+}
+
+export async function createOnnxPolicySession<TManifest extends OnnxPolicyManifest = OnnxPolicyManifest>(
+  options: CreateOnnxPolicySessionOptions<TManifest>
+): Promise<OnnxPolicySession<TManifest>> {
+  const fetcher = options.fetcher ?? fetch;
+  const resolved = await loadOnnxPolicyManifest(options.manifestUrl, options);
+  const response = await fetcher(resolved.modelUrl);
+  if (!response.ok) {
+    throw new Error(`Unable to load ONNX policy model from ${resolved.modelUrl.href} (${response.status}).`);
+  }
+  const modelBytes = await response.arrayBuffer();
+  const session = await options.runtime.InferenceSession.create(modelBytes, options.sessionOptions);
+  return {
+    ...resolved,
+    session,
+  };
+}
+
+export function onnxTensorToPolicyActionChunk(
+  tensor: ort.Tensor,
+  actionSize = tensor.dims.at(-1) ?? 1,
+  maxActions?: number
+): PolicyActionChunk {
+  const rawData = Array.from(tensor.data as ArrayLike<number>, (value) => Number(value));
+  const actionCount = Math.floor(rawData.length / actionSize);
+  const cappedActionCount = maxActions === undefined
+    ? actionCount
+    : Math.max(0, Math.min(actionCount, Math.floor(maxActions)));
+  const actions: number[][] = [];
+  for (let actionIndex = 0; actionIndex < cappedActionCount; actionIndex += 1) {
+    const start = actionIndex * actionSize;
+    actions.push(rawData.slice(start, start + actionSize));
+  }
+  return actions;
+}

package/src/policyImageTensors.ts

@@ -0,0 +1,150 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Helpers for turning browser camera captures into policy image tensors.
+ */
+
+export type PolicyImageTensorLayout = 'CHW' | 'HWC';
+export type PolicyImageTensorRange = readonly [number, number];
+/**
+ * Row order of a raw pixel buffer. WebGL `readRenderTargetPixels` returns rows
+ * bottom-to-top (`'bottom-left'`); `ImageData` is top-to-bottom (`'top-left'`).
+ */
+export type PolicyImageTensorSourceOrigin = 'top-left' | 'bottom-left';
+
+export interface PolicyImageTensorOptions {
+  width: number;
+  height: number;
+  channels?: 3 | 4;
+  layout?: PolicyImageTensorLayout;
+  range?: PolicyImageTensorRange;
+}
+
+export interface PolicyImageTensorPixelOptions extends PolicyImageTensorOptions {
+  /** Row order of the source buffer. Defaults to `'top-left'`. */
+  sourceOrigin?: PolicyImageTensorSourceOrigin;
+  /** Mirror horizontally while reading. */
+  flipX?: boolean;
+}
+
+export interface PolicyImageTensorResult {
+  data: Float32Array;
+  shape: [number, number, number];
+  width: number;
+  height: number;
+  channels: 3 | 4;
+  layout: PolicyImageTensorLayout;
+  range: PolicyImageTensorRange;
+}
+
+function resolveTensorOptions(options: PolicyImageTensorOptions): Required<PolicyImageTensorOptions> {
+  return {
+    channels: 3,
+    layout: 'CHW',
+    range: [0, 1],
+    ...options,
+  };
+}
+
+function normalizeChannel(value: number, range: PolicyImageTensorRange) {
+  const [min, max] = range;
+  if (min === 0 && max === 255) return value;
+  return min + (value / 255) * (max - min);
+}
+
+/**
+ * Convert a raw RGBA pixel buffer (4 bytes per pixel) directly into a policy
+ * image tensor. This is the fast path that skips canvas encoding entirely —
+ * feed it the `Uint8Array` returned by `readRenderTargetPixels` (which is
+ * bottom-left origin, so pass `sourceOrigin: 'bottom-left'`).
+ */
+export function pixelsToPolicyImageTensor(
+  pixels: Uint8Array | Uint8ClampedArray,
+  options: PolicyImageTensorPixelOptions
+): PolicyImageTensorResult {
+  const resolved = resolveTensorOptions(options);
+  const { width, height, channels, layout, range } = resolved;
+  const expected = width * height * 4;
+  if (pixels.length < expected) {
+    throw new Error(
+      `Pixel buffer of length ${pixels.length} is too small for ${width}x${height} RGBA data (${expected} bytes).`
+    );
+  }
+
+  const flipY = options.sourceOrigin === 'bottom-left';
+  const flipX = options.flipX ?? false;
+  const pixelCount = width * height;
+  const data = new Float32Array(pixelCount * channels);
+
+  for (let y = 0; y < height; y += 1) {
+    const sourceY = flipY ? height - y - 1 : y;
+    for (let x = 0; x < width; x += 1) {
+      const sourceX = flipX ? width - x - 1 : x;
+      const source = (sourceY * width + sourceX) * 4;
+      const target = y * width + x;
+      for (let channel = 0; channel < channels; channel += 1) {
+        const value = normalizeChannel(pixels[source + channel], range);
+        if (layout === 'CHW') {
+          data[channel * pixelCount + target] = value;
+        } else {
+          data[target * channels + channel] = value;
+        }
+      }
+    }
+  }
+
+  return {
+    data,
+    shape: layout === 'CHW' ? [channels, height, width] : [height, width, channels],
+    width,
+    height,
+    channels,
+    layout,
+    range,
+  };
+}
+
+export function imageDataToPolicyImageTensor(
+  imageData: ImageData,
+  options: PolicyImageTensorOptions
+): PolicyImageTensorResult {
+  const resolved = resolveTensorOptions(options);
+  if (imageData.width !== resolved.width || imageData.height !== resolved.height) {
+    throw new Error(
+      `ImageData size ${imageData.width}x${imageData.height} does not match tensor size ${resolved.width}x${resolved.height}.`
+    );
+  }
+  return pixelsToPolicyImageTensor(imageData.data, {
+    ...resolved,
+    sourceOrigin: 'top-left',
+  });
+}
+
+async function decodeImageSource(dataUrl: string) {
+  const image = new Image();
+  image.decoding = 'async';
+  image.src = dataUrl;
+  await image.decode();
+  return image;
+}
+
+export async function dataUrlToPolicyImageTensor(
+  dataUrl: string,
+  options: PolicyImageTensorOptions
+): Promise<PolicyImageTensorResult> {
+  const resolved = resolveTensorOptions(options);
+  const image = await decodeImageSource(dataUrl);
+  const canvas = document.createElement('canvas');
+  canvas.width = resolved.width;
+  canvas.height = resolved.height;
+  const context = canvas.getContext('2d', { willReadFrequently: true });
+  if (!context) {
+    throw new Error('Unable to create a 2D canvas context for policy image tensor conversion.');
+  }
+  context.drawImage(image, 0, 0, resolved.width, resolved.height);
+  return imageDataToPolicyImageTensor(
+    context.getImageData(0, 0, resolved.width, resolved.height),
+    resolved
+  );
+}

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. */
@@ -1458,6 +1473,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,