Allowing regularization of marker offsets in default MarkerFitter los…

…s function

dart/biomechanics/MarkerFitter.cpp

@@ -153,26 +153,36 @@ MarkerFitterState::MarkerFitterState(
 
     Eigen::VectorXs jointPoses = skeleton->getJointWorldPositions(joints);
     jointErrorsAtTimesteps.col(i) = jointPoses - jointCenters.col(i);
-    for (int j = 0; j < joints.size(); j++)
+    if (jointWeights.size() > 0)
     {
-      jointErrorsAtTimesteps.col(i).segment<3>(j * 3) *= jointWeights(j);
+      for (int j = 0; j < joints.size(); j++)
+      {
+        jointErrorsAtTimesteps.col(i).segment<3>(j * 3) *= jointWeights(j);
+      }
     }
 
     // Compute the axis errors at each timestep
 
-    for (int j = 0; j < joints.size(); j++)
+    if (jointAxis.size() > 0)
     {
-      Eigen::Vector3s jointPos = jointPoses.segment<3>(j * 3);
-      Eigen::Vector3s axisCenter = jointAxis.block<3, 1>(j * 6, i);
-      Eigen::Vector3s axisDir
-          = jointAxis.block<3, 1>(j * 6 + 3, i).normalized();
-
-      Eigen::Vector3s diff = jointPos - axisCenter;
-      // Subtract out the component of `diff` that's parallel to the axisDir
-      diff -= diff.dot(axisDir) * axisDir;
-      // Now our measured diff is only the distance perpendicular to axisDir (ie
-      // the shortest path to the axis)
-      axisErrorsAtTimesteps.block<3, 1>(j * 3, i) = diff * axisWeights(j);
+      for (int j = 0; j < joints.size(); j++)
+      {
+        Eigen::Vector3s jointPos = jointPoses.segment<3>(j * 3);
+        Eigen::Vector3s axisCenter = jointAxis.block<3, 1>(j * 6, i);
+        Eigen::Vector3s axisDir
+            = jointAxis.block<3, 1>(j * 6 + 3, i).normalized();
+
+        Eigen::Vector3s diff = jointPos - axisCenter;
+        // Subtract out the component of `diff` that's parallel to the axisDir
+        diff -= diff.dot(axisDir) * axisDir;
+        // Now our measured diff is only the distance perpendicular to axisDir
+        // (ie the shortest path to the axis)
+        axisErrorsAtTimesteps.block<3, 1>(j * 3, i) = diff;
+        if (axisWeights.size() > 0)
+        {
+          axisErrorsAtTimesteps.block<3, 1>(j * 3, i) *= axisWeights(j);
+        }
+      }
     }
   }
 
