add wrapper for joint-to-joint error function (#806)

Summary:
Pull Request resolved: #806

See title

Reviewed By: jeongseok-meta

Differential Revision: D86779635

fbshipit-source-id: dd93fc2064a229064db9937f9d02d82547511def

Author: cstollmeta

pymomentum/solver2/solver2_error_functions.cpp

@@ -12,6 +12,7 @@
 #include <momentum/character_solver/collision_error_function.h>
 #include <momentum/character_solver/distance_error_function.h>
 #include <momentum/character_solver/fixed_axis_error_function.h>
+#include <momentum/character_solver/joint_to_joint_distance_error_function.h>
 #include <momentum/character_solver/limit_error_function.h>
 #include <momentum/character_solver/model_parameters_error_function.h>
 #include <momentum/character_solver/normal_error_function.h>
@@ -1267,6 +1268,180 @@ or maintaining the width of body parts.)")
           "Returns the number of constraints.");
 }
 
+void defJointToJointDistanceErrorFunction(py::module_& m) {
+  py::class_<mm::JointToJointDistanceConstraintT<float>>(m, "JointToJointDistanceConstraint")
+      .def(
+          "__repr__",
+          [](const mm::JointToJointDistanceConstraintT<float>& self) {
+            return fmt::format(
+                "JointToJointDistanceConstraint(joint1={}, joint2={}, offset1=[{:.3f}, {:.3f}, {:.3f}], offset2=[{:.3f}, {:.3f}, {:.3f}], weight={}, target_distance={})",
+                self.joint1,
+                self.joint2,
+                self.offset1.x(),
+                self.offset1.y(),
+                self.offset1.z(),
+                self.offset2.x(),
+                self.offset2.y(),
+                self.offset2.z(),
+                self.weight,
+                self.targetDistance);
+          })
+      .def_readonly(
+          "joint1",
+          &mm::JointToJointDistanceConstraintT<float>::joint1,
+          "The index of the first joint")
+      .def_readonly(
+          "joint2",
+          &mm::JointToJointDistanceConstraintT<float>::joint2,
+          "The index of the second joint")
+      .def_readonly(
+          "offset1",
+          &mm::JointToJointDistanceConstraintT<float>::offset1,
+          "The offset from joint1 in the local coordinate system of joint1")
+      .def_readonly(
+          "offset2",
+          &mm::JointToJointDistanceConstraintT<float>::offset2,
+          "The offset from joint2 in the local coordinate system of joint2")
+      .def_readonly(
+          "weight",
+          &mm::JointToJointDistanceConstraintT<float>::weight,
+          "The weight of the constraint")
+      .def_readonly(
+          "target_distance",
+          &mm::JointToJointDistanceConstraintT<float>::targetDistance,
+          "The target distance between the two points");
+
+  py::class_<
+      mm::JointToJointDistanceErrorFunctionT<float>,
+      mm::SkeletonErrorFunction,
+      std::shared_ptr<mm::JointToJointDistanceErrorFunctionT<float>>>(
+      m,
+      "JointToJointDistanceErrorFunction",
+      R"(Error function that penalizes deviation from a target distance between two points attached to different joints.
+
+This is useful for enforcing distance constraints between different parts of a character,
+such as maintaining a fixed distance between hands or ensuring two joints stay a certain distance apart.)")
+      .def(
+          "__repr__",
+          [](const mm::JointToJointDistanceErrorFunctionT<float>& self) {
+            return fmt::format(
+                "JointToJointDistanceErrorFunction(weight={}, num_constraints={})",
+                self.getWeight(),
+                self.getConstraints().size());
+          })
+      .def(
+          py::init<>(
+              [](const mm::Character& character,
+                 float weight) -> std::shared_ptr<mm::JointToJointDistanceErrorFunctionT<float>> {
+                validateWeight(weight, "weight");
+                auto result =
+                    std::make_shared<mm::JointToJointDistanceErrorFunctionT<float>>(character);
+                result->setWeight(weight);
+                return result;
+              }),
+          R"(Initialize a JointToJointDistanceErrorFunction.
+
+            :param character: The character to use.
+            :param weight: The weight applied to the error function.)",
+          py::keep_alive<1, 2>(),
+          py::arg("character"),
+          py::kw_only(),
+          py::arg("weight") = 1.0f)
+      .def(
+          "add_constraint",
+          [](mm::JointToJointDistanceErrorFunctionT<float>& self,
+             size_t joint1,
+             const Eigen::Vector3f& offset1,
+             size_t joint2,
+             const Eigen::Vector3f& offset2,
+             float targetDistance,
+             float weight) {
+            validateJointIndex(joint1, "joint1", self.getSkeleton());
+            validateJointIndex(joint2, "joint2", self.getSkeleton());
+            validateWeight(weight, "weight");
+            self.addConstraint(joint1, offset1, joint2, offset2, targetDistance, weight);
+          },
+          R"(Adds a joint-to-joint distance constraint to the error function.
+
+        :param joint1: The index of the first joint.
+        :param offset1: The offset from joint1 in the local coordinate system of joint1.
+        :param joint2: The index of the second joint.
+        :param offset2: The offset from joint2 in the local coordinate system of joint2.
+        :param target_distance: The desired distance between the two points in world space.
+        :param weight: The weight of the constraint.)",
+          py::arg("joint1"),
+          py::arg("offset1"),
+          py::arg("joint2"),
+          py::arg("offset2"),
+          py::arg("target_distance"),
+          py::arg("weight") = 1.0f)
+      .def(
+          "add_constraints",
+          [](mm::JointToJointDistanceErrorFunctionT<float>& self,
+             const py::array_t<int>& joint1,
+             const py::array_t<float>& offset1,
+             const py::array_t<int>& joint2,
+             const py::array_t<float>& offset2,
+             const py::array_t<float>& targetDistance,
+             const std::optional<py::array_t<float>>& weight) {
+            ArrayShapeValidator validator;
+            const int nConsIdx = -1;
+            validator.validate(joint1, "joint1", {nConsIdx}, {"n_cons"});
+            validateJointIndex(joint1, "joint1", self.getSkeleton());
+            validator.validate(offset1, "offset1", {nConsIdx, 3}, {"n_cons", "xyz"});
+            validator.validate(joint2, "joint2", {nConsIdx}, {"n_cons"});
+            validateJointIndex(joint2, "joint2", self.getSkeleton());
+            validator.validate(offset2, "offset2", {nConsIdx, 3}, {"n_cons", "xyz"});
+            validator.validate(targetDistance, "target_distance", {nConsIdx}, {"n_cons"});
+            validator.validate(weight, "weight", {nConsIdx}, {"n_cons"});
+            validateWeights(weight, "weight");
+
+            auto joint1Acc = joint1.unchecked<1>();
+            auto offset1Acc = offset1.unchecked<2>();
+            auto joint2Acc = joint2.unchecked<1>();
+            auto offset2Acc = offset2.unchecked<2>();
+            auto targetDistanceAcc = targetDistance.unchecked<1>();
+            auto weightAcc =
+                weight.has_value() ? std::make_optional(weight->unchecked<1>()) : std::nullopt;
+
+            py::gil_scoped_release release;
+
+            for (py::ssize_t i = 0; i < joint1.shape(0); ++i) {
+              self.addConstraint(
+                  joint1Acc(i),
+                  Eigen::Vector3f(offset1Acc(i, 0), offset1Acc(i, 1), offset1Acc(i, 2)),
+                  joint2Acc(i),
+                  Eigen::Vector3f(offset2Acc(i, 0), offset2Acc(i, 1), offset2Acc(i, 2)),
+                  targetDistanceAcc(i),
+                  weightAcc.has_value() ? (*weightAcc)(i) : 1.0f);
+            }
+          },
+          R"(Adds multiple joint-to-joint distance constraints to the error function.
+
+        :param joint1: A numpy array of indices for the first joints.
+        :param offset1: A numpy array of shape (n, 3) for offsets from joint1 in local coordinates.
+        :param joint2: A numpy array of indices for the second joints.
+        :param offset2: A numpy array of shape (n, 3) for offsets from joint2 in local coordinates.
+        :param target_distance: A numpy array of desired distances between point pairs.
+        :param weight: A numpy array of weights for the constraints.)",
+          py::arg("joint1"),
+          py::arg("offset1"),
+          py::arg("joint2"),
+          py::arg("offset2"),
+          py::arg("target_distance"),
+          py::arg("weight") = std::nullopt)
+      .def(
+          "clear_constraints",
+          &mm::JointToJointDistanceErrorFunctionT<float>::clearConstraints,
+          "Clears all joint-to-joint distance constraints from the error function.")
+      .def_property_readonly(
+          "constraints",
+          [](const mm::JointToJointDistanceErrorFunctionT<float>& self) {
+            return self.getConstraints();
+          },
+          "Returns the list of joint-to-joint distance constraints.");
+}
+
 } // namespace
 
 void addErrorFunctions(py::module_& m) {
@@ -2241,6 +2416,9 @@ rotation matrix to a target rotation.)")
 
   // Vertex-to-vertex distance error function
   defVertexVertexDistanceErrorFunction(m);
