mujoco-react 10.4.0 → 10.6.0

package/src/hooks/usePolicyCameraTensors.ts

@@ -0,0 +1,215 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Capture policy observation tensors directly from Three/MuJoCo cameras,
+ * skipping the data-URL/PNG round-trip. Sessions are created once per camera
+ * and reused every step, so live inference and dataset recording read straight
+ * from the GPU into Float32 tensors.
+ */
+
+import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
+import { useMujoco } from '../core/MujocoSimProvider';
+import { createPolicyCameraFrameCapturePlanFromApi } from '../policyCameraFrames';
+import type { MountedPolicyCameraFrameCaptureOptions } from './usePolicyCameraFrames';
+import type {
+  CameraFrameCaptureSession,
+  CameraFrameCaptureTensorOptions,
+  CameraFrameTensorResult,
+} from '../rendering/cameraFrameCapture';
+import type { FrameCaptureStatus } from '../types';
+
+export interface PolicyCameraTensorStream extends CameraFrameCaptureTensorOptions {
+  /** Payload key this stream's tensor is stored under. */
+  key: string;
+  /** Additional payload keys that should reference the same tensor. */
+  aliases?: readonly string[];
+}
+
+export interface PolicyCameraTensorsOptions {
+  streams: PolicyCameraTensorStream[];
+  /** Also expose tensors under `observation.images.<key>` aliases. Defaults to `false`. */
+  includeObservationImageAliases?: boolean;
+}
+
+export interface PolicyCameraTensorsResult {
+  tensors: Record<string, CameraFrameTensorResult>;
+  sourceSummary: string;
+  capturedAt: number;
+}
+
+export interface PolicyCameraTensorsAPI {
+  status: FrameCaptureStatus;
+  error: Error | null;
+  isCapturing: boolean;
+  /** Synchronously render and convert every stream into a policy image tensor. */
+  capture: () => PolicyCameraTensorsResult;
+  reset: () => void;
+}
+
+export type MountedPolicyCameraTensorOptions = MountedPolicyCameraFrameCaptureOptions & {
+  tensor?: Pick<
+    CameraFrameCaptureTensorOptions,
+    'width' | 'height' | 'channels' | 'layout' | 'range'
+  >;
+};
+
+type SessionEntry = {
+  session: CameraFrameCaptureSession;
+  signature: string;
+};
+
+function sessionSignature(stream: PolicyCameraTensorStream): string {
+  return JSON.stringify({
+    width: stream.width,
+    height: stream.height,
+    channels: stream.channels,
+    renderIsolation: stream.renderIsolation ?? false,
+    cameraName: stream.cameraName,
+    siteName: stream.siteName,
+    bodyName: stream.bodyName,
+  });
+}
+
+function addTensorAliases(
+  tensors: Record<string, CameraFrameTensorResult>,
+  stream: PolicyCameraTensorStream,
+  tensor: CameraFrameTensorResult,
+  includeObservationImageAliases: boolean
+) {
+  const keys = new Set<string>([stream.key, ...(stream.aliases ?? [])]);
+  if (includeObservationImageAliases) {
+    for (const base of [stream.key, ...(stream.aliases ?? [])]) {
+      keys.add(`observation.images.${base}`);
+    }
+  }
+  for (const key of keys) tensors[key] = tensor;
+}
+
+export function usePolicyCameraTensors(
+  options: PolicyCameraTensorsOptions
+): PolicyCameraTensorsAPI {
+  const mujoco = useMujoco();
+  const [status, setStatus] = useState<FrameCaptureStatus>('idle');
+  const [error, setError] = useState<Error | null>(null);
+  const sessionsRef = useRef<Map<string, SessionEntry>>(new Map());
+
+  const disposeSessions = useCallback(() => {
+    for (const { session } of sessionsRef.current.values()) session.dispose();
+    sessionsRef.current.clear();
+  }, []);
+
+  useEffect(() => disposeSessions, [disposeSessions]);
+
+  const reset = useCallback(() => {
+    setStatus('idle');
+    setError(null);
+  }, []);
+
+  const capture = useCallback((): PolicyCameraTensorsResult => {
+    const api = mujoco.api;
+    if (!api) {
+      throw new Error('MuJoCo scene is not ready for policy camera tensor capture.');
+    }
+
+    setStatus('capturing');
+    setError(null);
+
+    try {
+      const sessions = sessionsRef.current;
+      const seen = new Set<string>();
+      const tensors: Record<string, CameraFrameTensorResult> = {};
+      const sourceParts: string[] = [];
+
+      for (const stream of options.streams) {
+        seen.add(stream.key);
+        const resolved: CameraFrameCaptureTensorOptions = {
+          ...api.resolveCameraCaptureOptions(stream),
+          channels: stream.channels,
+          layout: stream.layout,
+          range: stream.range,
+        };
+
+        const signature = sessionSignature(stream);
+        let entry = sessions.get(stream.key);
+        if (!entry || entry.signature !== signature) {
+          entry?.session.dispose();
+          entry = {
+            session: api.createCameraFrameCaptureSession(resolved),
+            signature,
+          };
+          sessions.set(stream.key, entry);
+        }
+
+        const tensor = entry.session.captureTensor(resolved);
+        addTensorAliases(
+          tensors,
+          stream,
+          tensor,
+          options.includeObservationImageAliases ?? false
+        );
+        sourceParts.push(`${stream.key}:${tensor.source.kind}`);
+      }
+
+      // Drop sessions for streams that are no longer requested.
+      for (const key of [...sessions.keys()]) {
+        if (!seen.has(key)) {
+          sessions.get(key)?.session.dispose();
+          sessions.delete(key);
+        }
+      }
+
+      setStatus('captured');
+      return {
+        tensors,
+        sourceSummary: sourceParts.join(' + ') || 'not used by policy',
+        capturedAt: Date.now(),
+      };
+    } catch (nextError) {
+      const captureError =
+        nextError instanceof Error
+          ? nextError
+          : new Error('Unable to capture policy camera tensors.');
+      setError(captureError);
+      setStatus('error');
+      throw captureError;
+    }
+  }, [mujoco.api, options.includeObservationImageAliases, options.streams]);
+
+  return {
+    status,
+    error,
+    isCapturing: status === 'capturing',
+    capture,
+    reset,
+  };
+}
+
+export function usePolicyCameraTensorsFromMountedStreams(
+  options: MountedPolicyCameraTensorOptions
+): PolicyCameraTensorsAPI {
+  const mujoco = useMujoco();
+  const tensorOptions = options.tensor;
+  const mountedOptions = useMemo<PolicyCameraTensorsOptions>(() => {
+    const api = mujoco.api;
+    if (!api) {
+      return {
+        streams: [],
+        includeObservationImageAliases: options.includeObservationImageAliases ?? false,
+      };
+    }
+
+    const plan = createPolicyCameraFrameCapturePlanFromApi(api, options);
+    return {
+      streams: plan.streams.map(({ key, aliases, ...stream }) => ({
+        ...stream,
+        ...tensorOptions,
+        key,
+        aliases,
+      })),
+      includeObservationImageAliases: plan.includeObservationImageAliases ?? false,
+    };
+  }, [mujoco.api, options, tensorOptions]);
+
+  return usePolicyCameraTensors(mountedOptions);
+}

