mujoco-react 10.0.1 → 10.2.0

package/src/rendering/imageProjection.ts

@@ -0,0 +1,186 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Project detector/image coordinates from a camera view into the rendered MuJoCo scene.
+ */
+
+import * as THREE from 'three';
+import type {
+  CameraFrameCaptureOptions,
+  CameraFrameCaptureSource,
+  ImagePointCoordinateSpace,
+  ImagePointProjectionOptions,
+  ImagePointProjectionResult,
+} from '../types';
+import { CAPTURE_EXCLUDE_KEY } from './cameraFrameCapture';
+
+const _raycaster = new THREE.Raycaster();
+const _ndc = new THREE.Vector2();
+
+function toVector3(
+  value: CameraFrameCaptureOptions['position'] | undefined,
+  fallback: 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 createProjectionCamera(
+  fallbackCamera: THREE.Camera,
+  options: CameraFrameCaptureOptions,
+  width: number,
+  height: number
+) {
+  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 getProjectionSource(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' };
+}
+
+function imageSize(
+  renderer: THREE.WebGLRenderer,
+  options: ImagePointProjectionOptions
+): [number, number] {
+  return [
+    Math.max(1, Math.floor(options.imageWidth ?? options.width ?? renderer.domElement.width)),
+    Math.max(1, Math.floor(options.imageHeight ?? options.height ?? renderer.domElement.height)),
+  ];
+}
+
+export function imagePointToNdc(
+  x: number,
+  y: number,
+  coordinateSpace: ImagePointCoordinateSpace = 'normalized',
+  width = 1,
+  height = 1
+): [number, number] {
+  if (coordinateSpace === 'ndc') return [x, y];
+  if (coordinateSpace === 'normalized-1000') {
+    return [(x / 1000) * 2 - 1, 1 - (y / 1000) * 2];
+  }
+  if (coordinateSpace === 'pixel') {
+    return [(x / width) * 2 - 1, 1 - (y / height) * 2];
+  }
+  return [x * 2 - 1, 1 - y * 2];
+}
+
+function isProjectionCandidate(object: THREE.Object3D, options: ImagePointProjectionOptions) {
+  if (!object.visible) return false;
+  if (object.userData[CAPTURE_EXCLUDE_KEY]) return false;
+
+  const geomGroup = object.userData.geomGroup;
+  const geomName = object.userData.geomName;
+  if (options.hiddenGeomNames && typeof geomName === 'string' && options.hiddenGeomNames.includes(geomName)) {
+    return false;
+  }
+  if (options.hiddenGeomGroups && typeof geomGroup === 'number' && options.hiddenGeomGroups.includes(geomGroup)) {
+    return false;
+  }
+  if (options.visibleGeomGroups && typeof geomGroup === 'number' && !options.visibleGeomGroups.includes(geomGroup)) {
+    return false;
+  }
+  return true;
+}
+
+function findBodyId(object: THREE.Object3D) {
+  let current: THREE.Object3D | null = object;
+  while (current && current.userData.bodyID === undefined && current.parent) {
+    current = current.parent;
+  }
+  return typeof current?.userData.bodyID === 'number' ? current.userData.bodyID : -1;
+}
+
+export function projectImagePointTo3D(
+  renderer: THREE.WebGLRenderer,
+  scene: THREE.Scene,
+  fallbackCamera: THREE.Camera,
+  options: ImagePointProjectionOptions
+): ImagePointProjectionResult | null {
+  const [width, height] = imageSize(renderer, options);
+  const [ndcX, ndcY] = imagePointToNdc(
+    options.x,
+    options.y,
+    options.coordinateSpace,
+    width,
+    height
+  );
+  const projectionCamera = createProjectionCamera(fallbackCamera, options, width, height);
+
+  scene.updateMatrixWorld(true);
+  _ndc.set(ndcX, ndcY);
+  _raycaster.setFromCamera(_ndc, projectionCamera);
+  _raycaster.far = options.maxDistance ?? Infinity;
+
+  const objects: THREE.Object3D[] = [];
+  scene.traverse((object) => {
+    if ((object as THREE.Mesh).isMesh && isProjectionCandidate(object, options)) {
+      objects.push(object);
+    }
+  });
+
+  const [hit] = _raycaster.intersectObjects(objects, true);
+  if (!hit) return null;
+
+  return {
+    point: hit.point.clone(),
+    bodyId: findBodyId(hit.object),
+    geomId: typeof hit.object.userData.geomID === 'number' ? hit.object.userData.geomID : -1,
+    distance: hit.distance,
+    ndc: [ndcX, ndcY],
+    imageSize: [width, height],
+    source: getProjectionSource(options),
+  };
+}

package/src/types.ts

@@ -266,8 +266,18 @@ export interface MujocoModel {
   geom_friction: Float64Array;
 
   // Material
+  mat_texid: Int32Array;
+  mat_texrepeat: Float32Array;
+  mat_texuniform: Uint8Array;
   mat_rgba: Float32Array;
 
+  // Texture
+  tex_adr: Int32Array;
+  tex_data: Uint8Array;
+  tex_height: Int32Array;
+  tex_nchannel: Int32Array;
+  tex_width: Int32Array;
+
   // Mesh
   mesh_vertadr: Int32Array;
   mesh_vertnum: Int32Array;
@@ -434,6 +444,8 @@ export interface MujocoModule {
 
 export interface SceneObject {
   name: string;
+  /** MuJoCo geom name. Defaults to `${name}_geom` for generated objects. */
+  geomName?: string;
   type: 'box' | 'sphere' | 'cylinder';
   size: [number, number, number];
   position: [number, number, number];
@@ -735,6 +747,42 @@ export interface RayHit {
   distance: number;
 }
 
+export type ImagePointCoordinateSpace =
+  | 'normalized'
+  | 'normalized-1000'
+  | 'pixel'
+  | 'ndc';
+
+export interface ImagePointProjectionOptions extends CameraFrameCaptureOptions {
+  /** X coordinate in the selected coordinate space. Defaults to normalized 0..1. */
+  x: number;
+  /** Y coordinate in the selected coordinate space. Defaults to normalized 0..1 with origin at top-left. */
+  y: number;
+  /**
+   * Coordinate convention for x/y:
+   * - normalized: 0..1 image coordinates, top-left origin
+   * - normalized-1000: 0..1000 detector coordinates, top-left origin
+   * - pixel: pixel coordinates, top-left origin
+   * - ndc: Three.js normalized device coordinates, -1..1
+   */
+  coordinateSpace?: ImagePointCoordinateSpace;
+  /** Image width for pixel coordinates. Falls back to `width` or renderer canvas width. */
+  imageWidth?: number;
+  /** Image height for pixel coordinates. Falls back to `height` or renderer canvas height. */
+  imageHeight?: number;
+  /** Ignore hits farther than this distance from the camera ray origin. */
+  maxDistance?: number;
+}
+
+export interface ImagePointProjectionResult extends RayHit {
+  /** NDC coordinates used for raycasting. */
+  ndc: [number, number];
+  /** Image dimensions used when interpreting pixel coordinates. */
+  imageSize: [number, number];
+  /** Camera pose provenance, matching camera-frame capture results. */
+  source: CameraFrameCaptureSource;
+}
+
 // ---- Model Options (spec 5.3) ----
 
 export interface ModelOptions {
@@ -787,17 +835,141 @@ export interface PolicyInferenceInput extends PolicyObservationInput {
   observation: PolicyVector;
 }
 
+export type PolicyActionChunk = readonly PolicyVector[];
+export type PolicyInferenceOutput = PolicyVector | PolicyActionChunk;
+export type PolicyInferenceResult = PolicyInferenceOutput | Promise<PolicyInferenceOutput>;
+
 export interface PolicyActionInput extends PolicyInferenceInput {
   action: PolicyVector;
 }
 
+export interface PolicyAPI {
+  readonly isRunning: boolean;
+  start: () => void;
+  stop: () => void;
+  clearQueue: () => void;
+  reset: () => void;
+  readonly inFlight: boolean;
+  readonly queuedActions: number;
+  readonly lastObservation: PolicyVector | null;
+  readonly lastAction: PolicyVector | null;
+  readonly lastError: unknown;
+}
+
 export interface PolicyConfig {
   frequency: number;
   enabled?: boolean;
+  /** Start async inference while this many queued actions remain. Defaults to 0. */
+  prefetchThreshold?: number;
+  /**
+   * How async action chunks update the queue.
+   * - append preserves legacy FIFO behavior.
+   * - replace is useful for receding-horizon policies where a fresh chunk should supersede stale queued actions.
+   */
+  queueStrategy?: 'append' | 'replace';
+  /**
+   * Clear queued actions and ignore in-flight async results when `stop()` is called.
+   * Defaults to false so callers can choose pause/resume behavior explicitly.
+   */
+  clearQueueOnStop?: boolean;
   onObservation: (input: PolicyObservationInput) => PolicyVector;
   /** Run policy inference. Omit to pass observations directly to `onAction` for custom inline controllers. */
-  infer?: (input: PolicyInferenceInput) => PolicyVector;
+  infer?: (input: PolicyInferenceInput) => PolicyInferenceResult;
   onAction: (input: PolicyActionInput) => void;
+  /** Called when async inference rejects. */
+  onError?: (error: unknown) => void;
+}
+
+export interface RemotePolicyRequestInput extends PolicyInferenceInput {
+  /** True for the first request after hook construction or `reset()`. */
+  reset: boolean;
+  /** Zero-based request index since construction or `reset()`. */
+  requestIndex: number;
+  /** Aborts when the request is no longer needed, e.g. after pause/reset. */
+  signal: AbortSignal;
+}
+
+export interface RemotePolicyRequestInfo extends RemotePolicyRequestInput {
+  body: unknown;
+  requestStartedAt: number;
+}
+
+export interface RemotePolicyResponseInfo extends RemotePolicyRequestInfo {
+  response: Response;
+  responseBody: unknown;
+  responseFinishedAt: number;
+  requestMs: number;
+}
+
+export type RemotePolicyStatus = 'idle' | 'requesting' | 'ready' | 'error' | 'aborted';
+
+export interface RemotePolicyConfig extends Omit<PolicyConfig, 'infer'> {
+  endpoint: string | URL;
+  method?: string;
+  headers?: HeadersInit;
+  credentials?: RequestCredentials;
+  /** Additional external cancellation signal for remote inference requests. */
+  signal?: AbortSignal;
+  /**
+   * Abort the active HTTP request when `stop()` or `reset()` is called.
+   * Defaults to true so paused policies stop consuming server work.
+   */
+  abortOnStop?: boolean;
+  fetcher?: typeof fetch;
+  requestInit?: Omit<RequestInit, 'body' | 'headers' | 'method' | 'credentials' | 'signal'>;
+  buildRequest?: (input: RemotePolicyRequestInput) => unknown | Promise<unknown>;
+  readResponse?: (response: Response) => unknown | Promise<unknown>;
+  parseResponse?: (
+    responseBody: unknown,
+    info: RemotePolicyResponseInfo
+  ) => PolicyInferenceResult;
+  onRequest?: (info: RemotePolicyRequestInfo) => void;
+  onResponse?: (info: RemotePolicyResponseInfo) => void;
+}
+
+export interface RemotePolicyAPI extends PolicyAPI {
+  abort: (reason?: unknown) => void;
+  readonly remoteStatus: RemotePolicyStatus;
+  readonly requestCount: number;
+  readonly responseCount: number;
+  readonly lastRequestBody: unknown;
+  readonly lastResponseBody: unknown;
+  readonly lastHttpStatus: number | null;
+  readonly lastRequestMs: number | null;
+}
+
+export interface PolicyCameraFrameStream extends CameraFrameCaptureOptions {
+  /** Image key used in policy payloads, e.g. `image`, `front`, or `wrist_cam`. */
+  key: string;
+  /** Additional payload keys that should receive the same data URL. */
+  aliases?: readonly string[];
+}
+
+export interface PolicyCameraFrameCaptureOptions {
+  streams: readonly PolicyCameraFrameStream[];
+  /**
+   * Include `observation.images.${key}` for every captured stream.
+   * Defaults to true because LeRobot-style policies usually use these names.
+   */
+  includeObservationImageAliases?: boolean;
+}
+
+export interface PolicyCameraFrameCaptureResult {
+  frames: Record<string, CameraFrameCaptureResult>;
+  images: Record<string, string>;
+  /** Human-readable source summary for UI/debug telemetry. */
+  sourceSummary: string;
+  capturedAt: number;
+}
+
+export interface PolicyCameraFrameCaptureAPI {
+  status: FrameCaptureStatus;
+  error: Error | null;
+  isCapturing: boolean;
+  capture: (
+    options?: Partial<PolicyCameraFrameCaptureOptions>
+  ) => Promise<PolicyCameraFrameCaptureResult>;
+  reset: () => void;
 }
 
 // ---- Observation Builder ----
@@ -1268,6 +1440,7 @@ export interface MujocoSimAPI {
     cameraPos: THREE.Vector3,
     lookAt: THREE.Vector3
   ): { point: THREE.Vector3; bodyId: number; geomId: number } | null;
+  projectImagePointTo3D(options: ImagePointProjectionOptions): ImagePointProjectionResult | null;
 
   // Domain randomization (spec 10.3)
   setBodyMass(name: Bodies, mass: number): void;
@@ -1343,6 +1516,10 @@ export interface CameraFrameCaptureOptions {
   lookAt?: CameraFrameCaptureVector3;
   quaternion?: CameraFrameCaptureQuaternion;
   up?: CameraFrameCaptureVector3;
+  /** Local-space offset applied after resolving a mounted MuJoCo camera/site/body pose. */
+  positionOffset?: CameraFrameCaptureVector3;
+  /** Local-space rotation applied after resolving a mounted MuJoCo camera/site/body pose. Array values use Three.js order: [x, y, z, w]. */
+  quaternionOffset?: CameraFrameCaptureQuaternion;
   width?: number;
   height?: number;
   type?: string;
@@ -1352,6 +1529,16 @@ export interface CameraFrameCaptureOptions {
   far?: number;
   /** Provenance for the camera pose used by the capture. Usually set by the MuJoCo provider. */
   source?: CameraFrameCaptureSource;
+  /** Hide rendered Three objects whose MuJoCo geom group is in this list. */
+  hiddenGeomGroups?: readonly number[];
+  /** When provided, only rendered Three objects whose MuJoCo geom group is in this list are visible. */
+  visibleGeomGroups?: readonly number[];
+  /** Hide rendered Three objects whose MuJoCo geom name is in this list. */
+  hiddenGeomNames?: readonly string[];
+  /** Optional clear color for this capture only. Defaults to the renderer's current clear color. */
+  background?: THREE.ColorRepresentation;
+  /** Optional clear alpha for this capture only. Defaults to the renderer's current clear alpha. */
+  backgroundAlpha?: number;
 }
 
 export type CameraFrameCaptureSource =