Add per-step snapshots and PartOfPlan relation for motion plans

.changeset/motion-plan-replayer-1c-snapshots.md

@@ -0,0 +1,5 @@
+---
+'@viamrobotics/motion-tools': patch
+---
+
+Add per-step snapshot generation (`parsedPlanToSnapshots`) and the `PartOfPlan` ECS relation for the Motion Plan Replayer, turning frame descriptors into renderable snapshots grouped under a plan entity.

src/lib/plugins/MotionPlanReplayer/__tests__/plan-to-snapshots.spec.ts

@@ -0,0 +1,123 @@
+import { describe, expect, it } from 'vitest'
+
+import { parsePlan } from '../parse-plan'
+import { parsedPlanToSnapshots } from '../plan-to-snapshots'
+
+// arm chain: waist (joint, Z-axis) → base (link, z=100mm, capsule geometry)
+const REQUEST = {
+	frame_system: {
+		frames: {
+			arm: {
+				frame_type: 'model',
+				frame: { name: 'arm', model: { joints: [{ id: 'waist' }] } },
+			},
+			'arm:waist': {
+				frame_type: 'named',
+				frame: {
+					inner_frame: {
+						frame_type: 'rotational',
+						frame: { axis: { X: 0, Y: 0, Z: 1 } },
+					},
+				},
+			},
+			'arm:base': {
+				frame_type: 'named',
+				frame: {
+					inner_frame: {
+						frame_type: 'static',
+						frame: {
+							translation: { X: 0, Y: 0, Z: 100 },
+							orientation: { type: 'quaternion', value: { W: 1, X: 0, Y: 0, Z: 0 } },
+							geometry: {
+								type: 'capsule',
+								r: 30,
+								l: 100,
+								translation: { X: 0, Y: 0, Z: 50 },
+								orientation: { type: 'quaternion', value: { W: 1, X: 0, Y: 0, Z: 0 } },
+								Label: 'base',
+							},
+						},
+					},
+				},
+			},
+		},
+		parents: { 'arm:waist': 'world', 'arm:base': 'arm:waist' },
+	},
+	goals: [],
+	start_state: {},
+}
+
+const RESULT = { trajectory: [{ arm: [0] }, { arm: [1.5708] }] }
+const CONTENT = JSON.stringify(REQUEST) + JSON.stringify(RESULT)
+
+const snapshotsFromContent = (content: string) => parsedPlanToSnapshots(parsePlan(content))
+
+describe('parsedPlanToSnapshots', () => {
+	it('returns one Snapshot per trajectory step', () => {
+		const snapshots = snapshotsFromContent(CONTENT)
+		expect(snapshots).toHaveLength(2)
+	})
+
+	it('joint and link both appear as transforms', () => {
+		const snapshots = snapshotsFromContent(CONTENT)
+		// waist (joint) and base (link) each emit a transform
+		expect(snapshots[0]!.transforms).toHaveLength(2)
+		const names = snapshots[0]!.transforms.map((t) => t.referenceFrame)
+		expect(names).toContain('arm:waist')
+		expect(names).toContain('arm:base')
+	})
+
+	it('joint pose changes per step; link local pose stays constant', () => {
+		const snapshots = snapshotsFromContent(CONTENT)
+		const waist0 = snapshots[0]!.transforms.find((t) => t.referenceFrame === 'arm:waist')!
+		const waist1 = snapshots[1]!.transforms.find((t) => t.referenceFrame === 'arm:waist')!
+		const base0 = snapshots[0]!.transforms.find((t) => t.referenceFrame === 'arm:base')!
+		const base1 = snapshots[1]!.transforms.find((t) => t.referenceFrame === 'arm:base')!
+
+		// joint theta changes: step 0 → 0°, step 1 → 90°
+		expect(waist0.poseInObserverFrame!.pose!.theta).toBeCloseTo(0, 1)
+		expect(waist1.poseInObserverFrame!.pose!.theta).toBeCloseTo(90, 1)
+
+		// link local pose is constant — WorldMatrix handles world-space composition
+		expect(base0.poseInObserverFrame!.pose!.z).toBeCloseTo(100, 1)
+		expect(base1.poseInObserverFrame!.pose!.z).toBeCloseTo(100, 1)
+		expect(base0.poseInObserverFrame!.pose!.theta).toBeCloseTo(0, 1)
+		expect(base1.poseInObserverFrame!.pose!.theta).toBeCloseTo(0, 1)
+	})
+
+	it('link is parented to its joint', () => {
+		const snapshots = snapshotsFromContent(CONTENT)
+		const base = snapshots[0]!.transforms.find((t) => t.referenceFrame === 'arm:base')!
+		expect(base.poseInObserverFrame!.referenceFrame).toBe('arm:waist')
+	})
+
+	it('joint is parented to its frame_system parent', () => {
+		const snapshots = snapshotsFromContent(CONTENT)
+		const waist = snapshots[0]!.transforms.find((t) => t.referenceFrame === 'arm:waist')!
+		expect(waist.poseInObserverFrame!.referenceFrame).toBe('world')
+	})
+
+	it('joint has no physicalObject', () => {
+		const snapshots = snapshotsFromContent(CONTENT)
+		const waist = snapshots[0]!.transforms.find((t) => t.referenceFrame === 'arm:waist')!
+		expect(waist.physicalObject).toBeUndefined()
+	})
+
+	it('link has same UUID across steps', () => {
+		const snapshots = snapshotsFromContent(CONTENT)
+		const base0 = snapshots[0]!.transforms.find((t) => t.referenceFrame === 'arm:base')!
+		const base1 = snapshots[1]!.transforms.find((t) => t.referenceFrame === 'arm:base')!
+		expect(base0.uuid).toStrictEqual(base1.uuid)
+	})
+
+	it('link has physicalObject when geometry is present', () => {
+		const snapshots = snapshotsFromContent(CONTENT)
+		const base = snapshots[0]!.transforms.find((t) => t.referenceFrame === 'arm:base')!
+		expect(base.physicalObject).toBeDefined()
+		expect(base.physicalObject!.geometryType.case).toBe('capsule')
+	})
+
+	it('returns empty array for plan with no trajectory', () => {
+		expect(parsedPlanToSnapshots(parsePlan(JSON.stringify(REQUEST)))).toHaveLength(0)
+	})
+})

