npm - mujoco-react - Versions diffs - 8.10.0 → 9.0.0 - Mend

mujoco-react 8.10.0 → 9.0.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Files changed (40) hide show

package/README.md +81 -44
package/dist/chunk-33CV6HSV.js +400 -0
package/dist/chunk-33CV6HSV.js.map +1 -0
package/dist/index.d.ts +92 -24
package/dist/index.js +338 -54
package/dist/index.js.map +1 -1
package/dist/spark.d.ts +24 -3
package/dist/spark.js +91 -6
package/dist/spark.js.map +1 -1
package/dist/{types-FFW7ykBu.d.ts → types-izZlUweI.d.ts} +109 -16
package/package.json +1 -1
package/src/components/Body.tsx +3 -1
package/src/components/DragInteraction.tsx +1 -1
package/src/components/IkGizmo.tsx +2 -2
package/src/components/SceneRenderer.tsx +1 -1
package/src/components/TrajectoryPlayer.tsx +4 -1
package/src/components/VisualScenario.tsx +343 -6
package/src/core/MujocoCanvas.tsx +8 -1
package/src/core/MujocoPhysics.tsx +10 -4
package/src/core/MujocoSimProvider.tsx +15 -12
package/src/core/SceneLoader.ts +182 -3
package/src/core/createController.tsx +2 -2
package/src/hooks/useBodyState.ts +1 -1
package/src/hooks/useContacts.ts +1 -1
package/src/hooks/useCtrlNoise.ts +1 -1
package/src/hooks/useFrameCapture.ts +206 -0
package/src/hooks/useGamepad.ts +1 -1
package/src/hooks/useGravityCompensation.ts +1 -1
package/src/hooks/useIkController.ts +22 -13
package/src/hooks/useJointState.ts +1 -1
package/src/hooks/useKeyboardTeleop.ts +1 -1
package/src/hooks/usePolicy.ts +13 -9
package/src/hooks/useSensor.ts +1 -1
package/src/hooks/useTrajectoryPlayer.ts +4 -4
package/src/hooks/useTrajectoryRecorder.ts +1 -1
package/src/index.ts +35 -0
package/src/spark.tsx +138 -4
package/src/types.ts +128 -21
package/dist/chunk-KGFRKPLS.js +0 -186
package/dist/chunk-KGFRKPLS.js.map +0 -1

package/src/core/SceneLoader.ts CHANGED Viewed

@@ -477,8 +477,9 @@ function sceneObjectToXml(obj: SceneObject): string {
   const solref = obj.solref ? ` solref="${obj.solref}"` : '';
   const solimp = obj.solimp ? ` solimp="${obj.solimp}"` : '';
   const condim = obj.condim ? ` condim="${obj.condim}"` : '';
+  const group = obj.group !== undefined ? ` group="${obj.group}"` : '';
   // Always set contype/conaffinity=1 so objects collide regardless of model defaults
-  return `<body name="${obj.name}" pos="${pos}">${joint}<geom type="${obj.type}" size="${size}" rgba="${rgba}" contype="1" conaffinity="1"${mass}${friction}${solref}${solimp}${condim}/></body>`;
+  return `<body name="${obj.name}" pos="${pos}">${joint}<geom type="${obj.type}" size="${size}" rgba="${rgba}" contype="1" conaffinity="1"${mass}${friction}${solref}${solimp}${condim}${group}/></body>`;
 }
 /** Create virtual directory structure for a file path. */
@@ -527,6 +528,22 @@ function localFilePath(file: LocalMujocoFile): string {
   return normalizeVfsPath(file.webkitRelativePath || file.name);
 }
