pymomentum-cpu 0.0.0__cp312-cp312-manylinux_2_39_x86_64.whl

include/momentum/common/memory.h

@@ -0,0 +1,71 @@
+/*
+ * 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 <memory>
+
+// Define smart pointers for a type
+#ifndef MOMENTUM_DEFINE_POINTERS
+#define MOMENTUM_DEFINE_POINTERS(x)               \
+  using x##_p = ::std::shared_ptr<x>;             \
+  using x##_u = ::std::unique_ptr<x>;             \
+  using x##_w = ::std::weak_ptr<x>;               \
+  using x##_const_p = ::std::shared_ptr<const x>; \
+  using x##_const_u = ::std::unique_ptr<const x>; \
+  using x##_const_w = ::std::weak_ptr<const x>;
+#endif // MOMENTUM_DEFINE_POINTERS
+
+// Forward-declare a struct, define smart pointers
+#ifndef MOMENTUM_FWD_DECLARE_STRUCT
+#define MOMENTUM_FWD_DECLARE_STRUCT(x) \
+  struct x;                            \
+  MOMENTUM_DEFINE_POINTERS(x);
+#endif // MOMENTUM_FWD_DECLARE_STRUCT
+
+// Forward-declare a class, define smart pointers
+#ifndef MOMENTUM_FWD_DECLARE_CLASS
+// NOLINTBEGIN(bugprone-macro-parentheses)
+#define MOMENTUM_FWD_DECLARE_CLASS(x) \
+  class x;                            \
+  MOMENTUM_DEFINE_POINTERS(x)
+// NOLINTEND(bugprone-macro-parentheses)
+#endif // MOMENTUM_FWD_DECLARE_CLASS
+
+// Forward-declare a templated class, define smart pointers for the class and its variants
+// Use the standard backward-compatible naming scheme:
+//    template <typename T> class BarT;
+//    using Bar = BarT<float>;
+//    using Bard = BarT<double>;
+#ifndef MOMENTUM_FWD_DECLARE_TEMPLATE_CLASS
+// NOLINTBEGIN(bugprone-macro-parentheses)
+#define MOMENTUM_FWD_DECLARE_TEMPLATE_CLASS(x) \
+  template <typename T>                        \
+  class x##T;                                  \
+  using x = x##T<float>;                       \
+  using x##d = x##T<double>;                   \
+  MOMENTUM_DEFINE_POINTERS(x)                  \
+  MOMENTUM_DEFINE_POINTERS(x##d)
+// NOLINTEND(bugprone-macro-parentheses)
+#endif // MOMENTUM_FWD_DECLARE_TEMPLATE_CLASS
+
+// Forward-declare a templated struct, define smart pointers for the struct and its variants
+// Use the standard backward-compatible naming scheme:
+//    template <typename T> struct BarT;
+//    using Bar = BarT<float>;
+//    using Bard = BarT<double>;
+#ifndef MOMENTUM_FWD_DECLARE_TEMPLATE_STRUCT
+// NOLINTBEGIN(bugprone-macro-parentheses)
+#define MOMENTUM_FWD_DECLARE_TEMPLATE_STRUCT(x) \
+  template <typename T>                         \
+  struct x##T;                                  \
+  using x = x##T<float>;                        \
+  using x##d = x##T<double>;                    \
+  MOMENTUM_DEFINE_POINTERS(x)                   \
+  MOMENTUM_DEFINE_POINTERS(x##d)
+// NOLINTEND(bugprone-macro-parentheses)
+#endif // MOMENTUM_FWD_DECLARE_TEMPLATE_STRUCT

include/momentum/common/profile.h

@@ -0,0 +1,79 @@
+/*
+ * 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
+
+#if defined(MOMENTUM_WITH_XR_PROFILER)
+
+#include <arvr/libraries/profile_redirect/annotate.hpp>
+
+#define MT_PROFILE_FUNCTION() XR_PROFILE_FUNCTION()
+#define MT_PROFILE_FUNCTION_CATEGORY(CATEGORY) XR_PROFILE_FUNCTION_CATEGORY(CATEGORY)
+#define MT_PROFILE_EVENT(NAME) XR_PROFILE_EVENT(NAME)
+#define MT_PROFILE_EVENT_DYNAMIC(NAME) XR_PROFILE_EVENT_DYNAMIC(NAME)
+#define MT_PROFILE_CATEGORY(NAME, CATEGORY) XR_PROFILE_CATEGORY(NAME, CATEGORY)
+#define MT_PROFILE_PREPARE_PUSH_POP() XR_PROFILE_PREPARE_PUSH_POP()
+#define MT_PROFILE_PUSH(NAME) XR_PROFILE_PUSH(NAME)
+#define MT_PROFILE_POP() XR_PROFILE_POP()
+#define MT_PROFILE_THREAD(THREAD_NAME) XR_PROFILE_THREAD(THREAD_NAME)
+#define MT_PROFILE_UPDATE() XR_PROFILE_UPDATE()
+#define MT_PROFILE_BEGIN_FRAME() XR_PROFILE_BEGIN_FRAME()
+#define MT_PROFILE_END_FRAME() XR_PROFILE_END_FRAME()
+#define MT_PROFILE_METADATA(NAME, DATA) XR_PROFILE_METADATA(NAME, DATA)
+
+#elif defined(MOMENTUM_WITH_TRACY_PROFILER)
+
+#include <tracy/Tracy.hpp>
+#include <tracy/TracyC.h>
+#include <stack>
+
+#define _MT_PROFILE_CONCATENATE_DETAIL(x, y) x##y
+#define _MT_PROFILE_CONCATENATE(x, y) _MT_PROFILE_CONCATENATE_DETAIL(x, y)
+#define _MT_PROFILE_MAKE_UNIQUE(x) _MT_PROFILE_CONCATENATE(x, __LINE__)
+
+#define MT_PROFILE_FUNCTION() ZoneScoped
+#define MT_PROFILE_FUNCTION_CATEGORY(CATEGORY)
+#define MT_PROFILE_EVENT(NAME) ZoneNamedN(_MT_PROFILE_MAKE_UNIQUE(__tracy), NAME, true)
+#define MT_PROFILE_EVENT_DYNAMIC(NAME) ZoneTransientN(_MT_PROFILE_MAKE_UNIQUE(__tracy), NAME, true)
+#define MT_PROFILE_CATEGORY(NAME, CATEGORY)
+#define MT_PROFILE_PREPARE_PUSH_POP() std::stack<TracyCZoneCtx> ___tracy_xr_stack
+#define MT_PROFILE_PUSH(NAME)                                                                    \
+  static const struct ___tracy_source_location_data TracyConcat(                                 \
+      __tracy_source_location, TracyLine) = {NAME, __func__, TracyFile, (uint32_t)TracyLine, 0}; \
+  ___tracy_xr_stack.push(                                                                        \
+      ___tracy_emit_zone_begin(&TracyConcat(__tracy_source_location, TracyLine), true));
+#define MT_PROFILE_POP()                  \
+  TracyCZoneEnd(___tracy_xr_stack.top()); \
+  ___tracy_xr_stack.pop()
+#define MT_PROFILE_THREAD(THREAD_NAME)           \
+  {                                              \
+    /* support both std::string and c strings */ \
+    std::string threadNameStr(THREAD_NAME);      \
+    TracyCSetThreadName(threadNameStr.c_str());  \
+  }
+#define MT_PROFILE_UPDATE()
+#define MT_PROFILE_BEGIN_FRAME() FrameMark
+#define MT_PROFILE_END_FRAME()
+#define MT_PROFILE_METADATA(NAME, DATA)
+
+#else
+
+#define MT_PROFILE_FUNCTION()
+#define MT_PROFILE_FUNCTION_CATEGORY(CATEGORY)
+#define MT_PROFILE_EVENT(NAME)
+#define MT_PROFILE_EVENT_DYNAMIC(NAME)
+#define MT_PROFILE_CATEGORY(NAME, CATEGORY)
+#define MT_PROFILE_PREPARE_PUSH_POP()
+#define MT_PROFILE_PUSH(NAME)
+#define MT_PROFILE_POP()
+#define MT_PROFILE_THREAD(THREAD_NAME)
+#define MT_PROFILE_UPDATE()
+#define MT_PROFILE_BEGIN_FRAME()
+#define MT_PROFILE_END_FRAME()
+#define MT_PROFILE_METADATA(NAME, DATA)
+
+#endif

