mujoco-react 10.0.1 → 10.1.0

package/src/policyControls.ts

@@ -0,0 +1,87 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Helpers for applying policy action vectors to MuJoCo controls.
+ */
+
+import type { MujocoData, MujocoModel, PolicyVector } from './types';
+
+export interface ApplyPolicyActionToControlsOptions {
+  /**
+   * First actuator/control index to write. Defaults to 0.
+   */
+  actuatorOffset?: number;
+  /**
+   * Maximum number of controls to write. Defaults to the action length.
+   */
+  actionSize?: number;
+  /**
+   * Clamp each action value to `model.actuator_ctrlrange` before writing.
+   * Defaults to true because most learned policies should not exceed actuator limits.
+   */
+  clamp?: boolean;
+  /**
+   * Leave the current control unchanged when an action entry is not finite.
+   * Defaults to true so a bad policy response cannot write NaN into the simulation.
+   */
+  skipInvalid?: boolean;
+}
+
+export interface ApplyPolicyActionToControlsResult {
+  /**
+   * Values actually written to `data.ctrl`, after offset, truncation, and clamping.
+   */
+  applied: number[];
+  /**
+   * Actuator indices that were not written because the corresponding action value
+   * was not finite and `skipInvalid` was enabled.
+   */
+  skipped: number[];
+  actuatorOffset: number;
+}
+
+export function clampPolicyActionValue(
+  model: MujocoModel,
+  actuatorIndex: number,
+  value: number
+) {
+  const ranges = model.actuator_ctrlrange;
+  const min = ranges?.[actuatorIndex * 2] ?? -Infinity;
+  const max = ranges?.[actuatorIndex * 2 + 1] ?? Infinity;
+  return Math.max(min, Math.min(max, value));
+}
+
+export function applyPolicyActionToControls(
+  model: MujocoModel,
+  data: MujocoData,
+  action: PolicyVector,
+  options: ApplyPolicyActionToControlsOptions = {}
+): ApplyPolicyActionToControlsResult {
+  const actuatorOffset = options.actuatorOffset ?? 0;
+  const actionSize = options.actionSize ?? action.length;
+  const shouldClamp = options.clamp ?? true;
+  const shouldSkipInvalid = options.skipInvalid ?? true;
+  const count = Math.max(
+    0,
+    Math.min(actionSize, action.length, data.ctrl.length - actuatorOffset, model.nu - actuatorOffset)
+  );
+  const applied: number[] = [];
+  const skipped: number[] = [];
+
+  for (let index = 0; index < count; index += 1) {
+    const actuatorIndex = actuatorOffset + index;
+    const value = Number(action[index]);
+    if (shouldSkipInvalid && !Number.isFinite(value)) {
+      skipped.push(actuatorIndex);
+      continue;
+    }
+    const nextValue = shouldClamp
+      ? clampPolicyActionValue(model, actuatorIndex, value)
+      : value;
+    data.ctrl[actuatorIndex] = nextValue;
+    applied.push(nextValue);
+  }
+
+  return { applied, skipped, actuatorOffset };
+}

package/src/policyObservation.ts

