mujoco-react 10.0.1 → 10.1.0

package/src/hooks/useRemotePolicy.ts

@@ -0,0 +1,329 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * useRemotePolicy — HTTP JSON inference wrapper around usePolicy.
+ */
+
+import { useMemo, useRef } from 'react';
+import { usePolicy } from './usePolicy';
+import type {
+  PolicyInferenceOutput,
+  PolicyVector,
+  RemotePolicyAPI,
+  RemotePolicyConfig,
+  RemotePolicyRequestInput,
+  RemotePolicyResponseInfo,
+  RemotePolicyStatus,
+} from '../types';
+
+function now() {
+  return typeof performance !== 'undefined' ? performance.now() : Date.now();
+}
+
+function isAbortError(error: unknown) {
+  return (
+    typeof DOMException !== 'undefined' &&
+    error instanceof DOMException && error.name === 'AbortError'
+  ) || (
+    error instanceof Error && error.name === 'AbortError'
+  );
+}
+
+function createAbortError(message: string) {
+  if (typeof DOMException !== 'undefined') {
+    return new DOMException(message, 'AbortError');
+  }
+  const error = new Error(message);
+  error.name = 'AbortError';
+  return error;
+}
+
+function abortController(controller: AbortController | null, reason?: unknown) {
+  if (!controller || controller.signal.aborted) return;
+  if (reason !== undefined) {
+    controller.abort(reason);
+  } else {
+    controller.abort();
+  }
+}
+
+function createMergedAbortSignal(
+  localSignal: AbortSignal,
+  externalSignal: AbortSignal | undefined
+) {
+  if (!externalSignal) return localSignal;
+  if (externalSignal.aborted) {
+    const controller = new AbortController();
+    abortController(controller, externalSignal.reason);
+    return controller.signal;
+  }
+  if (typeof AbortSignal !== 'undefined' && typeof AbortSignal.any === 'function') {
+    return AbortSignal.any([localSignal, externalSignal]);
+  }
+
+  const controller = new AbortController();
+  const abortFromLocal = () => abortController(controller, localSignal.reason);
+  const abortFromExternal = () => abortController(controller, externalSignal.reason);
+  localSignal.addEventListener('abort', abortFromLocal, { once: true });
+  externalSignal.addEventListener('abort', abortFromExternal, { once: true });
+  return controller.signal;
+}
+
+function vectorToArray(vector: PolicyVector) {
+  return Array.from(vector, (value) => Number(value));
+}
+
+function isPolicyVectorArray(value: unknown): value is PolicyVector[] {
+  return (
+    Array.isArray(value) &&
+    value.every((entry) => Array.isArray(entry) || ArrayBuffer.isView(entry))
+  );
+}
+
+function isPolicyVector(value: unknown): value is PolicyVector {
+  return Array.isArray(value) || ArrayBuffer.isView(value);
+}
+
+function defaultBuildRemotePolicyRequest(input: RemotePolicyRequestInput) {
+  const observation = vectorToArray(input.observation);
+  return {
+    observation,
+    state: observation,
+    time: input.data.time,
+    reset: input.reset,
+  };
+}
+
+async function defaultReadRemotePolicyResponse(response: Response) {
+  const text = await response.text();
+  if (text.length === 0) return null;
+  return JSON.parse(text);
+}
+
+function defaultParseRemotePolicyResponse(responseBody: unknown): PolicyInferenceOutput {
+  if (responseBody && typeof responseBody === 'object') {
+    const body = responseBody as {
+      action?: unknown;
+      actions?: unknown;
+      error?: unknown;
+    };
+
+    if (typeof body.error === 'string' && body.error.length > 0) {
+      throw new Error(body.error);
+    }
+
+    if (isPolicyVectorArray(body.actions) && body.actions.length > 0) {
+      return body.actions;
+    }
+
+    if (isPolicyVector(body.action)) {
+      return body.action;
+    }
+  }
+
+  if (isPolicyVectorArray(responseBody) && responseBody.length > 0) {
+    return responseBody;
+  }
+
+  if (isPolicyVector(responseBody)) {
+    return responseBody;
+  }
+
+  throw new Error('Remote policy response must include `action` or `actions`.');
+}
+
+function createHttpError(response: Response, responseBody: unknown) {
+  const suffix =
+    responseBody && typeof responseBody === 'object' && 'error' in responseBody
+      ? `: ${String((responseBody as { error?: unknown }).error)}`
+      : '';
+  return new Error(`Remote policy request failed with HTTP ${response.status}${suffix}`);
+}
+
+/**
+ * Run a policy whose inference step lives behind an HTTP JSON endpoint.
+ *
+ * The hook keeps `usePolicy` responsible for timing, queueing, and pause/reset
+ * behavior. This wrapper only builds requests, parses responses, and exposes
+ * request metadata for HUDs and debugging.
+ */
+export function useRemotePolicy(config: RemotePolicyConfig): RemotePolicyAPI {
+  const configRef = useRef(config);
+  configRef.current = config;
+  const requestCountRef = useRef(0);
+  const responseCountRef = useRef(0);
+  const remoteStatusRef = useRef<RemotePolicyStatus>('idle');
+  const lastRequestBodyRef = useRef<unknown>(null);
+  const lastResponseBodyRef = useRef<unknown>(null);
+  const lastHttpStatusRef = useRef<number | null>(null);
+  const lastRequestMsRef = useRef<number | null>(null);
+  const abortControllerRef = useRef<AbortController | null>(null);
+  const remoteEpochRef = useRef(0);
+
+  const policy = usePolicy({
+    ...config,
+    infer: async ({ observation, model, data }) => {
+      const cfg = configRef.current;
+      abortController(abortControllerRef.current, createAbortError('Remote policy request was superseded.'));
+      const controller = new AbortController();
+      abortControllerRef.current = controller;
+      const signal = createMergedAbortSignal(controller.signal, cfg.signal);
+      const remoteEpoch = remoteEpochRef.current;
+      const requestIndex = requestCountRef.current;
+      const requestInput: RemotePolicyRequestInput = {
+        observation,
+        model,
+        data,
+        reset: requestIndex === 0,
+        requestIndex,
+        signal,
+      };
+      requestCountRef.current += 1;
+
+      const requestStartedAt = now();
+      const body = await (
+        cfg.buildRequest?.(requestInput) ?? defaultBuildRemotePolicyRequest(requestInput)
+      );
+      signal.throwIfAborted();
+      if (remoteEpoch !== remoteEpochRef.current) {
+        throw createAbortError('Remote policy request was reset.');
+      }
+      lastRequestBodyRef.current = body;
+      remoteStatusRef.current = 'requesting';
+      cfg.onRequest?.({
+        ...requestInput,
+        body,
+        requestStartedAt,
+      });
+
+      let response: Response | null = null;
+      let responseBody: unknown = null;
+      try {
+        const headers = new Headers(cfg.headers);
+        if (!headers.has('content-type')) {
+          headers.set('content-type', 'application/json');
+        }
+
+        const fetcher = cfg.fetcher ?? fetch;
+        response = await fetcher(String(cfg.endpoint), {
+          ...cfg.requestInit,
+          method: cfg.method ?? 'POST',
+          credentials: cfg.credentials,
+          headers,
+          signal,
+          body: typeof body === 'string' ? body : JSON.stringify(body),
+        });
+        if (remoteEpoch === remoteEpochRef.current) {
+          lastHttpStatusRef.current = response.status;
+        }
+        responseBody = await (
+          cfg.readResponse?.(response) ?? defaultReadRemotePolicyResponse(response)
+        );
+        signal.throwIfAborted();
+        if (remoteEpoch !== remoteEpochRef.current) {
+          throw createAbortError('Remote policy request was reset.');
+        }
+        lastResponseBodyRef.current = responseBody;
+
+        if (!response.ok) {
+          throw createHttpError(response, responseBody);
+        }
+
+        const responseFinishedAt = now();
+        const info: RemotePolicyResponseInfo = {
+          ...requestInput,
+          body,
+          requestStartedAt,
+          response,
+          responseBody,
+          responseFinishedAt,
+          requestMs: responseFinishedAt - requestStartedAt,
+        };
+        if (remoteEpoch === remoteEpochRef.current) {
+          lastRequestMsRef.current = info.requestMs;
+          responseCountRef.current += 1;
+        }
+        cfg.onResponse?.(info);
+        const output = await (
+          cfg.parseResponse?.(responseBody, info) ??
+          defaultParseRemotePolicyResponse(responseBody)
+        );
+        if (remoteEpoch === remoteEpochRef.current) {
+          remoteStatusRef.current = 'ready';
+        }
+        return output;
+      } catch (error) {
+        if (response && remoteEpoch === remoteEpochRef.current) {
+          lastHttpStatusRef.current = response.status;
+        }
+        if (isAbortError(error) || signal.aborted) {
+          if (remoteEpoch === remoteEpochRef.current) {
+            remoteStatusRef.current = 'aborted';
+          }
+          throw error;
+        }
+        if (remoteEpoch === remoteEpochRef.current) {
+          remoteStatusRef.current = 'error';
+        }
+        throw error;
+      } finally {
+        if (abortControllerRef.current === controller) {
+          abortControllerRef.current = null;
+        }
+      }
+    },
+  });
+
+  return useMemo(() => {
+    const abort = (reason?: unknown) => {
+      abortController(abortControllerRef.current, reason);
+      if (abortControllerRef.current) {
+        remoteStatusRef.current = 'aborted';
+      }
+    };
+    const resetRemoteState = () => {
+      remoteEpochRef.current += 1;
+      abort(createAbortError('Remote policy request was reset.'));
+      requestCountRef.current = 0;
+      responseCountRef.current = 0;
+      remoteStatusRef.current = 'idle';
+      lastRequestBodyRef.current = null;
+      lastResponseBodyRef.current = null;
+      lastHttpStatusRef.current = null;
+      lastRequestMsRef.current = null;
+    };
+
+    return {
+      get isRunning() { return policy.isRunning; },
+      start: policy.start,
+      stop: () => {
+        if (configRef.current.abortOnStop ?? true) {
+          abort(createAbortError('Remote policy request was stopped.'));
+        }
+        policy.stop();
+        if (configRef.current.clearQueueOnStop) {
+          resetRemoteState();
+        }
+      },
+      clearQueue: policy.clearQueue,
+      abort,
+      reset: () => {
+        resetRemoteState();
+        policy.reset();
+      },
+      get inFlight() { return policy.inFlight; },
+      get queuedActions() { return policy.queuedActions; },
+      get lastObservation() { return policy.lastObservation; },
+      get lastAction() { return policy.lastAction; },
+      get lastError() { return policy.lastError; },
+      get remoteStatus() { return remoteStatusRef.current; },
+      get requestCount() { return requestCountRef.current; },
+      get responseCount() { return responseCountRef.current; },
+      get lastRequestBody() { return lastRequestBodyRef.current; },
+      get lastResponseBody() { return lastResponseBodyRef.current; },
+      get lastHttpStatus() { return lastHttpStatusRef.current; },
+      get lastRequestMs() { return lastRequestMsRef.current; },
+    };
+  }, [policy]);
+}

