mujoco-react 0.1.0

package/src/index.ts

@@ -0,0 +1,108 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+// Core
+export { MujocoProvider, useMujoco } from './core/MujocoProvider';
+export { MujocoCanvas } from './core/MujocoCanvas';
+export { MujocoSimProvider, useMujocoSim, useBeforePhysicsStep, useAfterPhysicsStep } from './core/MujocoSimProvider';
+export {
+  loadScene,
+  getName,
+  findSiteByName,
+  findActuatorByName,
+  findKeyframeByName,
+  findBodyByName,
+  findJointByName,
+  findGeomByName,
+  findSensorByName,
+  findTendonByName,
+} from './core/SceneLoader';
+
+// Components
+export { SceneRenderer } from './components/SceneRenderer';
+export { IkGizmo } from './components/IkGizmo';
+export { ContactMarkers } from './components/ContactMarkers';
+export { DragInteraction } from './components/DragInteraction';
+export { SceneLights } from './components/SceneLights';
+export { Debug } from './components/Debug';
+export { TendonRenderer } from './components/TendonRenderer';
+export { FlexRenderer } from './components/FlexRenderer';
+export { ContactListener } from './components/ContactListener';
+export { TrajectoryPlayer } from './components/TrajectoryPlayer';
+export { SelectionHighlight } from './components/SelectionHighlight';
+
+// Hooks
+export { useActuators } from './hooks/useActuators';
+export { useSitePosition } from './hooks/useSitePosition';
+export { useGravityCompensation } from './hooks/useGravityCompensation';
+export { useSensor, useSensors } from './hooks/useSensor';
+export { useJointState } from './hooks/useJointState';
+export { useBodyState } from './hooks/useBodyState';
+export { useCtrl } from './hooks/useCtrl';
+export { useContacts, useContactEvents } from './hooks/useContacts';
+export { useKeyboardTeleop } from './hooks/useKeyboardTeleop';
+export { usePolicy } from './hooks/usePolicy';
+export { useTrajectoryPlayer } from './hooks/useTrajectoryPlayer';
+export { useTrajectoryRecorder } from './hooks/useTrajectoryRecorder';
+export { useGamepad } from './hooks/useGamepad';
+export { useVideoRecorder } from './hooks/useVideoRecorder';
+export { useCtrlNoise } from './hooks/useCtrlNoise';
+
+// Types
+export type {
+  // Scene config
+  SceneConfig,
+  SceneObject,
+  XmlPatch,
+  SceneMarker,
+  PhysicsConfig,
+  // IK
+  IKSolveFn,
+  // Callbacks
+  PhysicsStepCallback,
+  // State management
+  StateSnapshot,
+  // Model introspection
+  BodyInfo,
+  JointInfo,
+  GeomInfo,
+  SiteInfo,
+  ActuatorInfo,
+  SensorInfo,
+  // Contacts
+  ContactInfo,
+  // Raycast
+  RayHit,
+  // Model options
+  ModelOptions,
+  // Trajectory
+  TrajectoryFrame,
+  TrajectoryData,
+  // Keyboard teleop
+  KeyBinding,
+  KeyboardTeleopConfig,
+  // Policy
+  PolicyConfig,
+  // Component props
+  IkGizmoProps,
+  DragInteractionProps,
+  DebugProps,
+  SceneLightsProps,
+  TrajectoryPlayerProps,
+  SelectionHighlightProps,
+  ContactListenerProps,
+  // API
+  MujocoSimAPI,
+  MujocoCanvasProps,
+  MujocoContextValue,
+  // Hook return types
+  SitePositionResult,
+  SensorResult,
+  BodyStateResult,
+  JointStateResult,
+} from './types';
+
+// Re-export MuJoCo types for convenience
+export type { MujocoModule, MujocoModel, MujocoData } from './types';

package/src/rendering/CapsuleGeometry.ts