package/src/index.ts

@@ -34,6 +34,8 @@ export type { ControllerOptions, ControllerComponent } from './core/createContro
 
 // IK controller hook
 export { useIkController } from './hooks/useIkController';
+export { GenericIK } from './core/GenericIK';
+export type { GenericIKOptions } from './core/GenericIK';
 
 // Components
 export { Body } from './components/Body';
@@ -74,6 +76,10 @@ export type {
   SplatCollisionProxyPreviewVector3,
   UseSplatCollisionProxyGeomsOptions,
 } from './components/SplatCollisionProxyPreview';
+export { CameraView, useCameraViewport } from './components/CameraView';
+export type { CameraViewProps, CameraViewportOptions } from './components/CameraView';
+export { useCameraStream } from './hooks/useCameraStream';
+export type { CameraStreamOptions } from './hooks/useCameraStream';
 export { Debug } from './components/Debug';
 export { TendonRenderer } from './components/TendonRenderer';
 export { FlexRenderer } from './components/FlexRenderer';
@@ -91,6 +97,13 @@ export { useBodyState } from './hooks/useBodyState';
 export { useBodyPose, useGeomPose, useSitePose } from './hooks/usePose';
 export type { PoseReadout, PoseResourceKind } from './hooks/usePose';
 export { useCtrl } from './hooks/useCtrl';
+export { controlGroup, useControlGroup } from './hooks/useControlGroup';
+export type {
+  ControlGroup,
+  ControlGroupHandle,
+  ControlGroupSetOptions,
+  UseControlGroupOptions,
+} from './hooks/useControlGroup';
 export { useControlWriter } from './hooks/useControlWriter';
 export type {
   ControlWriterConflict,
@@ -130,6 +143,17 @@ export type {
   MountedPolicyCameraFrameCaptureAPI,
   MountedPolicyCameraFrameCaptureOptions,
 } from './hooks/usePolicyCameraFrames';
+export {
+  usePolicyCameraTensors,
+  usePolicyCameraTensorsFromMountedStreams,
+} from './hooks/usePolicyCameraTensors';
+export type {
+  MountedPolicyCameraTensorOptions,
+  PolicyCameraTensorsAPI,
+  PolicyCameraTensorsOptions,
+  PolicyCameraTensorsResult,
+  PolicyCameraTensorStream,
+} from './hooks/usePolicyCameraTensors';
 export { useCameraSequenceRecorder } from './hooks/useCameraSequenceRecorder';
 export { useMountedCameraSequenceRecorder } from './hooks/useMountedCameraSequenceRecorder';
 export type {
@@ -145,9 +169,16 @@ export {
   CAPTURE_EXCLUDE_KEY,
   captureCameraFrame,
   captureCameraFrameBlob,
+  captureCameraFrameTensor,
   createCameraFrameCaptureSession,
   renderCameraFrameToCanvas,
 } from './rendering/cameraFrameCapture';
+export type {
+  CameraFrameCaptureSession,
+  CameraFrameCaptureTensorOptions,
+  CameraFramePixelsResult,
+  CameraFrameTensorResult,
+} from './rendering/cameraFrameCapture';
 export {
   imagePointToNdc,
   projectImagePointTo3D,
@@ -172,6 +203,11 @@ export {
   readNamedObservation,
   sitePositionField,
 } from './policyObservation';
+export {
+  dataUrlToPolicyImageTensor,
+  imageDataToPolicyImageTensor,
+  pixelsToPolicyImageTensor,
+} from './policyImageTensors';
 export type {
   CreatePolicyCameraFrameCapturePlanOptions,
   PolicyCameraFrameCapturePlan,
@@ -191,6 +227,14 @@ export type {
   NamedObservationOptions,
   NamedObservationResult,
 } from './policyObservation';
+export type {
+  PolicyImageTensorLayout,
+  PolicyImageTensorOptions,
+  PolicyImageTensorPixelOptions,
+  PolicyImageTensorRange,
+  PolicyImageTensorResult,
+  PolicyImageTensorSourceOrigin,
+} from './policyImageTensors';
 export {
   createMountedCameraFrameSequenceManifest,
   createMountedCameraFrameSequenceReadiness,

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