AutoBone iteration abstraction

server/core/src/main/java/dev/slimevr/autobone/AutoBoneStep.kt

@@ -1,62 +1,13 @@
 package dev.slimevr.autobone
 
-import dev.slimevr.config.AutoBoneConfig
-import dev.slimevr.config.ConfigManager
-import dev.slimevr.poseframeformat.PoseFrames
-import dev.slimevr.poseframeformat.player.TrackerFramesPlayer
-import dev.slimevr.tracking.processor.HumanPoseManager
-import java.util.function.Consumer
-
 class AutoBoneStep(
-	val config: AutoBoneConfig,
-	val targetHmdHeight: Float,
-	val frames: PoseFrames,
-	val epochCallback: Consumer<AutoBone.Epoch>?,
-	serverConfig: ConfigManager,
-	var curEpoch: Int = 0,
-	var curAdjustRate: Float = 0f,
-	var cursor1: Int = 0,
-	var cursor2: Int = 0,
-	var currentHmdHeight: Float = 0f,
+	var hmdHeight: Float = 1f,
+	val targetHmdHeight: Float = 1f,
+	var adjustRate: Float = 0f,
 ) {
-	var maxFrameCount = frames.maxFrameCount
-
-	val framePlayer1 = TrackerFramesPlayer(frames)
-	val framePlayer2 = TrackerFramesPlayer(frames)
-
-	val trackers1 = framePlayer1.trackers.toList()
-	val trackers2 = framePlayer2.trackers.toList()
-
-	val skeleton1 = HumanPoseManager(trackers1)
-	val skeleton2 = HumanPoseManager(trackers2)
 
 	val errorStats = StatsCalculator()
 
-	init {
-		// Load server configs into the skeleton
-		skeleton1.loadFromConfig(serverConfig)
-		skeleton2.loadFromConfig(serverConfig)
-		// Disable leg tweaks and IK solver, these will mess with the resulting positions
-		skeleton1.setLegTweaksEnabled(false)
-		skeleton2.setLegTweaksEnabled(false)
-	}
-
-	fun setCursors(cursor1: Int, cursor2: Int, updatePlayerCursors: Boolean) {
-		this.cursor1 = cursor1
-		this.cursor2 = cursor2
-
-		if (updatePlayerCursors) {
-			updatePlayerCursors()
-		}
-	}
-
-	fun updatePlayerCursors() {
-		framePlayer1.setCursors(cursor1)
-		framePlayer2.setCursors(cursor2)
-		skeleton1.update()
-		skeleton2.update()
-	}
-
 	val heightOffset: Float
-		get() = targetHmdHeight - currentHmdHeight
+		get() = targetHmdHeight - hmdHeight
 }

server/core/src/main/java/dev/slimevr/autobone/BoneContribution.kt

@@ -0,0 +1,84 @@
+package dev.slimevr.autobone
+
+import dev.slimevr.autobone.AutoBone.Companion.MIN_SLIDE_DIST
+import dev.slimevr.autobone.AutoBone.Companion.SYMM_CONFIGS
+import dev.slimevr.tracking.processor.BoneType
+import dev.slimevr.tracking.processor.HumanPoseManager
+import dev.slimevr.tracking.processor.config.SkeletonConfigOffsets
+import io.github.axisangles.ktmath.Vector3
+
+object BoneContribution {
+	/**
+	 * Computes the local tail position of the bone after rotation.
+	 */
+	fun getBoneLocalTail(
+		skeleton: HumanPoseManager,
+		boneType: BoneType,
+	): Vector3 {
+		val bone = skeleton.getBone(boneType)
+		return bone.getTailPosition() - bone.getPosition()
+	}
+
+	/**
+	 * Computes the direction of the bone tail's movement between skeletons 1 and 2.
+	 */
+	fun getBoneLocalTailDir(
+		skeleton1: HumanPoseManager,
+		skeleton2: HumanPoseManager,
+		boneType: BoneType,
+	): Vector3? {
+		val boneOff = getBoneLocalTail(skeleton2, boneType) - getBoneLocalTail(skeleton1, boneType)
+		val boneOffLen = boneOff.len()
+		// If the offset is approx 0, just return null so it can be easily ignored
+		return if (boneOffLen > MIN_SLIDE_DIST) boneOff / boneOffLen else null
+	}
+
+	/**
+	 * Predicts how much the provided config should be affecting the slide offsets
+	 * of the left and right ankles.
+	 */
+	fun getSlideDot(
+		skeleton1: HumanPoseManager,
+		skeleton2: HumanPoseManager,
+		config: SkeletonConfigOffsets,
+		slideL: Vector3?,
+		slideR: Vector3?,
+	): Float {
+		var slideDot = 0f
+		// Used for right offset if not a symmetric bone
+		var boneOffL: Vector3? = null
+
+		// Treat null as 0
+		if (slideL != null) {
+			boneOffL = getBoneLocalTailDir(skeleton1, skeleton2, config.affectedOffsets[0])
+
+			// Treat null as 0
+			if (boneOffL != null) {
+				slideDot += slideL.dot(boneOffL)
+			}
+		}
+
+		// Treat null as 0
+		if (slideR != null) {
+			// IMPORTANT: This assumption for acquiring BoneType only works if
+			// SkeletonConfigOffsets is set up to only affect one BoneType, make sure no
+			// changes to SkeletonConfigOffsets goes against this assumption, please!
+			val boneOffR = if (SYMM_CONFIGS.contains(config)) {
+				getBoneLocalTailDir(skeleton1, skeleton2, config.affectedOffsets[1])
+			} else if (slideL != null) {
+				// Use cached offset if slideL was used
+				boneOffL
+			} else {
+				// Compute offset if missing because of slideL
+				getBoneLocalTailDir(skeleton1, skeleton2, config.affectedOffsets[0])
+			}
+
+			// Treat null as 0
+			if (boneOffR != null) {
+				slideDot += slideR.dot(boneOffR)
+			}
+		}
+
+		return slideDot / 2f
+	}
+}