@@ -0,0 +1,35 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+*/
+
+
+import * as THREE from 'three';
+
+interface InternalBufferGeometry extends THREE.BufferGeometry {
+  type: string;
+}
+
+/**
+ * CapsuleGeometry
+ * Custom geometry for capsule shape.
+ */
+export class CapsuleGeometry extends THREE.BufferGeometry {
+  parameters: { radius: number; length: number; capSegments: number; radialSegments: number; };
+
+  constructor(radius = 1, length = 1, capSegments = 4, radialSegments = 8) {
+    super();
+    (this as unknown as InternalBufferGeometry).type = 'CapsuleGeometry';
+    this.parameters = { radius, length, capSegments, radialSegments };
+    const path = new THREE.Path();
+    path.absarc(0, -length / 2, radius, Math.PI * 1.5, 0, false);
+    path.absarc(0, length / 2, radius, 0, Math.PI * 0.5, false);
+    const latheGeometry = new THREE.LatheGeometry(path.getPoints(capSegments), radialSegments);
+    
+    const self = this as THREE.BufferGeometry;
+    self.setIndex(latheGeometry.getIndex());
+    self.setAttribute('position', latheGeometry.getAttribute('position'));
+    self.setAttribute('normal', latheGeometry.getAttribute('normal'));
+    self.setAttribute('uv', latheGeometry.getAttribute('uv'));
+  }
+}

package/src/rendering/GeomBuilder.ts

@@ -0,0 +1,140 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+*/
+
+
+import * as THREE from 'three';
+import { CapsuleGeometry } from './CapsuleGeometry';
+import { Reflector } from './Reflector';
+import { MujocoModel, MujocoModule } from '../types';
+
+/**
+ * GeomBuilder
+ * RESPONSIBILITY: Manufacturing visual objects.
+ * 
+ * This class knows how to read a single MuJoCo 'geom' (collision shape) definition
+ * and build the corresponding Three.js Mesh for it.
+ * It handles all the different shape types (Box, Sphere, Cylinder, generic Mesh, etc.).
+ */
+export class GeomBuilder {
+    private mujoco: MujocoModule; 
+
+    constructor(mujoco: MujocoModule) {
+        this.mujoco = mujoco;
+    }
+
+    /**
+     * 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).
+     */
+    create(mjModel: MujocoModel, g: number): THREE.Object3D | null {
+        // 1. Check if this geom is meant to be visible
+        // Group 3 in MuJoCo is conventionally used for invisible 'helper' geoms.
+        if (mjModel.geom_group[g] === 3) return null;
+
+        // 2. Read raw data from MuJoCo's WASM memory arrays
+        const type = mjModel.geom_type[g];
+        const size = mjModel.geom_size.subarray(g * 3, g * 3 + 3); // [x, y, z] size parameters
+        const pos = mjModel.geom_pos.subarray(g * 3, g * 3 + 3);   // [x, y, z] local position
+        const quat = mjModel.geom_quat.subarray(g * 4, g * 4 + 4); // [w, x, y, z] local rotation
+
+        // 3. Determine material color
+        // 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);
+        let opacity = 1.0;
+
+        if (matId >= 0) {
+            // Use shared material
+            const rgba = mjModel.mat_rgba.subarray(matId * 4, matId * 4 + 4);
+            color.setRGB(rgba[0], rgba[1], rgba[2]);
+            opacity = rgba[3];
+        } else {
+            // Use geom-specific color
+            const rgba = mjModel.geom_rgba.subarray(g * 4, g * 4 + 4);
+            color.setRGB(rgba[0], rgba[1], rgba[2]);
+            opacity = rgba[3];
+        }
+
+        // 4. Build the Geometry based on type
+        const MG = this.mujoco.mjtGeom; // Short alias for MuJoCo Geometry Types enum
+        let geo: THREE.BufferGeometry | null = null;
+
+        // The '.value ?? MG.XYZ' pattern handles slightly different versions of the mujoco-js bindings.
+        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.
+            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);
+        } else if (type === getVal(MG.mjGEOM_CAPSULE)) {
+            // Capsules in MuJoCo are Z-axis aligned by default.
+            // Our custom CapsuleGeometry might need rotation to match.
+            geo = new CapsuleGeometry(size[0], size[1] * 2, 24, 12);
+            geo.rotateX(Math.PI / 2); 
+        } else if (type === getVal(MG.mjGEOM_BOX)) {
+            // MuJoCo defines box size as "half-extents" (center to edge). Three.js uses full width.
+            geo = new THREE.BoxGeometry(size[0] * 2, size[1] * 2, size[2] * 2);
+        } else if (type === getVal(MG.mjGEOM_CYLINDER)) {
+            geo = new THREE.CylinderGeometry(size[0], size[0], size[1] * 2, 24);
+            geo.rotateX(Math.PI / 2);
+        } else if (type === getVal(MG.mjGEOM_MESH)) {
+            // Arbitrary 3D meshes (like the robot parts).
+            // We must read the vertex and face data directly from MuJoCo's buffers.
+            const mId = mjModel.geom_dataid[g];
+            const vAdr = mjModel.mesh_vertadr[mId];
+            const vNum = mjModel.mesh_vertnum[mId];
+            const fAdr = mjModel.mesh_faceadr[mId];
+            const fNum = mjModel.mesh_facenum[mId];
+
+            geo = new THREE.BufferGeometry();
+            // 'position' attribute = vertices
+            geo.setAttribute('position', new THREE.Float32BufferAttribute(mjModel.mesh_vert.subarray(vAdr * 3, (vAdr + vNum) * 3), 3));
+            // 'index' = faces (triangles connecting vertices)
+            geo.setIndex(Array.from(mjModel.mesh_face.subarray(fAdr * 3, (fAdr + fNum) * 3)));
+            geo.computeVertexNormals(); // Auto-calculate smooth lighting normals
+        }
+
+        // 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;
+            }
+
+            // Apply the local position offset and rotation specified in the MJCF XML
+            mesh.position.set(pos[0], pos[1], pos[2]);
+            // MuJoCo quaternions are [w, x, y, z], Three.js are [x, y, z, w]
+            mesh.quaternion.set(quat[1], quat[2], quat[3], quat[0]);
+
+            // 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;
+
+            return mesh;
+        }
+
+        return null;
+    }
+}

