Move root with pelvis tracker if it's active

nicoco007 · nicoco007 · commit 96c1a379e271 · 2024-10-07T21:55:50.000-04:00
diff --git a/Source/CustomAvatar/Avatar/AvatarIK.cs b/Source/CustomAvatar/Avatar/AvatarIK.cs
@@ -59,6 +59,9 @@ public bool isLocomotionEnabled
         private Pose _defaultRootPose;
         private Pose _previousParentPose;
 
+        private bool _hasPelvisTarget;
+        private bool _hasBothLegTargets;
+
         #region Behaviour Lifecycle
 
         protected void Awake()
@@ -139,7 +142,6 @@ protected void Start()
 
             _input.inputChanged += OnInputChanged;
 
-            UpdateLocomotion();
             UpdateSolverTargets();
 
             foreach (BeatSaberDynamicBone::DynamicBone dynamicBone in _dynamicBones)
@@ -179,6 +181,8 @@ protected void OnDestroy()
 
         private void OnPreUpdate()
         {
+            ApplyRootPose();
+
             foreach (BeatSaberDynamicBone::DynamicBone dynamicBone in _dynamicBones)
             {
                 dynamicBone.Update();
@@ -195,6 +199,39 @@ private void OnPostUpdate()
             }
         }
 
+        // adapted from VRIKRootController to work when IK is rotated by parent
+        private void ApplyRootPose()
+        {
+            // don't move the root if locomotion is disabled and both feet aren't being tracked
+            // (i.e. keep previous behaviour of sticking to the origin when locomotion is disabled and we're missing one or more FBT trackers)
+            if (!_isLocomotionEnabled && !_hasBothLegTargets)
+            {
+                return;
+            }
+
+            if (!_hasPelvisTarget)
+            {
+                return;
+            }
+
+            Transform pelvisTarget = _vrik.solver.spine.pelvisTarget;
+            Transform root = _vrik.references.root;
+            Transform parent = root.parent;
+            bool hasParent = parent != null;
+
+            Vector3 up = hasParent ? parent.up : Vector3.up;
+            root.rotation = Quaternion.LookRotation(Vector3.ProjectOnPlane(pelvisTarget.rotation * _avatar.prefab.pelvisRootForward, up), up);
+
+            var position = Vector3.ProjectOnPlane(pelvisTarget.position - root.TransformVector(_avatar.prefab.pelvisRootOffset), up);
+
+            if (hasParent)
+            {
+                position = parent.InverseTransformPoint(position);
+            }
+
+            root.localPosition = new Vector3(position.x, root.localPosition.y, position.z);
+        }
+
         private void ApplyPlatformMotion()
         {
             Transform parent = _vrik.references.root.parent;
@@ -221,10 +258,16 @@ private void UpdateLocomotion()
                 return;
             }
 
-            _vrik.solver.locomotion.weight = _isLocomotionEnabled ? vrikManager.solver_locomotion_weight : 0;
+            // don't enable locomotion if FBT is applied
+            bool shouldEnableLocomotion = _isLocomotionEnabled && !(_hasPelvisTarget && _hasBothLegTargets);
 
-            if (!_isLocomotionEnabled && _vrik.references.root != null)
+            if (shouldEnableLocomotion)
+            {
+                _vrik.solver.locomotion.weight = vrikManager.solver_locomotion_weight;
+            }
+            else
             {
+                _vrik.solver.locomotion.weight = 0;
                 _vrik.references.root.SetLocalPose(_defaultRootPose);
             }
         }