server/core/src/main/java/dev/slimevr/autobone/PoseFrameIterator.kt

@@ -0,0 +1,90 @@
+package dev.slimevr.autobone
+
+import kotlin.random.Random
+
+object PoseFrameIterator {
+	fun <T> iterateFrames(
+		step: PoseFrameStep<T>,
+	) {
+		check(step.frames.frameHolders.isNotEmpty()) { "Recording has no trackers." }
+		check(step.maxFrameCount > 0) { "Recording has no frames." }
+
+		// Epoch loop, each epoch is one full iteration over the full dataset
+		for (epoch in (if (step.config.calcInitError) -1 else 0) until step.config.numEpochs) {
+			// Set the current epoch to process
+			step.epoch = epoch
+			// Process the epoch
+			epoch(step)
+		}
+	}
+
+	private fun randomIndices(count: Int, random: Random): IntArray {
+		val randIndices = IntArray(count)
+
+		var zeroPos = -1
+		for (i in 0 until count) {
+			var index = random.nextInt(count)
+			if (i > 0) {
+				while (index == zeroPos || randIndices[index] > 0) {
+					index = random.nextInt(count)
+				}
+			} else {
+				zeroPos = index
+			}
+			randIndices[index] = i
+		}
+
+		return randIndices
+	}
+
+	private fun <T> epoch(step: PoseFrameStep<T>) {
+		val config = step.config
+		val frameCount = step.maxFrameCount
+
+		// Perform any setup that needs to be done before the current epoch
+		step.preEpoch?.accept(step)
+
+		val randIndices = if (config.randomizeFrameOrder) {
+			randomIndices(step.maxFrameCount, step.random)
+		} else {
+			null
+		}
+
+		// Iterate over the frames using a cursor and an offset for comparing
+		// frames a certain number of frames apart
+		var cursorOffset = config.minDataDistance
+		while (cursorOffset <= config.maxDataDistance &&
+			cursorOffset < frameCount
+		) {
+			var frameCursor = 0
+			while (frameCursor < frameCount - cursorOffset) {
+				val frameCursor2 = frameCursor + cursorOffset
+
+				// Then set the frame cursors and apply them to both skeletons
+				if (config.randomizeFrameOrder && randIndices != null) {
+					step
+						.setCursors(
+							randIndices[frameCursor],
+							randIndices[frameCursor2],
+							updatePlayerCursors = true,
+						)
+				} else {
+					step.setCursors(
+						frameCursor,
+						frameCursor2,
+						updatePlayerCursors = true,
+					)
+				}
+
+				// Process the iteration
+				step.onStep.accept(step)
+
+				// Move on to the next iteration
+				frameCursor += config.cursorIncrement
+			}
+			cursorOffset++
+		}
+
+		step.postEpoch?.accept(step)
+	}
+}

server/core/src/main/java/dev/slimevr/autobone/PoseFrameStep.kt