src/lib/plugins/MotionPlanReplayer/__tests__/relations.spec.ts

@@ -0,0 +1,38 @@
+import { createWorld, type World } from 'koota'
+import { afterEach, describe, expect, it } from 'vitest'
+
+import { relations, traits } from '$lib/ecs'
+
+import { PartOfPlan } from '../relations'
+
+describe('PartOfPlan', () => {
+	let world: World
+	afterEach(() => world?.destroy())
+
+	it('destroys every member when the plan root is destroyed', () => {
+		world = createWorld()
+		const planRoot = world.spawn(traits.Name('plan'))
+		const member = world.spawn(PartOfPlan(planRoot))
+		const nestedParent = world.spawn(PartOfPlan(planRoot))
+		const nestedChild = world.spawn(relations.ChildOf(nestedParent), PartOfPlan(planRoot))
+
+		planRoot.destroy()
+
+		expect(planRoot.isAlive()).toBe(false)
+		expect(member.isAlive()).toBe(false)
+		expect(nestedParent.isAlive()).toBe(false)
+		expect(nestedChild.isAlive()).toBe(false)
+	})
+
+	it('leaves entities outside the plan alone', () => {
+		world = createWorld()
+		const planRoot = world.spawn(traits.Name('plan'))
+		const member = world.spawn(PartOfPlan(planRoot))
+		const unrelated = world.spawn(traits.Name('unrelated'))
+
+		planRoot.destroy()
+
+		expect(member.isAlive()).toBe(false)
+		expect(unrelated.isAlive()).toBe(true)
+	})
+})

src/lib/plugins/MotionPlanReplayer/plan-to-snapshots.ts

@@ -0,0 +1,72 @@
+import { Quaternion, Vector3 } from 'three'
+import { UuidTool } from 'uuid-tool'
+
+import { Pose, PoseInFrame, Transform } from '$lib/buf/common/v1/common_pb'
+import { Snapshot } from '$lib/buf/draw/v1/snapshot_pb'
+import { quaternionToPose } from '$lib/transform'
+
+import type { ParsedPlan } from './parse-plan'
+
+import {
+	buildFrameDescriptors,
+	type FrameDescriptor,
+	type JointFrameDescriptor,
+} from './build-frame-descriptors'
+
+// Shared scratch objects — safe in single-threaded JS
+const tmpQ = new Quaternion()
+const tmpVec = new Vector3()
+
+const computeJointPose = (descriptor: JointFrameDescriptor, angleRad: number): Pose => {
+	tmpQ.setFromAxisAngle(
+		tmpVec.set(descriptor.axis.X, descriptor.axis.Y, descriptor.axis.Z).normalize(),
+		angleRad
+	)
+	const pose = new Pose()
+	quaternionToPose(tmpQ, pose)
+	return pose
+}
+
+const descriptorToTransform = (
+	descriptor: FrameDescriptor,
+	stepInputs: Record<string, number[]>
+): Transform => {
+	if (descriptor.kind === 'static') {
+		return new Transform({
+			referenceFrame: descriptor.name,
+			poseInObserverFrame: new PoseInFrame({
+				referenceFrame: descriptor.parent,
+				pose: descriptor.localPose,
+			}),
+			physicalObject: descriptor.geometry ?? undefined,
+			uuid: descriptor.uuid,
+		})
+	}
+
+	const angleRad = stepInputs[descriptor.componentName]?.[descriptor.jointIndex] ?? 0
+	return new Transform({
+		referenceFrame: descriptor.name,
+		poseInObserverFrame: new PoseInFrame({
+			referenceFrame: descriptor.parent,
+			pose: computeJointPose(descriptor, angleRad),
+		}),
+		uuid: descriptor.uuid,
+	})
+}
+
+const planToSnapshots = (
+	descriptors: FrameDescriptor[],
+	trajectory: Array<Record<string, number[]>>
+): Snapshot[] =>
+	trajectory.map((stepInputs) => {
+		const transforms = descriptors.map((d) => descriptorToTransform(d, stepInputs))
+		return new Snapshot({
+			transforms,
+			uuid: Uint8Array.from(UuidTool.toBytes(crypto.randomUUID())),
+		})
+	})
+
+export const parsedPlanToSnapshots = (plan: ParsedPlan): Snapshot[] => {
+	const descriptors = buildFrameDescriptors(plan)
+	return planToSnapshots(descriptors, plan.trajectory)
+}

src/lib/plugins/MotionPlanReplayer/relations.ts

@@ -0,0 +1,10 @@
+import { relation } from 'koota'
+
+/**
+ * Members are the source, the plan root is the target. Destroying the plan
+ * root cascades to destroy every member — the opposite direction from
+ * Selection's `PointsCapturedBy` (autoDestroy: 'target'), since here we want
+ * "destroy the group destroys its members", not "destroy a member destroys
+ * the group".
+ */
+export const PartOfPlan = relation({ autoDestroy: 'source' })