@@ -236,24 +279,31 @@ private void OnInputChanged()
 
         private void UpdateSolverTargets()
         {
-            if (_vrik == null || vrikManager == null) return;
+            if (_vrik == null || vrikManager == null)
+            {
+                return;
+            }
 
             _logger.LogTrace("Updating solver targets");
 
             UpdateSolverTarget(DeviceUse.Head, ref _vrik.solver.spine.headTarget, ref _vrik.solver.spine.positionWeight, ref _vrik.solver.spine.rotationWeight);
             UpdateSolverTarget(DeviceUse.LeftHand, ref _vrik.solver.leftArm.target, ref _vrik.solver.leftArm.positionWeight, ref _vrik.solver.leftArm.rotationWeight);
             UpdateSolverTarget(DeviceUse.RightHand, ref _vrik.solver.rightArm.target, ref _vrik.solver.rightArm.positionWeight, ref _vrik.solver.rightArm.rotationWeight);
-            UpdateSolverTarget(DeviceUse.LeftFoot, ref _vrik.solver.leftLeg.target, ref _vrik.solver.leftLeg.positionWeight, ref _vrik.solver.leftLeg.rotationWeight);
-            UpdateSolverTarget(DeviceUse.RightFoot, ref _vrik.solver.rightLeg.target, ref _vrik.solver.rightLeg.positionWeight, ref _vrik.solver.rightLeg.rotationWeight);
 
-            if (UpdateSolverTarget(DeviceUse.Waist, ref _vrik.solver.spine.pelvisTarget, ref _vrik.solver.spine.pelvisPositionWeight, ref _vrik.solver.spine.pelvisRotationWeight))
+            _hasBothLegTargets = true;
+            _hasBothLegTargets &= UpdateSolverTarget(DeviceUse.LeftFoot, ref _vrik.solver.leftLeg.target, ref _vrik.solver.leftLeg.positionWeight, ref _vrik.solver.leftLeg.rotationWeight);
+            _hasBothLegTargets &= UpdateSolverTarget(DeviceUse.RightFoot, ref _vrik.solver.rightLeg.target, ref _vrik.solver.rightLeg.positionWeight, ref _vrik.solver.rightLeg.rotationWeight);
+
+            if (_hasPelvisTarget = UpdateSolverTarget(DeviceUse.Waist, ref _vrik.solver.spine.pelvisTarget, ref _vrik.solver.spine.pelvisPositionWeight, ref _vrik.solver.spine.pelvisRotationWeight))
             {
                 _vrik.solver.plantFeet = false;
             }
             else
             {
                 _vrik.solver.plantFeet = vrikManager.solver_plantFeet;
             }
+
+            UpdateLocomotion();
         }
 
         private bool UpdateSolverTarget(DeviceUse deviceUse, ref Transform target, ref float positionWeight, ref float rotationWeight)
diff --git a/Source/CustomAvatar/Avatar/AvatarPrefab.cs b/Source/CustomAvatar/Avatar/AvatarPrefab.cs
@@ -82,6 +82,10 @@ public class AvatarPrefab : MonoBehaviour
 
         internal Pose rightFootCalibrationOffset { get; private set; }
 
+        internal Vector3 pelvisRootForward { get; private set; }
+
+        internal Vector3 pelvisRootOffset { get; private set; }
+
         private ILogger<AvatarPrefab> _logger;
 
         [Inject]
@@ -195,6 +199,9 @@ internal void Construct(ILoggerFactory loggerFactory, DiContainer container)
 
                     Destroy(targetObj);
                 }
+
+                pelvisRootForward = Quaternion.Inverse(vrikManager.references_pelvis.rotation) * vrikManager.references_root.forward;
+                pelvisRootOffset = vrikManager.references_root.InverseTransformPoint(vrikManager.references_pelvis.position);
             }
         }
 