+function dirname(path: string): string {
+  const normalized = normalizeVfsPath(path);
+  const idx = normalized.lastIndexOf('/');
+  return idx === -1 ? '' : normalized.slice(0, idx + 1);
+}
+function relativeVfsPath(fromDir: string, targetPath: string): string {
+  const from = normalizeVfsPath(fromDir).split('/').filter(Boolean);
+  const target = normalizeVfsPath(targetPath).split('/').filter(Boolean);
+  while (from.length && target.length && from[0] === target[0]) {
+    from.shift();
+    target.shift();
+  }
+  return [...from.map(() => '..'), ...target].join('/') || '.';
+}
 function inferSceneFile(files: readonly LocalMujocoFile[], options?: LoadFromFilesOptions): string {
   if (options?.sceneFile) return normalizeVfsPath(options.sceneFile);
@@ -551,6 +568,7 @@ export function createSceneConfigFromFiles(
     src: '',
     sceneFile: inferSceneFile(fileArray, options),
     files: fileArray,
+    environmentFiles: options.environmentFiles?.map(normalizeVfsPath),
     homeJoints: options.homeJoints,
     xmlPatches: options.xmlPatches,
     sceneObjects: options.sceneObjects,
@@ -558,6 +576,137 @@ export function createSceneConfigFromFiles(
   };
 }
+const ENVIRONMENT_MERGE_SECTIONS = [
+  'asset',
+  'worldbody',
+  'contact',
+  'equality',
+  'tendon',
+  'sensor',
+  'keyframe',
+  'custom',
+  'extension',
+] as const;
+function directChild(parent: Element, tagName: string): Element | null {
+  const lower = tagName.toLowerCase();
+  for (const child of Array.from(parent.children)) {
+    if (child.tagName.toLowerCase() === lower) return child;
+  }
+  return null;
+}
+function ensureTopLevelSection(doc: XMLDocument, tagName: string): Element {
+  const root = doc.documentElement;
+  const existing = directChild(root, tagName);
+  if (existing) return existing;
+  const section = doc.createElement(tagName);
+  if (tagName === 'asset') {
+    const worldbody = directChild(root, 'worldbody');
+    if (worldbody) root.insertBefore(section, worldbody);
+    else root.appendChild(section);
+  } else {
+    root.appendChild(section);
+  }
+  return section;
+}
+function readCompilerDirs(doc: XMLDocument) {
+  const compiler = directChild(doc.documentElement, 'compiler');
+  const assetDir = compiler?.getAttribute('assetdir') || '';
+  return {
+    meshDir: compiler?.getAttribute('meshdir') || assetDir,
+    textureDir: compiler?.getAttribute('texturedir') || assetDir,
+  };
+}
+function isExternalPath(path: string): boolean {
+  return /^[a-z]+:\/\//i.test(path) || path.startsWith('package://') || path.startsWith('/');
+}
+function fileReferencePrefix(el: Element, compilerDirs: ReturnType<typeof readCompilerDirs>): string {
+  const tag = el.tagName.toLowerCase();
+  if (tag === 'mesh') return compilerDirs.meshDir ? compilerDirs.meshDir + '/' : '';
+  if (tag === 'texture' || tag === 'hfield') return compilerDirs.textureDir ? compilerDirs.textureDir + '/' : '';
+  return '';
+}
+function rewriteFileReferencesForMerge(node: Element, sourceFile: string, targetFile: string, sourceDoc: XMLDocument) {
+  const sourceDir = dirname(sourceFile);
+  const targetDir = dirname(targetFile);
+  const compilerDirs = readCompilerDirs(sourceDoc);
+  node.querySelectorAll('[file], [filename]').forEach((el) => {
+    const attr = el.hasAttribute('file') ? 'file' : 'filename';
+    const value = el.getAttribute(attr);
+    if (!value || isExternalPath(value)) return;
+    const sourceRelativePath = normalizeVfsPath(fileReferencePrefix(el, compilerDirs) + value);
+    const resolvedPath = normalizeVfsPath(sourceDir + sourceRelativePath);
+    el.setAttribute(attr, relativeVfsPath(targetDir, resolvedPath));
+  });
+}
+function hasParseError(doc: XMLDocument): boolean {
+  return doc.getElementsByTagName('parsererror').length > 0;
+}
+function composeEnvironmentXml(
+  sceneXml: string,
+  config: SceneConfig,
+  parser: DOMParser,
+  environmentXmlByPath: Map<string, string>
+): string {
+  const environmentFiles = config.environmentFiles?.map(normalizeVfsPath) ?? [];
+  if (!environmentFiles.length) return sceneXml;
+  const sceneDoc = parser.parseFromString(sceneXml, 'text/xml');
+  if (hasParseError(sceneDoc)) {
+    console.warn(`Could not compose environments: failed to parse ${config.sceneFile}`);
+    return sceneXml;
+  }
+  for (const environmentFile of environmentFiles) {
+    const environmentXml = environmentXmlByPath.get(environmentFile);
+    if (!environmentXml) {
+      console.warn(`Environment XML not found: ${environmentFile}`);
+      continue;
+    }
+    const environmentDoc = parser.parseFromString(environmentXml, 'text/xml');
+    if (hasParseError(environmentDoc)) {
+      console.warn(`Skipping environment XML with parse errors: ${environmentFile}`);
+      continue;
+    }
+    for (const sectionName of ENVIRONMENT_MERGE_SECTIONS) {
+      const environmentSection = directChild(environmentDoc.documentElement, sectionName);
+      if (!environmentSection?.children.length) continue;
+      const targetSection = ensureTopLevelSection(sceneDoc, sectionName);
+      for (const child of Array.from(environmentSection.children)) {
+        const imported = sceneDoc.importNode(child, true) as Element;
+        rewriteFileReferencesForMerge(imported, environmentFile, config.sceneFile, environmentDoc);
+        targetSection.appendChild(imported);
+      }
+    }
+  }
+  return new XMLSerializer().serializeToString(sceneDoc);
+}
+function findTextByConfiguredPath(textByPath: Map<string, string>, configuredPath: string): string | undefined {
+  const normalized = normalizeVfsPath(configuredPath);
+  const direct = textByPath.get(normalized);
+  if (direct) return direct;
+  const suffix = '/' + normalized;
+  for (const [path, text] of textByPath) {
+    if (path.endsWith(suffix) || path === normalized.split('/').pop()) return text;
+  }
+  return undefined;
+}
 function applyXmlPatches(text: string, fname: string, config: SceneConfig): string {
   let result = text;
   for (const patch of config.xmlPatches ?? []) {
@@ -627,10 +776,25 @@ async function loadSceneFromFiles(
     if (isModelTextFile(path)) {
       const text = applyXmlPatches(await file.text(), path, config);
       textByPath.set(path, text);
-      mujoco.FS.writeFile(`/working/${path}`, text);
     } else {
       mujoco.FS.writeFile(`/working/${path}`, new Uint8Array(await file.arrayBuffer()));
+      written.add(path);
     }
+  }
+  const environmentXmlByPath = new Map<string, string>();
+  for (const environmentFile of config.environmentFiles?.map(normalizeVfsPath) ?? []) {
+    const environmentXml = findTextByConfiguredPath(textByPath, environmentFile);
+    if (environmentXml) environmentXmlByPath.set(environmentFile, environmentXml);
+  }
+  for (const [path, text] of textByPath) {
+    const composedText = path === config.sceneFile
+      ? composeEnvironmentXml(text, config, parser, environmentXmlByPath)
+      : text;
+    textByPath.set(path, composedText);
+    ensureDir(mujoco, path);
+    mujoco.FS.writeFile(`/working/${path}`, composedText);
     written.add(path);
   }
@@ -689,6 +853,18 @@ export async function loadScene(
   const baseUrl = config.src.endsWith('/') ? config.src : config.src + '/';
+  const environmentXmlByPath = new Map<string, string>();
+  const environmentFiles = config.environmentFiles?.map(normalizeVfsPath) ?? [];
+  for (const environmentFile of environmentFiles) {
+    onProgress?.(`Downloading ${environmentFile}...`);
+    const res = await fetch(baseUrl + environmentFile);
+    if (!res.ok) {
+      console.warn(`Failed to fetch environment XML ${environmentFile}: ${res.status} ${res.statusText}`);
+      continue;
+    }
+    environmentXmlByPath.set(environmentFile, applyXmlPatches(await res.text(), environmentFile, config));
+  }
   const downloaded = new Set<string>();
   const xmlQueue: string[] = [config.sceneFile];
   const assetFiles: string[] = [];
@@ -714,7 +890,10 @@ export async function loadScene(
       continue;
     }
-    const text = applyXmlPatches(await res.text(), fname, config);
+    const patchedText = applyXmlPatches(await res.text(), fname, config);
+    const text = fname === config.sceneFile
+      ? composeEnvironmentXml(patchedText, config, parser, environmentXmlByPath)
+      : patchedText;
     ensureDir(mujoco, fname);
     mujoco.FS.writeFile(`/working/${fname}`, text);

package/src/core/createController.tsx CHANGED Viewed

@@ -43,7 +43,7 @@ export type ControllerComponent<TConfig> = React.FC<{
  * const MyController = createController<{ speed: number }>(
  *   { name: 'my-controller', defaultConfig: { speed: 1.0 } },
  *   function MyControllerImpl({ config }) {
- *     useBeforePhysicsStep((_model, data) => {
+ *     useBeforePhysicsStep(({ data }) => {
  *       data.ctrl[0] = config.speed;
  *     });
  *     return null;
@@ -100,7 +100,7 @@ export function createController<TConfig>(
  *   { name: 'useMyController', defaultConfig: { gain: 1.0 } },
  *   function useMyControllerImpl(config) {
  *     // config is MyConfig | null — hooks must be called unconditionally
- *     useBeforePhysicsStep((_model, data) => {
+ *     useBeforePhysicsStep(({ data }) => {
  *       if (!config) return;
  *       data.ctrl[0] = config.gain * Math.sin(data.time);
  *     });

package/src/hooks/useBodyState.ts CHANGED Viewed

@@ -29,7 +29,7 @@ export function useBodyState(name: Bodies): BodyStateResult {
     bodyIdRef.current = findBodyByName(model, name);
   }, [name, status, mjModelRef]);
-  useAfterPhysicsStep((_model, data) => {
+  useAfterPhysicsStep(({ data }) => {
     const bid = bodyIdRef.current;
     if (bid < 0) return;

package/src/hooks/useContacts.ts CHANGED Viewed

@@ -60,7 +60,7 @@ export function useContacts(
     bodyResolvedRef.current = true;
   }, [bodyName, status, mjModelRef]);
-  useAfterPhysicsStep((model, data) => {
+  useAfterPhysicsStep(({ model, data }) => {
     // Resolve body id lazily once model exists, to avoid missing the first ready frame.
     if (bodyName && !bodyResolvedRef.current) {
       bodyIdRef.current = findBodyByName(model, bodyName);

package/src/hooks/useCtrlNoise.ts CHANGED Viewed

@@ -30,7 +30,7 @@ export function useCtrlNoise(config: CtrlNoiseConfig = {}) {
   configRef.current = config;
   const noiseRef = useRef<Float64Array | null>(null);
-  useBeforePhysicsStep((_model, data) => {
+  useBeforePhysicsStep(({ data }) => {
     const cfg = configRef.current;
     if (cfg.enabled === false) return;

package/src/hooks/useFrameCapture.ts ADDED Viewed

@@ -0,0 +1,206 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * useFrameCapture — still-frame capture for canvas-backed MuJoCo/R3F scenes.
+ */
+import { useCallback, useState } from 'react';
+import type React from 'react';
+export type FrameCaptureStatus = 'idle' | 'capturing' | 'captured' | 'error';
+export type FrameCaptureTarget =
+  | HTMLCanvasElement
+  | HTMLElement
+  | null
+  | undefined;
+export type FrameCaptureTargetRef =
+  React.RefObject<HTMLCanvasElement | HTMLElement | null>;
+export interface FrameCaptureOptions {
+  target?: FrameCaptureTarget | FrameCaptureTargetRef;
+  type?: string;
+  quality?: number;
+  waitForAnimationFrame?: boolean;
+}
+export interface FrameCaptureResult {
+  canvas: HTMLCanvasElement;
+  dataUrl: string;
+  type: string;
+}
+export interface FrameCaptureBlobResult {
+  canvas: HTMLCanvasElement;
+  blob: Blob;
+  type: string;
+}
+export interface FrameCaptureAPI {
+  status: FrameCaptureStatus;
+  error: Error | null;
+  isCapturing: boolean;
+  capture: (options?: FrameCaptureOptions) => Promise<FrameCaptureResult>;
+  captureBlob: (
+    options?: FrameCaptureOptions
+  ) => Promise<FrameCaptureBlobResult>;
+  reset: () => void;
+}
+function isTargetRef(
+  target: FrameCaptureOptions['target']
+): target is FrameCaptureTargetRef {
+  return Boolean(target && typeof target === 'object' && 'current' in target);
+}
+function resolveCanvasTarget(
+  target: FrameCaptureOptions['target']
+): HTMLCanvasElement {
+  const resolvedTarget = isTargetRef(target) ? target.current : target;
+  if (!resolvedTarget) {
+    throw new Error('No frame capture target is available.');
+  }
+  if (resolvedTarget instanceof HTMLCanvasElement) {
+    return resolvedTarget;
+  }
+  const canvas = resolvedTarget.querySelector('canvas');
+  if (!canvas) {
+    throw new Error('Frame capture target does not contain a canvas.');
+  }
+  return canvas;
+}
+function waitForNextAnimationFrame() {
+  return new Promise<void>((resolve) => {
+    requestAnimationFrame(() => resolve());
+  });
+}
+/**
+ * Capture the current canvas frame as a data URL.
+ *
+ * For WebGL scenes, create the renderer with `preserveDrawingBuffer: true`
+ * when you need deterministic captures after the frame has presented.
+ */
+export async function captureFrame(
+  options: FrameCaptureOptions
+): Promise<FrameCaptureResult> {
+  const type = options.type ?? 'image/png';
+  const canvas = resolveCanvasTarget(options.target);
+  if (options.waitForAnimationFrame ?? true) {
+    await waitForNextAnimationFrame();
+  }
+  return {
+    canvas,
+    dataUrl: canvas.toDataURL(type, options.quality),
+    type,
+  };
+}
+/**
+ * Capture the current canvas frame as a Blob.
+ */
+export async function captureFrameBlob(
+  options: FrameCaptureOptions
+): Promise<FrameCaptureBlobResult> {
+  const type = options.type ?? 'image/png';
+  const canvas = resolveCanvasTarget(options.target);
+  if (options.waitForAnimationFrame ?? true) {
+    await waitForNextAnimationFrame();
+  }
+  const blob = await new Promise<Blob>((resolve, reject) => {
+    canvas.toBlob(
+      (nextBlob) => {
+        if (nextBlob) {
+          resolve(nextBlob);
+        } else {
+          reject(new Error('Canvas frame capture did not produce a Blob.'));
+        }
+      },
+      type,
+      options.quality
+    );
+  });
+  return { canvas, blob, type };
+}
+/**
+ * React state wrapper around `captureFrame` and `captureFrameBlob`.
+ */
+export function useFrameCapture(
+  defaultOptions: FrameCaptureOptions = {}
+): FrameCaptureAPI {
+  const [status, setStatus] = useState<FrameCaptureStatus>('idle');
+  const [error, setError] = useState<Error | null>(null);
+  const reset = useCallback(() => {
+    setStatus('idle');
+    setError(null);
+  }, []);
+  const capture = useCallback(
+    async (options: FrameCaptureOptions = {}) => {
+      setStatus('capturing');
+      setError(null);
+      try {
+        const result = await captureFrame({ ...defaultOptions, ...options });
+        setStatus('captured');
+        return result;
+      } catch (nextError) {
+        const error =
+          nextError instanceof Error
+            ? nextError
+            : new Error('Unable to capture the current canvas frame.');
+        setError(error);
+        setStatus('error');
+        throw error;
+      }
+    },
+    [defaultOptions]
+  );
+  const captureBlob = useCallback(
+    async (options: FrameCaptureOptions = {}) => {
+      setStatus('capturing');
+      setError(null);
+      try {
+        const result = await captureFrameBlob({
+          ...defaultOptions,
+          ...options,
+        });
+        setStatus('captured');
+        return result;
+      } catch (nextError) {
+        const error =
+          nextError instanceof Error
+            ? nextError
+            : new Error('Unable to capture the current canvas frame.');
+        setError(error);
+        setStatus('error');
+        throw error;
+      }
+    },
+    [defaultOptions]
+  );
+  return {
+    status,
+    error,
+    isCapturing: status === 'capturing',
+    capture,
+    captureBlob,
+    reset,
+  };
+}

package/src/hooks/useGamepad.ts CHANGED Viewed

@@ -50,7 +50,7 @@ export function useGamepad(config: GamepadConfig) {
     }
   }, [config.axes, config.buttons, status, mjModelRef]);
-  useBeforePhysicsStep((_model, data) => {
+  useBeforePhysicsStep(({ data }) => {
     const cfg = configRef.current;
     if (cfg.enabled === false) return;

package/src/hooks/useGravityCompensation.ts CHANGED Viewed

@@ -13,7 +13,7 @@ import { useBeforePhysicsStep } from '../core/MujocoSimProvider';
  * hook (and DragInteraction) compose correctly — both add to a clean slate.
  */
 export function useGravityCompensation(enabled = true): void {
-  useBeforePhysicsStep((model, data) => {
+  useBeforePhysicsStep(({ model, data }) => {
     if (!enabled) return;
     for (let i = 0; i < model.nv; i++) {
       data.qfrc_applied[i] += data.qfrc_bias[i];

package/src/hooks/useIkController.ts CHANGED Viewed

@@ -10,7 +10,7 @@ import { createControllerHook } from '../core/createController';
 import { useMujocoContext, useBeforePhysicsStep } from '../core/MujocoSimProvider';
 import { GenericIK } from '../core/GenericIK';
 import { createContiguousControlGroup, findSiteByName, resolveControlGroup } from '../core/SceneLoader';
-import type { ControlGroupInfo, IkConfig, IkContextValue, IKSolveFn, MujocoData } from '../types';
+import type { ControlGroupInfo, IkConfig, IkContextValue, IKSolveFn, IkSolveInput, MujocoData } from '../types';
 // Preallocated temp for syncGizmoToSite
 const _syncMat4 = new THREE.Matrix4();
@@ -84,16 +84,16 @@ export const useIkController = createControllerHook<IkConfig, IkContextValue>(
     // IK solve function
     const ikSolveFn = useCallback(
-      (pos: THREE.Vector3, quat: THREE.Quaternion, currentQ: number[]): number[] | null => {
+      ({ position, quaternion, currentQ, context }: IkSolveInput): number[] | null => {
         if (!config) return null;
-        if (config.ikSolveFn) return config.ikSolveFn(pos, quat, currentQ);
+        if (config.ikSolveFn) return config.ikSolveFn({ position, quaternion, currentQ, context });
         const model = mjModelRef.current;
         const data = mjDataRef.current;
         const controlGroup = controlGroupRef.current;
         if (!model || !data || !controlGroup || siteIdRef.current === -1) return null;
         return genericIkRef.current.solve(
           model, data, siteIdRef.current, controlGroup.qposAdr,
-          pos, quat, currentQ,
+          position, quaternion, currentQ,
           { damping: config.damping, maxIterations: config.maxIterations },
         );
       },
@@ -128,7 +128,7 @@ export const useIkController = createControllerHook<IkConfig, IkContextValue>(
     });
     // IK solve in physics loop
-    useBeforePhysicsStep((model, data) => {
+    useBeforePhysicsStep(({ model, data }) => {
       if (!config || !ikEnabledRef.current) {
         ikCalculatingRef.current = false;
         return;
@@ -142,13 +142,22 @@ export const useIkController = createControllerHook<IkConfig, IkContextValue>(
       const currentQ = Array.from(controlGroup.readQpos(data));
       const solution = config.ikSolveFn
-        ? config.ikSolveFn(target.position, target.quaternion, currentQ, {
-            model,
-            data,
-            siteId: siteIdRef.current,
-            controlGroup,
+        ? config.ikSolveFn({
+            position: target.position,
+            quaternion: target.quaternion,
+            currentQ,
+            context: {
+              model,
+              data,
+              siteId: siteIdRef.current,
+              controlGroup,
+            },
           })
-        : ikSolveFnRef.current(target.position, target.quaternion, currentQ);
+        : ikSolveFnRef.current({
+            position: target.position,
+            quaternion: target.quaternion,
+            currentQ,
+          });
       if (solution) {
         controlGroup.writeCtrl(data, solution);
       }
@@ -192,8 +201,8 @@ export const useIkController = createControllerHook<IkConfig, IkContextValue>(
     }, [mjDataRef]);
     const solveIK = useCallback(
-      (pos: THREE.Vector3, quat: THREE.Quaternion, currentQ: number[]): number[] | null => {
-        return ikSolveFnRef.current(pos, quat, currentQ);
+      (input: IkSolveInput): number[] | null => {
+        return ikSolveFnRef.current(input);
       },
       [],
     );

package/src/hooks/useJointState.ts CHANGED Viewed

@@ -59,7 +59,7 @@ export function useJointState(name: Joints): JointStateResult {
     jointIdRef.current = -1;
   }, [name, status, mjModelRef]);
-  useAfterPhysicsStep((_model, data) => {
+  useAfterPhysicsStep(({ data }) => {
     if (jointIdRef.current < 0) return;
     const qa = qposAdrRef.current;
     const da = dofAdrRef.current;

package/src/hooks/useKeyboardTeleop.ts CHANGED Viewed

@@ -70,7 +70,7 @@ export function useKeyboardTeleop(config: KeyboardTeleopConfig) {
   }, []);
   // Apply bindings each physics frame
-  useBeforePhysicsStep((_model, data) => {
+  useBeforePhysicsStep(({ data }) => {
     if (!enabledRef.current) return;
     const bindings = bindingsRef.current;
     const cache = actuatorCacheRef.current;

package/src/hooks/usePolicy.ts CHANGED Viewed

@@ -6,7 +6,7 @@
  */
 import { useRef } from 'react';
-import { useMujocoContext, useBeforePhysicsStep } from '../core/MujocoSimProvider';
+import { useBeforePhysicsStep } from '../core/MujocoSimProvider';
 import type { PolicyConfig } from '../types';
 /**
@@ -20,14 +20,15 @@ import type { PolicyConfig } from '../types';
  * @returns { step, isRunning } control handles
  */
 export function usePolicy(config: PolicyConfig) {
-  const { mjModelRef } = useMujocoContext();
   const lastActionTimeRef = useRef(0);
+  const lastObservationRef = useRef<ReturnType<PolicyConfig['onObservation']> | null>(null);
   const lastActionRef = useRef<Float32Array | Float64Array | number[] | null>(null);
-  const isRunningRef = useRef(true);
+  const isRunningRef = useRef(config.enabled ?? true);
   const configRef = useRef(config);
   configRef.current = config;
+  isRunningRef.current = config.enabled ?? isRunningRef.current;
-  useBeforePhysicsStep((model, data) => {
+  useBeforePhysicsStep(({ model, data }) => {
     if (!isRunningRef.current) return;
     const cfg = configRef.current;
@@ -37,13 +38,15 @@ export function usePolicy(config: PolicyConfig) {
     // Check if it's time for a new action
     if (data.time - lastActionTimeRef.current >= interval) {
       // Build observation
-      const obs = cfg.onObservation(model, data);
+      const observation = cfg.onObservation({ model, data });
+      const action = cfg.infer ? cfg.infer({ observation, model, data }) : observation;
-      // Apply action (consumer does inference inline or uses cached result)
-      cfg.onAction(obs, model, data);
+      // Apply action. If `infer` is omitted, this preserves the legacy inline-controller path.
+      cfg.onAction({ action, observation, model, data });
       lastActionTimeRef.current = data.time;
-      lastActionRef.current = obs;
+      lastObservationRef.current = observation;
+      lastActionRef.current = action;
     }
   });
@@ -51,6 +54,7 @@ export function usePolicy(config: PolicyConfig) {
     get isRunning() { return isRunningRef.current; },
     start: () => { isRunningRef.current = true; },
     stop: () => { isRunningRef.current = false; },
-    get lastObservation() { return lastActionRef.current; },
+    get lastObservation() { return lastObservationRef.current; },
+    get lastAction() { return lastActionRef.current; },
   };
 }

package/src/hooks/useSensor.ts CHANGED Viewed

@@ -39,7 +39,7 @@ export function useSensor(name: Sensors): SensorHandle {
   }, [name, status, mjModelRef]);
   // Update every frame after physics step
-  useAfterPhysicsStep((_model, data) => {
+  useAfterPhysicsStep(({ data }) => {
     if (sensorIdRef.current < 0) return;
     const adr = sensorAdrRef.current;
     const dim = sensorDimRef.current;