pymomentum-cpu 0.0.0__cp312-cp312-manylinux_2_39_x86_64.whl

include/momentum/character_solver/constraint_error_function-inl.h

@@ -0,0 +1,321 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <momentum/character/skeleton_state.h>
+#include <momentum/common/checks.h>
+#include <momentum/common/profile.h>
+
+namespace momentum {
+
+template <typename T, class Data, size_t FuncDim, size_t NumVec, size_t NumPos>
+ConstraintErrorFunctionT<T, Data, FuncDim, NumVec, NumPos>::ConstraintErrorFunctionT(
+    const Skeleton& skel,
+    const ParameterTransform& pt,
+    const T& lossAlpha,
+    const T& lossC)
+    : SkeletonErrorFunctionT<T>(skel, pt),
+      loss_(lossAlpha, lossC),
+      jointGrad_(pt.numAllModelParameters()) {
+  static_assert(FuncDim > 0, "The error function cannot be empty.");
+  static_assert(NumVec > 0, "At least one vector is required in the constraint.");
+  static_assert(
+      NumVec >= NumPos, "The number of points cannot be more than the total number of vectors.");
+}
+
+template <typename T, class Data, size_t FuncDim, size_t NumVec, size_t NumPos>
+double ConstraintErrorFunctionT<T, Data, FuncDim, NumVec, NumPos>::getError(
+    const ModelParametersT<T>& /* params */,
+    const SkeletonStateT<T>& state) {
+  MT_PROFILE_FUNCTION();
+
+  // loop over all constraints and calculate the error
+  FuncType f;
+  double error = 0.0;
+  for (size_t iConstr = 0; iConstr < constraints_.size(); ++iConstr) {
+    const Data& constr = constraints_[iConstr];
+    if (constr.weight != 0) {
+      evalFunction(iConstr, state.jointState.at(constr.parent), f);
+      error += constr.weight * loss_.value(f.squaredNorm());
+    }
+  }
+  return this->weight_ * error;
+}
+
+template <typename T, class Data, size_t FuncDim, size_t NumVec, size_t NumPos>
+double ConstraintErrorFunctionT<T, Data, FuncDim, NumVec, NumPos>::getJacobianForSingleConstraint(
+    const JointStateListT<T>& jointStates,
+    const size_t iConstr,
+    Ref<Eigen::MatrixX<T>> jacobian,
+    Ref<Eigen::VectorX<T>> residual) {
+  const Data& constr = this->constraints_[iConstr];
+  if (constr.weight == 0) {
+    return 0;
+  }
+
+  FuncType f;
+  std::array<VType, NumVec> v;
+  std::array<DfdvType, NumVec> dfdv;
+  evalFunction(iConstr, jointStates.at(constr.parent), f, v, dfdv);
+  const T sqrError = f.squaredNorm();
+  const T w = constr.weight * this->weight_;
+  const double error = w * this->loss_.value(sqrError);
+  const T deriv = std::sqrt(w * this->loss_.deriv(sqrError));
+
+  // The input jacobian is a subblock from the full jacobian tailored for this error function, so
+  // the rows start from zero.
+  const size_t rowIndex = FuncDim * iConstr;
+  // The gradient of the loss function needs to be split between J and r, so both J'J and
+  // J'r are scaled correctly. The factor 2 from the square term is accounted for in the base
+  // class function.
+  residual.template middleRows<FuncDim>(rowIndex).noalias() = deriv * f;
+
+  // Small optimization for early termination
+  if (isApprox<T>(deriv, T(0), Eps<T>(1e-9, 1e-16))) {
+    return error;
+  }
+  const bool zeroDeriv =
+      !std::any_of(dfdv.begin(), dfdv.end(), [](const auto& vec) { return !vec.isZero(); });
+  if (zeroDeriv) {
+    return error;
+  }
+
+  // shorthands for code readability
+  const auto outerIndexPtr = this->parameterTransform_.transform.outerIndexPtr();
+  const auto innerIndexPtr = this->parameterTransform_.transform.innerIndexPtr();
+  const auto valuePtr = this->parameterTransform_.transform.valuePtr();
+  auto&& jac = jacobian.template middleRows<FuncDim>(rowIndex);
+
+  // Fill in the jacobian by walking up the joint hierarchy for each constraint.
+  // Need to consider both active joints and enabled model parameters.
+  size_t jntIndex = constr.parent;
+  while (jntIndex != kInvalidIndex) {
+    // check for valid index
+    MT_CHECK(jntIndex < this->skeleton_.joints.size());
+
+    const auto& jntState = jointStates[jntIndex];
+    const size_t jntParamIndex = jntIndex * kParametersPerJoint;
+
+    for (size_t jVec = 0; jVec < NumVec; ++jVec) {
+      if (dfdv[jVec].isZero()) {
+        continue;
+      }
+
+      Eigen::Vector3<T> offset;
+      if (jVec < NumPos) {
+        offset.noalias() = v[jVec] - jntState.translation();
+      } else {
+        offset = v[jVec];
+      }
+
+      // Translational dofs -- only affect POINT
+      if (jVec < NumPos) {
+        for (size_t d = 0; d < 3; ++d) {
+          if (!this->activeJointParams_[jntParamIndex + d]) {
+            continue;
+          }
+
+          const FuncType jc = deriv * dfdv[jVec] * jntState.getTranslationDerivative(d);
+          for (auto index = outerIndexPtr[jntParamIndex + d];
+               index < outerIndexPtr[jntParamIndex + d + 1];
+               ++index) {
+            if (this->enabledParameters_.test(innerIndexPtr[index])) {
+              jac.col(innerIndexPtr[index]).noalias() += jc * valuePtr[index];
+            }
+          }
+        }
+      }
+
+      // Rotational dofs -- affect both POINT and AXIS
+      for (size_t d = 0; d < 3; ++d) {
+        if (!this->activeJointParams_[jntParamIndex + 3 + d]) {
+          continue;
+        }
+
+        const FuncType jc = deriv * dfdv[jVec] * jntState.getRotationDerivative(d, offset);
+        for (auto index = outerIndexPtr[jntParamIndex + 3 + d];
+             index < outerIndexPtr[jntParamIndex + 3 + d + 1];
+             ++index) {
+          if (this->enabledParameters_.test(innerIndexPtr[index])) {
+            jac.col(innerIndexPtr[index]).noalias() += jc * valuePtr[index];
+          }
+        }
+      }
+
+      // Scale dof -- only affect POINT
+      if (jVec < NumPos) {
+        if (this->activeJointParams_[jntParamIndex + 6]) {
+          const FuncType jc = deriv * dfdv[jVec] * jntState.getScaleDerivative(offset);
+          for (auto index = outerIndexPtr[jntParamIndex + 6];
+               index < outerIndexPtr[jntParamIndex + 6 + 1];
+               ++index) {
+            if (this->enabledParameters_.test(innerIndexPtr[index])) {
+              jac.col(innerIndexPtr[index]).noalias() += jc * valuePtr[index];
+            }
+          }
+        }
+      }
+    }
+    // go up the hierarchy to the parent
+    jntIndex = this->skeleton_.joints[jntIndex].parent;
+  }
+
+  return error;
+}
+
+template <typename T, class Data, size_t FuncDim, size_t NumVec, size_t NumPos>
+double ConstraintErrorFunctionT<T, Data, FuncDim, NumVec, NumPos>::getJacobian(
+    const ModelParametersT<T>& /*params*/,
+    const SkeletonStateT<T>& state,
+    Ref<Eigen::MatrixX<T>> jacobian,
+    Ref<Eigen::VectorX<T>> residual,
+    int& usedRows) {
+  MT_PROFILE_FUNCTION();
+  usedRows = getJacobianSize();
+
+  double error = 0.0;
+  for (size_t iConstr = 0; iConstr < this->constraints_.size(); ++iConstr) {
+    error += getJacobianForSingleConstraint(state.jointState, iConstr, jacobian, residual);
+  }
+  return error;
+}
+
+template <typename T, class Data, size_t FuncDim, size_t NumVec, size_t NumPos>
+double ConstraintErrorFunctionT<T, Data, FuncDim, NumVec, NumPos>::getGradientForSingleConstraint(
+    const JointStateListT<T>& jointStates,
+    size_t iConstr,
+    Ref<VectorX<T>> /*gradient*/) {
+  const Data& constr = this->constraints_[iConstr];
+  if (constr.weight == 0) {
+    return 0;
+  }
+
+  FuncType f;
+  std::array<VType, NumVec> v;
+  std::array<DfdvType, NumVec> dfdv;
+  evalFunction(iConstr, jointStates.at(constr.parent), f, v, dfdv);
+  if (f.isZero()) {
+    return 0;
+  }
+
+  const T sqrError = f.squaredNorm();
+  const T w = constr.weight * this->weight_;
+  const double error = w * this->loss_.value(sqrError);
+  // The gradient is scaled by the gradient of the loss function; factor 2 falls out of the
+  // square.
+  const FuncType deriv = T(2) * w * this->loss_.deriv(sqrError) * f;
+
+  const auto outerIndexPtr = this->parameterTransform_.transform.outerIndexPtr();
+  const auto innerIndexPtr = this->parameterTransform_.transform.innerIndexPtr();
+  const auto valuePtr = this->parameterTransform_.transform.valuePtr();
+  size_t jntIndex = constr.parent;
+  while (jntIndex != kInvalidIndex) {
+    // check for valid index
+    MT_CHECK(jntIndex < this->skeleton_.joints.size());
+
+    const auto& jntState = jointStates[jntIndex];
+    const size_t jntParamIndex = jntIndex * kParametersPerJoint;
+
+    for (size_t jVec = 0; jVec < NumVec; ++jVec) {
+      if (dfdv[jVec].isZero()) {
+        continue;
+      }
+
+      Eigen::Vector3<T> offset;
+      if (jVec < NumPos) {
+        offset.noalias() = v[jVec] - jntState.translation();
+      } else {
+        offset = v[jVec];
+      }
+
+      // Translational dofs -- only affect POINT
+      if (jVec < NumPos) {
+        for (size_t d = 0; d < 3; d++) {
+          if (!this->activeJointParams_[jntParamIndex + d]) {
+            continue;
+          }
+
+          // joint space gradient
+          const T val = deriv.dot(dfdv[jVec] * jntState.getTranslationDerivative(d));
+          // multiply with parameter transform to get model parameter gradient
+          for (auto index = outerIndexPtr[jntParamIndex + d];
+               index < outerIndexPtr[jntParamIndex + d + 1];
+               ++index) {
+            if (this->enabledParameters_.test(innerIndexPtr[index])) {
+              jointGrad_[innerIndexPtr[index]] += val * valuePtr[index];
+            }
+          }
+        }
+      }
+
+      // Rotational dofs -- affect both POINT and AXIS
+      for (size_t d = 0; d < 3; d++) {
+        if (!this->activeJointParams_[jntParamIndex + 3 + d]) {
+          continue;
+        }
+
+        // joint space gradient
+        const T val = deriv.dot(dfdv[jVec] * jntState.getRotationDerivative(d, offset));
+        // multiply with parameter transform to get model parameter gradient
+        for (auto index = outerIndexPtr[jntParamIndex + 3 + d];
+             index < outerIndexPtr[jntParamIndex + 3 + d + 1];
+             ++index) {
+          if (this->enabledParameters_.test(innerIndexPtr[index])) {
+            jointGrad_[innerIndexPtr[index]] += val * valuePtr[index];
+          }
+        }
+      }
+
+      // Scale dof -- only affect POINT
+      if (jVec < NumPos) {
+        if (this->activeJointParams_[jntParamIndex + 6]) {
+          // joint space gradient
+          const T val = deriv.dot(dfdv[jVec] * jntState.getScaleDerivative(offset));
+          // multiply with parameter transform to get model parameter gradient
+          for (auto index = outerIndexPtr[jntParamIndex + 6];
+               index < outerIndexPtr[jntParamIndex + 6 + 1];
+               ++index) {
+            if (this->enabledParameters_.test(innerIndexPtr[index])) {
+              jointGrad_[innerIndexPtr[index]] += val * valuePtr[index];
+            }
+          }
+        }
+      }
+    }
+    // go up the skeleton hierarchy to the parent joint
+    jntIndex = this->skeleton_.joints[jntIndex].parent;
+  }
+
+  return error;
+}
+
+template <typename T, class Data, size_t FuncDim, size_t NumVec, size_t NumPos>
+double ConstraintErrorFunctionT<T, Data, FuncDim, NumVec, NumPos>::getGradient(
+    const ModelParametersT<T>& /*unused*/,
+    const SkeletonStateT<T>& state,
+    Ref<Eigen::VectorX<T>> gradient) {
+  MT_PROFILE_FUNCTION();
+
+  // initialize joint gradients storage
+  jointGrad_.setZero();
+
+  double error = 0.0;
+  for (size_t iConstr = 0; iConstr < this->constraints_.size(); ++iConstr) {
+    error += getGradientForSingleConstraint(state.jointState, iConstr, gradient);
+  }
+
+  gradient += jointGrad_;
+  return error;
+}
+
+template <typename T, class Data, size_t FuncDim, size_t NumVec, size_t NumPos>
+size_t ConstraintErrorFunctionT<T, Data, FuncDim, NumVec, NumPos>::getJacobianSize() const {
+  return FuncDim * constraints_.size();
+}
+
+} // namespace momentum