diff --git a/Source/CustomAvatar/Patches/IKSolverVR.Arm.cs b/Source/CustomAvatar/Patches/IKSolverVR.Arm.cs
@@ -0,0 +1,87 @@
+﻿//  Beat Saber Custom Avatars - Custom player models for body presence in Beat Saber.
+//  Copyright © 2018-2024  Nicolas Gnyra and Beat Saber Custom Avatars Contributors
+//
+//  This library is free software: you can redistribute it and/or
+//  modify it under the terms of the GNU Lesser General Public
+//  License as published by the Free Software Foundation, either
+//  version 3 of the License, or (at your option) any later version.
+//
+//  This program is distributed in the hope that it will be useful,
+//  but WITHOUT ANY WARRANTY; without even the implied warranty of
+//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+//  GNU Lesser General Public License for more details.
+//
+//  You should have received a copy of the GNU Lesser General Public License
+//  along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+extern alias BeatSaberFinalIK;
+
+using System;
+using System.Collections.Generic;
+using System.Reflection;
+using System.Reflection.Emit;
+using BeatSaberFinalIK::RootMotion.FinalIK;
+using CustomAvatar.Scripts;
+using CustomAvatar.Utilities;
+using HarmonyLib;
+
+namespace CustomAvatar.Patches
+{
+    /// <summary>
+    /// This patch enables the use of <see cref="IKSolverVR_Arm.shoulderPitchOffset"/>.
+    /// </summary>
+
+    [HarmonyPatch(typeof(IKSolverVR.Arm), nameof(IKSolverVR.Arm.Solve))]
+    internal static class IKSolverVR_Arm_PitchAngleOffset
+    {
+        private static readonly FieldInfo kPitchOffsetAngleField = AccessTools.DeclaredField(typeof(IKSolverVR_Arm), nameof(IKSolverVR_Arm.shoulderPitchOffset));
+
+        public static IEnumerable<CodeInstruction> Transpiler(IEnumerable<CodeInstruction> instructions)
+        {
+            return new CodeMatcher(instructions)
+                /* Quaternion.AngleAxis(isLeft ? pitchOffsetAngle : -pitchOffsetAngle, chestForward) */
+                .MatchForward(false,
+                    new CodeMatch(i => i.Equals(OpCodes.Ldc_R4, 30f)),
+                    new CodeMatch(i => i.Branches(out Label? _)),
+                    new CodeMatch(i => i.Equals(OpCodes.Ldc_R4, -30f)))
+                .SetAndAdvance(OpCodes.Ldarg_0, null)
+                .InsertAndAdvance(
+                    new CodeInstruction(OpCodes.Ldfld, kPitchOffsetAngleField),
+                    new CodeInstruction(OpCodes.Neg))
+                .Advance(1)
+                .SetAndAdvance(OpCodes.Ldarg_0, null)
+                .InsertAndAdvance(
+                    new CodeInstruction(OpCodes.Ldfld, kPitchOffsetAngleField))
+
+                /* pitch -= pitchOffsetAngle */
+                .MatchForward(false,
+                    new CodeMatch(i => i.LoadsLocal(11)),
+                    new CodeMatch(i => i.opcode == OpCodes.Ldc_R4 && (float)i.operand == -30f),
+                    new CodeMatch(OpCodes.Sub),
+                    new CodeMatch(i => i.SetsLocal(11)))
+                .ThrowIfInvalid("`pitch -= pitchOffsetAngle` not found")
+                .Advance(1)
+                .SetAndAdvance(OpCodes.Ldarg_0, null)
+                .InsertAndAdvance(
+                    new CodeInstruction(OpCodes.Ldfld, kPitchOffsetAngleField))
+
+                /* DamperValue(pitch, -45f - pitchOffsetAngle, 45f - pitchOffsetAngle) */
+                .MatchForward(false,
+                    new CodeMatch(i => i.LoadsLocal(11)),
+                    new CodeMatch(i => i.Equals(OpCodes.Ldc_R4, -15f)),
+                    new CodeMatch(i => i.Equals(OpCodes.Ldc_R4, 75f)))
+                .Advance(1)
+                .SetOperandAndAdvance(-45f)
+                .InsertAndAdvance(
+                    new CodeInstruction(OpCodes.Ldarg_0, null),
+                    new CodeInstruction(OpCodes.Ldfld, kPitchOffsetAngleField),
+                    new CodeInstruction(OpCodes.Sub))
+                .SetOperandAndAdvance(45f)
+                .InsertAndAdvance(
+                    new CodeInstruction(OpCodes.Ldarg_0, null),
+                    new CodeInstruction(OpCodes.Ldfld, kPitchOffsetAngleField),
+                    new CodeInstruction(OpCodes.Sub))
+                .InstructionEnumeration();
+        }
+    }
+}
diff --git a/Source/CustomAvatar/Patches/IKSolverVR.Spine.cs b/Source/CustomAvatar/Patches/IKSolverVR.Spine.cs
diff --git a/Source/CustomAvatar/Patches/IKSolverVR.cs b/Source/CustomAvatar/Patches/IKSolverVR.cs
@@ -16,8 +16,8 @@
 
 extern alias BeatSaberFinalIK;
 
-using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Reflection;
 using System.Reflection.Emit;
 using BeatSaberFinalIK::RootMotion.FinalIK;
