From 9bad5950000865f7fc7cfee51089c10b6d6a0f7b Mon Sep 17 00:00:00 2001
From: Collin Kees <Collin@kees.net>
Date: Tue, 15 Oct 2024 16:02:24 -0700
Subject: [PATCH] Add target and end effector offset to improve solver
 stability

---
 .../tracking/processor/skeleton/IKChain.kt    | 21 ++++++++++---------
 .../tracking/processor/skeleton/IKSolver.kt   | 10 ++++-----
 2 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/server/core/src/main/java/dev/slimevr/tracking/processor/skeleton/IKChain.kt b/server/core/src/main/java/dev/slimevr/tracking/processor/skeleton/IKChain.kt
index bd2c96e9b9..a564347647 100644
--- a/server/core/src/main/java/dev/slimevr/tracking/processor/skeleton/IKChain.kt
+++ b/server/core/src/main/java/dev/slimevr/tracking/processor/skeleton/IKChain.kt
@@ -5,7 +5,7 @@ import dev.slimevr.tracking.processor.Constraint.Companion.ConstraintType
 import dev.slimevr.tracking.trackers.Tracker
 import io.github.axisangles.ktmath.Quaternion
 import io.github.axisangles.ktmath.Vector3
-import kotlin.math.*
+import kotlin.math.pow
 
 /*
  * This class implements a chain of Bones
@@ -40,7 +40,7 @@ class IKChain(
 	 */
 	private fun prepBones() {
 		for (i in 0..<bones.size) {
-			if (bones[i].rotationConstraint.constraintType != ConstraintType.COMPLETE && bones[i].boneType.bodyPart !in IKSolver.LOCK_ROTATION) {
+			if (bones[i].rotationConstraint.constraintType != ConstraintType.COMPLETE) {
 				bones[i].setRotationRaw(rotations[i])
 			}
 		}
@@ -48,21 +48,18 @@ class IKChain(
 
 	fun backwardsCCDIK(useConstraints: Boolean) {
 		target = computedTailPosition?.getPosition() ?: getChildTargetAvg()
+		var offset = Vector3.NULL
 
 		for (i in bones.size - 1 downTo 0) {
 			val currentBone = bones[i]
 
-			if (currentBone.boneType.bodyPart in IKSolver.LOCK_ROTATION) {
-				rotations[i] = currentBone.getGlobalRotation()
-				continue
-			}
-
 			// Get the local position of the end effector and the target relative to the current node
-			val endEffectorLocal = (getEndEffectorsAvg() - currentBone.getPosition()).unit()
-			val targetLocal = (target - currentBone.getPosition()).unit()
+			val endEffectorLocal = ((getEndEffectorsAvg() - offset) - currentBone.getPosition()).unit()
+			val targetLocal = ((target - offset) - currentBone.getPosition()).unit()
 
 			// Compute the axis of rotation and angle for this bone
-			val scalar = IKSolver.DAMPENING_FACTOR * if (currentBone.rotationConstraint.allowModifications) 1f else IKSolver.STATIC_DAMPENING
+			var scalar = IKSolver.DAMPENING_FACTOR * if (currentBone.rotationConstraint.hasTrackerRotation) IKSolver.STATIC_DAMPENING else 1f
+			scalar *= ((bones.size - i).toFloat() / bones.size).pow(IKSolver.ANNEALING_EXPONENT)
 			var adjustment = Quaternion.fromTo(endEffectorLocal, targetLocal).pow(scalar).unit()
 
 			// Bones that are not supposed to be modified should tend towards their origin
@@ -74,6 +71,10 @@ class IKChain(
 			val correctedRot = (adjustment * rotation).unit()
 
 			rotations[i] = setBoneRotation(currentBone, correctedRot, useConstraints)
+
+			if (currentBone.rotationConstraint.hasTrackerRotation) {
+				offset += rotations[i].sandwich(Vector3.NEG_Y) * currentBone.length
+			}
 		}
 	}
 
diff --git a/server/core/src/main/java/dev/slimevr/tracking/processor/skeleton/IKSolver.kt b/server/core/src/main/java/dev/slimevr/tracking/processor/skeleton/IKSolver.kt
index cde6f695ed..e0d4f59389 100644
--- a/server/core/src/main/java/dev/slimevr/tracking/processor/skeleton/IKSolver.kt
+++ b/server/core/src/main/java/dev/slimevr/tracking/processor/skeleton/IKSolver.kt
@@ -7,19 +7,17 @@ import solarxr_protocol.datatypes.BodyPart
 /*
  * Implements CCDIK (Cyclic Coordinate Descent Inverse Kinematics) to allow
  * positional trackers such as vive/tundra trackers to be used in conjunction with
- * IMU trackers
+ * IMU only trackers
  */
 
 class IKSolver(private val root: Bone) {
 	companion object {
 		const val TOLERANCE_SQR = 1e-8 // == 0.01 cm
 		const val MAX_ITERATIONS = 20
+		const val ANNEALING_EXPONENT = 1
 		const val DAMPENING_FACTOR = 0.5f
-		const val STATIC_DAMPENING = 0.1f
-		const val CORRECTION_FACTOR = 0.1f
-
-		// Short limbs positioned on the end of the skeleton are prone to over rotation
-		val LOCK_ROTATION = setOf(BodyPart.LEFT_HAND, BodyPart.RIGHT_HAND, BodyPart.LEFT_FOOT, BodyPart.RIGHT_FOOT)
+		const val STATIC_DAMPENING = 0.5f
+		const val CORRECTION_FACTOR = 0.01f
 	}
 
 	var enabled = true