@@ -462,8 +472,10 @@ MarkerFitter::MarkerFitter(
     mAnthropometrics(nullptr),
     mAnthropometricWeight(0.001),
     mMinVarianceCutoff(3.0),
-    mMinSphereFitScore(3e-5),
-    mMinAxisFitScore(6e-5),
+    mMinSphereFitScore(6e-5),
+    mMinAxisFitScore(1.2e-4),
+    mRegularizeTrackingMarkerOffsets(0.05),
+    mRegularizeAnatomicalMarkerOffsets(1.0),
     mDebugJointVariability(false)
 {
   mSkeletonBallJoints = mSkeleton->convertSkeletonToBallJoints();
@@ -497,6 +509,7 @@ MarkerFitter::MarkerFitter(
 
   // Default to a least-squares loss over just the marker errors
   mLossAndGrad = [this](MarkerFitterState* state) {
+    int numTimesteps = state->posesAtTimesteps.cols();
     // 1. Compute loss as a simple squared norm of marker and joint errors
     s_t loss = state->markerErrorsAtTimesteps.squaredNorm()
                + state->jointErrorsAtTimesteps.squaredNorm()
@@ -505,6 +518,20 @@ MarkerFitter::MarkerFitter(
     state->markerErrorsAtTimestepsGrad = 2 * state->markerErrorsAtTimesteps;
     state->jointErrorsAtTimestepsGrad = 2 * state->jointErrorsAtTimesteps;
     state->axisErrorsAtTimestepsGrad = 2 * state->axisErrorsAtTimesteps;
+
+    // Regularize tracking vs anatomical differently
+    state->markerOffsetsGrad = 2 * numTimesteps * state->markerOffsets;
+    for (int i = 0; i < this->mMarkerIsTracking.size(); i++)
+    {
+      s_t multiple
+          = (this->mMarkerIsTracking[i]
+                 ? this->mRegularizeTrackingMarkerOffsets
+                 : this->mRegularizeAnatomicalMarkerOffsets);
+      loss += numTimesteps * state->markerOffsets.col(i).squaredNorm()
+              * multiple;
+      state->markerOffsetsGrad.col(i) *= multiple;
+    }
+
     // 3. If we've got an anthropometrics prior, use it
     if (this->mAnthropometrics)
     {
@@ -2455,6 +2482,22 @@ void MarkerFitter::setMinAxisFitScore(s_t score)
   mMinAxisFitScore = score;
 }
 
+//==============================================================================
+/// This sets the value weight used to regularize tracking marker offsets from
+/// where the model thinks they should be
+void MarkerFitter::setRegularizeTrackingMarkerOffsets(s_t weight)
+{
+  mRegularizeTrackingMarkerOffsets = weight;
+}
+
+//==============================================================================
+/// This sets the value weight used to regularize anatomical marker offsets from
+/// where the model thinks they should be
+void MarkerFitter::setRegularizeAnatomicalMarkerOffsets(s_t weight)
+{
+  mRegularizeAnatomicalMarkerOffsets = weight;
+}
+
 //==============================================================================
 /// If set to true, we print the pair observation counts and data for
 /// computing joint variability.

dart/biomechanics/MarkerFitter.hpp

@@ -434,6 +434,14 @@ class MarkerFitter
   /// This sets the value used to compute axis fit weights
   void setMinAxisFitScore(s_t score);
 
+  /// This sets the value weight used to regularize tracking marker offsets from
+  /// where the model thinks they should be
+  void setRegularizeTrackingMarkerOffsets(s_t weight);
+
+  /// This sets the value weight used to regularize anatomical marker offsets
+  /// from where the model thinks they should be
+  void setRegularizeAnatomicalMarkerOffsets(s_t weight);
+
   /// If set to true, we print the pair observation counts and data for
   /// computing joint variability.
   void setDebugJointVariability(bool debug);
@@ -779,6 +787,8 @@ class MarkerFitter
   s_t mMinSphereFitScore;
   s_t mMinAxisFitScore;
   bool mDebugJointVariability;
+  s_t mRegularizeTrackingMarkerOffsets;
+  s_t mRegularizeAnatomicalMarkerOffsets;
 
   // These are IPOPT settings
   double mTolerance;

python/_nimblephysics/biomechanics/MarkerFitter.cpp

@@ -260,6 +260,15 @@ void MarkerFitter(py::module& m)
           "setMinAxisFitScore",
           &dart::biomechanics::MarkerFitter::setMinAxisFitScore,
           ::py::arg("score"))
+      .def(
+          "setRegularizeAnatomicalMarkerOffsets",
+          &dart::biomechanics::MarkerFitter::
+              setRegularizeAnatomicalMarkerOffsets,
+          ::py::arg("weight"))
+      .def(
+          "setRegularizeTrackingMarkerOffsets",
+          &dart::biomechanics::MarkerFitter::setRegularizeTrackingMarkerOffsets,
+          ::py::arg("weight"))
       .def(
           "setDebugJointVariability",
           &dart::biomechanics::MarkerFitter::setDebugJointVariability,

unittests/unit/test_MarkerFitterKinematics.cpp

@@ -2702,7 +2702,12 @@ TEST(MarkerFitter, FULL_KINEMATIC_STACK_LAI_ARNOLD_2)
   finalKinematicsReport.printReport(5);
 
   // Target markers
-  fitter.debugTrajectoryAndMarkersToGUI(finalKinematicInit, subsetTimesteps);
+  std::shared_ptr<server::GUIWebsocketServer> server
+      = std::make_shared<server::GUIWebsocketServer>();
+  server->serve(8070);
+  fitter.debugTrajectoryAndMarkersToGUI(
+      server, finalKinematicInit, subsetTimesteps);
+  server->blockWhileServing();
 }
 #endif
 
@@ -2867,7 +2872,12 @@ TEST(MarkerFitter, FULL_KINEMATIC_STACK_SPRINTER)
   finalKinematicsReport.printReport(5);
 
   // Target markers
-  fitter.debugTrajectoryAndMarkersToGUI(finalKinematicInit, subsetTimesteps);
+  std::shared_ptr<server::GUIWebsocketServer> server
+      = std::make_shared<server::GUIWebsocketServer>();
+  server->serve(8070);
+  fitter.debugTrajectoryAndMarkersToGUI(
+      server, finalKinematicInit, subsetTimesteps);
+  server->blockWhileServing();
 }
 #endif
 
@@ -3000,7 +3010,7 @@ TEST(MarkerFitter, FULL_KINEMATIC_RAJAGOPAL)
   // Bilevel optimization
   fitter.setIterationLimit(100);
   std::shared_ptr<BilevelFitResult> bilevelFit
-      = fitter.optimizeBilevel(subsetTimesteps, reinit, 150);
+      = fitter.optimizeBilevel(subsetTimesteps, reinit, 50);
 
   // Fine-tune IK and re-fit all the points
   MarkerInitialization finalKinematicInit = fitter.completeBilevelResult(
@@ -3036,7 +3046,12 @@ TEST(MarkerFitter, FULL_KINEMATIC_RAJAGOPAL)
   finalKinematicsReport.printReport(5);
 
   // Target markers
-  fitter.debugTrajectoryAndMarkersToGUI(finalKinematicInit, subsetTimesteps);
+  std::shared_ptr<server::GUIWebsocketServer> server
+      = std::make_shared<server::GUIWebsocketServer>();
+  server->serve(8070);
+  fitter.debugTrajectoryAndMarkersToGUI(
+      server, finalKinematicInit, subsetTimesteps);
+  server->blockWhileServing();
 }
 #endif