mujoco-react 8.4.1 → 8.5.0

package/README.md

@@ -100,6 +100,58 @@ function ResetButton() {
 }
 ```
 
+## Map Controls to Joints
+
+Use control groups when a robot's actuator order does not match a simple `qpos[0..n]` layout:
+
+```tsx
+import { useRef } from "react";
+import { resolveControlGroup, useBeforePhysicsStep } from "mujoco-react";
+import type { ControlGroupInfo } from "mujoco-react";
+
+function HoldTcpPose() {
+  const armRef = useRef<ControlGroupInfo | null>(null);
+
+  useBeforePhysicsStep((model, data) => {
+    armRef.current ??= resolveControlGroup(model, { siteName: "tcp" });
+    if (!armRef.current) return;
+
+    armRef.current.writeCtrl(data, armRef.current.readQpos(data));
+  });
+
+  return null;
+}
+```
+
+`resolveControlGroup()` accepts `{ siteName }`, `{ bodyName }`, `{ joints }`, or `{ actuators }`. Selectors can be a name, ordered name array, regex, or predicate.
+
+## Build Observations
+
+Build policy-ready observation vectors from common MuJoCo state without hard-coding offsets:
+
+```tsx
+import { buildObservation, useBeforePhysicsStep } from "mujoco-react";
+
+function PolicyDriver() {
+  useBeforePhysicsStep((model, data) => {
+    const obs = buildObservation(model, data, {
+      qpos: true,
+      qvel: true,
+      ctrl: true,
+      sensors: ["imu_gyro", "imu_accel"],
+      sites: ["tcp"],
+      projectedGravity: "torso",
+    });
+
+    runPolicy(obs.values, obs.layout);
+  });
+
+  return null;
+}
+```
+
+Use `output: "float64"` when a downstream model expects double precision. Named resources are skipped when absent, so `obs.layout` is the source of truth for the current model.
+
 ## WebSocket Control
 
 Stream actuator commands over a WebSocket and send simulation state back. Use a schema validator such as Zod at this boundary because socket messages are untrusted app input (`npm install zod` for this example):
@@ -109,32 +161,28 @@ import { useEffect, useRef } from "react";
 import { z } from "zod";
 import { useMujoco, useBeforePhysicsStep, useAfterPhysicsStep } from "mujoco-react";
 
-const CtrlCommand = z.object({
-  type: z.literal("ctrl_command"),
-  ctrl: z.array(z.number()),
-});
-
-type CtrlCommand = z.infer<typeof CtrlCommand>;
-
-function parseSocketMessage(data: string): CtrlCommand | null {
+const CtrlCommand = z.preprocess((data) => {
   try {
-    const parsed = CtrlCommand.safeParse(JSON.parse(data));
-    return parsed.success ? parsed.data : null;
+    return typeof data === "string" ? JSON.parse(data) : data;
   } catch {
-    return null;
+    return undefined;
   }
-}
+}, z.object({
+  type: z.literal("ctrl_command"),
+  ctrl: z.array(z.number()),
+}));
 
 function useWebSocketControls(url: string) {
   const wsRef = useRef<WebSocket | null>(null);
-  const latestCommandRef = useRef<CtrlCommand | null>(null);
+  const latestCtrlRef = useRef<number[] | null>(null);
 
   useEffect(() => {
     const ws = new WebSocket(url);
     wsRef.current = ws;
 
     ws.onmessage = (evt) => {
-      latestCommandRef.current = parseSocketMessage(evt.data);
+      const command = CtrlCommand.safeParse(evt.data);
+      if (command.success) latestCtrlRef.current = command.data.ctrl;
     };
 
     return () => ws.close();
@@ -142,10 +190,10 @@ function useWebSocketControls(url: string) {
 
   // Apply incoming actuator controls each physics step.
   useBeforePhysicsStep((model, data) => {
-    const command = latestCommandRef.current;
-    if (!command) return;
-    for (let i = 0; i < Math.min(command.ctrl.length, model.nu); i++) {
-      data.ctrl[i] = command.ctrl[i];
+    const ctrl = latestCtrlRef.current;
+    if (!ctrl) return;
+    for (let i = 0; i < Math.min(ctrl.length, model.nu); i++) {
+      data.ctrl[i] = ctrl[i];
     }
   });
 
@@ -659,13 +707,19 @@ useGamepad({
 Framework-agnostic decimation loop for RL policies:
 
 ```tsx