@@ -0,0 +1,172 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Named policy observation builders with layout and units metadata.
+ */
+
+import type {
+  Bodies,
+  Geoms,
+  MujocoData,
+  MujocoModel,
+  ObservationOutput,
+  Sites,
+} from './types';
+import { findBodyByName, findGeomByName, findSiteByName } from './core/SceneLoader';
+
+export type NamedObservationMissing = 'skip' | 'zeros' | 'throw';
+
+export interface NamedObservationInput {
+  model: MujocoModel;
+  data: MujocoData;
+}
+
+export interface NamedObservationField {
+  name: string;
+  size: number;
+  units?: string;
+  read: (input: NamedObservationInput) => ArrayLike<number> | number | null | undefined;
+}
+
+export interface NamedObservationLayoutItem {
+  name: string;
+  start: number;
+  size: number;
+  units?: string;
+}
+
+export interface NamedObservationOptions {
+  fields: readonly NamedObservationField[];
+  output?: ObservationOutput;
+  missing?: NamedObservationMissing;
+}
+
+export interface NamedObservationResult {
+  values: Float32Array | Float64Array;
+  layout: NamedObservationLayoutItem[];
+}
+
+function pushValues(target: number[], value: ArrayLike<number> | number, size: number) {
+  if (typeof value === 'number') {
+    target.push(value);
+    for (let index = 1; index < size; index += 1) target.push(0);
+    return;
+  }
+
+  for (let index = 0; index < size; index += 1) {
+    target.push(Number(value[index] ?? 0));
+  }
+}
+
+export function readNamedObservation(
+  model: MujocoModel,
+  data: MujocoData,
+  options: NamedObservationOptions
+): NamedObservationResult {
+  const values: number[] = [];
+  const layout: NamedObservationLayoutItem[] = [];
+  const missing = options.missing ?? 'skip';
+
+  for (const field of options.fields) {
+    const start = values.length;
+    const value = field.read({ model, data });
+    if (value === null || value === undefined) {
+      if (missing === 'skip') continue;
+      if (missing === 'throw') {
+        throw new Error(`Unable to read named observation field "${field.name}".`);
+      }
+      for (let index = 0; index < field.size; index += 1) values.push(0);
+    } else {
+      pushValues(values, value, field.size);
+    }
+    layout.push({
+      name: field.name,
+      start,
+      size: field.size,
+      units: field.units,
+    });
+  }
+
+  return {
+    values: options.output === 'float64'
+      ? new Float64Array(values)
+      : new Float32Array(values),
+    layout,
+  };
+}
+
+export function createNamedObservationBuilder(options: NamedObservationOptions) {
+  return (model: MujocoModel, data: MujocoData) => (
+    readNamedObservation(model, data, options)
+  );
+}
+
+export function qposField(name: string, index: number, units = 'qpos'): NamedObservationField {
+  return {
+    name,
+    size: 1,
+    units,
+    read: ({ data }) => data.qpos[index],
+  };
+}
+
+export function qvelField(name: string, index: number, units = 'qvel'): NamedObservationField {
+  return {
+    name,
+    size: 1,
+    units,
+    read: ({ data }) => data.qvel[index],
+  };
+}
+
+export function ctrlField(name: string, index: number, units = 'ctrl'): NamedObservationField {
+  return {
+    name,
+    size: 1,
+    units,
+    read: ({ data }) => data.ctrl[index],
+  };
+}
+
+export function bodyPositionField(name: Bodies, units = 'world_position'): NamedObservationField {
+  return {
+    name: `body:${name}:xpos`,
+    size: 3,
+    units,
+    read: ({ model, data }) => {
+      const bodyId = findBodyByName(model, name);
+      if (bodyId < 0) return null;
+      const offset = bodyId * 3;
+      return data.xpos.subarray(offset, offset + 3);
+    },
+  };
+}
+
+export function geomPositionField(name: Geoms, units = 'world_position'): NamedObservationField {
+  return {
+    name: `geom:${name}:xpos`,
+    size: 3,
+    units,
+    read: ({ model, data }) => {
+      const geomId = findGeomByName(model, name);
+      if (geomId < 0) return null;
+      const offset = geomId * 3;
+      return data.geom_xpos.subarray(offset, offset + 3);
+    },
+  };
+}
+
+export function sitePositionField(name: Sites, units = 'world_position'): NamedObservationField {
+  return {
+    name: `site:${name}:xpos`,
+    size: 3,
+    units,
+    read: ({ model, data }) => {
+      const siteId = findSiteByName(model, name);
+      if (siteId < 0) return null;
+      const offset = siteId * 3;
+      return data.site_xpos.subarray(offset, offset + 3);
+    },
+  };
+}

package/src/rendering/GeomBuilder.ts

@@ -6,7 +6,7 @@
 
 import * as THREE from 'three';
 import { CapsuleGeometry } from './CapsuleGeometry';