include/momentum/character_solver/constraint_error_function.h

@@ -0,0 +1,241 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <momentum/character/character.h>
+#include <momentum/character/types.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+#include <momentum/math/generalized_loss.h>
+#include <momentum/math/types.h>
+
+#include <optional>
+
+namespace momentum {
+
+/// Base structure of constraint data
+struct ConstraintData {
+  /// Parent joint index this constraint is under
+  size_t parent = kInvalidIndex;
+  /// Weight of the constraint
+  float weight = 0.0;
+  /// Name of the constraint
+  std::string name = {};
+
+  ConstraintData(size_t pIndex, float w, const std::string& n = "")
+      : parent(pIndex), weight(w), name(n) {}
+};
+
+/// A list of ConstraintData
+using ConstraintDataList = std::vector<ConstraintData>;
+
+/// An optional of a reference. Because optional<T&> is invalid, we need to use reference_wrapper to
+/// make the reference of T as an optional.
+template <typename T>
+using optional_ref = std::optional<std::reference_wrapper<T>>;
+
+/// The ConstraintErrorFunction is a base class of a general form of constraint errors l = w *
+/// loss(f^2), where w is the weight (could be a product of different weighting terms), loss() is
+/// the generalized loss function (see math/generalized_loss.h), and f is a difference vector we
+/// want to minimize.
+///
+/// f takes the form of f(v, target), where v = T(q)*source. T is the global transformation of the
+/// parent joint of the source, and target is the desired value of source in global space. f
+/// computes the differences between v and target.
+///
+/// Based on the above, we have
+/// Jacobian: df/dq = df/dv * dv/dT * dT/dq, and
+/// Gradient: dl/dq = dl/df * Jac
+/// Both dl/df and dT/dq are boiler plate code that we can implement in the base class, so a derived
+/// class only needs to implement f and df/dT. However, dT/dq is not efficient to compute, and
+/// Momentum instead implements dv/dq, for any 3-vector v. Therefore, we will compute df/dq = df/dv
+/// * dv/dq, and implement dv/dq in the base class. So a derived class now needs to implement f and
+/// df/dv.
+///
+/// This should work for a point (eg. PositionErrorFunction), or an axis (eg.
+/// FixedAxisErrorFunction), but it can also work for a rotation matrix, or a 3x4 transformation
+/// matrix, by applying the transformation one axis/point at a time.The number of 3-vectors to be
+/// transformed in a constraint is NumVec.
+template <
+    typename T, // float or double
+    class Data, // derived types from ConstraintData
+    size_t FuncDim = 3, // dimension of f
+    size_t NumVec =
+        1, // how many 3-vector v in one constraint, eg. a point is 1, and a rotation matrix is 3
+    size_t NumPos =
+        1> // we assume a constraint can be a function of both points and axes, and points come
+           // before axes in the NumVec of "v"s. This specifies how many "v"s are points. For
+           // example, it's 1 for a point constraint, and 0 for a rotation matrix.
+class ConstraintErrorFunctionT : public SkeletonErrorFunctionT<T> {
+ public:
+  static constexpr size_t kFuncDim = FuncDim;
+  static constexpr size_t kNumVec = NumVec;
+  static constexpr size_t kNumPos = NumPos;
+
+  using FuncType = Vector<T, FuncDim>; // vector type for f
+  using VType = Vector3<T>; // vector type for v
+  using DfdvType = Eigen::Matrix<T, FuncDim, 3>; // type for dfdv - it's effectively a vector if f
+                                                 // is a scalar (FuncDim=1, eg. PlaneErrorFunction)
+  /// Constructor
+  ///
+  /// @param[in] skel: character skeleton
+  /// @param[in] pt: parameter transformation
+  /// @param[in] lossAlpha: alpha parameter for the loss function
+  /// @param[in] lossC: c parameter for the loss function
+  explicit ConstraintErrorFunctionT(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1));
+
+  /// A convenience constructor where character contains info of the skeleton and parameter
+  /// transform.
+  ///
+  /// @param[in] character: character definition
+  /// @param[in] lossAlpha: alpha parameter for the loss function
+  /// @param[in] lossC: c parameter for the loss function
+  explicit ConstraintErrorFunctionT(
+      const Character& character,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : ConstraintErrorFunctionT<T, Data, FuncDim, NumVec, NumPos>(
+            character.skeleton,
+            character.parameterTransform,
+            lossAlpha,
+            lossC) {}
+
+  // The functions below should just work for most constraints. When we have new constraints that
+  // don't fit these implementations, we can remove the "final" annotation to allow override.
+
+  /// Computes the error function value l = w * loss(f^2). It gets f from the derived class, and
+  /// implements the rest.
+  ///
+  /// @param[in] params: current model parameters
+  /// @param[in] state: curren global skeleton joint states computed from the model parameters
+  ///
+  /// @return the error function value l
+  [[nodiscard]] double getError(const ModelParametersT<T>& params, const SkeletonStateT<T>& state)
+      final;
+
+  /// The gradient of the error function: dl/dq = dl/d[f^2] * 2f * df/dv * dv/dq. It gets df/dv from
+  /// the derived class, and implements the rest.
+  ///
+  /// @param[in] params: current model parameters
+  /// @param[in] state: curren global skeleton joint states computed from the model parameters
+  /// @param[out] gradient: the gradient vector to accumulate into
+  ///
+  /// @return the error function value l
+  double getGradient(
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      Ref<VectorX<T>> gradient) final;
+
+  /// For least-square problems, we assume l is the square of a vector function F. The jacobian is
+  /// then dF/dq. (A factor 2 is implemented in the solver.) With l2 loss, we have F = sqrt(w) * f,
+  /// and the jacobian is sqrt(w) * df/dv * dv/dq. With the generalized loss, the jacobian becomes
+  /// sqrt(w * d[loss]/d[f^2]) * df/dv * dv/dq. It gets df/dv from the derived class, and implements
+  /// the rest.
+  ///
+  /// @param[in] params: current model parameters
+  /// @param[in] state: curren global skeleton joint states computed from the model parameters
+  /// @param[out] jacobian: the output jacobian matrix
+  /// @param[out] residual: the output function residual (ie. f scaled by the loss gradient)
+  /// @param[out] usedRows: number of rows in the jacobian/residual used by this error function
+  ///
+  /// @return the error function l
+  double getJacobian(
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      Ref<MatrixX<T>> jacobian,
+      Ref<VectorX<T>> residual,
+      int& usedRows) final;
+
+  /// The number of rows in the jacobian is the dimension of f, FuncDim, times the number of
+  /// constraints.
+  ///
+  /// @return number of rows in the jacobian
+  [[nodiscard]] size_t getJacobianSize() const final;
+
+  /// Adds a constraint to the list
+  ///
+  /// @param[in] constr: the constraint to be added
+  void addConstraint(const Data& constr) {
+    constraints_.push_back(constr);
+  }
+
+  /// Appends a list of constraints
+  ///
+  /// @param[in] constrs: a list of constraints to be added
+  void addConstraints(gsl::span<const Data> constrs) {
+    constraints_.insert(constraints_.end(), constrs.begin(), constrs.end());
+  }
+
+  /// Replace the current list of constraints with the input
+  ///
+  /// @param[in] constrs: the new list of constraints
+  void setConstraints(gsl::span<const Data> constrs) {
+    constraints_.assign(constrs.begin(), constrs.end());
+  }
+
+  /// @return the current list of constraints immutable
+  [[nodiscard]] const std::vector<Data>& getConstraints() const {
+    return constraints_;
+  }
+
+  [[nodiscard]] size_t numConstraints() const {
+    return constraints_.size();
+  }
+
+  /// Clear the current list of constraints
+  void clearConstraints() {
+    constraints_.clear();
+  }
+
+ protected:
+  /// List of constraints
+  std::vector<Data> constraints_;
+  /// The generalized loss function that transforms f^2
+  const GeneralizedLossT<T> loss_;
+  /// Intermediate storage of the gradient from this error function. We can allocate the space in
+  /// the constructor to save some dynamic allocation.
+  VectorX<T> jointGrad_;
+
+  /// The only function a derived class needs to implement.
+  /// f is needed both for errors and derivatives, v and dfdv are needed only for computing
+  /// derivatives, and therefore optional. An implementation should check if they are provided
+  /// before computing and setting their values. And don't assume they are zero-initialized. Each
+  /// returned v and df/dv corresponds to a source 3-vector (eg. a position constraint, or each of
+  /// the three axis of a rotation constraint).
+  ///
+  /// @param[in] constrIndex: index of the constraint to evaluate
+  /// @param[in] state: JointState of the parent joint with transformation T
+  /// @param[out] f: output the value of f of dimension FuncDim
+  /// @param[out] v: if valid, output the vector v=T*source; there could be NumVec of vs
+  /// @param[out] dfdv: if valid, output the matrix df/dv of dimension FuncDim x 3 per v
+  virtual void evalFunction(
+      size_t constrIndex,
+      const JointStateT<T>& state,
+      FuncType& f,
+      optional_ref<std::array<VType, NumVec>> v = {},
+      optional_ref<std::array<DfdvType, NumVec>> dfdv = {}) const = 0;
+
+ private:
+  double getJacobianForSingleConstraint(
+      const JointStateListT<T>& jointStates,
+      size_t iConstr,
+      Ref<Eigen::MatrixX<T>> jacobian,
+      Ref<Eigen::VectorX<T>> residual);
+
+  double getGradientForSingleConstraint(
+      const JointStateListT<T>& jointStates,
+      size_t iConstr,
+      Ref<VectorX<T>> gradient);
+};
+
+} // namespace momentum
+
+#include "momentum/character_solver/constraint_error_function-inl.h"

