SuperHQ Motion experiment: locomotion and action checkpoints, runtime, clips, evaluation
6adbbb1 verified | import { Bone, Group, Quaternion, Skeleton, Vector3 } from 'three'; | |
| // Three.js integration; the model and dense inference runtime are engine-independent. | |
| export function validateRig(rig) { | |
| const n = rig.names?.length; | |
| if (!Number.isInteger(n) || n < 1 || new Set(rig.names).size !== n || | |
| rig.parents?.length !== n || rig.offsets?.length !== n || rig.end_sites?.length !== n) { | |
| throw new Error('Rig requires unique names and equally sized hierarchy, offset and end-site arrays.'); | |
| } | |
| if (rig.units !== 'metres' || rig.up_axis !== 'Y' || rig.handedness !== 'right') { | |
| throw new Error('Convert the rig to right-handed Y-up metres before binding.'); | |
| } | |
| for (let i = 0; i < n; i++) { | |
| const p = rig.parents[i]; | |
| if (!Number.isInteger(p) || (i === 0 ? p !== -1 : p < 0 || p >= i)) { | |
| throw new Error(`Invalid parent for joint ${rig.names[i]}.`); | |
| } | |
| if (typeof rig.end_sites[i] !== 'boolean' || rig.offsets[i]?.length !== 3 || | |
| !rig.offsets[i].every(Number.isFinite)) throw new Error(`Invalid joint ${rig.names[i]}.`); | |
| } | |
| return rig; | |
| } | |
| export function createSourceRig(rig) { | |
| validateRig(rig); | |
| const group = new Group(); | |
| group.name = 'SuperHQMotion'; | |
| const bones = rig.names.map((name, i) => { | |
| const bone = new Bone(); | |
| bone.name = name; | |
| bone.position.fromArray(rig.offsets[i]); | |
| return bone; | |
| }); | |
| bones.forEach((bone, i) => (rig.parents[i] < 0 ? group : bones[rig.parents[i]]).add(bone)); | |
| group.updateMatrixWorld(true); | |
| const skeleton = new Skeleton(bones); | |
| return { group, skeleton, adapter: bindSkeleton(skeleton, rig) }; | |
| } | |
| /** Apply local rotations to a compatible Three.js Skeleton. Retargeting is a separate operation. */ | |
| export function bindSkeleton(skeleton, rig, nameMap = {}) { | |
| validateRig(rig); | |
| if (!(skeleton instanceof Skeleton)) throw new Error('Expected a Three.js Skeleton.'); | |
| const byName = new Map(); | |
| for (const bone of skeleton.bones) { | |
| if (byName.has(bone.name)) throw new Error(`Duplicate target bone name: ${bone.name}`); | |
| byName.set(bone.name, bone); | |
| } | |
| const mapped = rig.names.map((name, i) => { | |
| const bone = byName.get(nameMap[name] ?? name); | |
| if (!bone && !rig.end_sites[i]) throw new Error(`Missing articulated joint: ${name}`); | |
| return bone; | |
| }); | |
| if (new Set(mapped.filter(Boolean)).size !== mapped.filter(Boolean).length) { | |
| throw new Error('Multiple source joints map to the same target bone.'); | |
| } | |
| mapped.forEach((bone, i) => { | |
| if (!bone) return; | |
| if (rig.parents[i] >= 0 && bone.parent !== mapped[rig.parents[i]]) { | |
| throw new Error(`Hierarchy mismatch: ${rig.names[i]}`); | |
| } | |
| // This initial adapter requires the canonical bind pose. It does not silently retarget proportions. | |
| if (bone.position.distanceTo(new Vector3(...rig.offsets[i])) > 1e-4) { | |
| throw new Error(`Rest offset mismatch: ${rig.names[i]}. Retarget before binding.`); | |
| } | |
| }); | |
| const root = mapped[0]; | |
| const work = new Quaternion(); | |
| return { | |
| skeleton, | |
| rig, | |
| applyPose(pose) { | |
| if (pose.rootPosition?.length !== 3 || !pose.rootPosition.every(Number.isFinite) || | |
| pose.localQuaternions?.length !== rig.names.length) throw new Error('Invalid pose dimensions.'); | |
| // Validate the whole pose before mutating any bone. | |
| for (const q of pose.localQuaternions) { | |
| if (q?.length !== 4 || !q.every(Number.isFinite) || Math.hypot(...q) < 1e-8) { | |
| throw new Error('Nonfinite or degenerate pose rotation.'); | |
| } | |
| } | |
| root.position.fromArray(pose.rootPosition); | |
| mapped.forEach((bone, i) => { | |
| if (bone) bone.quaternion.copy(work.fromArray(pose.localQuaternions[i]).normalize()); | |
| }); | |
| root.updateWorldMatrix(true, true); | |
| }, | |
| }; | |
| } | |
| /** Playback of saved source or prediction sequences. This does not run inference. */ | |
| export function samplePoseSequence(clip, seconds) { | |
| if (!['recorded_training_motion', 'saved_model_rollout'].includes(clip.kind) || clip.partition !== 'train' || | |
| !Number.isFinite(seconds) || seconds < 0 || !(clip.dt > 0) || clip.root_positions.length < 1) { | |
| throw new Error('Expected a training pose sequence and a finite nonnegative time.'); | |
| } | |
| const t = Math.min(seconds / clip.dt, clip.root_positions.length - 1); | |
| const a = Math.floor(t), b = Math.min(a + 1, clip.root_positions.length - 1), alpha = t - a; | |
| const rootPosition = new Vector3(...clip.root_positions[a]).lerp(new Vector3(...clip.root_positions[b]), alpha).toArray(); | |
| const localQuaternions = clip.local_quaternions[a].map((q, i) => new Quaternion(...q) | |
| .slerp(new Quaternion(...clip.local_quaternions[b][i]), alpha).toArray()); | |
| return { rootPosition, localQuaternions }; | |
| } | |
| export function sampleRecordedMotion(clip, seconds) { | |
| if (clip.kind !== 'recorded_training_motion') throw new Error('Expected recorded source motion.'); | |
| return samplePoseSequence(clip, seconds); | |
| } | |
| /** Rest-aligned world-rotation retargeting. This does not solve foot contacts. */ | |
| export function createRetargeter(skeleton, rig, {nameMap, directionMap, restCorrectionMap, scale=1, rotation=[0,0,0,1], offset=[0,0,0]}={}) { | |
| validateRig(rig); | |
| const scaleVector=typeof scale==='number'?[scale,scale,scale]:scale; | |
| if(!(skeleton instanceof Skeleton)||!nameMap||scaleVector?.length!==3||!scaleVector.every(v=>Number.isFinite(v)&&v>0)|| | |
| rotation?.length!==4||!rotation.every(Number.isFinite)||Math.hypot(...rotation)<1e-8|| | |
| offset?.length!==3||!offset.every(Number.isFinite))throw Error('Invalid retargeting configuration.'); | |
| const byName=new Map(); | |
| for(const bone of skeleton.bones){if(byName.has(bone.name))throw Error('Duplicate target bone name.');byName.set(bone.name,bone);} | |
| const alignment=new Quaternion(...rotation).normalize(),inverseAlignment=alignment.clone().invert(); | |
| const entries=Object.entries(nameMap).map(([source,target])=>{ | |
| const index=rig.names.indexOf(source),bone=byName.get(target); | |
| if(index<0||rig.end_sites[index]||!bone)throw Error(`Invalid retarget mapping: ${source} -> ${target}`); | |
| bone.updateWorldMatrix(true,false); | |
| return {index,bone,rest:bone.getWorldQuaternion(new Quaternion())}; | |
| }); | |
| if(!entries.some(e=>e.index===0)||new Set(entries.map(e=>e.bone)).size!==entries.length)throw Error('Map the source root and use unique target bones.'); | |
| function depth(bone){let n=0;for(let p=bone.parent;p;p=p.parent)n++;return n;} | |
| entries.sort((a,b)=>depth(a.bone)-depth(b.bone)); | |
| const root=entries.find(e=>e.index===0).bone; | |
| if(entries.some(e=>e.bone!==root&&!isBelow(e.bone,root)))throw Error('Mapped target bones must descend from the mapped root.'); | |
| function isBelow(bone,parent){for(let p=bone.parent;p;p=p.parent)if(p===parent)return true;return false;} | |
| if(directionMap||restCorrectionMap){ | |
| const sourceRest=[]; | |
| rig.offsets.forEach((v,i)=>sourceRest.push(new Vector3(...v).add(rig.parents[i]<0?new Vector3():sourceRest[rig.parents[i]]))); | |
| const corrections=new Map(); | |
| for(const entry of [...entries].sort((a,b)=>a.index-b.index)){ | |
| const spec=directionMap?.[rig.names[entry.index]],explicit=restCorrectionMap?.[rig.names[entry.index]];let correction; | |
| if(explicit){ | |
| if(explicit.length!==4||!explicit.every(Number.isFinite)||Math.hypot(...explicit)<1e-8)throw Error('Invalid rest correction quaternion.'); | |
| correction=new Quaternion(...explicit).normalize(); | |
| }else if(spec){ | |
| const child=rig.names.indexOf(spec.sourceChild),targetChild=byName.get(spec.targetChild); | |
| if(child<0||!targetChild||targetChild.parent!==entry.bone)throw Error('Direction calibration requires a source descendant and direct target child.'); | |
| let ancestor=rig.parents[child];while(ancestor>=0&&ancestor!==entry.index)ancestor=rig.parents[ancestor]; | |
| if(ancestor!==entry.index)throw Error('Source calibration child must descend from its mapped bone.'); | |
| const from=targetChild.getWorldPosition(new Vector3()).sub(entry.bone.getWorldPosition(new Vector3())); | |
| const to=sourceRest[child].clone().sub(sourceRest[entry.index]).applyQuaternion(alignment); | |
| if(from.length()<1e-8||to.length()<1e-8)throw Error('Degenerate calibration direction.'); | |
| correction=new Quaternion().setFromUnitVectors(from.normalize(),to.normalize()); | |
| }else{ | |
| let ancestor=rig.parents[entry.index];while(ancestor>=0&&!corrections.has(ancestor))ancestor=rig.parents[ancestor]; | |
| correction=(corrections.get(ancestor)??new Quaternion()).clone(); | |
| } | |
| corrections.set(entry.index,correction);entry.rest.premultiply(correction); | |
| } | |
| } | |
| return {skeleton,rig, | |
| applyPose(pose) { | |
| if(pose.rootPosition?.length!==3||!pose.rootPosition.every(Number.isFinite)||pose.localQuaternions?.length!==rig.names.length|| | |
| pose.localQuaternions.some(q=>q?.length!==4||!q.every(Number.isFinite)||Math.hypot(...q)<1e-8))throw Error('Invalid source pose.'); | |
| const sourceWorld=[]; | |
| pose.localQuaternions.forEach((q,i)=>{ | |
| const local=new Quaternion(...q).normalize(); | |
| sourceWorld.push(rig.parents[i]<0?local:local.premultiply(sourceWorld[rig.parents[i]])); | |
| }); | |
| const rootWorld=new Vector3(...pose.rootPosition).multiply(new Vector3(...scaleVector)).applyQuaternion(alignment).add(new Vector3(...offset)); | |
| if(root.parent){root.parent.updateWorldMatrix(true,false);root.position.copy(root.parent.worldToLocal(rootWorld));} | |
| else root.position.copy(rootWorld); | |
| for(const {index,bone,rest} of entries){ | |
| const world=alignment.clone().multiply(sourceWorld[index]).multiply(inverseAlignment).multiply(rest); | |
| if(bone.parent){bone.parent.updateWorldMatrix(true,false);world.premultiply(bone.parent.getWorldQuaternion(new Quaternion()).invert());} | |
| bone.quaternion.copy(world).normalize();bone.updateWorldMatrix(false,false); | |
| } | |
| root.updateWorldMatrix(true,true); | |
| }, | |
| }; | |
| } | |