@@ -27,66 +27,65 @@
 
 namespace CustomAvatar.Patches
 {
+    /// <summary>
+    /// This patch makes the constructor of <see cref="IKSolverVR"/> instantiate our <see cref="IKSolverVR_Arm"/> in the <see cref="IKSolverVR.leftArm"/> and <see cref="IKSolverVR.rightArm"/> fields.
+    /// </summary>
     [HarmonyPatch(typeof(IKSolverVR), MethodType.Constructor)]
     internal static class IKSolverVR_Constructor
     {
-        public static void Postfix(IKSolverVR __instance)
+        private static readonly ConstructorInfo kIKSolverVRArmConstructor = typeof(IKSolverVR.Arm).GetConstructors().Single();
+        private static readonly ConstructorInfo kNewIKSolverVRArmConstructor = typeof(IKSolverVR_Arm).GetConstructors().Single();
+
+        public static IEnumerable<CodeInstruction> Transpiler(IEnumerable<CodeInstruction> instructions)
         {
-            __instance.leftArm = new IKSolverVR_Arm();
-            __instance.rightArm = new IKSolverVR_Arm();
+            foreach (CodeInstruction instruction in instructions)
+            {
+                if (instruction.opcode == OpCodes.Newobj && (ConstructorInfo)instruction.operand == kIKSolverVRArmConstructor)
+                {
+                    instruction.operand = kNewIKSolverVRArmConstructor;
+                }
+
+                yield return instruction;
+            }
         }
     }
 
-    [HarmonyPatch(typeof(IKSolverVR.Arm), nameof(IKSolverVR.Arm.Solve))]
-    internal static class IKSolverVR_Arm_PitchAngleOffset
+    /// <summary>
+    /// This patch prevents locomotion from fighting against the position we set in <see cref="Avatar.AvatarIK"/> when the pelvis target exists.
+    /// </summary>
+    [HarmonyPatch(typeof(IKSolverVR), nameof(IKSolverVR.Solve))]
+    internal static class IKSolverVR_Solve
     {
-        private static readonly FieldInfo kPitchOffsetAngleField = AccessTools.DeclaredField(typeof(IKSolverVR_Arm), nameof(IKSolverVR_Arm.shoulderPitchOffset));
+        private static readonly MethodInfo kRootBonePropertyGetter = AccessTools.DeclaredPropertyGetter(typeof(IKSolverVR), nameof(IKSolverVR.rootBone));
+        private static readonly FieldInfo kVirtualBoneSolverPositionField = AccessTools.DeclaredField(typeof(IKSolverVR.VirtualBone), nameof(IKSolverVR.VirtualBone.solverPosition));
+        private static readonly FieldInfo kSpineField = AccessTools.DeclaredField(typeof(IKSolverVR), nameof(IKSolverVR.spine));
+        private static readonly FieldInfo kSpinePelvisTargetField = AccessTools.DeclaredField(typeof(IKSolverVR.Spine), nameof(IKSolverVR.Spine.pelvisTarget));
+        private static readonly MethodInfo kUnityObjectEqualsMethod = AccessTools.DeclaredMethod(typeof(UnityEngine.Object), "op_Equality");
 
-        public static IEnumerable<CodeInstruction> Transpiler(IEnumerable<CodeInstruction> instructions)
+        public static IEnumerable<CodeInstruction> Transpiler(IEnumerable<CodeInstruction> instructions, ILGenerator ilGenerator)
         {
-            return new CodeMatcher(instructions)
-                /* Quaternion.AngleAxis(isLeft ? pitchOffsetAngle : -pitchOffsetAngle, chestForward) */
-                .MatchForward(false,
-                    new CodeMatch(i => i.Equals(OpCodes.Ldc_R4, 30f)),
-                    new CodeMatch(i => i.Branches(out Label? _)),
-                    new CodeMatch(i => i.Equals(OpCodes.Ldc_R4, -30f)))
-                .SetAndAdvance(OpCodes.Ldarg_0, null)
-                .InsertAndAdvance(
-                    new CodeInstruction(OpCodes.Ldfld, kPitchOffsetAngleField),
-                    new CodeInstruction(OpCodes.Neg))
-                .Advance(1)
-                .SetAndAdvance(OpCodes.Ldarg_0, null)
-                .InsertAndAdvance(
-                    new CodeInstruction(OpCodes.Ldfld, kPitchOffsetAngleField))
-
-                /* pitch -= pitchOffsetAngle */
-                .MatchForward(false,
-                    new CodeMatch(i => i.LoadsLocal(11)),
-                    new CodeMatch(i => i.opcode == OpCodes.Ldc_R4 && (float)i.operand == -30f),
-                    new CodeMatch(OpCodes.Sub),
-                    new CodeMatch(i => i.SetsLocal(11)))
-                .ThrowIfInvalid("`pitch -= pitchOffsetAngle` not found")
-                .Advance(1)
-                .SetAndAdvance(OpCodes.Ldarg_0, null)
-                .InsertAndAdvance(
-                    new CodeInstruction(OpCodes.Ldfld, kPitchOffsetAngleField))
-
-                /* DamperValue(pitch, -45f - pitchOffsetAngle, 45f - pitchOffsetAngle) */
-                .MatchForward(false,
-                    new CodeMatch(i => i.LoadsLocal(11)),
-                    new CodeMatch(i => i.Equals(OpCodes.Ldc_R4, -15f)),
-                    new CodeMatch(i => i.Equals(OpCodes.Ldc_R4, 75f)))
-                .Advance(1)
-                .SetOperandAndAdvance(-45f)
-                .InsertAndAdvance(
-                    new CodeInstruction(OpCodes.Ldarg_0, null),
-                    new CodeInstruction(OpCodes.Ldfld, kPitchOffsetAngleField),
-                    new CodeInstruction(OpCodes.Sub))
-                .SetOperandAndAdvance(45f)
+            return new CodeMatcher(instructions, ilGenerator)
+                .MatchForward(
+                    false,
+                    new CodeMatch(OpCodes.Ldarg_0),
+                    new CodeMatch(i => i.Calls(kRootBonePropertyGetter)),
+                    new CodeMatch(i => i.LoadsLocal(14)),
+                    new CodeMatch(i => i.StoresField(kVirtualBoneSolverPositionField)))
+                .CreateLabelWithOffsets(4, out Label label)
                 .InsertAndAdvance(
-                    new CodeInstruction(OpCodes.Ldarg_0, null),
-                    new CodeInstruction(OpCodes.Ldfld, kPitchOffsetAngleField),
-                    new CodeInstruction(OpCodes.Sub))
+                    new CodeInstruction(OpCodes.Ldarg_0),
+                    new CodeInstruction(OpCodes.Ldfld, kSpineField),
+                    new CodeInstruction(OpCodes.Ldfld, kSpinePelvisTargetField),
+                    new CodeInstruction(OpCodes.Ldnull),
+                    new CodeInstruction(OpCodes.Call, kUnityObjectEqualsMethod),
+                    new CodeInstruction(OpCodes.Brfalse_S, label))
+                .MatchForward(
+                    false,
+                    new CodeMatch(OpCodes.Ldarg_0),
+                    new CodeMatch(i => i.Calls(kRootBonePropertyGetter)),
+                    new CodeMatch(i => i.LoadsField(kVirtualBoneSolverPositionField)))
+                .RemoveInstructions(3)
+                .InsertAndAdvance(new CodeInstruction(OpCodes.Ldloc_S, 14))
                 .InstructionEnumeration();
         }
     }

Original file line number	Diff line number	Diff line change
`@@ -82,6 +82,10 @@ public class AvatarPrefab : MonoBehaviour`
`82`	`82`
`83`	`83`	`internal Pose rightFootCalibrationOffset { get; private set; }`
`84`	`84`
	`85`	`+ internal Vector3 pelvisRootForward { get; private set; }`
	`86`	`+`
	`87`	`+ internal Vector3 pelvisRootOffset { get; private set; }`
	`88`	`+`
`85`	`89`	`private ILogger<AvatarPrefab> _logger;`
`86`	`90`
`87`	`91`	`[Inject]`
`@@ -195,6 +199,9 @@ internal void Construct(ILoggerFactory loggerFactory, DiContainer container)`
`195`	`199`
`196`	`200`	`Destroy(targetObj);`
`197`	`201`	`}`
	`202`	`+`
	`203`	`+ pelvisRootForward = Quaternion.Inverse(vrikManager.references_pelvis.rotation) * vrikManager.references_root.forward;`
	`204`	`+ pelvisRootOffset = vrikManager.references_root.InverseTransformPoint(vrikManager.references_pelvis.position);`
`198`	`205`	`}`
`199`	`206`	`}`
`200`	`207`