extract the root motion from a three.js AnimationClip

extract-root-motion.ts

// Adapted from https://github.com/donmccurdy/three.js/blob/feat/extract-root-motion/examples/jsm/utils/SkeletonUtils.js#L606


import {AnimationClip, Bone, BufferAttribute, Object3D, Quaternion, Vector3, Vector4} from "three";

function getRootBone(rootNode: Object3D) {
  const rootBones: Bone[] = [];

  rootNode.traverse(function (object) {
    // @ts-ignore
    if (object.isBone && !(object.parent && object.parent.isBone)) {
      rootBones.push(object as Bone);
    }
  });

  if (rootBones.length !== 1) {
    throw new Error(`THREE.SkeletonUtils: Expected 1 root bone, found ${rootBones.length}.`);
  }

  return rootBones[0];
}


export type RootMotionData = {
  startPosition: Vector3,
  endPosition: Vector3,
  totalDistance: number,
  averageVelocity: Vector3,
}

export function extractRootMotionPosition(rootNode: Object3D, targetClip: AnimationClip, trackName?: string): RootMotionData {
  // if no explicit name for the track containing the root motion was specified, we try it to infer it from the name of the root bone
  if (trackName == undefined) {
    const rootBone = getRootBone(rootNode);
    trackName = (rootBone.name || rootBone.uuid) + '.position';
  }

  const track = targetClip.tracks.find((track) => track.name === trackName);
  if (track == undefined) {
    throw new Error(`THREE.SkeletonUtils: Track ${trackName} not found.`);
  }

  const startTime = track.times[0];
  const endTime = track.times[track.times.length - 1];
  const duration = endTime - startTime;
  const values = new BufferAttribute(track.values, 3);

  const itemsPerKeyframe = track.getValueSize() / 3; // 1 (value), or 3 (inTan,value,outTan)
  const itemOffset = (itemsPerKeyframe === 3) ? 1 : 0;

  const startPosition = new Vector3().fromBufferAttribute(values, itemOffset);
  const endPositionIndex = values.count - itemsPerKeyframe + itemOffset;
  const endPosition = new Vector3().fromBufferAttribute(values, endPositionIndex);

  const averageVelocity = new Vector3();

  averageVelocity
    .copy(endPosition)
    .sub(startPosition)
    .divideScalar(duration);
  
  const position = new Vector3();

  for (let i = 0, il = track.times.length; i < il; i++) {
    const dt = track.times[i] - startTime;

    position.fromBufferAttribute(values, i + itemOffset);
    position.x -= averageVelocity.x * dt;
    position.y -= averageVelocity.y * dt;
    position.z -= averageVelocity.z * dt;
    values.setXYZ(i + itemOffset, position.x, position.y, position.z);
  }

  const resultObject = {
    startPosition,
    endPosition,
    totalDistance: startPosition.distanceTo(endPosition),
    averageVelocity,
  }

  return resultObject;
}