+
+  // Joint-to-joint distance error function
+  defJointToJointDistanceErrorFunction(m);
 }
 
 } // namespace pymomentum

pymomentum/test/test_solver2.py

@@ -1626,6 +1626,165 @@ def test_vertex_vertex_distance_constraint(self) -> None:
         self.assertIn("vertex_index1=0", constraint_repr)
         self.assertIn("vertex_index2=2", constraint_repr)
 
+    def test_joint_to_joint_distance_constraint(self) -> None:
+        """Test JointToJointDistanceErrorFunction to ensure joints are pulled to target distance."""
+
+        # Create a test character
+        character = pym_geometry.create_test_character(num_joints=4)
+
+        n_params = character.parameter_transform.size
+
+        # Ensure repeatability in the rng:
+        torch.manual_seed(0)
+        model_params_init = torch.zeros(n_params, dtype=torch.float32)
+
+        # Choose two joints to constrain - use joints that are initially far apart
+        joint_index1 = 0
+        joint_index2 = character.skeleton.size - 1  # Last joint
+        offset1 = np.array([0.5, 0.0, 0.0], dtype=np.float32)
+        offset2 = np.array([0.0, 0.5, 0.0], dtype=np.float32)
+        target_distance = 1.5  # Target distance between the two points
+        weight = 1.0
+
+        # Get initial positions of the points
+        skel_state_init = pym_geometry.model_parameters_to_skeleton_state(
+            character, model_params_init
+        )
+        initial_point1 = pym_skel_state.transform_points(
+            skel_state_init[joint_index1],
+            torch.from_numpy(offset1),
+        )
+        initial_point2 = pym_skel_state.transform_points(
+            skel_state_init[joint_index2],
+            torch.from_numpy(offset2),
+        )
+        initial_distance = torch.norm(initial_point2 - initial_point1).item()
+
+        # Create JointToJointDistanceErrorFunction
+        joint_distance_error = pym_solver2.JointToJointDistanceErrorFunction(character)
+
+        # Test basic properties
+        self.assertEqual(len(joint_distance_error.constraints), 0)
+
+        # Add a single constraint
+        joint_distance_error.add_constraint(
+            joint1=joint_index1,
+            offset1=offset1,
+            joint2=joint_index2,
+            offset2=offset2,
+            target_distance=target_distance,
+            weight=weight,
+        )
+
+        # Verify constraint was added
+        self.assertEqual(len(joint_distance_error.constraints), 1)
+
+        constraint = joint_distance_error.constraints[0]
+        self.assertEqual(constraint.joint1, joint_index1)
+        self.assertEqual(constraint.joint2, joint_index2)
+        self.assertTrue(np.allclose(constraint.offset1, offset1))
+        self.assertTrue(np.allclose(constraint.offset2, offset2))
+        self.assertAlmostEqual(constraint.weight, weight)
+        self.assertAlmostEqual(constraint.target_distance, target_distance)
+
+        # Create solver function with the joint distance error
+        solver_function = pym_solver2.SkeletonSolverFunction(
+            character, [joint_distance_error]
+        )
+
+        # Set solver options
+        solver_options = pym_solver2.GaussNewtonSolverOptions()
+        solver_options.max_iterations = 100
+        solver_options.regularization = 1e-5
+
+        # Create and run the solver
+        solver = pym_solver2.GaussNewtonSolver(solver_function, solver_options)
+        model_params_final = solver.solve(model_params_init.numpy())
+
+        # Convert final model parameters to skeleton state
+        skel_state_final = pym_geometry.model_parameters_to_skeleton_state(
+            character, torch.from_numpy(model_params_final)
+        )
+
+        # Compute final positions of the points
+        final_point1 = pym_skel_state.transform_points(
+            skel_state_final[joint_index1],
+            torch.from_numpy(offset1),
+        )
+        final_point2 = pym_skel_state.transform_points(
+            skel_state_final[joint_index2],
+            torch.from_numpy(offset2),
+        )
+        final_distance = torch.norm(final_point2 - final_point1).item()
+
+        # Assert that the final distance is close to the target distance
+        self.assertAlmostEqual(
+            final_distance,
+            target_distance,
+            delta=1e-3,
+            msg=f"Final distance {final_distance} does not match target {target_distance}",
+        )
+
+        # Verify that the distance actually changed from the initial distance
+        self.assertNotAlmostEqual(
+            initial_distance,
+            final_distance,
+            delta=1e-1,
+            msg=f"Distance did not change significantly from initial {initial_distance} to final {final_distance}",
+        )
+
+        # Test multiple constraints using add_constraints
+        joint_distance_error.clear_constraints()
+        self.assertEqual(len(joint_distance_error.constraints), 0)
+
+        # Add multiple constraints
+        joint_indices1 = np.array([0, 1], dtype=np.int32)
+        offsets1 = np.array([[0.5, 0.0, 0.0], [0.0, 0.5, 0.0]], dtype=np.float32)
+        joint_indices2 = np.array([2, 3], dtype=np.int32)
+        offsets2 = np.array([[0.0, 0.0, 0.5], [0.5, 0.5, 0.0]], dtype=np.float32)
+        target_distances = np.array([0.8, 1.2], dtype=np.float32)
+        weights = np.array([1.0, 2.0], dtype=np.float32)
+
+        joint_distance_error.add_constraints(
+            joint1=joint_indices1,
+            offset1=offsets1,
+            joint2=joint_indices2,
+            offset2=offsets2,
+            target_distance=target_distances,
+            weight=weights,
+        )
+
+        # Verify multiple constraints were added
+        self.assertEqual(len(joint_distance_error.constraints), 2)
+        constraints = joint_distance_error.constraints
+
+        # Check first constraint
+        self.assertEqual(constraints[0].joint1, 0)
+        self.assertEqual(constraints[0].joint2, 2)
+        self.assertTrue(np.allclose(constraints[0].offset1, [0.5, 0.0, 0.0]))
+        self.assertTrue(np.allclose(constraints[0].offset2, [0.0, 0.0, 0.5]))
+        self.assertAlmostEqual(constraints[0].weight, 1.0)
+        self.assertAlmostEqual(constraints[0].target_distance, 0.8)
+
+        # Check second constraint
+        self.assertEqual(constraints[1].joint1, 1)
+        self.assertEqual(constraints[1].joint2, 3)
+        self.assertTrue(np.allclose(constraints[1].offset1, [0.0, 0.5, 0.0]))
+        self.assertTrue(np.allclose(constraints[1].offset2, [0.5, 0.5, 0.0]))
+        self.assertAlmostEqual(constraints[1].weight, 2.0)
+        self.assertAlmostEqual(constraints[1].target_distance, 1.2)
+
+        # Test string representation
+        repr_str = repr(joint_distance_error)
+        self.assertIn("JointToJointDistanceErrorFunction", repr_str)
+        self.assertIn("num_constraints=2", repr_str)
+
+        # Test constraint string representation
+        constraint_repr = repr(constraints[0])
+        self.assertIn("JointToJointDistanceConstraint", constraint_repr)
+        self.assertIn("joint1=0", constraint_repr)
+        self.assertIn("joint2=2", constraint_repr)
+
     def test_weight_validation(self) -> None:
         """Test that error functions throw ValueError when negative weights are passed."""