From 9967340eee3137166b593db068493ba21755eed8 Mon Sep 17 00:00:00 2001 From: Benjamin Bolte Date: Sun, 10 Nov 2024 00:04:46 -0800 Subject: [PATCH] stuff sort of working --- .../src/components/files/MJCFRenderer.tsx | 425 +++++++++--------- .../src/components/files/mujoco/mujoco.ts | 196 +++++++- 2 files changed, 387 insertions(+), 234 deletions(-) diff --git a/frontend/src/components/files/MJCFRenderer.tsx b/frontend/src/components/files/MJCFRenderer.tsx index 2cb72eab..bc33aa1f 100644 --- a/frontend/src/components/files/MJCFRenderer.tsx +++ b/frontend/src/components/files/MJCFRenderer.tsx @@ -1,4 +1,11 @@ import { useEffect, useRef, useState } from "react"; +import { + FaChevronLeft, + FaChevronRight, + FaPause, + FaPlay, + FaUndo, +} from "react-icons/fa"; import humanoid from "@/components/files/demo/humanoid.xml"; import { @@ -6,6 +13,9 @@ import { cleanupMujoco, initializeMujoco, initializeThreeJS, + setupModelGeometry, + setupScene, + updateBodyTransforms, } from "@/components/files/mujoco/mujoco"; import Spinner from "@/components/ui/Spinner"; import { humanReadableError } from "@/hooks/useAlertQueue"; @@ -36,194 +46,33 @@ const MJCFRenderer = ({ useControls = true }: Props) => { // State management const isSimulatingRef = useRef(false); const mujocoTimeRef = useRef(0); - const tmpVecRef = useRef(new THREE.Vector3()); - const tmpQuatRef = useRef(new THREE.Quaternion()); const [isMujocoReady, setIsMujocoReady] = useState(false); const [error, setError] = useState(null); // Constants const DEFAULT_TIMESTEP = 0.01; - const setupModelGeometry = () => { - const { sceneRef, modelRef, mujocoRef } = refs; - if (!sceneRef.current || !modelRef.current || !mujocoRef.current) return; - const model = modelRef.current; - const mj = mujocoRef.current; - - // Create body groups first - const bodies: { [key: number]: THREE.Group } = {}; - for (let b = 0; b < model.nbody; b++) { - bodies[b] = new THREE.Group(); - bodies[b].name = `body_${b}`; - bodies[b].userData.bodyId = b; - - // Add to parent body or scene - if (b === 0 || !bodies[0]) { - sceneRef.current.add(bodies[b]); - } else { - bodies[0].add(bodies[b]); - } - } - - try { - for (let i = 0; i < model.ngeom; i++) { - // Get geom properties from the model - const geomType = model.geom_type[i]; - const geomSize = model.geom_size.subarray(i * 3, i * 3 + 3); - const geomPos = model.geom_pos.subarray(i * 3, i * 3 + 3); - - // Create corresponding Three.js geometry - let geometry: THREE.BufferGeometry; - switch (geomType) { - case mj.mjtGeom.mjGEOM_PLANE.value: - geometry = new THREE.PlaneGeometry( - geomSize[0] * 2, - geomSize[1] * 2, - ); - break; - case mj.mjtGeom.mjGEOM_SPHERE.value: - geometry = new THREE.SphereGeometry(geomSize[0], 32, 32); - break; - case mj.mjtGeom.mjGEOM_CAPSULE.value: - // Capsule is a cylinder with hemispheres at the ends - geometry = new THREE.CapsuleGeometry( - geomSize[0], - geomSize[1] * 2, - 4, - 32, - ); - break; - case mj.mjtGeom.mjGEOM_ELLIPSOID.value: - // Create a sphere and scale it to make an ellipsoid - geometry = new THREE.SphereGeometry(1, 32, 32); - geometry.scale(geomSize[0], geomSize[1], geomSize[2]); - break; - case mj.mjtGeom.mjGEOM_CYLINDER.value: - geometry = new THREE.CylinderGeometry( - geomSize[0], - geomSize[0], - geomSize[1] * 2, - 32, - ); - break; - case mj.mjtGeom.mjGEOM_BOX.value: - geometry = new THREE.BoxGeometry( - geomSize[0] * 2, - geomSize[1] * 2, - geomSize[2] * 2, - ); - break; - case mj.mjtGeom.mjGEOM_MESH.value: - // For mesh, you'll need to load the actual mesh data - console.warn("Mesh geometry requires additional mesh data loading"); - geometry = new THREE.BoxGeometry(1, 1, 1); // Placeholder - break; - case mj.mjtGeom.mjGEOM_LINE.value: - geometry = new THREE.BufferGeometry().setFromPoints([ - new THREE.Vector3(0, 0, 0), - new THREE.Vector3(geomSize[0], 0, 0), - ]); - break; - default: - console.warn(`Unsupported geom type: ${geomType}`); - continue; - } - - // Create material (you can customize this based on geom properties) - const material = new THREE.MeshPhongMaterial({ color: 0x808080 }); - - // Create mesh and add to corresponding body - const mesh = new THREE.Mesh(geometry, material); - mesh.userData.geomIndex = i; - - const bodyId = model.geom_bodyid[i]; - if (bodies[bodyId]) { - // Set local position and orientation relative to body - const pos = model.geom_pos.subarray(i * 3, i * 3 + 3); - mesh.position.set(pos[0], pos[2], -pos[1]); - - const quat = model.geom_quat.subarray(i * 4, i * 4 + 4); - mesh.quaternion.set(-quat[1], -quat[3], quat[2], -quat[0]); - - bodies[bodyId].add(mesh); - } - } - - // Update initial body positions and orientations - for (let b = 0; b < model.nbody; b++) { - if (bodies[b]) { - const pos = refs.simulationRef.current!.xpos.subarray( - b * 3, - b * 3 + 3, - ); - bodies[b].position.set(pos[0], pos[2], -pos[1]); - - const quat = refs.simulationRef.current!.xquat.subarray( - b * 4, - b * 4 + 4, - ); - bodies[b].quaternion.set(-quat[1], -quat[3], quat[2], -quat[0]); - - bodies[b].updateWorldMatrix(true, true); - } - } + // Add new state for sidebar visibility + const [showControls, setShowControls] = useState(true); - setIsMujocoReady(true); - } catch (error) { - console.error(error); - } - }; + // Add state for simulation control + const [isSimulating, setIsSimulating] = useState(false); - const setupScene = () => { - if (!refs.sceneRef.current || !refs.cameraRef.current) return; + // Add new state for joint controls + const [joints, setJoints] = useState<{ name: string; value: number }[]>([]); - // Set up scene properties - refs.sceneRef.current.background = new THREE.Color(0.15, 0.25, 0.35); - refs.sceneRef.current.fog = new THREE.Fog( - refs.sceneRef.current.background, - 30, - 50, - ); + // Add function to update joint positions + const updateJointPosition = (index: number, value: number) => { + if (refs.stateRef.current && refs.simulationRef.current) { + const qpos = refs.stateRef.current?.qpos || []; + qpos[index] = value; + refs.stateRef.current.setQpos(qpos); + refs.simulationRef.current.forward(); + updateBodyTransforms(refs); - // Add ambient light - const ambientLight = new THREE.AmbientLight(0xffffff, 0.1); - ambientLight.name = "AmbientLight"; - refs.sceneRef.current.add(ambientLight); - - // Set up camera with wider view - refs.cameraRef.current.position.set(4.0, 3.4, 3.4); - refs.cameraRef.current.far = 100; - refs.cameraRef.current.updateProjectionMatrix(); - - // Update OrbitControls settings - if (refs.controlsRef.current) { - refs.controlsRef.current.target.set(0, 0.7, 0); - refs.controlsRef.current.panSpeed = 2; - refs.controlsRef.current.zoomSpeed = 1; - refs.controlsRef.current.enableDamping = true; - refs.controlsRef.current.dampingFactor = 0.1; - refs.controlsRef.current.screenSpacePanning = true; - refs.controlsRef.current.update(); - } - }; - - const updateBodyTransforms = () => { - const { modelRef, simulationRef, sceneRef } = refs; - if (!modelRef.current || !simulationRef.current || !sceneRef.current) - return; - - // Update body transforms - for (let b = 0; b < modelRef.current.nbody; b++) { - const body = sceneRef.current.getObjectByName(`body_${b}`); - if (body) { - const pos = simulationRef.current.xpos.subarray(b * 3, b * 3 + 3); - body.position.set(pos[0], pos[2], -pos[1]); - - const quat = simulationRef.current.xquat.subarray(b * 4, b * 4 + 4); - body.quaternion.set(-quat[1], -quat[3], quat[2], -quat[0]); - - body.updateWorldMatrix(true, true); - } + setJoints((prev) => + prev.map((joint, i) => (i === index ? { ...joint, value } : joint)), + ); } }; @@ -248,7 +97,7 @@ const MJCFRenderer = ({ useControls = true }: Props) => { } // Update body transforms after physics step - updateBodyTransforms(); + updateBodyTransforms(refs); } // Update controls if enabled @@ -266,54 +115,41 @@ const MJCFRenderer = ({ useControls = true }: Props) => { const stopPhysicsSimulation = () => { isSimulatingRef.current = false; + setIsSimulating(false); if (animationFrameRef.current) { cancelAnimationFrame(animationFrameRef.current); } }; - useEffect(() => { - if (!containerRef.current) return; - - (async () => { - try { - if (!containerRef.current) throw new Error("Container ref is null"); - - // Initialize MuJoCo with the humanoid model - const humanoidXML = await fetch(humanoid).then((res) => res.text()); - const { mj, model, state, simulation } = await initializeMujoco({ - modelXML: humanoidXML, - refs, - }); - - refs.mujocoRef.current = mj; - refs.modelRef.current = model; - refs.stateRef.current = state; - refs.simulationRef.current = simulation; - - // Initialize Three.js scene - const { renderer, scene, camera, controls } = initializeThreeJS( - containerRef.current, - ); - refs.rendererRef.current = renderer; - refs.sceneRef.current = scene; - refs.cameraRef.current = camera; - refs.controlsRef.current = controls; + // Add a function to start physics simulation + const startPhysicsSimulation = () => { + isSimulatingRef.current = true; + setIsSimulating(true); + animate(performance.now()); + }; - // Add model-specific setup (bodies, geometries, etc.) - setupModelGeometry(); + // Toggle simulation function + const toggleSimulation = () => { + if (isSimulating) { + stopPhysicsSimulation(); + } else { + startPhysicsSimulation(); + } + }; - // Add new scene setup - setupScene(); + // Add function to restart simulation + const restartSimulation = () => { + if (refs.stateRef.current && refs.simulationRef.current) { + refs.simulationRef.current.resetData(); + refs.simulationRef.current.forward(); + updateBodyTransforms(refs); + } + }; - // Start animation loop - animate(performance.now()); - } catch (error) { - console.error(error); - setError(error as Error); - } - })(); + useEffect(() => { + if (!containerRef.current) return; - // Handle window resize + // Define handleResize first const handleResize = () => { if ( refs.rendererRef.current && @@ -330,28 +166,171 @@ const MJCFRenderer = ({ useControls = true }: Props) => { } }; - window.addEventListener("resize", handleResize); - - return () => { + // Then define cleanup + const cleanup = () => { window.removeEventListener("resize", handleResize); + if (animationFrameRef.current) { + cancelAnimationFrame(animationFrameRef.current); + } cleanupMujoco(refs); stopPhysicsSimulation(); + Object.keys(refs).forEach((key) => { + refs[key as keyof MujocoRefs].current = null; + }); }; - }, [useControls, containerRef.current]); + + // Rest of the effect remains the same... + cleanup(); + + (async () => { + try { + if (!containerRef.current) throw new Error("Container ref is null"); + + // Initialize MuJoCo with the humanoid model + const humanoidXML = await fetch(humanoid).then((res) => res.text()); + const { mj, model, state, simulation } = await initializeMujoco({ + modelXML: humanoidXML, + refs, + }); + + // Only set refs if component is still mounted + if (containerRef.current) { + refs.mujocoRef.current = mj; + refs.modelRef.current = model; + refs.stateRef.current = state; + refs.simulationRef.current = simulation; + + // Initialize Three.js scene + const { renderer, scene, camera, controls } = initializeThreeJS( + containerRef.current, + ); + refs.rendererRef.current = renderer; + refs.sceneRef.current = scene; + refs.cameraRef.current = camera; + refs.controlsRef.current = controls; + + // Add model-specific setup (bodies, geometries, etc.) + setupModelGeometry(refs); + setupScene(refs); + setIsMujocoReady(true); + + // Add joint information after MuJoCo initialization + if (refs.modelRef.current) { + const jointNames = []; + const numJoints = refs.modelRef.current.nq; + + for (let i = 0; i < numJoints; i++) { + const name = refs.simulationRef.current.id2name( + mj.mjtObj.mjOBJ_JOINT.value, + i, + ); + const qpos = refs.stateRef.current?.qpos || []; + jointNames.push({ name, value: qpos[i] || 0 }); + } + setJoints(jointNames); + } + + // Start animation loop + animate(performance.now()); + } + } catch (error) { + console.error(error); + setError(error as Error); + } + })(); + + window.addEventListener("resize", handleResize); + return cleanup; + }, []); // Empty dependency array return ( -
+
+ {!isMujocoReady && (
)} + {error && (
{humanReadableError(error)}
)} + + {useControls && showControls && ( +
+
+
+
+ + + + + {/* Add joint controls */} +
+

Joint Controls

+ {joints.map((joint, index) => ( +
+ + + updateJointPosition(index, parseFloat(e.target.value)) + } + className="w-full" + /> +
+ ))} +
+
+
+
+
+ )} + + {useControls && !showControls && ( + + )}
); }; diff --git a/frontend/src/components/files/mujoco/mujoco.ts b/frontend/src/components/files/mujoco/mujoco.ts index fca8b311..e15779c3 100644 --- a/frontend/src/components/files/mujoco/mujoco.ts +++ b/frontend/src/components/files/mujoco/mujoco.ts @@ -1,3 +1,4 @@ +// @ts-nocheck import load_mujoco, { mujoco } from "@/lib/mujoco/mujoco_wasm"; import * as THREE from "three"; import { OrbitControls } from "three/examples/jsm/controls/OrbitControls"; @@ -34,7 +35,6 @@ export const initializeMujoco = async ({ const mj = await load_mujoco(); // Set up file system and load XML model - // @ts-ignore if (!mj.FS.analyzePath(MODEL_DIR).exists) { mj.FS.mkdir(MODEL_DIR); } @@ -114,16 +114,16 @@ export const initializeThreeJS = ( scene.add(dirLight); // Add floor - // const floorGeometry = new THREE.PlaneGeometry(10, 10); - // const floorMaterial = new THREE.MeshStandardMaterial({ - // color: 0xe0e0e0, - // roughness: 0.7, - // metalness: 0.1, - // }); - // const floor = new THREE.Mesh(floorGeometry, floorMaterial); - // floor.rotation.x = -Math.PI / 2; - // floor.position.y = 0; - // scene.add(floor); + const floorGeometry = new THREE.PlaneGeometry(10, 10); + const floorMaterial = new THREE.MeshStandardMaterial({ + color: 0xe0e0e0, + roughness: 0.7, + metalness: 0.1, + }); + const floor = new THREE.Mesh(floorGeometry, floorMaterial); + floor.rotation.x = -Math.PI / 2; + floor.position.y = 0; + scene.add(floor); return { renderer, scene, camera, controls }; }; @@ -172,3 +172,177 @@ export const cleanupMujoco = (refs: MujocoRefs) => { refs.cameraRef.current = null; refs.controlsRef.current = null; }; + +export const setupModelGeometry = (refs: MujocoRefs) => { + const { sceneRef, modelRef, mujocoRef } = refs; + if (!sceneRef.current || !modelRef.current || !mujocoRef.current) return; + const model = modelRef.current; + const mj = mujocoRef.current; + + // Create body groups first + const bodies: { [key: number]: THREE.Group } = {}; + for (let b = 0; b < model.nbody; b++) { + bodies[b] = new THREE.Group(); + bodies[b].name = `body_${b}`; + bodies[b].userData.bodyId = b; + + // Add to parent body or scene + if (b === 0 || !bodies[0]) { + sceneRef.current.add(bodies[b]); + } else { + bodies[0].add(bodies[b]); + } + } + + try { + for (let i = 0; i < model.ngeom; i++) { + // Get geom properties from the model + const geomType = model.geom_type[i]; + const geomSize = model.geom_size.subarray(i * 3, i * 3 + 3); + const geomPos = model.geom_pos.subarray(i * 3, i * 3 + 3); + + // Create corresponding Three.js geometry + let geometry: THREE.BufferGeometry; + switch (geomType) { + case mj.mjtGeom.mjGEOM_PLANE.value: + geometry = new THREE.PlaneGeometry(geomSize[0] * 2, geomSize[1] * 2); + break; + case mj.mjtGeom.mjGEOM_SPHERE.value: + geometry = new THREE.SphereGeometry(geomSize[0], 32, 32); + break; + case mj.mjtGeom.mjGEOM_CAPSULE.value: + // Capsule is a cylinder with hemispheres at the ends + geometry = new THREE.CapsuleGeometry( + geomSize[0], + geomSize[1] * 2, + 4, + 32, + ); + break; + case mj.mjtGeom.mjGEOM_ELLIPSOID.value: + // Create a sphere and scale it to make an ellipsoid + geometry = new THREE.SphereGeometry(1, 32, 32); + geometry.scale(geomSize[0], geomSize[1], geomSize[2]); + break; + case mj.mjtGeom.mjGEOM_CYLINDER.value: + geometry = new THREE.CylinderGeometry( + geomSize[0], + geomSize[0], + geomSize[1] * 2, + 32, + ); + break; + case mj.mjtGeom.mjGEOM_BOX.value: + geometry = new THREE.BoxGeometry( + geomSize[0] * 2, + geomSize[1] * 2, + geomSize[2] * 2, + ); + break; + case mj.mjtGeom.mjGEOM_MESH.value: + // For mesh, you'll need to load the actual mesh data + console.warn("Mesh geometry requires additional mesh data loading"); + geometry = new THREE.BoxGeometry(1, 1, 1); // Placeholder + break; + case mj.mjtGeom.mjGEOM_LINE.value: + geometry = new THREE.BufferGeometry().setFromPoints([ + new THREE.Vector3(0, 0, 0), + new THREE.Vector3(geomSize[0], 0, 0), + ]); + break; + default: + console.warn(`Unsupported geom type: ${geomType}`); + continue; + } + + // Create material (you can customize this based on geom properties) + const material = new THREE.MeshPhongMaterial({ color: 0x808080 }); + + // Create mesh and add to corresponding body + const mesh = new THREE.Mesh(geometry, material); + mesh.userData.geomIndex = i; + + const bodyId = model.geom_bodyid[i]; + if (bodies[bodyId]) { + // Set local position and orientation relative to body + const pos = model.geom_pos.subarray(i * 3, i * 3 + 3); + mesh.position.set(pos[0], pos[2], -pos[1]); + + const quat = model.geom_quat.subarray(i * 4, i * 4 + 4); + mesh.quaternion.set(-quat[1], -quat[3], quat[2], -quat[0]); + + bodies[bodyId].add(mesh); + } + } + + // Update initial body positions and orientations + for (let b = 0; b < model.nbody; b++) { + if (bodies[b]) { + const pos = refs.simulationRef.current!.xpos.subarray(b * 3, b * 3 + 3); + bodies[b].position.set(pos[0], pos[2], -pos[1]); + + const quat = refs.simulationRef.current!.xquat.subarray( + b * 4, + b * 4 + 4, + ); + bodies[b].quaternion.set(-quat[1], -quat[3], quat[2], -quat[0]); + + bodies[b].updateWorldMatrix(true, true); + } + } + } catch (error) { + console.error(error); + } +}; + +export const setupScene = (refs: MujocoRefs) => { + if (!refs.sceneRef.current || !refs.cameraRef.current) return; + + // Set up scene properties + refs.sceneRef.current.background = new THREE.Color(0.15, 0.25, 0.35); + refs.sceneRef.current.fog = new THREE.Fog( + refs.sceneRef.current.background, + 30, + 50, + ); + + // Add ambient light + const ambientLight = new THREE.AmbientLight(0xffffff, 0.1); + ambientLight.name = "AmbientLight"; + refs.sceneRef.current.add(ambientLight); + + // Set up camera with wider view + refs.cameraRef.current.position.set(4.0, 3.4, 3.4); + refs.cameraRef.current.far = 100; + refs.cameraRef.current.updateProjectionMatrix(); + + // Update OrbitControls settings + if (refs.controlsRef.current) { + refs.controlsRef.current.target.set(0, 0.7, 0); + refs.controlsRef.current.panSpeed = 2; + refs.controlsRef.current.zoomSpeed = 1; + refs.controlsRef.current.enableDamping = true; + refs.controlsRef.current.dampingFactor = 0.1; + refs.controlsRef.current.screenSpacePanning = true; + refs.controlsRef.current.update(); + } +}; + +export const updateBodyTransforms = (refs: MujocoRefs) => { + const { modelRef, simulationRef, sceneRef } = refs; + if (!modelRef.current || !simulationRef.current || !sceneRef.current) return; + + // Update body transforms + for (let b = 0; b < modelRef.current.nbody; b++) { + const body = sceneRef.current.getObjectByName(`body_${b}`); + if (body) { + const pos = simulationRef.current.xpos.subarray(b * 3, b * 3 + 3); + body.position.set(pos[0], pos[2], -pos[1]); + + const quat = simulationRef.current.xquat.subarray(b * 4, b * 4 + 4); + body.quaternion.set(-quat[1], -quat[3], quat[2], -quat[0]); + + body.updateWorldMatrix(true, true); + } + } +};