include/momentum/character_solver/distance_error_function.h

@@ -0,0 +1,74 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include "momentum/character/locator.h"
+#include "momentum/character_solver/skeleton_error_function.h"
+
+namespace momentum {
+
+// The distance constraint is defined as ||(p_joint - origin)^2 - target||^2
+template <typename T>
+struct DistanceConstraintDataT {
+  Eigen::Vector3<T> origin = Eigen::Vector3<T>::Zero(); // origin in world space
+  T target{}; // distance target in world space
+  size_t parent{}; // parent joint of the constraint
+  Eigen::Vector3<T> offset; // relative offset to the parent
+  T weight{}; // constraint weight
+  // comment for now
+  static DistanceConstraintDataT<T> createFromLocator(const momentum::Locator& locator);
+};
+
+template <typename T>
+class DistanceErrorFunctionT : public momentum::SkeletonErrorFunctionT<T> {
+ public:
+  DistanceErrorFunctionT(const momentum::Skeleton& skel, const momentum::ParameterTransform& pt);
+
+  [[nodiscard]] double getError(
+      const momentum::ModelParametersT<T>& params,
+      const momentum::SkeletonStateT<T>& state) final;
+  double getGradient(
+      const momentum::ModelParametersT<T>& params,
+      const momentum::SkeletonStateT<T>& state,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) final;
+  double getJacobian(
+      const momentum::ModelParametersT<T>& params,
+      const momentum::SkeletonStateT<T>& state,
+      Eigen::Ref<Eigen::MatrixX<T>> jacobian,
+      Eigen::Ref<Eigen::VectorX<T>> residual,
+      int& usedRows) final;
+  [[nodiscard]] size_t getJacobianSize() const final;
+
+  void addConstraint(const DistanceConstraintDataT<T>& constr) {
+    constraints_.push_back(constr);
+  }
+  void clearConstraints() {
+    constraints_.clear();
+  }
+  void setConstraints(std::vector<DistanceConstraintDataT<T>> constr) {
+    constraints_ = std::move(constr);
+  }
+  [[nodiscard]] bool empty() const {
+    return constraints_.empty();
+  }
+  [[nodiscard]] size_t numConstraints() const {
+    return constraints_.size();
+  }
+
+  [[nodiscard]] const std::vector<DistanceConstraintDataT<T>>& getConstraints() const {
+    return constraints_;
+  }
+
+ protected:
+  // TODO: what should we use here?
+  const T kDistanceWeight = 1.0f;
+
+  std::vector<DistanceConstraintDataT<T>> constraints_;
+};
+
+} // namespace momentum