lerobot/glove_visualizer

1// Generic MJCF rig builder, adapted from the microduck sandbox's2// src/game/duck.js. Loads kinematics.json (produced by3// tools/mjcf_to_kinematics.py) and builds an Object3D tree with one Group per4// body; joints become per-body local rotations driven by name.5//6// Differences from the duck version: the jaw hinge and its SITTING_POSE are7// gone (G1 has neither), meshes are keyed by the original glTF node name8// because three strips '.' from node names, and a missing GLB is fatal rather9// than optional since every G1 mesh is packed into one.10 11import * as THREE from "three";12import { GLTFLoader } from "three/addons/loaders/GLTFLoader.js";13import { STLLoader } from "three/addons/loaders/STLLoader.js";14import { mergeVertices, toCreasedNormals } from "three/addons/utils/BufferGeometryUtils.js";15 16const CREASE = Math.PI / 5; // 36 deg17 18export async function loadKinematics(url) {19  const r = await fetch(url);20  if (!r.ok) throw new Error(`kinematics fetch ${r.status} for ${url}`);21  return r.json();22}23 24// three's GLTFLoader runs node names through PropertyBinding.sanitizeNodeName,25// which strips '.' and ':' - so a node exported as "pelvis.STL" arrives as26// "pelvisSTL". The untouched original survives on userData.name, so register27// every alias we might be asked for.28function registerAliases(map, mesh, entry) {29  for (const key of [mesh.userData?.name, mesh.name, mesh.geometry?.name]) {30    if (key && !map.has(key)) map.set(key, entry);31  }32}33 34function prepare(raw) {35  // STL repeats every vertex per facet and the packed GLB is already welded;36  // normals are rebuilt here either way so creases stay sharp.37  raw.deleteAttribute("normal");38  const welded = raw.index ? raw : mergeVertices(raw, 1e-4);39  // toCreasedNormals hashes on a 0.01-unit grid. Meshes are in metres, so40  // scale a clone to mm (10 um grid) and back.41  const scaled = welded.clone();42  scaled.scale(1000, 1000, 1000);43  const display = toCreasedNormals(scaled, CREASE);44  display.scale(1e-3, 1e-3, 1e-3);45  return { display, welded };46}47 48export async function loadGlbGeometries(url) {49  const gltf = await new GLTFLoader().loadAsync(url);50  const map = new Map();51  gltf.scene.traverse((o) => {52    if (!o.isMesh || !o.geometry) return;53    registerAliases(map, o, prepare(o.geometry));54  });55  return map;56}57 58export async function buildRig(k, opts = {}) {59  // placer carries the world pose in three.js space (no axis conversion)...60  const placer = new THREE.Group();61  placer.name = "robot_placer";62  // ...and root applies the MJCF +Z up -> three.js +Y up convention fix, so63  // everything below it stays in MJCF coordinates.64  const root = new THREE.Group();65  root.name = "robot_root";66  root.rotation.x = -Math.PI / 2;67  placer.add(root);68 69  const bodies = new Map();70  const joints = new Map();71  const geomByName = opts.glbUrl ? await loadGlbGeometries(opts.glbUrl) : new Map();72 73  const stlCache = new Map();74  const loadMesh = (name) => {75    const entry = geomByName.get(name);76    if (entry) return Promise.resolve(entry);77    if (!stlCache.has(name)) {78      stlCache.set(79        name,80        new STLLoader().loadAsync(`${k.mesh_dir}/${name}`).then(prepare),81      );82    }83    return stlCache.get(name);84  };85 86  for (const b of k.bodies) {87    const g = new THREE.Group();88    g.name = b.name;89    g.position.set(b.pos[0], b.pos[1], b.pos[2]);90    g.quaternion.set(b.quat[1], b.quat[2], b.quat[3], b.quat[0]);91    bodies.set(b.name, g);92  }93  for (const b of k.bodies) {94    const g = bodies.get(b.name);95    if (b.parent && bodies.has(b.parent)) bodies.get(b.parent).add(g);96    else root.add(g);97  }98  for (const b of k.bodies) {99    if (!b.joint || (b.joint.type && b.joint.type !== "hinge")) continue;100    const g = bodies.get(b.name);101    joints.set(b.joint.name, {102      body: g,103      axis: new THREE.Vector3(...b.joint.axis).normalize(),104      baseQuat: g.quaternion.clone(),105      range: b.joint.range ?? null,106    });107  }108 109  // Identical PBR props share one GPU material instance.110  const matCache = new Map();111  const matFor = (rgba) => {112    const key = rgba.join(",");113    const cached = matCache.get(key);114    if (cached) return cached;115    const m = new THREE.MeshStandardMaterial({116      color: new THREE.Color(rgba[0], rgba[1], rgba[2]),117      roughness: opts.roughness ?? 0.55,118      metalness: opts.metalness ?? 0.25,119      transparent: (rgba[3] ?? 1) < 1,120      opacity: rgba[3] ?? 1,121    });122    matCache.set(key, m);123    return m;124  };125 126  const pending = [];127  const missing = new Set();128  for (const b of k.bodies) {129    const g = bodies.get(b.name);130    if (!g) continue;131    for (const geom of b.geoms) {132      if (geom.type && geom.type !== "mesh") continue;133      if (!geom.mesh) continue;134      pending.push(135        loadMesh(geom.mesh)136          .then(({ display }) => {137            const m = new THREE.Mesh(display, matFor(geom.color ?? [0.8, 0.8, 0.8, 1]));138            m.userData.meshName = geom.mesh;139            m.castShadow = true;140            m.receiveShadow = true;141            if (geom.pos) m.position.set(...geom.pos);142            if (geom.quat) m.quaternion.set(geom.quat[1], geom.quat[2], geom.quat[3], geom.quat[0]);143            g.add(m);144          })145          .catch(() => missing.add(geom.mesh)),146      );147    }148  }149  await Promise.all(pending);150  if (missing.size) console.warn(`[rig] ${missing.size} mesh(es) failed:`, [...missing]);151 152  return { placer, root, bodies, joints, missing, pose: Object.fromEntries([...joints.keys()].map(name => [name, 0])) };153}154 155const _q = new THREE.Quaternion();156 157// Set one named joint, clamped to its MJCF range when known.158export function setJoint(rig, name, angle) {159  const j = rig.joints.get(name);160  if (!j || !Number.isFinite(angle)) return false;161  let a = angle;162  if (j.range) a = Math.min(j.range[1], Math.max(j.range[0], a));163  j.body.quaternion.copy(j.baseQuat).multiply(_q.setFromAxisAngle(j.axis, a));164  rig.pose[name] = a;165  return true;166}167 168export function applyPose(rig, pose) {169  for (const [name, ang] of Object.entries(pose)) setJoint(rig, name, ang);170}171 172// MJCF is Z-up, three.js is Y-up. The rig's root already applies Rx(-90) to173// everything below it, so a base pose given in MJCF world coords must be174// mapped into three.js space before it lands on the placer: positions become175// (x, z, -y) and the orientation is conjugated by the same rotation.176const AXIS_FIX = new THREE.Quaternion().setFromAxisAngle(new THREE.Vector3(1, 0, 0), -Math.PI / 2);177const AXIS_FIX_INV = AXIS_FIX.clone().invert();178const _qb = new THREE.Quaternion();179 180export function setBasePose(rig, pos, quat) {181  rig.placer.position.set(pos[0], pos[2], -pos[1]);182  if (!quat) return;183  _qb.set(quat[1], quat[2], quat[3], quat[0]);184  rig.placer.quaternion.copy(AXIS_FIX).multiply(_qb).multiply(AXIS_FIX_INV);185}186 187/** three.js world point -> MJCF world point, the inverse of the mapping above. */188export function toMujoco(v) {189  return [v.x, -v.z, v.y];190}191 192/** MJCF world point -> three.js world point. */193export function toThree(p, out) {194  return out.set(p[0], p[2], -p[1]);195}196 197/** Walk up from a hit mesh to the rig body that owns it. */198export function bodyOf(rig, object) {199  for (let o = object; o; o = o.parent) {200    if (rig.bodies.get(o.name) === o) return o.name;201  }202  return null;203}204