// Copyright 2025 DeepMind Technologies Limited // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. import * as THREE from "three" import { OrbitControls } from "three/examples/jsm/controls/OrbitControls.js" import loadMujoco from "../dist/mujoco.js" declare function loadMujoco(): Promise; let mujoco: any; const modelXml = ` `; // Backport of CapsuleGeometry class introduced in THREE.js r139 class CapsuleGeometry extends THREE.BufferGeometry { readonly parameters: { readonly radius: number, readonly length: number, readonly capSegments: number, readonly radialSegments: number }; constructor(radius = 1, length = 1, capSegments = 4, radialSegments = 8) { 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); super(); this.setIndex(latheGeometry.getIndex()); this.setAttribute('position', latheGeometry.getAttribute('position')); this.setAttribute('normal', latheGeometry.getAttribute('normal')); this.setAttribute('uv', latheGeometry.getAttribute('uv')); this.type = 'CapsuleGeometry'; this.parameters = { radius, length, capSegments, radialSegments, }; } } class App { // TODO(matijak): We can use better types here by doing the following: // https://emscripten.org/docs/porting/connecting_cpp_and_javascript/embind.html#typescript-definitions mjModel: any; mjData: any; mjvOption: any; mjvPerturb: any; mjvCamera: any; mjvScene: any; paused = false; frameId: number|null = null; maxGeoms: number = 2 ** 15; scene: THREE.Scene; renderer: THREE.WebGLRenderer; camera: THREE.PerspectiveCamera; controls: OrbitControls; meshes: THREE.Mesh[] = []; bufferGeometryCache = new Map(); constructor() { this.mjvPerturb = new mujoco.MjvPerturb(); this.mjvOption = new mujoco.MjvOption(); this.mjvCamera = new mujoco.MjvCamera(); this.scene = new THREE.Scene(); this.renderer = new THREE.WebGLRenderer(); this.renderer.setSize(window.innerWidth, window.innerHeight); this.renderer.shadowMap.enabled = true; this.renderer.shadowMap.type = THREE.PCFSoftShadowMap; document.body.appendChild(this.renderer.domElement); this.camera = new THREE.PerspectiveCamera( 45, window.innerWidth / window.innerHeight, .1, 1000); this.camera.up.set(0, 0, 1); // Mujoco uses z-up this.camera.position.set(-2, 0, 2); this.controls = new OrbitControls(this.camera, this.renderer.domElement); } dispose() { // Release all C++ objects if (this.mjvScene) { this.mjvScene.delete(); } if (this.mjvCamera) { this.mjvCamera.delete(); } if (this.mjvPerturb) { this.mjvPerturb.delete(); } if (this.mjvOption) { this.mjvOption.delete(); } if (this.mjData) { this.mjData.delete(); } if (this.mjModel) { this.mjModel.delete(); } // Release all the THREE.js objects this.meshes.forEach((mesh) => { if (mesh.material) { if (Array.isArray(mesh.material)) { mesh.material.forEach(material => material.dispose()); } else { mesh.material.dispose(); } } if (mesh.geometry) { mesh.geometry.dispose(); } }); this.bufferGeometryCache.clear(); if (this.controls) { this.controls.dispose(); } if (this.renderer) { this.renderer.dispose(); } // Stop the animation loop since we've disposed of all the data if (this.frameId) { cancelAnimationFrame(this.frameId); this.frameId = null; } } loadModel(xmlContent: string) { this.mjModel = mujoco.MjModel.from_xml_string(xmlContent); if (!app.mjModel) { throw new Error('Failed to load model'); } this.mjData = new mujoco.MjData(this.mjModel); if (!this.mjData) { throw new Error('Failed to load data'); } this.initScene(); } pauseButton() { this.paused = !this.paused; const button = document.getElementById('pause-button'); if (button) { button.textContent = this.paused ? 'Resume' : 'Pause'; } } // TODO(matijak): Fix the bug where contact cylinders are wrong if the // simulation is reset while they are being visualized resetButton() { if (this.mjModel && this.mjData) { console.log('Resetting model and data'); mujoco.mj_resetData(this.mjModel, this.mjData); mujoco.mj_forward(this.mjModel, this.mjData); this.clearScene(); this.initScene(); } } contactButton() { const index = mujoco.mjtVisFlag.mjVIS_CONTACTPOINT.value; const value = this.mjvOption.flags[index]; this.mjvOption.flags[index] = !value; const button = document.getElementById('contact-button'); if (button) { button.textContent = this.mjvOption.flags[index] ? 'Hide Contacts' : 'Show Contacts'; } this.clearScene(); this.initScene(); } initScene() { this.mjvScene = new mujoco.MjvScene(this.mjModel, this.maxGeoms); const pointLight = new THREE.PointLight(0xffffff, .4); pointLight.position.set(0, 0, 2); pointLight.castShadow = true; pointLight.shadow.mapSize.set(2048, 2048); this.scene.add(pointLight); const ambientLight = new THREE.AmbientLight(0xffffff, .2); this.scene.add(ambientLight); const spotLight = new THREE.SpotLight(0xffffff, .2); spotLight.position.set(0, 0, 2); spotLight.target.position.set(0, 0, 0); spotLight.castShadow = true; spotLight.shadow.mapSize.set(2048, 2048); this.scene.add(spotLight); this.scene.add(spotLight.target); } clearScene() { // clear cached meshes this.meshes.forEach((mesh) => { if (mesh.material) { if (Array.isArray(mesh.material)) { mesh.material.forEach(material => material.dispose()); } else { mesh.material.dispose(); } } if (mesh.geometry) { mesh.geometry.dispose(); } }); this.bufferGeometryCache.clear(); this.meshes.length = 0; while (this.scene.children.length > 0) { this.scene.remove(this.scene.children[0]); } this.mjvScene.delete(); } getBufferGeometry(mjvGeom: any): [boolean, THREE.BufferGeometry] { if (!(mjvGeom instanceof mujoco.MjvGeom)) { throw new Error('mjvGeom is not an instance of mujoco.MjvGeom'); } // Lookup the geometry and return it if found const key = JSON.stringify([mjvGeom.type, mjvGeom.size, mjvGeom.dataid]); const found = this.bufferGeometryCache.get(key); if (found) { return [false, found]; } // Create geometry let geom: THREE.BufferGeometry; if (mjvGeom.type === mujoco.mjtGeom.mjGEOM_PLANE.value) { geom = new THREE.PlaneGeometry( 2 * (mjvGeom.size[0] ? mjvGeom.size[0] : 10000), 2 * (mjvGeom.size[1] ? mjvGeom.size[1] : 10000)); const uv = geom.getAttribute('uv'); for (let i = 0; i < uv.count; ++i) { uv.setY(i, 1 - uv.getY(i)); } } else if (mjvGeom.type === mujoco.mjtGeom.mjGEOM_SPHERE.value) { geom = new THREE.SphereGeometry(mjvGeom.size[0]); } else if (mjvGeom.type === mujoco.mjtGeom.mjGEOM_CAPSULE.value) { geom = new CapsuleGeometry(mjvGeom.size[0], 2 * mjvGeom.size[2], 32, 16); geom.rotateX(0.5 * Math.PI); } else if (mjvGeom.type === mujoco.mjtGeom.mjGEOM_BOX.value) { geom = new THREE.BoxGeometry( 2 * mjvGeom.size[0], 2 * mjvGeom.size[1], 2 * mjvGeom.size[2]); } else if (mjvGeom.type === mujoco.mjtGeom.mjGEOM_CYLINDER.value) { geom = new THREE.CylinderGeometry( mjvGeom.size[0], mjvGeom.size[1], 2 * mjvGeom.size[2], 32); geom.rotateX(0.5 * Math.PI); } else if (mjvGeom.type === mujoco.mjtGeom.mjGEOM_ELLIPSOID.value) { geom = new THREE.SphereGeometry(1); geom.scale(mjvGeom.size[0], mjvGeom.size[1], mjvGeom.size[2]); } else { console.log('Unsupported geom type: ', mjvGeom.type); geom = new THREE.BufferGeometry(); } this.bufferGeometryCache.set(key, geom); return [true, geom]; } update() { if (!this.mjModel || !this.mjData) { return; } app.controls.update(); // Simulate physics for 1/60 sec if (!app.paused) { let sim_start = app.mjData.time; while (app.mjData.time - sim_start < 1. / 60.) { mujoco.mj_step(app.mjModel, app.mjData); } } // Update the mujoco scene mujoco.mjv_updateScene( this.mjModel, this.mjData, this.mjvOption, this.mjvPerturb, this.mjvCamera, mujoco.mjtCatBit.mjCAT_ALL.value, this.mjvScene); const geoms = this.mjvScene.geoms; for (let i = 0; i < geoms.size(); i++) { const mjvGeom = geoms.get(i); let mesh: THREE.Mesh; if (i < this.meshes.length) { mesh = this.meshes[i]; } else { const mjvGeom = geoms.get(i); const [added, geom] = this.getBufferGeometry(mjvGeom); // Create material let material = new THREE.MeshPhongMaterial(); material.color.setRGB( mjvGeom.rgba[0], mjvGeom.rgba[1], mjvGeom.rgba[2]); material.opacity = mjvGeom.rgba[3]; material.transparent = mjvGeom.rgba[3] !== 0; // Create mesh mesh = new THREE.Mesh(geom, material); mesh.castShadow = true; mesh.receiveShadow = true; this.meshes.push(mesh); this.scene.add(mesh); } mesh.matrixAutoUpdate = false; const sz = 1; mesh.matrix.set( mjvGeom.mat[0], mjvGeom.mat[1], mjvGeom.mat[2] * sz, mjvGeom.pos[0], mjvGeom.mat[3], mjvGeom.mat[4], mjvGeom.mat[5] * sz, mjvGeom.pos[1], mjvGeom.mat[6], mjvGeom.mat[7], mjvGeom.mat[8] * sz, mjvGeom.pos[2], 0, 0, 0, 1); mesh.matrixWorldNeedsUpdate = true; mjvGeom.delete(); } geoms.delete(); } render() { this.renderer.render(this.scene, this.camera); } run() { const animate = () => { try { this.update(); this.render(); } catch (error) { console.error('Simulation error:', error); } // Request next frame this.frameId = requestAnimationFrame(animate); }; // Request first frame this.frameId = requestAnimationFrame(animate); } } function setupWindowEvents() { // Add an event listener to clean up when the page is unloaded // Tip: put "window.dispatchEvent(new Event('unload'))" in the console to test window.addEventListener('unload', () => { app.dispose(); }); window.addEventListener('keydown', (event) => { if (event.code === 'Backspace') { app.resetButton(); } }); window.addEventListener('keydown', (event) => { if (event.code === 'Space') { app.pauseButton(); } }); window.addEventListener('keydown', (event) => { if (event.key === 'c') { app.contactButton(); } }); } let app: App; async function main() { try { mujoco = await loadMujoco(); app = new App(); setupWindowEvents(); // Note: all elements will be destroyed with the page const pauseButtonElement = document.getElementById('pause-button'); if (pauseButtonElement) { pauseButtonElement.onclick = () => app.pauseButton(); } const resetButtonElement = document.getElementById('reset-button'); if (resetButtonElement) { resetButtonElement.onclick = () => app.resetButton(); } const contactButtonElement = document.getElementById('contact-button'); if (contactButtonElement) { contactButtonElement.onclick = () => app.contactButton(); } app.loadModel(modelXml); app.run(); } catch (error) { console.error('Initialization error: ', error); app.dispose(); } } main();