-import { Reflector } from './Reflector';
+import { getName } from '../core/SceneLoader';
 import { MujocoModel, MujocoModule } from '../types';
 
 /**
@@ -19,11 +19,64 @@ import { MujocoModel, MujocoModule } from '../types';
  */
 export class GeomBuilder {
     private mujoco: MujocoModule; 
+    private textureCache = new Map<number, THREE.Texture>();
 
     constructor(mujoco: MujocoModule) {
         this.mujoco = mujoco;
     }
 
+    private getMaterialTexture(mjModel: MujocoModel, matId: number): THREE.Texture | null {
+        if (matId < 0 || !mjModel.mat_texid || !mjModel.tex_data) return null;
+
+        const materialCount = Math.max(1, Math.floor(mjModel.mat_rgba.length / 4));
+        const textureRoles = Math.max(1, Math.floor(mjModel.mat_texid.length / materialCount));
+        let texId = -1;
+        for (let role = 0; role < textureRoles; role += 1) {
+            const candidate = mjModel.mat_texid[matId * textureRoles + role];
+            if (candidate >= 0) {
+                texId = candidate;
+                break;
+            }
+        }
+        if (texId < 0) return null;
+
+        const cached = this.textureCache.get(texId);
+        if (cached) return cached;
+
+        const width = Number(mjModel.tex_width[texId]);
+        const height = Number(mjModel.tex_height[texId]);
+        const channels = Number(mjModel.tex_nchannel[texId]);
+        const offset = Number(mjModel.tex_adr[texId]);
+        if (width <= 0 || height <= 0 || channels <= 0 || offset < 0) return null;
+
+        const source = mjModel.tex_data.subarray(offset, offset + width * height * channels);
+        const rgba = new Uint8Array(width * height * 4);
+        for (let i = 0, j = 0; i < width * height; i += 1, j += channels) {
+            const r = source[j] ?? 255;
+            const g = channels > 1 ? source[j + 1] : r;
+            const b = channels > 2 ? source[j + 2] : r;
+            const a = channels > 3 ? source[j + 3] : 255;
+            const out = i * 4;
+            rgba[out] = r;
+            rgba[out + 1] = g;
+            rgba[out + 2] = b;
+            rgba[out + 3] = a;
+        }
+
+        const texture = new THREE.DataTexture(rgba, width, height, THREE.RGBAFormat);
+        texture.colorSpace = THREE.LinearSRGBColorSpace;
+        texture.wrapS = THREE.RepeatWrapping;
+        texture.wrapT = THREE.RepeatWrapping;
+        texture.flipY = true;
+        const repeatOffset = matId * 2;
+        const repeatS = mjModel.mat_texrepeat?.[repeatOffset] ?? 1;
+        const repeatT = mjModel.mat_texrepeat?.[repeatOffset + 1] ?? 1;
+        texture.repeat.set(repeatS || 1, repeatT || 1);
+        texture.needsUpdate = true;
+        this.textureCache.set(texId, texture);
+        return texture;
+    }
+
     /**
      * Creates a Three.js Object3D (usually a Mesh) for a specific geometry in the MuJoCo model.
      * Returns null if the geometry shouldn't be rendered (e.g., invisible collision triggers).
@@ -43,6 +96,7 @@ export class GeomBuilder {
         // Sometimes color is on the geom itself, sometimes it uses a shared material definition.
         const matId = mjModel.geom_matid[g];
         const color = new THREE.Color(0xffffff);
+        const map = this.getMaterialTexture(mjModel, matId);
         let opacity = 1.0;
 
         if (matId >= 0) {
@@ -65,8 +119,7 @@ export class GeomBuilder {
         const getVal = (v: unknown) => (v as { value: number })?.value ?? v;
 
         if (type === getVal(MG.mjGEOM_PLANE)) {
-            // Planes are infinite in MuJoCo, but we need a finite mesh for Three.js. 
-            // Fallback reduced to 5m to match grid as requested.
+            // Planes are infinite in MuJoCo, but Three needs finite UVs for textured captures.
             geo = new THREE.PlaneGeometry(size[0] * 2 || 5, size[1] * 2 || 5);
         } else if (type === getVal(MG.mjGEOM_SPHERE)) {
             geo = new THREE.SphereGeometry(size[0], 24, 24);
@@ -100,28 +153,22 @@ export class GeomBuilder {
 
         // 5. Construct the final Mesh
         if (geo) {
-            let mesh;
-            // Special handling for the floor plane to make it shiny
-            if (type === getVal(MG.mjGEOM_PLANE)) {
-                mesh = new Reflector(geo, {
-                    clipBias: 0.003,
-                    textureWidth: 1024, textureHeight: 1024,
-                    color,
-                    mixStrength: 0.25
-                });
-            } else {
-                // Standard physical material for everything else
-                mesh = new THREE.Mesh(geo, new THREE.MeshStandardMaterial({
-                    color,
-                    transparent: opacity < 1,
-                    opacity,
-                    roughness: 0.6,
-                    metalness: 0.2
-                }));
-                // Enable shadows
-                mesh.castShadow = true;
-                mesh.receiveShadow = true;
+            const isPlane = type === getVal(MG.mjGEOM_PLANE);
+            const materialMap = isPlane && map ? map.clone() : map;
+            if (isPlane && materialMap) {
+                materialMap.repeat.multiplyScalar(2.5);
+                materialMap.needsUpdate = true;
             }
+            const mesh = new THREE.Mesh(geo, new THREE.MeshStandardMaterial({
+                color,
+                map: materialMap,
+                transparent: opacity < 1,
+                opacity,
+                roughness: 0.6,
+                metalness: 0
+            }));
+            mesh.castShadow = type !== getVal(MG.mjGEOM_PLANE);
+            mesh.receiveShadow = true;
 
             // Apply the local position offset and rotation specified in the MJCF XML
             mesh.position.set(pos[0], pos[1], pos[2]);
@@ -131,6 +178,8 @@ export class GeomBuilder {
             // Tag the mesh with its MuJoCo body and geom IDs for interaction (picking/dragging)
             mesh.userData.bodyID = mjModel.geom_bodyid[g];
             mesh.userData.geomID = g;
+            mesh.userData.geomGroup = mjModel.geom_group[g];
+            mesh.userData.geomName = getName(mjModel, mjModel.name_geomadr[g]);
 
             return mesh;
         }

package/src/rendering/cameraFrameCapture.ts

@@ -41,6 +41,8 @@ export interface CameraFrameCaptureSession {
 
 export const CAMERA_FRAME_CAPTURE_RENDER_USER_DATA_KEY =
   'mujocoReactCameraFrameCaptureRender';
+export const CAMERA_FRAME_CAPTURE_PRE_RENDER_USER_DATA_KEY =
+  'mujocoReactCameraFrameCapturePreRender';
 export const CAPTURE_EXCLUDE_KEY =
   'mujoco.capture.exclude';
 
@@ -84,6 +86,8 @@ type VisibilityState = {
   visible: boolean;
 };
 
+type CameraFrameCapturePreRender = () => void;
+
 function toVector3(
   value: CameraFrameCaptureVector3 | undefined,
   fallback: THREE.Vector3
@@ -260,6 +264,39 @@ function hideExcludedCaptureObjects(scene: THREE.Scene): VisibilityState[] {
   return hidden;
 }
 
+function hideCaptureGeomGroups(
+  scene: THREE.Scene,
+  options: CameraFrameCaptureOptions
+): VisibilityState[] {
+  const hidden: VisibilityState[] = [];
+  const hiddenGroups = options.hiddenGeomGroups
+    ? new Set(options.hiddenGeomGroups)
+    : null;
+  const visibleGroups = options.visibleGeomGroups
+    ? new Set(options.visibleGeomGroups)
+    : null;
+  const hiddenNames = options.hiddenGeomNames
+    ? new Set(options.hiddenGeomNames)
+    : null;
+  if (!hiddenGroups && !visibleGroups && !hiddenNames) return hidden;
+
+  scene.traverse((object) => {
+    if (!object.visible) return;
+    const geomGroup = object.userData.geomGroup;
+    const geomName = object.userData.geomName;
+    if (typeof geomGroup !== 'number' && typeof geomName !== 'string') return;
+    if (
+      hiddenNames?.has(geomName) ||
+      hiddenGroups?.has(geomGroup) ||
+      (typeof geomGroup === 'number' && visibleGroups && !visibleGroups.has(geomGroup))
+    ) {
+      hidden.push({ object, visible: object.visible });
+      object.visible = false;
+    }
+  });
+  return hidden;
+}
+
 function restoreObjectVisibility(hidden: VisibilityState[]) {
   for (const { object, visible } of hidden) {
     object.visible = visible;
@@ -332,6 +369,17 @@ function getCaptureRenderer(
   return renderers[0] ?? null;
 }
 
+function runCapturePreRenderHooks(scene: THREE.Scene) {
+  const callbacks: CameraFrameCapturePreRender[] = [];
+  scene.traverse((object) => {
+    const callback = object.userData[
+      CAMERA_FRAME_CAPTURE_PRE_RENDER_USER_DATA_KEY
+    ] as CameraFrameCapturePreRender | undefined;
+    if (typeof callback === 'function') callbacks.push(callback);
+  });
+  for (const callback of callbacks) callback();
+}
+
 export function createCameraFrameCaptureSession(
   renderer: THREE.WebGLRenderer,
   scene: THREE.Scene,
@@ -382,8 +430,12 @@ export function createCameraFrameCaptureSession(
 
   function renderPreparedCapture(captureOptions: CameraFrameCaptureOptions) {
     const previousState = saveRendererState(renderer);
-    const hidden = hideExcludedCaptureObjects(scene);
+    const hidden = [
+      ...hideExcludedCaptureObjects(scene),
+      ...hideCaptureGeomGroups(scene, captureOptions),
+    ];
 
+    runCapturePreRenderHooks(scene);
     scene.updateMatrixWorld(true);
     try {
       renderer.xr.enabled = false;
@@ -391,6 +443,14 @@ export function createCameraFrameCaptureSession(
       renderer.setViewport(0, 0, width, height);
       renderer.setScissor(0, 0, width, height);
       renderer.setScissorTest(false);
+      if (captureOptions.background !== undefined) {
+        renderer.setClearColor(
+          new THREE.Color(captureOptions.background),
+          captureOptions.backgroundAlpha ?? previousState.clearAlpha
+        );
+      } else if (captureOptions.backgroundAlpha !== undefined) {
+        renderer.setClearColor(previousState.clearColor, captureOptions.backgroundAlpha);
+      }
       renderer.clear();
       renderer.render(scene, camera);
       readRenderTargetToCanvas(
@@ -423,13 +483,25 @@ export function createCameraFrameCaptureSession(
 
   async function captureAsync(nextOptions: CameraFrameCaptureOptions = {}) {
     const captureOptions = resolveCaptureOptions(nextOptions);
+    runCapturePreRenderHooks(scene);
     scene.updateMatrixWorld(true);
     const captureRenderer = getCaptureRenderer(scene);
     if (captureRenderer) {
       const previousState = saveRendererState(renderer);
-      const hidden = hideExcludedCaptureObjects(scene);
+      const hidden = [
+        ...hideExcludedCaptureObjects(scene),
+        ...hideCaptureGeomGroups(scene, captureOptions),
+      ];
       try {
         renderer.xr.enabled = false;
+        if (captureOptions.background !== undefined) {
+          renderer.setClearColor(
+            new THREE.Color(captureOptions.background),
+            captureOptions.backgroundAlpha ?? previousState.clearAlpha
+          );
+        } else if (captureOptions.backgroundAlpha !== undefined) {
+          renderer.setClearColor(previousState.clearColor, captureOptions.backgroundAlpha);
+        }
         const captureResult = await captureRenderer({
           renderer,
           scene,