package/src/index.ts

@@ -88,12 +88,30 @@ export { useGravityCompensation } from './hooks/useGravityCompensation';
 export { useSensor, useSensors } from './hooks/useSensor';
 export { useJointState } from './hooks/useJointState';
 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 { useControlWriter } from './hooks/useControlWriter';
+export type {
+  ControlWriterConflict,
+  ControlWriterHandle,
+  ControlWriterOptions,
+  ControlWriterWriteOptions,
+} from './hooks/useControlWriter';
 export { useContacts, useContactEvents } from './hooks/useContacts';
+export { useContactHistory } from './hooks/useContactHistory';
+export type {
+  ContactHistoryEntry,
+  ContactHistoryHandle,
+  ContactHistoryOptions,
+} from './hooks/useContactHistory';
 export { useKeyboardTeleop } from './hooks/useKeyboardTeleop';
 export { useKeyboardIkTarget } from './hooks/useKeyboardIkTarget';
 export { usePolicy } from './hooks/usePolicy';
+export { useRemotePolicy } from './hooks/useRemotePolicy';
 export { useObservation } from './hooks/useObservation';
+export { useNamedObservation } from './hooks/useNamedObservation';
+export type { NamedObservationHandle } from './hooks/useNamedObservation';
 export { useTrajectoryPlayer } from './hooks/useTrajectoryPlayer';
 export { useTrajectoryRecorder } from './hooks/useTrajectoryRecorder';
 export { useGamepad } from './hooks/useGamepad';