-const { step, isRunning } = usePolicy({
+const obs = useObservation({ qpos: true, qvel: true, projectedGravity: "torso" });
+
+const policy = usePolicy({
   frequency: 50,
-  onObservation: (model, data) => buildObs(model, data),
+  onObservation: () => obs.readValues(),
   onAction: (action, model, data) => applyAction(action, data),
 });
 ```
 
+### `buildObservation(model, data, config)` / `useObservation(config)`
+
+Build a flat `Float32Array` or `Float64Array` plus a layout map from qpos, qvel, ctrl, actuator activations, sensordata, named sensors, named site positions, and projected gravity.
+
 ### `useTrajectoryRecorder(config)` / `useTrajectoryPlayer(trajectory, config)`
 
 Record and play back simulation trajectories:
@@ -770,7 +824,11 @@ The full API object available via `ref` or `useMujoco()` (when `isReady`):
 | `setQpos(values)` / `getQpos()` | Direct qpos access |
 | `setQvel(values)` / `getQvel()` | Direct qvel access |
 | `setCtrl(nameOrValues, value?)` | Set control by name or batch |
-| `getCtrl(name?)` | Get control values |
+| `getCtrl()` | Get all control values |
+| `getControlMap()` | Map directly actuated scalar joints to qpos/dof/ctrl addresses |
+| `getActuatedJoints()` | List scalar hinge/slide joints with matching actuators |
+| `resolveControlGroup(selector)` | Resolve qpos/ctrl mapping from a site, body, joint selector, or actuator selector |
+| `buildObservation(model, data, config)` | Build policy-ready observation vectors with layout metadata |
 | `applyKeyframe(nameOrIndex)` | Apply a keyframe |
 | `getKeyframeNames()` / `getKeyframeCount()` | Keyframe introspection |
 
@@ -861,7 +919,6 @@ Features planned but not yet implemented:
 | **User-uploaded model loading** | P2 | `loadFromFiles(FileList)` -- detect meshdir, write to VFS |
 | **URDF loading** | P2 | Load URDF models via MuJoCo's built-in URDF compiler |
 | **XML mutation / recompile** | P1 | `addBody()`, `removeBody()`, `recompile()` for runtime XML editing |
-| **Observation builder utilities** | P2 | Helpers for projected gravity, joint positions/velocities for RL |
 | **Physics interpolation** | P1 | Smooth rendering between physics ticks for very high refresh displays |
 | **Instanced geom rendering** | P2 | `<InstancedGeomRenderer />` for particle/granular sims |
 | **Web Worker physics** | P2 | Run `mj_step` off main thread via SharedArrayBuffer |

package/dist/index.d.ts

@@ -477,6 +477,44 @@ interface PolicyConfig {
     onObservation: (model: MujocoModel, data: MujocoData) => Float32Array | Float64Array | number[];
     onAction: (action: Float32Array | Float64Array | number[], model: MujocoModel, data: MujocoData) => void;
 }
+type ObservationOutput = 'float32' | 'float64';
+interface ObservationConfig {
+    /** Include scalar simulation time. */
+    time?: boolean;
+    /** Include all qpos values. */
+    qpos?: boolean;
+    /** Include all qvel values. */
+    qvel?: boolean;
+    /** Include all ctrl values. */
+    ctrl?: boolean;
+    /** Include all actuator activation values. */
+    act?: boolean;
+    /** Include all raw sensordata values. */
+    sensordata?: boolean;
+    /** Include named sensor values in the configured order. */
+    sensors?: readonly Sensors[];
+    /** Include named site world positions in the configured order. */
+    sites?: readonly Sites[];
+    /** Include world gravity projected into each named body's local frame. */
+    projectedGravity?: Bodies | readonly Bodies[];
+    /** Output array type. Defaults to Float32Array. */
+    output?: ObservationOutput;
+}
+interface ObservationLayoutItem {
+    name: string;
+    start: number;
+    size: number;
+}
+interface ObservationResult {
+    values: Float32Array | Float64Array;
+    layout: ObservationLayoutItem[];
+}
+interface ObservationHandle {
+    /** Read a fresh observation from the current live MuJoCo model/data refs. */
+    read(): ObservationResult;
+    /** Read just the vector values for policy inference. */
+    readValues(): Float32Array | Float64Array;
+}
 interface DebugProps {
     showGeoms?: boolean;
     showSites?: boolean;
@@ -831,6 +869,19 @@ interface LoadResult {
  */
 declare function loadScene(mujoco: MujocoModule, config: SceneConfig, onProgress?: (msg: string) => void): Promise<LoadResult>;
 
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+/**
+ * Build a flat observation vector plus a layout map from live MuJoCo state.
+ *
+ * Missing named resources are skipped. The returned layout is the source of
+ * truth for the vector produced from the current model.
+ */
+declare function buildObservation(model: MujocoModel, data: MujocoData, config: ObservationConfig): ObservationResult;
+
 /**
  * @license
  * SPDX-License-Identifier: Apache-2.0
@@ -1182,6 +1233,19 @@ declare function usePolicy(config: PolicyConfig): {
     readonly lastObservation: Float64Array<ArrayBufferLike> | Float32Array<ArrayBufferLike> | number[] | null;
 };
 
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+/**
+ * Live observation reader for policy loops and telemetry.
+ *
+ * The handle is stable; call `read()` inside callbacks to sample the latest
+ * MuJoCo model/data state without forcing React renders.
+ */
+declare function useObservation(config: ObservationConfig): ObservationHandle;
+
 /**
  * @license
  * SPDX-License-Identifier: Apache-2.0
@@ -1386,4 +1450,4 @@ interface CameraAnimationAPI {
  */
 declare function useCameraAnimation(): CameraAnimationAPI;
 
-export { type ActuatedJointInfo, type ActuatorInfo, type Actuators, type Bodies, Body, type BodyInfo, type BodyProps, type BodyStateResult, type CameraAnimationAPI, type ContactInfo, ContactListener, type ContactListenerProps, ContactMarkers, type ControlGroupInfo, type ControlGroupSelector, type ControlJointInfo, type ControllerComponent, type ControllerOptions, type CtrlHandle, Debug, type DebugProps, DragInteraction, type DragInteractionProps, FlexRenderer, type GeomInfo, type Geoms, type IKSolveFn, type IkConfig, type IkContextValue, IkGizmo, type IkGizmoProps, type JointInfo, type JointStateResult, type Joints, type KeyBinding, type KeyboardTeleopConfig, type Keyframes, type ModelOptions, MujocoCanvas, type MujocoCanvasProps, type MujocoContact, type MujocoContactArray, type MujocoContextValue, type MujocoData, type MujocoLoader, type MujocoLoaderOptions, type MujocoModel, type MujocoModule, MujocoPhysics, type MujocoPhysicsProps, MujocoProvider, type MujocoProviderProps, type MujocoSimAPI, MujocoSimProvider, type MujocoWasmVariant, type PhysicsConfig, type PhysicsStepCallback, type PlaybackState, type PolicyConfig, type RayHit, type Register, type ResourceSelector, type SceneConfig, SceneLights, type SceneLightsProps, type SceneMarker, type SceneObject, type SensorHandle, type SensorInfo, type SensorResult, type Sensors, type SiteInfo, type SitePositionResult, type Sites, type StateSnapshot, TendonRenderer, type TrajectoryData, type TrajectoryFrame, type TrajectoryInput, TrajectoryPlayer, type TrajectoryPlayerProps, type XmlPatch, createContiguousControlGroup, createController, createControllerHook, findActuatorByName, findBodyByName, findGeomByName, findJointByName, findKeyframeByName, findSensorByName, findSiteByName, findTendonByName, getActuatedJoints, getContact, getControlMap, getName, loadScene, resolveControlGroup, useActuators, useAfterPhysicsStep, useBeforePhysicsStep, useBodyMeshes, useBodyState, useCameraAnimation, useContactEvents, useContacts, useCtrl, useCtrlNoise, useGamepad, useGravityCompensation, useIkController, useJointState, useKeyboardTeleop, useMujoco, useMujocoWasm, usePolicy, useSceneLights, useSelectionHighlight, useSensor, useSensors, useSitePosition, useTrajectoryPlayer, useTrajectoryRecorder, useVideoRecorder };
+export { type ActuatedJointInfo, type ActuatorInfo, type Actuators, type Bodies, Body, type BodyInfo, type BodyProps, type BodyStateResult, type CameraAnimationAPI, type ContactInfo, ContactListener, type ContactListenerProps, ContactMarkers, type ControlGroupInfo, type ControlGroupSelector, type ControlJointInfo, type ControllerComponent, type ControllerOptions, type CtrlHandle, Debug, type DebugProps, DragInteraction, type DragInteractionProps, FlexRenderer, type GeomInfo, type Geoms, type IKSolveFn, type IkConfig, type IkContextValue, IkGizmo, type IkGizmoProps, type JointInfo, type JointStateResult, type Joints, type KeyBinding, type KeyboardTeleopConfig, type Keyframes, type ModelOptions, MujocoCanvas, type MujocoCanvasProps, type MujocoContact, type MujocoContactArray, type MujocoContextValue, type MujocoData, type MujocoLoader, type MujocoLoaderOptions, type MujocoModel, type MujocoModule, MujocoPhysics, type MujocoPhysicsProps, MujocoProvider, type MujocoProviderProps, type MujocoSimAPI, MujocoSimProvider, type MujocoWasmVariant, type ObservationConfig, type ObservationHandle, type ObservationLayoutItem, type ObservationOutput, type ObservationResult, type PhysicsConfig, type PhysicsStepCallback, type PlaybackState, type PolicyConfig, type RayHit, type Register, type ResourceSelector, type SceneConfig, SceneLights, type SceneLightsProps, type SceneMarker, type SceneObject, type SensorHandle, type SensorInfo, type SensorResult, type Sensors, type SiteInfo, type SitePositionResult, type Sites, type StateSnapshot, TendonRenderer, type TrajectoryData, type TrajectoryFrame, type TrajectoryInput, TrajectoryPlayer, type TrajectoryPlayerProps, type XmlPatch, buildObservation, createContiguousControlGroup, createController, createControllerHook, findActuatorByName, findBodyByName, findGeomByName, findJointByName, findKeyframeByName, findSensorByName, findSiteByName, findTendonByName, getActuatedJoints, getContact, getControlMap, getName, loadScene, resolveControlGroup, useActuators, useAfterPhysicsStep, useBeforePhysicsStep, useBodyMeshes, useBodyState, useCameraAnimation, useContactEvents, useContacts, useCtrl, useCtrlNoise, useGamepad, useGravityCompensation, useIkController, useJointState, useKeyboardTeleop, useMujoco, useMujocoWasm, useObservation, usePolicy, useSceneLights, useSelectionHighlight, useSensor, useSensors, useSitePosition, useTrajectoryPlayer, useTrajectoryRecorder, useVideoRecorder };

package/dist/index.js

@@ -1953,6 +1953,81 @@ var MujocoPhysics = forwardRef(
     );
   }
 );
+
+// src/core/ObservationBuilder.ts
+function append(values, layout, name, chunk) {
+  const start = values.length;
+  for (let i = 0; i < chunk.length; i++) values.push(chunk[i] ?? 0);
+  layout.push({ name, start, size: chunk.length });
+}
+function appendScalar(values, layout, name, value) {
+  const start = values.length;
+  values.push(value);
+  layout.push({ name, start, size: 1 });
+}
+function namedBodyList(names) {
+  if (!names) return [];
+  return typeof names === "string" ? [names] : names;
+}
+function normalizedGravity(model) {
+  const gx = model.opt.gravity[0] ?? 0;
+  const gy = model.opt.gravity[1] ?? 0;
+  const gz = model.opt.gravity[2] ?? -9.81;
+  const length = Math.hypot(gx, gy, gz);
+  if (length === 0) return [0, 0, 0];
+  return [gx / length, gy / length, gz / length];
+}
+function rotateWorldVectorToBody(data, bodyId, world) {
+  const adr = bodyId * 4;
+  const w = data.xquat[adr] ?? 1;
+  const x = -(data.xquat[adr + 1] ?? 0);
+  const y = -(data.xquat[adr + 2] ?? 0);
+  const z = -(data.xquat[adr + 3] ?? 0);
+  const vx = world[0];
+  const vy = world[1];
+  const vz = world[2];
+  const tx = 2 * (y * vz - z * vy);
+  const ty = 2 * (z * vx - x * vz);
+  const tz = 2 * (x * vy - y * vx);
+  return [
+    vx + w * tx + y * tz - z * ty,
+    vy + w * ty + z * tx - x * tz,
+    vz + w * tz + x * ty - y * tx
+  ];
+}
+function buildObservation(model, data, config) {
+  const values = [];
+  const layout = [];
+  if (config.time) appendScalar(values, layout, "time", data.time);
+  if (config.qpos) append(values, layout, "qpos", data.qpos);
+  if (config.qvel) append(values, layout, "qvel", data.qvel);
+  if (config.ctrl) append(values, layout, "ctrl", data.ctrl);
+  if (config.act) append(values, layout, "act", data.act);
+  if (config.sensordata) append(values, layout, "sensordata", data.sensordata);
+  for (const name of config.sensors ?? []) {
+    const sensorId = findSensorByName(model, name);
+    if (sensorId < 0) continue;
+    const start = model.sensor_adr[sensorId] ?? 0;
+    const dim = model.sensor_dim[sensorId] ?? 0;
+    append(values, layout, `sensor:${name}`, data.sensordata.subarray(start, start + dim));
+  }
+  for (const name of config.sites ?? []) {
+    const siteId = findSiteByName(model, name);
+    if (siteId < 0) continue;
+    const start = siteId * 3;
+    append(values, layout, `site:${name}:xpos`, data.site_xpos.subarray(start, start + 3));
+  }
+  const gravity = normalizedGravity(model);
+  for (const name of namedBodyList(config.projectedGravity)) {
+    const bodyId = findBodyByName(model, name);
+    if (bodyId < 0) continue;
+    append(values, layout, `projectedGravity:${name}`, rotateWorldVectorToBody(data, bodyId, gravity));
+  }
+  return {
+    values: config.output === "float64" ? new Float64Array(values) : new Float32Array(values),
+    layout
+  };
+}
 function shallowEqual(a, b) {
   const keysA = Object.keys(a);
   const keysB = Object.keys(b);
@@ -4086,6 +4161,26 @@ function usePolicy(config) {
     }
   };
 }
+var EMPTY_OBSERVATION = {
+  values: new Float32Array(0),
+  layout: []
+};
+function useObservation(config) {
+  const { mjModelRef, mjDataRef } = useMujocoContext();
+  const configRef = useRef(config);
+  configRef.current = config;
+  return useMemo(() => ({
+    read() {
+      const model = mjModelRef.current;
+      const data = mjDataRef.current;
+      if (!model || !data) return EMPTY_OBSERVATION;
+      return buildObservation(model, data, configRef.current);
+    },
+    readValues() {
+      return this.read().values;
+    }
+  }), [mjDataRef, mjModelRef]);
+}
 function useTrajectoryRecorder(options = {}) {
   const { mjModelRef } = useMujocoContext();
   const recordingRef = useRef(false);
@@ -4573,6 +4668,6 @@ function useCameraAnimation() {
  * useCameraAnimation — composable camera animation hook.
  */
 
-export { Body, ContactListener, ContactMarkers, Debug, DragInteraction, FlexRenderer, IkGizmo, MujocoCanvas, MujocoPhysics, MujocoProvider, MujocoSimProvider, SceneLights, TendonRenderer, TrajectoryPlayer, createContiguousControlGroup, createController, createControllerHook, findActuatorByName, findBodyByName, findGeomByName, findJointByName, findKeyframeByName, findSensorByName, findSiteByName, findTendonByName, getActuatedJoints, getContact, getControlMap, getName, loadScene, resolveControlGroup, useActuators, useAfterPhysicsStep, useBeforePhysicsStep, useBodyMeshes, useBodyState, useCameraAnimation, useContactEvents, useContacts, useCtrl, useCtrlNoise, useGamepad, useGravityCompensation, useIkController, useJointState, useKeyboardTeleop, useMujoco, useMujocoWasm, usePolicy, useSceneLights, useSelectionHighlight, useSensor, useSensors, useSitePosition, useTrajectoryPlayer, useTrajectoryRecorder, useVideoRecorder };
+export { Body, ContactListener, ContactMarkers, Debug, DragInteraction, FlexRenderer, IkGizmo, MujocoCanvas, MujocoPhysics, MujocoProvider, MujocoSimProvider, SceneLights, TendonRenderer, TrajectoryPlayer, buildObservation, createContiguousControlGroup, createController, createControllerHook, findActuatorByName, findBodyByName, findGeomByName, findJointByName, findKeyframeByName, findSensorByName, findSiteByName, findTendonByName, getActuatedJoints, getContact, getControlMap, getName, loadScene, resolveControlGroup, useActuators, useAfterPhysicsStep, useBeforePhysicsStep, useBodyMeshes, useBodyState, useCameraAnimation, useContactEvents, useContacts, useCtrl, useCtrlNoise, useGamepad, useGravityCompensation, useIkController, useJointState, useKeyboardTeleop, useMujoco, useMujocoWasm, useObservation, usePolicy, useSceneLights, useSelectionHighlight, useSensor, useSensors, useSitePosition, useTrajectoryPlayer, useTrajectoryRecorder, useVideoRecorder };
 //# sourceMappingURL=index.js.map
 //# sourceMappingURL=index.js.map