@@ -0,0 +1,70 @@
+package dev.slimevr.autobone
+
+import dev.slimevr.config.AutoBoneConfig
+import dev.slimevr.config.ConfigManager
+import dev.slimevr.poseframeformat.PoseFrames
+import dev.slimevr.poseframeformat.player.TrackerFramesPlayer
+import dev.slimevr.tracking.processor.HumanPoseManager
+import java.util.function.Consumer
+import kotlin.random.Random
+
+class PoseFrameStep<T>(
+	val config: AutoBoneConfig,
+	/** The config to initialize skeletons. */
+	serverConfig: ConfigManager? = null,
+	val frames: PoseFrames,
+	/** The consumer run before each epoch. */
+	val preEpoch: Consumer<PoseFrameStep<T>>? = null,
+	/** The consumer run for each step. */
+	val onStep: Consumer<PoseFrameStep<T>>,
+	/** The consumer run after each epoch. */
+	val postEpoch: Consumer<PoseFrameStep<T>>? = null,
+	/** The current epoch. */
+	var epoch: Int = 0,
+	/** The current frame cursor position in [frames] for skeleton1. */
+	var cursor1: Int = 0,
+	/** The current frame cursor position in [frames] for skeleton2. */
+	var cursor2: Int = 0,
+	randomSeed: Long = 0,
+	val data: T,
+) {
+	var maxFrameCount = frames.maxFrameCount
+
+	val framePlayer1 = TrackerFramesPlayer(frames)
+	val framePlayer2 = TrackerFramesPlayer(frames)
+
+	val trackers1 = framePlayer1.trackers.toList()
+	val trackers2 = framePlayer2.trackers.toList()
+
+	val skeleton1 = HumanPoseManager(trackers1)
+	val skeleton2 = HumanPoseManager(trackers2)
+
+	val random = Random(randomSeed)
+
+	init {
+		// Load server configs into the skeleton
+		if (serverConfig != null) {
+			skeleton1.loadFromConfig(serverConfig)
+			skeleton2.loadFromConfig(serverConfig)
+		}
+		// Disable leg tweaks and IK solver, these will mess with the resulting positions
+		skeleton1.setLegTweaksEnabled(false)
+		skeleton2.setLegTweaksEnabled(false)
+	}
+
+	fun setCursors(cursor1: Int, cursor2: Int, updatePlayerCursors: Boolean) {
+		this.cursor1 = cursor1
+		this.cursor2 = cursor2
+
+		if (updatePlayerCursors) {
+			updatePlayerCursors()
+		}
+	}
+
+	fun updatePlayerCursors() {
+		framePlayer1.setCursors(cursor1)
+		framePlayer2.setCursors(cursor2)
+		skeleton1.update()
+		skeleton2.update()
+	}
+}

server/core/src/main/java/dev/slimevr/autobone/errors/BodyProportionError.kt

@@ -1,6 +1,7 @@
 package dev.slimevr.autobone.errors
 
 import dev.slimevr.autobone.AutoBoneStep
+import dev.slimevr.autobone.PoseFrameStep
 import dev.slimevr.autobone.errors.proportions.ProportionLimiter
 import dev.slimevr.autobone.errors.proportions.RangeProportionLimiter
 import dev.slimevr.tracking.processor.HumanPoseManager
@@ -11,8 +12,8 @@ import kotlin.math.*
 // The distance from average human proportions
 class BodyProportionError : IAutoBoneError {
 	@Throws(AutoBoneException::class)
-	override fun getStepError(trainingStep: AutoBoneStep): Float = getBodyProportionError(
-		trainingStep.skeleton1,
+	override fun getStepError(step: PoseFrameStep<AutoBoneStep>): Float = getBodyProportionError(
+		step.skeleton1,
 		// Skeletons are now normalized to reduce bias, so height is always 1
 		1f,
 	)

server/core/src/main/java/dev/slimevr/autobone/errors/FootHeightOffsetError.kt

@@ -1,6 +1,7 @@
 package dev.slimevr.autobone.errors
 
 import dev.slimevr.autobone.AutoBoneStep
+import dev.slimevr.autobone.PoseFrameStep
 import dev.slimevr.tracking.processor.skeleton.HumanSkeleton
 import dev.slimevr.tracking.trackers.Tracker
 import dev.slimevr.tracking.trackers.TrackerRole
@@ -9,9 +10,9 @@ import kotlin.math.*
 // The offset between the height both feet at one instant and over time
 class FootHeightOffsetError : IAutoBoneError {
 	@Throws(AutoBoneException::class)
-	override fun getStepError(trainingStep: AutoBoneStep): Float = getSlideError(
-		trainingStep.skeleton1.skeleton,
-		trainingStep.skeleton2.skeleton,
+	override fun getStepError(step: PoseFrameStep<AutoBoneStep>): Float = getSlideError(
+		step.skeleton1.skeleton,
+		step.skeleton2.skeleton,
 	)
 
 	companion object {

server/core/src/main/java/dev/slimevr/autobone/errors/HeightError.kt

@@ -1,14 +1,15 @@
 package dev.slimevr.autobone.errors
 
 import dev.slimevr.autobone.AutoBoneStep
+import dev.slimevr.autobone.PoseFrameStep
 import kotlin.math.*
 
 // The difference from the current height to the target height
 class HeightError : IAutoBoneError {
 	@Throws(AutoBoneException::class)
-	override fun getStepError(trainingStep: AutoBoneStep): Float = getHeightError(
-		trainingStep.currentHmdHeight,
-		trainingStep.targetHmdHeight,
+	override fun getStepError(step: PoseFrameStep<AutoBoneStep>): Float = getHeightError(
+		step.data.hmdHeight,
+		step.data.targetHmdHeight,
 	)
 
 	fun getHeightError(currentHeight: Float, targetHeight: Float): Float = abs(targetHeight - currentHeight)

server/core/src/main/java/dev/slimevr/autobone/errors/IAutoBoneError.kt

@@ -1,8 +1,9 @@
 package dev.slimevr.autobone.errors
 
 import dev.slimevr.autobone.AutoBoneStep
+import dev.slimevr.autobone.PoseFrameStep
 
 interface IAutoBoneError {
 	@Throws(AutoBoneException::class)
-	fun getStepError(trainingStep: AutoBoneStep): Float
+	fun getStepError(step: PoseFrameStep<AutoBoneStep>): Float
 }