Spaces:
Running
Running
| // Generic MJCF rig builder, adapted from the microduck sandbox's | |
| // src/game/duck.js. Loads kinematics.json (produced by | |
| // tools/mjcf_to_kinematics.py) and builds an Object3D tree with one Group per | |
| // body; joints become per-body local rotations driven by name. | |
| // | |
| // Differences from the duck version: the jaw hinge and its SITTING_POSE are | |
| // gone (G1 has neither), meshes are keyed by the original glTF node name | |
| // because three strips '.' from node names, and a missing GLB is fatal rather | |
| // than optional since every G1 mesh is packed into one. | |
| import * as THREE from "three"; | |
| import { GLTFLoader } from "three/addons/loaders/GLTFLoader.js"; | |
| import { STLLoader } from "three/addons/loaders/STLLoader.js"; | |
| import { mergeVertices, toCreasedNormals } from "three/addons/utils/BufferGeometryUtils.js"; | |
| const CREASE = Math.PI / 5; // 36 deg | |
| export async function loadKinematics(url) { | |
| const r = await fetch(url); | |
| if (!r.ok) throw new Error(`kinematics fetch ${r.status} for ${url}`); | |
| return r.json(); | |
| } | |
| // three's GLTFLoader runs node names through PropertyBinding.sanitizeNodeName, | |
| // which strips '.' and ':' - so a node exported as "pelvis.STL" arrives as | |
| // "pelvisSTL". The untouched original survives on userData.name, so register | |
| // every alias we might be asked for. | |
| function registerAliases(map, mesh, entry) { | |
| for (const key of [mesh.userData?.name, mesh.name, mesh.geometry?.name]) { | |
| if (key && !map.has(key)) map.set(key, entry); | |
| } | |
| } | |
| function prepare(raw) { | |
| // STL repeats every vertex per facet and the packed GLB is already welded; | |
| // normals are rebuilt here either way so creases stay sharp. | |
| raw.deleteAttribute("normal"); | |
| const welded = raw.index ? raw : mergeVertices(raw, 1e-4); | |
| // toCreasedNormals hashes on a 0.01-unit grid. Meshes are in metres, so | |
| // scale a clone to mm (10 um grid) and back. | |
| const scaled = welded.clone(); | |
| scaled.scale(1000, 1000, 1000); | |
| const display = toCreasedNormals(scaled, CREASE); | |
| display.scale(1e-3, 1e-3, 1e-3); | |
| return { display, welded }; | |
| } | |
| export async function loadGlbGeometries(url) { | |
| const gltf = await new GLTFLoader().loadAsync(url); | |
| const map = new Map(); | |
| gltf.scene.traverse((o) => { | |
| if (!o.isMesh || !o.geometry) return; | |
| registerAliases(map, o, prepare(o.geometry)); | |
| }); | |
| return map; | |
| } | |
| export async function buildRig(k, opts = {}) { | |
| // placer carries the world pose in three.js space (no axis conversion)... | |
| const placer = new THREE.Group(); | |
| placer.name = "robot_placer"; | |
| // ...and root applies the MJCF +Z up -> three.js +Y up convention fix, so | |
| // everything below it stays in MJCF coordinates. | |
| const root = new THREE.Group(); | |
| root.name = "robot_root"; | |
| root.rotation.x = -Math.PI / 2; | |
| placer.add(root); | |
| const bodies = new Map(); | |
| const joints = new Map(); | |
| const geomByName = opts.glbUrl ? await loadGlbGeometries(opts.glbUrl) : new Map(); | |
| const stlCache = new Map(); | |
| const loadMesh = (name) => { | |
| const entry = geomByName.get(name); | |
| if (entry) return Promise.resolve(entry); | |
| if (!stlCache.has(name)) { | |
| stlCache.set( | |
| name, | |
| new STLLoader().loadAsync(`${k.mesh_dir}/${name}`).then(prepare), | |
| ); | |
| } | |
| return stlCache.get(name); | |
| }; | |
| for (const b of k.bodies) { | |
| const g = new THREE.Group(); | |
| g.name = b.name; | |
| g.position.set(b.pos[0], b.pos[1], b.pos[2]); | |
| g.quaternion.set(b.quat[1], b.quat[2], b.quat[3], b.quat[0]); | |
| bodies.set(b.name, g); | |
| } | |
| for (const b of k.bodies) { | |
| const g = bodies.get(b.name); | |
| if (b.parent && bodies.has(b.parent)) bodies.get(b.parent).add(g); | |
| else root.add(g); | |
| } | |
| for (const b of k.bodies) { | |
| if (!b.joint || (b.joint.type && b.joint.type !== "hinge")) continue; | |
| const g = bodies.get(b.name); | |
| joints.set(b.joint.name, { | |
| body: g, | |
| axis: new THREE.Vector3(...b.joint.axis).normalize(), | |
| baseQuat: g.quaternion.clone(), | |
| range: b.joint.range ?? null, | |
| }); | |
| } | |
| // Identical PBR props share one GPU material instance. | |
| const matCache = new Map(); | |
| const matFor = (rgba) => { | |
| const key = rgba.join(","); | |
| const cached = matCache.get(key); | |
| if (cached) return cached; | |
| const m = new THREE.MeshStandardMaterial({ | |
| color: new THREE.Color(rgba[0], rgba[1], rgba[2]), | |
| roughness: opts.roughness ?? 0.55, | |
| metalness: opts.metalness ?? 0.25, | |
| transparent: (rgba[3] ?? 1) < 1, | |
| opacity: rgba[3] ?? 1, | |
| }); | |
| matCache.set(key, m); | |
| return m; | |
| }; | |
| const pending = []; | |
| const missing = new Set(); | |
| for (const b of k.bodies) { | |
| const g = bodies.get(b.name); | |
| if (!g) continue; | |
| for (const geom of b.geoms) { | |
| if (geom.type && geom.type !== "mesh") continue; | |
| if (!geom.mesh) continue; | |
| pending.push( | |
| loadMesh(geom.mesh) | |
| .then(({ display }) => { | |
| const m = new THREE.Mesh(display, matFor(geom.color ?? [0.8, 0.8, 0.8, 1])); | |
| m.userData.meshName = geom.mesh; | |
| m.castShadow = true; | |
| m.receiveShadow = true; | |
| if (geom.pos) m.position.set(...geom.pos); | |
| if (geom.quat) m.quaternion.set(geom.quat[1], geom.quat[2], geom.quat[3], geom.quat[0]); | |
| g.add(m); | |
| }) | |
| .catch(() => missing.add(geom.mesh)), | |
| ); | |
| } | |
| } | |
| await Promise.all(pending); | |
| if (missing.size) console.warn(`[rig] ${missing.size} mesh(es) failed:`, [...missing]); | |
| return { placer, root, bodies, joints, missing, pose: Object.fromEntries([...joints.keys()].map(name => [name, 0])) }; | |
| } | |
| const _q = new THREE.Quaternion(); | |
| // Set one named joint, clamped to its MJCF range when known. | |
| export function setJoint(rig, name, angle) { | |
| const j = rig.joints.get(name); | |
| if (!j || !Number.isFinite(angle)) return false; | |
| let a = angle; | |
| if (j.range) a = Math.min(j.range[1], Math.max(j.range[0], a)); | |
| j.body.quaternion.copy(j.baseQuat).multiply(_q.setFromAxisAngle(j.axis, a)); | |
| rig.pose[name] = a; | |
| return true; | |
| } | |
| export function applyPose(rig, pose) { | |
| for (const [name, ang] of Object.entries(pose)) setJoint(rig, name, ang); | |
| } | |
| // MJCF is Z-up, three.js is Y-up. The rig's root already applies Rx(-90) to | |
| // everything below it, so a base pose given in MJCF world coords must be | |
| // mapped into three.js space before it lands on the placer: positions become | |
| // (x, z, -y) and the orientation is conjugated by the same rotation. | |
| const AXIS_FIX = new THREE.Quaternion().setFromAxisAngle(new THREE.Vector3(1, 0, 0), -Math.PI / 2); | |
| const AXIS_FIX_INV = AXIS_FIX.clone().invert(); | |
| const _qb = new THREE.Quaternion(); | |
| export function setBasePose(rig, pos, quat) { | |
| rig.placer.position.set(pos[0], pos[2], -pos[1]); | |
| if (!quat) return; | |
| _qb.set(quat[1], quat[2], quat[3], quat[0]); | |
| rig.placer.quaternion.copy(AXIS_FIX).multiply(_qb).multiply(AXIS_FIX_INV); | |
| } | |
| /** three.js world point -> MJCF world point, the inverse of the mapping above. */ | |
| export function toMujoco(v) { | |
| return [v.x, -v.z, v.y]; | |
| } | |
| /** MJCF world point -> three.js world point. */ | |
| export function toThree(p, out) { | |
| return out.set(p[0], p[2], -p[1]); | |
| } | |
| /** Walk up from a hit mesh to the rig body that owns it. */ | |
| export function bodyOf(rig, object) { | |
| for (let o = object; o; o = o.parent) { | |
| if (rig.bodies.get(o.name) === o) return o.name; | |
| } | |
| return null; | |
| } | |