include/momentum/common/progress_bar.h

@@ -0,0 +1,37 @@
+/*
+ * 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 <indicators/progress_bar.hpp>
+
+#include <string>
+
+namespace momentum {
+
+// A simple progress bar that prints hash marks (e.g., "Name [===>  ] 60%")
+class ProgressBar {
+ public:
+  /// @param prefix Displayed prefix
+  /// @param numOperations Total operations (determines progress ratio)
+  ProgressBar(const std::string& prefix, size_t numOperations);
+
+  /// Increments the progress by the given count (default 1)
+  void increment(size_t count = 1);
+
+  /// Sets the progress to the given count
+  void set(size_t count);
+
+  /// Returns the current progress in the range [0, 100]
+  [[nodiscard]] size_t getCurrentProgress();
+
+ private:
+  static constexpr size_t kMaxWidth = 80;
+  indicators::ProgressBar bar_;
+};
+
+} // namespace momentum

include/momentum/common/string.h

@@ -0,0 +1,52 @@
+/*
+ * 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 <string>
+#include <string_view>
+#include <vector>
+
+namespace momentum {
+
+/// Returns a copy of the string from which all the characters in whiteChars
+/// at the beginning or at the end of the string have been removed.
+/// @tparam StringType: The type of the string to trim (e.g., std::string, std::string_view)
+/// @param text: The string to trim
+/// @param whiteChars: A series of 1-byte characters to remove
+/// @return The trimmed string
+template <typename StringType>
+[[nodiscard]] StringType trim(const StringType& text, const char* whiteChars = " \t");
+
+/// Returns a copy of the string from which all the characters in whiteChars
+/// at the beginning or at the end of the string have been removed.
+/// @param text: The string to trim (null-terminated)
+/// @param whiteChars: A series of 1-byte characters to remove
+/// @return The trimmed string
+[[nodiscard]] std::string trim(const char* text, const char* whiteChars = " \t");
+
+/// Tokenize a string using the specified delimiters.
+///
+/// @param inputString The string to tokenize.
+/// @param delimiters The string containing delimiter characters.
+/// @param trim Whether to trim spaces around the tokens (default: true).
+/// @return A vector of tokenized strings.
+std::vector<std::string> tokenize(
+    const std::string& inputString,
+    const std::string& delimiters = " \t\r\n",
+    bool trim = true);
+
+/// Tokenize a string_view using the specified delimiters.
+///
+/// @param inputString The string_view to tokenize.
+/// @param delimiters The string_view containing delimiter characters.
+/// @param trim Whether to trim spaces around the tokens (default: true).
+/// @return A vector of tokenized string_views.
+std::vector<std::string_view>
+tokenize(std::string_view inputString, std::string_view delimiters = " \t\r\n", bool trim = true);
+
+} // namespace momentum

include/momentum/diff_ik/ceres_utility.h

@@ -0,0 +1,73 @@
+/*
+ * 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/parameter_transform.h>
+#include <momentum/character/skeleton_state.h>
+
+#include <ceres/jet.h>
+#include <Eigen/Core>
+
+namespace momentum {
+
+template <typename T, typename T2>
+Eigen::Vector3<T>
+getRotationDerivative(const JointStateT<T2>& js, const size_t index, const Eigen::Vector3<T>& ref) {
+  return js.rotationAxis.col(index).cross(ref);
+}
+
+template <typename T, typename T2>
+Eigen::Vector3<T> getScaleDerivative(const JointStateT<T2>& js, const Eigen::Vector3<T>& ref) {
+  (void)js;
+  return ref * ln2<T2>();
+}
+
+template <typename T, int N>
+Eigen::Matrix<ceres::Jet<T, N>, N, 1> buildJetVec(const Eigen::Matrix<T, N, 1>& v) {
+  Eigen::Matrix<ceres::Jet<T, N>, N, 1> result;
+  for (int k = 0; k < N; ++k) {
+    result(k).a = v(k);
+    result(k).v.setZero();
+    result(k).v(k) = 1;
+  }
+
+  return result;
+}
+
+template <typename T, int Rows, int Cols>
+Eigen::Matrix<ceres::Jet<T, Cols * Rows>, Rows, Cols> buildJetMat(
+    const Eigen::Matrix<T, Rows, Cols>& m) {
+  Eigen::Matrix<ceres::Jet<T, Cols * Rows>, Rows, Cols> result;
+  for (int iRow = 0; iRow < Rows; ++iRow) {
+    for (int jCol = 0; jCol < Cols; ++jCol) {
+      result(iRow, jCol).a = m(iRow, jCol);
+      result(iRow, jCol).v.setZero();
+      result(iRow, jCol).v(iRow * Cols + jCol) = 1;
+    }
+  }
+  return result;
+}
+
+template <typename T, typename T2>
+T times_parameterTransform_times_v(
+    const T& value,
+    const size_t jointParamIdx,
+    const ParameterTransform& parameterTransform,
+    Eigen::Ref<const Eigen::VectorX<T2>> v) {
+  T result = T();
+  for (auto index = parameterTransform.transform.outerIndexPtr()[jointParamIdx];
+       index < parameterTransform.transform.outerIndexPtr()[jointParamIdx + 1];
+       ++index) {
+    result += value * static_cast<T2>(parameterTransform.transform.valuePtr()[index]) *
+        v[parameterTransform.transform.innerIndexPtr()[index]];
+  }
+
+  return result;
+}
+
+} // namespace momentum

include/momentum/diff_ik/fully_differentiable_body_ik.h

@@ -0,0 +1,58 @@
+/*
+ * 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_solver/skeleton_error_function.h>
+#include <momentum/character_solver/skeleton_solver_function.h>
+#include <momentum/diff_ik/fwd.h>
+
+#include <unordered_map>
+#include <vector>
+
+namespace momentum {
+
+// This is going to be used within Pytorch's automatic differentiation package,
+// so we will assume the existence of a scalar "loss" function that is the thing
+// we are minimizing in Pytorch.  To clarify:
+//   "Error"/Error function: the thing that momentum's IK solver is minimizing
+//      to compute the body pose.
+//   "Loss"/Loss function: the thing that Pytorch is minimizing.
+// Here, we will be computing the derivative of the loss function with respect
+// to all the various error function inputs, like "weight" or "target".
+
+template <typename T>
+struct ErrorFunctionDerivativesT {
+  std::shared_ptr<const SkeletonErrorFunctionT<T>> errorFunction = nullptr;
+
+  // Gradient of the energy wrt the global weight.
+  T gradWeight = 0;
+
+  // Derivative of the energy wrt the various inputs.
+  std::unordered_map<std::string, Eigen::VectorX<T>> gradInputs;
+};
+
+// Takes in the derivative of the loss wrt the model parameters and
+// computes the derivative of the loss wrt the error function inputs,
+// assuming the model parameters are always the ones that (locally)
+// minimize the momentum IK error.
+//
+// The gradient_norm here can be used to determine if the derivatives
+// are valid; the computation only really works when the gradient
+// (wrt the subset of parameters that are allowed to vary) is close to 0.
+template <typename T>
+std::vector<ErrorFunctionDerivativesT<T>> d_modelParams_d_inputs(
+    const Skeleton& skeleton,
+    const ParameterTransformT<T>& parameterTransform,
+    const ParameterSet& activeParams,
+    const ModelParametersT<T>& parameters,
+    SkeletonSolverFunctionT<T>& solverFunction,
+    Eigen::Ref<const Eigen::VectorX<T>> dLoss_dModelParams,
+    T* gradientRmse = nullptr);
+
+} // namespace momentum

include/momentum/diff_ik/fully_differentiable_distance_error_function.h

@@ -0,0 +1,69 @@
+/*
+ * 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/distance_error_function.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+#include <momentum/diff_ik/fully_differentiable_skeleton_error_function.h>
+#include <momentum/diff_ik/fwd.h>
+
+namespace momentum {
+
+MOMENTUM_FWD_DECLARE_TEMPLATE_CLASS(DistanceErrorFunction);
+MOMENTUM_FWD_DECLARE_TEMPLATE_STRUCT(DistanceConstraintData);
+
+template <typename T>
+class FullyDifferentiableDistanceErrorFunctionT
+    : public DistanceErrorFunctionT<T>,
+      public virtual momentum::FullyDifferentiableSkeletonErrorFunctionT<T> {
+ public:
+  static constexpr const char* kOrigins = "origins";
+  static constexpr const char* kParents = "parents";
+  static constexpr const char* kOffsets = "offsets";
+  static constexpr const char* kWeights = "weights";
+  static constexpr const char* kTargets = "targets";
+
+  FullyDifferentiableDistanceErrorFunctionT(
+      const momentum::Skeleton& skel,
+      const momentum::ParameterTransform& pt);
+
+  [[nodiscard]] std::vector<std::string> inputs() const override;
+  [[nodiscard]] Eigen::Index getInputSize(const std::string& name) const override;
+  void getInputImp(const std::string& name, Eigen::Ref<Eigen::VectorX<T>> result) const override;
+  void setInputImp(const std::string& name, Eigen::Ref<const Eigen::VectorX<T>> value) override;
+  [[nodiscard]] const char* name() const override {
+    return "DistanceErrorFunction";
+  }
+
+  Eigen::VectorX<T> d_gradient_d_input_dot(
+      const std::string& inputName,
+      const momentum::ModelParametersT<T>& modelParams,
+      const momentum::SkeletonStateT<T>& state,
+      Eigen::Ref<const Eigen::VectorX<T>> inputVec) override;
+
+  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
+
+ private:
+  // Compute dError/dModelParams^T * v for some v.
+  template <typename JetType>
+  JetType constraintGradient_dot(
+      const momentum::ModelParametersT<T>& modelParams,
+      const momentum::SkeletonStateT<T>& skelState,
+      const Eigen::Vector3<JetType>& origin_cons,
+      int parentJointIndex,
+      const Eigen::Vector3<JetType>& offset,
+      const JetType& weight_cons,
+      const JetType& target_cons,
+      Eigen::Ref<const Eigen::VectorX<T>> vec);
+
+  // weights for the error functions
+  using DistanceErrorFunctionT<T>::kDistanceWeight;
+  using DistanceErrorFunctionT<T>::constraints_;
+};
+
+} // namespace momentum

include/momentum/diff_ik/fully_differentiable_motion_error_function.h

@@ -0,0 +1,46 @@
+/*
+ * 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_solver/model_parameters_error_function.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+#include <momentum/diff_ik/fully_differentiable_skeleton_error_function.h>
+#include <momentum/diff_ik/fwd.h>
+
+#include <Eigen/Core>
+
+namespace momentum {
+
+template <typename T>
+class FullyDifferentiableMotionErrorFunctionT : public FullyDifferentiableSkeletonErrorFunctionT<T>,
+                                                public ModelParametersErrorFunctionT<T> {
+ public:
+  FullyDifferentiableMotionErrorFunctionT(const Skeleton& skel, const ParameterTransform& pt);
+  ~FullyDifferentiableMotionErrorFunctionT() override;
+
+  [[nodiscard]] std::vector<std::string> inputs() const final;
+  [[nodiscard]] Eigen::Index getInputSize(const std::string& name) const final;
+  [[nodiscard]] const char* name() const final {
+    return "ModelParametersErrorFunction";
+  }
+
+  static constexpr const char* kTargetParameters = "target_parameters";
+  static constexpr const char* kTargetWeights = "target_weights";
+
+  Eigen::VectorX<T> d_gradient_d_input_dot(
+      const std::string& inputName,
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      Eigen::Ref<const Eigen::VectorX<T>> inputVec) final;
+
+ protected:
+  void getInputImp(const std::string& name, Eigen::Ref<Eigen::VectorX<T>> value) const final;
+  void setInputImp(const std::string& name, Eigen::Ref<const Eigen::VectorX<T>> value) final;
+};
+
+} // namespace momentum

include/momentum/diff_ik/fully_differentiable_orientation_error_function.h

@@ -0,0 +1,109 @@
+/*
+ * 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_solver/skeleton_error_function.h>
+#include <momentum/diff_ik/fully_differentiable_skeleton_error_function.h>
+#include <momentum/diff_ik/fwd.h>
+
+#include <vector>
+
+namespace momentum {
+
+template <typename T>
+struct OrientationConstraintT {
+  Quaternion<T> offset; ///< Constraint rotation offset in local joint space
+  Quaternion<T> target; ///< Target orientation in global space
+  size_t parent; ///< Parent joint index this constraint is under
+  float weight; ///< Weight of the constraint
+  std::string name; ///< Name of the constraint
+
+  explicit OrientationConstraintT(
+      const Quaternion<T>& offset,
+      const Quaternion<T>& target,
+      size_t parent,
+      float w,
+      const std::string& n = "")
+      : offset(offset), target(target), parent(parent), weight(w), name(n) {
+    // Empty
+  }
+};
+
+template <typename T>
+class FullyDifferentiableOrientationErrorFunctionT
+    : public FullyDifferentiableSkeletonErrorFunctionT<T>,
+      public SkeletonErrorFunctionT<T> {
+ public:
+  static constexpr T kOrientationWeight = 1e-1f;
+  static constexpr const char* kParents = "parents";
+  static constexpr const char* kOffsets = "offsets";
+  static constexpr const char* kWeights = "weights";
+  static constexpr const char* kTargets = "targets";
+
+  FullyDifferentiableOrientationErrorFunctionT(const Skeleton& skel, const ParameterTransform& pt);
+  ~FullyDifferentiableOrientationErrorFunctionT() override = default;
+
+  [[nodiscard]] std::vector<std::string> inputs() const final;
+  [[nodiscard]] Eigen::Index getInputSize(const std::string& name) const final;
+  [[nodiscard]] const char* name() const final {
+    return "OrientationErrorFunction";
+  }
+
+  Eigen::VectorX<T> d_gradient_d_input_dot(
+      const std::string& inputName,
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      Eigen::Ref<const Eigen::VectorX<T>> inputVec) final;
+
+  double getError(const ModelParametersT<T>& params, const SkeletonStateT<T>& state) final;
+
+  double getGradient(
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      Ref<Eigen::VectorX<T>> gradient) final;
+
+  double getJacobian(
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      Ref<Eigen::MatrixX<T>> jacobian,
+      Ref<Eigen::VectorX<T>> residual,
+      int& usedRows) final;
+
+  [[nodiscard]] size_t getJacobianSize() const final;
+
+  void addConstraint(const OrientationConstraintT<T>& constr);
+
+ private:
+  void getInputImp(const std::string& name, Eigen::Ref<Eigen::VectorX<T>> value) const final;
+  void setInputImp(const std::string& name, Eigen::Ref<const Eigen::VectorX<T>> value) final;
+
+  T calculateOrientationGradient(
+      const SkeletonStateT<T>& state,
+      const OrientationConstraintT<T>& constr,
+      Eigen::VectorX<T>& jGrad) const;
+
+  template <typename JetType>
+  JetType calculateOrientationGradient_dot(
+      const SkeletonStateT<T>& state,
+      size_t constrParent,
+      const JetType& constrWeight,
+      const Eigen::Quaternion<JetType>& constrOrientationOffset,
+      const Eigen::Quaternion<JetType>& constrOrientationTarget,
+      Eigen::Ref<const Eigen::VectorX<T>> vec) const;
+
+  [[nodiscard]] T calculateOrientationJacobian(
+      const SkeletonStateT<T>& state,
+      const OrientationConstraintT<T>& constr,
+      Ref<Eigen::MatrixX<T>> jac,
+      Ref<Eigen::VectorX<T>> res) const;
+
+  std::vector<OrientationConstraintT<T>> constraints_;
+  Eigen::VectorX<T> jointGrad_;
+};
+
+} // namespace momentum

include/momentum/diff_ik/fully_differentiable_pose_prior_error_function.h

@@ -0,0 +1,76 @@
+/*
+ * 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_solver/pose_prior_error_function.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+#include <momentum/diff_ik/fully_differentiable_skeleton_error_function.h>
+#include <momentum/diff_ik/fwd.h>
+
+#include <Eigen/Core>
+
+namespace momentum {
+
+template <typename T>
+class FullyDifferentiablePosePriorErrorFunctionT
+    : public FullyDifferentiableSkeletonErrorFunctionT<T>,
+      public PosePriorErrorFunctionT<T> {
+ public:
+  FullyDifferentiablePosePriorErrorFunctionT(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      std::vector<std::string> names);
+  ~FullyDifferentiablePosePriorErrorFunctionT() override;
+
+  void setPosePrior(
+      const VectorX<T>& pi,
+      const MatrixX<T>& mmu,
+      const std::vector<MatrixX<T>>& W,
+      const VectorX<T>& sigma);
+
+  static constexpr const char* kPi = "pi";
+  static constexpr const char* kMu = "mu";
+  static constexpr const char* kW = "W";
+  static constexpr const char* kSigma = "sigma";
+  static constexpr const char* kParameterIndices = "parameter_indices";
+
+  [[nodiscard]] std::vector<std::string> inputs() const final;
+  [[nodiscard]] Eigen::Index getInputSize(const std::string& name) const final;
+  [[nodiscard]] const char* name() const final {
+    return "MarkerErrorFunction";
+  }
+
+  Eigen::VectorX<T> d_gradient_d_input_dot(
+      const std::string& inputName,
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      Eigen::Ref<const Eigen::VectorX<T>> inputVec) final;
+
+ protected:
+  void getInputImp(const std::string& name, Eigen::Ref<Eigen::VectorX<T>> value) const final;
+  void setInputImp(const std::string& name, Eigen::Ref<const Eigen::VectorX<T>> value) final;
+
+  // pi_ are the mixture weights
+  Eigen::VectorX<T> pi_;
+
+  // The columns of mu are the means for the different mixtures
+  Eigen::MatrixX<T> mmu_;
+
+  // Each matrix W is the basis for a mixture:
+  std::vector<Eigen::MatrixX<T>> W_;
+
+  // The covariance matrix is (sigma^2 * I + W^T * W)^{-1}
+  // Note that this is sigma, not sigma squared, because the scaling should
+  // be more reasonable (sigma is the same order of magnitude as the mu and
+  // W arrays).
+  Eigen::VectorX<T> sigma_;
+
+  const std::vector<std::string> names_;
+};
+
+} // namespace momentum