package/src/rendering/Reflector.ts

@@ -0,0 +1,225 @@
+/**
+ * @license
+ * SPDX-License-Identifier: Apache-2.0
+*/
+
+
+import * as THREE from 'three';
+
+/**
+ * Options for configuring the Reflector.
+ */
+export interface ReflectorOptions {
+    color?: THREE.ColorRepresentation;
+    textureWidth?: number;
+    textureHeight?: number;
+    clipBias?: number;
+    multisample?: number;
+    texture?: THREE.Texture;
+    mixStrength?: number; // How strong the reflection is (0.0 - 1.0)
+}
+
+interface ReflectorMesh extends THREE.Mesh {
+    type: string;
+    material: THREE.MeshPhysicalMaterial;
+    // tslint:disable-next-line:no-any
+    onBeforeRender: (renderer: any, scene: any, camera: any) => void;
+}
+
+/**
+ * Reflector
+ * Creates a reflective surface.
+ */
+export class Reflector extends THREE.Mesh {
+    isReflector = true;
+    camera: THREE.PerspectiveCamera;
+    private reflectorPlane = new THREE.Plane();
+    private normal = new THREE.Vector3();
+    private reflectorWorldPosition = new THREE.Vector3();
+    private cameraWorldPosition = new THREE.Vector3();
+    private rotationMatrix = new THREE.Matrix4();
+    private lookAtPosition = new THREE.Vector3(0, 0, -1);
+    private clipPlane = new THREE.Vector4();
+    private view = new THREE.Vector3();
+    private target = new THREE.Vector3();
+    private q = new THREE.Vector4();
+    private textureMatrix = new THREE.Matrix4();
+    private virtualCamera: THREE.PerspectiveCamera;
+    private renderTarget: THREE.WebGLRenderTarget;
+
+    constructor(geometry: THREE.BufferGeometry, options: ReflectorOptions = {}) {
+        super(geometry);
+
+        (this as unknown as ReflectorMesh).type = 'Reflector';
+        this.camera = new THREE.PerspectiveCamera();
+
+        const color = (options.color !== undefined) ? new THREE.Color(options.color) : new THREE.Color(0x7F7F7F);
+        const textureWidth = options.textureWidth || 512;
+        const textureHeight = options.textureHeight || 512;
+        const clipBias = options.clipBias || 0;
+        const multisample = (options.multisample !== undefined) ? options.multisample : 4;
+        const blendTexture = options.texture || undefined;
+        const mixStrength = (options.mixStrength !== undefined) ? options.mixStrength : 0.25; 
+
+        this.virtualCamera = this.camera;
+
+        this.renderTarget = new THREE.WebGLRenderTarget(textureWidth, textureHeight, { 
+            samples: multisample, 
+            type: THREE.HalfFloatType 
+        });
+
+        (this as unknown as ReflectorMesh).material = new THREE.MeshPhysicalMaterial({
+            map: blendTexture,
+            color,
+            roughness: 0.5, 
+            metalness: 0.1, 
+        });
+
+        (this as unknown as ReflectorMesh).material.onBeforeCompile = (shader) => {
+            shader.uniforms.tDiffuse = { value: this.renderTarget.texture };
+            shader.uniforms.textureMatrix = { value: this.textureMatrix };
+            shader.uniforms.mixStrength = { value: mixStrength };
+
+            // Vertex Shader: Set Vertex Positions to the Unwrapped UV Positions
+            const bodyStart = shader.vertexShader.indexOf('void main() {');
+            shader.vertexShader =
+                'uniform mat4 textureMatrix;\n' +
+                'varying vec4 vUvReflection;\n' +
+                shader.vertexShader.slice(0, bodyStart) +
+                shader.vertexShader.slice(bodyStart, -1) +
+                '   vUvReflection = textureMatrix * vec4( position, 1.0 );\n' +
+                '}';
+
+            // Fragment Shader: Mix reflection with base material
+            const fragmentBodyStart = shader.fragmentShader.indexOf('void main() {');
+            shader.fragmentShader =
+                'uniform sampler2D tDiffuse;\n' +
+                'uniform float mixStrength;\n' +
+                'varying vec4 vUvReflection;\n' +
+                shader.fragmentShader.slice(0, fragmentBodyStart) +
+                shader.fragmentShader.slice(fragmentBodyStart, -1) +
+                '   vec4 reflectionColor = texture2DProj( tDiffuse, vUvReflection );\n' +
+                '   gl_FragColor = vec4( mix( gl_FragColor.rgb, reflectionColor.rgb, mixStrength ), gl_FragColor.a );\n' +
+                '}';
+        };
+
+        (this as THREE.Object3D).receiveShadow = true;
+
+        (this as unknown as ReflectorMesh).onBeforeRender = (renderer: THREE.WebGLRenderer, scene: THREE.Scene, camera: THREE.Camera) => {
+            this.reflectorWorldPosition.setFromMatrixPosition((this as THREE.Object3D).matrixWorld);
+            this.cameraWorldPosition.setFromMatrixPosition(camera.matrixWorld);
+
+            this.rotationMatrix.extractRotation((this as THREE.Object3D).matrixWorld);
+
+            this.normal.set(0, 0, 1);
+            this.normal.applyMatrix4(this.rotationMatrix);
+
+            this.view.subVectors(this.reflectorWorldPosition, this.cameraWorldPosition);
+
+            // Avoid rendering when reflector is facing away
+            if (this.view.dot(this.normal) > 0) return;
+
+            this.view.reflect(this.normal).negate();
+            this.view.add(this.reflectorWorldPosition);
+
+            this.rotationMatrix.extractRotation(camera.matrixWorld);
+
+            this.lookAtPosition.set(0, 0, -1);
+            this.lookAtPosition.applyMatrix4(this.rotationMatrix);
+            this.lookAtPosition.add(this.cameraWorldPosition);
+
+            this.target.subVectors(this.reflectorWorldPosition, this.lookAtPosition);
+            this.target.reflect(this.normal).negate();
+            this.target.add(this.reflectorWorldPosition);
+
+            this.virtualCamera.position.copy(this.view);
+            this.virtualCamera.up.set(0, 1, 0);
+            this.virtualCamera.up.applyMatrix4(this.rotationMatrix);
+            this.virtualCamera.up.reflect(this.normal);
+            this.virtualCamera.lookAt(this.target);
+
+            this.virtualCamera.far = (camera as THREE.PerspectiveCamera).far; 
+
+            this.virtualCamera.updateMatrixWorld();
+            this.virtualCamera.projectionMatrix.copy((camera as THREE.PerspectiveCamera).projectionMatrix);
+
+            // Update the texture matrix
+            this.textureMatrix.set(
+                0.5, 0.0, 0.0, 0.5,
+                0.0, 0.5, 0.0, 0.5,
+                0.0, 0.0, 0.5, 0.5,
+                0.0, 0.0, 0.0, 1.0
+            );
+            this.textureMatrix.multiply(this.virtualCamera.projectionMatrix);
+            this.textureMatrix.multiply(this.virtualCamera.matrixWorldInverse);
+            this.textureMatrix.multiply((this as THREE.Object3D).matrixWorld);
+
+            // Now update projection matrix with new clip plane
+            this.reflectorPlane.setFromNormalAndCoplanarPoint(this.normal, this.reflectorWorldPosition);
+            this.reflectorPlane.applyMatrix4(this.virtualCamera.matrixWorldInverse);
+
+            this.clipPlane.set(this.reflectorPlane.normal.x, this.reflectorPlane.normal.y, this.reflectorPlane.normal.z, this.reflectorPlane.constant);
+
+            const projectionMatrix = this.virtualCamera.projectionMatrix;
+
+            this.q.x = (Math.sign(this.clipPlane.x) + projectionMatrix.elements[8]) / projectionMatrix.elements[0];
+            this.q.y = (Math.sign(this.clipPlane.y) + projectionMatrix.elements[9]) / projectionMatrix.elements[5];
+            this.q.z = -1.0;
+            this.q.w = (1.0 + projectionMatrix.elements[10]) / projectionMatrix.elements[14];
+
+            // Calculate the scaled plane vector
+            this.clipPlane.multiplyScalar(2.0 / this.clipPlane.dot(this.q));
+
+            // Replacing the third row of the projection matrix
+            projectionMatrix.elements[2] = this.clipPlane.x;
+            projectionMatrix.elements[6] = this.clipPlane.y;
+            projectionMatrix.elements[10] = this.clipPlane.z + 1.0 - clipBias;
+            projectionMatrix.elements[14] = this.clipPlane.w;
+
+            // Render
+            (this as THREE.Object3D).visible = false;
+
+            const currentRenderTarget = renderer.getRenderTarget();
+            const currentXrEnabled = renderer.xr.enabled;
+            const currentShadowAutoUpdate = renderer.shadowMap.autoUpdate;
+
+            renderer.xr.enabled = false; 
+            renderer.shadowMap.autoUpdate = false; 
+
+            renderer.setRenderTarget(this.renderTarget);
+
+            renderer.state.buffers.depth.setMask(true); 
+
+            if (renderer.autoClear === false) renderer.clear();
+            renderer.render(scene, this.virtualCamera);
+
+            renderer.xr.enabled = currentXrEnabled;
+            renderer.shadowMap.autoUpdate = currentShadowAutoUpdate;
+
+            renderer.setRenderTarget(currentRenderTarget);
+
+            // Restore viewport
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            const viewport = (camera as any).viewport;
+            if (viewport !== undefined) {
+                renderer.state.viewport(viewport);
+            }
+
+            (this as THREE.Object3D).visible = true;
+        };
+    }
+
+    getRenderTarget() {
+        return this.renderTarget;
+    }
+
+    dispose() {
+        this.renderTarget.dispose();
+        const mesh = this as THREE.Mesh;
+        if (Array.isArray(mesh.material)) {
+            mesh.material.forEach((m) => m.dispose());
+        } else {
+            mesh.material.dispose();
+        }
+    }
+}