@@ -104,6 +122,14 @@ export {
   useFrameCapture,
 } from './hooks/useFrameCapture';
 export { useCameraFrameCapture } from './hooks/useCameraFrameCapture';
+export {
+  usePolicyCameraFrames,
+  usePolicyCameraFramesFromMountedStreams,
+} from './hooks/usePolicyCameraFrames';
+export type {
+  MountedPolicyCameraFrameCaptureAPI,
+  MountedPolicyCameraFrameCaptureOptions,
+} from './hooks/usePolicyCameraFrames';
 export { useCameraSequenceRecorder } from './hooks/useCameraSequenceRecorder';
 export { useMountedCameraSequenceRecorder } from './hooks/useMountedCameraSequenceRecorder';
 export type {
@@ -114,12 +140,53 @@ export type {
   MountedCameraSequenceRecordResult,
 } from './hooks/useMountedCameraSequenceRecorder';
 export {
+  CAMERA_FRAME_CAPTURE_RENDER_USER_DATA_KEY,
+  CAMERA_FRAME_CAPTURE_PRE_RENDER_USER_DATA_KEY,
   CAPTURE_EXCLUDE_KEY,
   captureCameraFrame,
   captureCameraFrameBlob,
   createCameraFrameCaptureSession,
   renderCameraFrameToCanvas,
 } from './rendering/cameraFrameCapture';
+export {
+  capturePolicyCameraFrames,
+  capturePolicyCameraFramesFromMountedStreams,
+  createPolicyCameraFrameCapturePlan,
+  createPolicyCameraFrameCapturePlanFromApi,
+} from './policyCameraFrames';
+export {
+  applyPolicyActionToControls,
+  clampPolicyActionValue,
+} from './policyControls';
+export {
+  bodyPositionField,
+  createNamedObservationBuilder,
+  ctrlField,
+  geomPositionField,
+  qposField,
+  qvelField,
+  readNamedObservation,
+  sitePositionField,
+} from './policyObservation';
+export type {
+  CreatePolicyCameraFrameCapturePlanOptions,
+  PolicyCameraFrameCapturePlan,
+  PolicyCameraFrameCaptureTarget,
+  PolicyCameraFramePlanTarget,
+  PolicyCameraFrameStreamOptions,
+} from './policyCameraFrames';
+export type {
+  ApplyPolicyActionToControlsOptions,
+  ApplyPolicyActionToControlsResult,
+} from './policyControls';
+export type {
+  NamedObservationField,
+  NamedObservationInput,
+  NamedObservationLayoutItem,
+  NamedObservationMissing,
+  NamedObservationOptions,
+  NamedObservationResult,
+} from './policyObservation';
 export {
   createMountedCameraFrameSequenceManifest,
   createMountedCameraFrameSequenceReadiness,
@@ -216,11 +283,25 @@ export type {
   KeyboardIkTargetBinding,
   KeyboardIkTargetConfig,
   // Policy
+  PolicyAPI,
   PolicyConfig,
   PolicyVector,
+  PolicyActionChunk,
+  PolicyInferenceOutput,
   PolicyObservationInput,
   PolicyInferenceInput,
+  PolicyInferenceResult,
   PolicyActionInput,
+  RemotePolicyAPI,
+  RemotePolicyConfig,
+  RemotePolicyRequestInput,
+  RemotePolicyRequestInfo,
+  RemotePolicyResponseInfo,
+  RemotePolicyStatus,
+  PolicyCameraFrameCaptureAPI,
+  PolicyCameraFrameCaptureOptions,
+  PolicyCameraFrameCaptureResult,
+  PolicyCameraFrameStream,
   // Observations
   ObservationConfig,
   ObservationHandle,

package/src/policyCameraFrames.ts

@@ -0,0 +1,213 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Helpers for turning Three/MuJoCo camera captures into policy image payloads.
+ */
+
+import type {
+  CameraFrameCaptureResult,
+  MujocoSimAPI,
+  PolicyCameraFrameCaptureOptions,
+  PolicyCameraFrameCaptureResult,
+  PolicyCameraFrameStream,
+} from './types';
+import {
+  createMountedCameraFrameSequencePlan,
+  type MountedCameraFrameSequenceCameraOptions,
+  type MountedCameraFrameSequenceDefaults,
+  type MountedCameraFrameSequencePlan,
+  type ResolveMountedCameraFrameSourceOptions,
+} from './rendering/cameraFrameSource';
+
+export type PolicyCameraFrameCaptureTarget = Pick<MujocoSimAPI, 'captureCameraFrame'>;
+export type PolicyCameraFramePlanTarget = Pick<MujocoSimAPI, 'getCameras' | 'getSites' | 'getBodies'>;
+
+export type PolicyCameraFrameStreamOptions =
+  Partial<Omit<PolicyCameraFrameStream, 'key'>> &
+  MountedCameraFrameSequenceCameraOptions & {
+  /** Additional policy payload keys that should receive this stream's data URL. */
+  aliases?: readonly string[];
+};
+
+export interface CreatePolicyCameraFrameCapturePlanOptions
+  extends Omit<
+    ResolveMountedCameraFrameSourceOptions,
+    'cameras' | 'sites' | 'bodies'
+  > {
+  cameraKeys: readonly string[];
+  cameras?: ResolveMountedCameraFrameSourceOptions['cameras'];
+  sites?: ResolveMountedCameraFrameSourceOptions['sites'];
+  bodies?: ResolveMountedCameraFrameSourceOptions['bodies'];
+  defaults?: MountedCameraFrameSequenceDefaults;
+  streamOptions?: Record<string, PolicyCameraFrameStreamOptions>;
+  includeObservationImageAliases?: boolean;
+  requireAll?: boolean;
+}
+
+export interface PolicyCameraFrameCapturePlan
+  extends PolicyCameraFrameCaptureOptions {
+  cameraKeys: string[];
+  streams: PolicyCameraFrameStream[];
+  mountedPlan: MountedCameraFrameSequencePlan;
+  missingKeys: string[];
+}
+
+function addPolicyImageAliases(
+  images: Record<string, string>,
+  stream: PolicyCameraFrameStream,
+  frame: CameraFrameCaptureResult,
+  includeObservationImageAliases: boolean
+) {
+  const keys = new Set<string>();
+  keys.add(stream.key);
+  for (const alias of stream.aliases ?? []) keys.add(alias);
+  if (includeObservationImageAliases) {
+    keys.add(`observation.images.${stream.key}`);
+    for (const alias of stream.aliases ?? []) {
+      keys.add(`observation.images.${alias}`);
+    }
+  }
+  for (const key of keys) {
+    images[key] = frame.dataUrl;
+  }
+}
+
+function describeFrameSource(key: string, frame: CameraFrameCaptureResult) {
+  return `${key}:${frame.source.kind}`;
+}
+
+function hasExplicitPolicyCameraSource(
+  options: PolicyCameraFrameStreamOptions | undefined
+) {
+  return Boolean(
+    options?.camera ||
+    options?.position ||
+    options?.quaternion ||
+    options?.source
+  );
+}
+
+export function createPolicyCameraFrameCapturePlan(
+  options: CreatePolicyCameraFrameCapturePlanOptions
+): PolicyCameraFrameCapturePlan {
+  const {
+    cameraKeys,
+    defaults,
+    streamOptions,
+    includeObservationImageAliases,
+    requireAll,
+    ...sourceOptions
+  } = options;
+  const mountedPlan = createMountedCameraFrameSequencePlan(cameraKeys, {
+    ...sourceOptions,
+    defaults,
+    cameraOptions: streamOptions as
+      | Record<string, MountedCameraFrameSequenceCameraOptions>
+      | undefined,
+  });
+  const streams: PolicyCameraFrameStream[] = [];
+  const missingKeys = new Set(mountedPlan.missingKeys);
+
+  for (const key of cameraKeys) {
+    const perStreamOptions = streamOptions?.[key];
+    if (hasExplicitPolicyCameraSource(perStreamOptions)) {
+      missingKeys.delete(key);
+      streams.push({
+        ...defaults,
+        ...perStreamOptions,
+        key,
+        aliases: perStreamOptions?.aliases,
+      });
+      continue;
+    }
+
+    const mountedCamera = mountedPlan.cameras.find((camera) => camera.key === key);
+    if (!mountedCamera) continue;
+    const { key: _mountedKey, ...captureOptions } = mountedCamera;
+    streams.push({
+      ...captureOptions,
+      key,
+      aliases: perStreamOptions?.aliases,
+    });
+  }
+
+  const result: PolicyCameraFrameCapturePlan = {
+    cameraKeys: [...cameraKeys],
+    streams,
+    includeObservationImageAliases,
+    mountedPlan,
+    missingKeys: [...missingKeys],
+  };
+
+  if (requireAll && result.missingKeys.length > 0) {
+    throw new Error(
+      `Unable to resolve policy camera stream${
+        result.missingKeys.length === 1 ? '' : 's'
+      } for ${result.missingKeys.join(', ')}.`
+    );
+  }
+
+  return result;
+}
+
+export function createPolicyCameraFrameCapturePlanFromApi(
+  api: PolicyCameraFramePlanTarget,
+  options: Omit<CreatePolicyCameraFrameCapturePlanOptions, 'cameras' | 'sites' | 'bodies'>
+): PolicyCameraFrameCapturePlan {
+  return createPolicyCameraFrameCapturePlan({
+    ...options,
+    cameras: api.getCameras(),
+    sites: api.getSites(),
+    bodies: api.getBodies(),
+  });
+}
+
+export async function capturePolicyCameraFrames(
+  target: PolicyCameraFrameCaptureTarget,
+  options: PolicyCameraFrameCaptureOptions
+): Promise<PolicyCameraFrameCaptureResult> {
+  const includeObservationImageAliases =
+    options.includeObservationImageAliases ?? true;
+
+  const entries = await Promise.all(
+    options.streams.map(async ({ key, aliases, ...captureOptions }) => {
+      const frame = await target.captureCameraFrame(captureOptions);
+      return [key, { frame, aliases }] as const;
+    })
+  );
+
+  const frames: Record<string, CameraFrameCaptureResult> = {};
+  const images: Record<string, string> = {};
+  const sourceParts: string[] = [];
+
+  for (const [key, { frame, aliases }] of entries) {
+    const stream = { key, aliases };
+    frames[key] = frame;
+    addPolicyImageAliases(
+      images,
+      stream,
+      frame,
+      includeObservationImageAliases
+    );
+    sourceParts.push(describeFrameSource(key, frame));
+  }
+
+  return {
+    frames,
+    images,
+    sourceSummary: sourceParts.length > 0
+      ? sourceParts.join(' + ')
+      : 'not used by policy',
+    capturedAt: Date.now(),
+  };
+}
+
+export async function capturePolicyCameraFramesFromMountedStreams(
+  target: PolicyCameraFrameCaptureTarget & PolicyCameraFramePlanTarget,
+  options: Omit<CreatePolicyCameraFrameCapturePlanOptions, 'cameras' | 'sites' | 'bodies'>
+): Promise<PolicyCameraFrameCaptureResult & { plan: PolicyCameraFrameCapturePlan }> {
+  const plan = createPolicyCameraFrameCapturePlanFromApi(target, options);
+  const result = await capturePolicyCameraFrames(target, plan);
+  return { ...result, plan };
+}