pymomentum-cpu 0.1.93.post0__cp312-cp312-macosx_14_0_arm64.whl

include/momentum/character_solver/gauss_newton_solver_qr.h

@@ -0,0 +1,64 @@
+/*
+ * 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/character_solver/fwd.h>
+#include <momentum/character_solver/skeleton_solver_function.h>
+#include <momentum/math/online_householder_qr.h>
+#include <momentum/solver/solver.h>
+
+#include <Eigen/Core>
+
+namespace momentum {
+
+/// Gauss-Newton solver with QR decomposition specific options
+struct GaussNewtonSolverQROptions : SolverOptions {
+  /// Damping parameter added to Hessian diagonal for numerical stability; see
+  /// https://en.wikipedia.org/wiki/Levenberg%E2%80%93Marquardt_algorithm
+  float regularization = 0.05f;
+
+  /// Flag to enable line search during optimization.
+  bool doLineSearch = false;
+
+  GaussNewtonSolverQROptions() = default;
+
+  /* implicit */ GaussNewtonSolverQROptions(const SolverOptions& baseOptions)
+      : SolverOptions(baseOptions) {
+    // Empty
+  }
+};
+
+template <typename T>
+class GaussNewtonSolverQRT : public SolverT<T> {
+ public:
+  GaussNewtonSolverQRT(const SolverOptions& options, SkeletonSolverFunctionT<T>* solver);
+  ~GaussNewtonSolverQRT() override;
+
+  [[nodiscard]] std::string_view getName() const override;
+
+  void setOptions(const SolverOptions& options) final;
+
+ protected:
+  void doIteration() final;
+  void initializeSolver() final;
+
+ private:
+  std::unique_ptr<SkeletonStateT<T>> skeletonState_;
+  std::unique_ptr<MeshStateT<T>> meshState_;
+
+  ResizeableMatrix<T> jacobian_;
+  ResizeableMatrix<T> residual_;
+
+  OnlineHouseholderQR<T> qrSolver_;
+
+  float regularization_;
+  bool doLineSearch_;
+};
+
+} // namespace momentum

include/momentum/character_solver/height_error_function.h

@@ -0,0 +1,176 @@
+/*
+ * 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/mesh_state.h>
+#include <momentum/character_solver/fwd.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+#include <momentum/math/fwd.h>
+
+#include <Eigen/Dense>
+
+#include <queue>
+#include <vector>
+
+namespace momentum {
+
+/// Error function for character height constraints
+///
+/// This error function measures the height of a character mesh by projecting
+/// all vertices onto a specified "up" direction and computing the difference
+/// between the maximum and minimum projections.
+///
+/// IMPORTANT: Unlike most error functions, this one only depends on specific
+/// active parameters that are automatically determined in the constructor.
+/// Specifically, it uses blend shape, face expression, and scale parameters only
+/// (not pose parameters). The error function achieves this by zeroing out all
+/// inactive parameters before computing the skeleton and mesh state. This allows
+/// us to constrain height independent of pose, e.g. the solver doesn't try to fix
+/// height by having the character bend over.
+template <typename T>
+class HeightErrorFunctionT : public SkeletonErrorFunctionT<T> {
+ public:
+  /// Construct a height error function
+  /// @param character The character to measure height for
+  /// @param targetHeight The target height for the character (required)
+  /// @param upDirection The direction to measure height along (defaults to Y-axis)
+  /// @param k Number of vertices to average for min/max height calculation (defaults to 1)
+  ///
+  /// This error function automatically uses blend shape, face expression, and
+  /// scale parameters only. All other parameters (pose, etc.) are kept inactive.
+  explicit HeightErrorFunctionT(
+      const Character& character,
+      T targetHeight,
+      const Eigen::Vector3<T>& upDirection = Eigen::Vector3<T>::UnitY(),
+      size_t k = 1);
+  ~HeightErrorFunctionT() override;
+
+  HeightErrorFunctionT(const HeightErrorFunctionT& other) = delete;
+  HeightErrorFunctionT(HeightErrorFunctionT&& other) noexcept = delete;
+  HeightErrorFunctionT& operator=(const HeightErrorFunctionT& other) = delete;
+  HeightErrorFunctionT& operator=(HeightErrorFunctionT&& other) = delete;
+
+  [[nodiscard]] double getError(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState) final;
+
+  double getGradient(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) final;
+
+  double getJacobian(
+      const ModelParametersT<T>& modelParameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      Eigen::Ref<Eigen::MatrixX<T>> jacobian,
+      Eigen::Ref<Eigen::VectorX<T>> residual,
+      int& usedRows) final;
+
+  [[nodiscard]] size_t getJacobianSize() const final;
+
+  /// Set the target height
+  /// @param height The desired height
+  void setTargetHeight(T height);
+
+  /// Get the target height
+  [[nodiscard]] T getTargetHeight() const {
+    return targetHeight_;
+  }
+
+  /// Set the up direction for height measurement
+  /// @param upDirection The direction to measure height along (will be normalized)
+  void setUpDirection(const Eigen::Vector3<T>& upDirection);
+
+  /// Get the up direction
+  [[nodiscard]] const Eigen::Vector3<T>& getUpDirection() const {
+    return upDirection_;
+  }
+
+  /// Get the character
+  [[nodiscard]] const Character* getCharacter() const override {
+    return &character_;
+  }
+
+  /// Override to indicate this function does NOT require mesh state
+  /// (we maintain our own meshes)
+  [[nodiscard]] bool needsMesh() const override {
+    return false;
+  }
+
+ private:
+  /// Create a copy of model parameters with inactive parameters zeroed out
+  /// @param modelParameters The input parameters
+  /// @return Modified parameters with only active parameters
+  [[nodiscard]] ModelParametersT<T> applyActiveParameters(
+      const ModelParametersT<T>& modelParameters) const;
+
+  /// Result of height calculation
+  struct HeightResult {
+    T height{};
+    std::vector<size_t> minVertexIndices;
+    std::vector<T> minVertexWeights;
+    std::vector<size_t> maxVertexIndices;
+    std::vector<T> maxVertexWeights;
+  };
+
+  /// Calculate the current height of the mesh
+  /// @return Height result containing the height and the min/max vertex indices
+  [[nodiscard]] HeightResult calculateHeight() const;
+
+  /// Calculate jacobian contribution from a vertex
+  template <typename Derived>
+  void calculateVertexJacobian(
+      size_t vertexIndex,
+      const Eigen::Vector3<T>& jacobianDirection,
+      Eigen::MatrixBase<Derived>& jacobian) const;
+
+  /// Calculate gradient contribution from a vertex
+  void calculateVertexGradient(
+      size_t vertexIndex,
+      const Eigen::Vector3<T>& gradientDirection,
+      Eigen::Ref<Eigen::VectorX<T>> gradient) const;
+
+  /// Calculate world space position derivative for blend shape parameters
+  void calculateDWorldPos(
+      size_t vertexIndex,
+      const Eigen::Vector3<T>& d_restPos,
+      Eigen::Vector3<T>& d_worldPos) const;
+
+  const Character& character_;
+
+  T targetHeight_;
+  Eigen::Vector3<T> upDirection_;
+  size_t k_;
+
+  /// The set of active model parameters (blend shapes, face expressions, and scale)
+  ///
+  /// This is automatically set up in the constructor and determines which parameters
+  /// affect the height measurement. Inactive parameters are zeroed out before
+  /// computing the skeleton and mesh state.
+  ParameterSet activeModelParams_;
+
+  /// Internal skeleton state for skinning (computed from active parameters only)
+  ///
+  /// This is updated in getError/getGradient/getJacobian by zeroing out all
+  /// inactive parameters from the input model parameters. This ensures derivatives
+  /// w.r.t. inactive parameters are automatically zero.
+  mutable SkeletonStateT<T> skeletonState_;
+
+  /// Internal mesh state (updated using skeletonState_ and active parameters only)
+  ///
+  /// We maintain our own mesh state because we use skeletonState_ (computed from
+  /// active parameters only) rather than the state parameter passed to
+  /// getError/getGradient/getJacobian. This is marked mutable so we can update
+  /// it in const methods like getError.
+  mutable MeshStateT<T> meshState_;
+};
+
+} // namespace momentum

include/momentum/character_solver/joint_to_joint_distance_error_function.h

@@ -0,0 +1,111 @@
+/*
+ * 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/fwd.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+#include <momentum/math/types.h>
+
+namespace momentum {
+
+/// Constraint that enforces a distance relationship between two points attached to different
+/// joints.
+///
+/// Each point is specified by a joint index and an offset in the local coordinate system of that
+/// joint.
+template <typename T>
+struct JointToJointDistanceConstraintT {
+  /// Index of the first joint.
+  size_t joint1 = kInvalidIndex;
+
+  /// Offset from joint1 in the local coordinate system of joint1.
+  Vector3<T> offset1 = Vector3<T>::Zero();
+
+  /// Index of the second joint.
+  size_t joint2 = kInvalidIndex;
+
+  /// Offset from joint2 in the local coordinate system of joint2.
+  Vector3<T> offset2 = Vector3<T>::Zero();
+
+  /// Target distance between the two points (in world space).
+  T targetDistance = T(0);
+
+  /// Weight for this constraint.
+  T weight = T(1);
+
+  template <typename T2>
+  JointToJointDistanceConstraintT<T2> cast() const {
+    return {
+        this->joint1,
+        this->offset1.template cast<T2>(),
+        this->joint2,
+        this->offset2.template cast<T2>(),
+        static_cast<T2>(this->targetDistance),
+        static_cast<T2>(this->weight)};
+  }
+};
+
+/// 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.
+template <typename T>
+class JointToJointDistanceErrorFunctionT : public SkeletonErrorFunctionT<T> {
+ public:
+  explicit JointToJointDistanceErrorFunctionT(const Skeleton& skel, const ParameterTransform& pt);
+
+  explicit JointToJointDistanceErrorFunctionT(const Character& character);
+
+  [[nodiscard]] double getError(
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState) final;
+
+  double getGradient(
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      Ref<VectorX<T>> gradient) override;
+
+  double getJacobian(
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      Ref<MatrixX<T>> jacobian,
+      Ref<VectorX<T>> residual,
+      int& usedRows) override;
+
+  [[nodiscard]] size_t getJacobianSize() const final;
+
+  /// Add a constraint between two joints with local offsets.
+  void addConstraint(
+      size_t joint1,
+      const Vector3<T>& offset1,
+      size_t joint2,
+      const Vector3<T>& offset2,
+      T targetDistance,
+      T weight = T(1));
+
+  /// Clear all constraints.
+  void clearConstraints();
+
+  /// Get all constraints.
+  [[nodiscard]] const std::vector<JointToJointDistanceConstraintT<T>>& getConstraints() const {
+    return constraints_;
+  }
+
+  /// Default weight for distance constraints.
+  static constexpr T kDistanceWeight = 1e-2f;
+
+ private:
+  std::vector<JointToJointDistanceConstraintT<T>> constraints_;
+};
+
+} // namespace momentum

include/momentum/character_solver/limit_error_function.h

@@ -0,0 +1,57 @@
+/*
+ * 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_limits.h>
+#include <momentum/character_solver/fwd.h>
+#include <momentum/character_solver/limit_error_function.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+
+namespace momentum {
+
+template <typename T>
+class LimitErrorFunctionT : public SkeletonErrorFunctionT<T> {
+ public:
+  LimitErrorFunctionT(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      const ParameterLimits& pl = ParameterLimits());
+  explicit LimitErrorFunctionT(const Character& character);
+  LimitErrorFunctionT(const Character& character, const ParameterLimits& pl);
+
+  [[nodiscard]] double getError(
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState) final;
+
+  double getGradient(
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      Ref<Eigen::VectorX<T>> gradient) final;
+
+  double getJacobian(
+      const ModelParametersT<T>& params,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      Ref<Eigen::MatrixX<T>> jacobian,
+      Ref<Eigen::VectorX<T>> residual,
+      int& usedRows) final;
+
+  [[nodiscard]] size_t getJacobianSize() const final;
+
+  void setLimits(const ParameterLimits& lm);
+  void setLimits(const Character& character);
+
+ private:
+  // weights for the error functions
+  static constexpr float kLimitWeight = 1e+1;
+  ParameterLimits limits_;
+};
+
+} // namespace momentum

include/momentum/character_solver/model_parameters_error_function.h

@@ -0,0 +1,64 @@
+/*
+ * 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/character_solver/fwd.h>
+#include <momentum/character_solver/skeleton_error_function.h>
+
+namespace momentum {
+
+template <typename T>
+class ModelParametersErrorFunctionT : public SkeletonErrorFunctionT<T> {
+ public:
+  ModelParametersErrorFunctionT(const Skeleton& skel, const ParameterTransform& pt);
+  explicit ModelParametersErrorFunctionT(const Character& character);
+
+  // Create a ModelParametersError that only targets the specified parameters:
+  ModelParametersErrorFunctionT(const Character& character, const ParameterSet& active);
+
+  [[nodiscard]] double getError(
+      const ModelParametersT<T>& parameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState) final;
+  double getGradient(
+      const ModelParametersT<T>& parameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      Ref<Eigen::VectorX<T>> gradient) final;
+  double getJacobian(
+      const ModelParametersT<T>& parameters,
+      const SkeletonStateT<T>& state,
+      const MeshStateT<T>& meshState,
+      Ref<Eigen::MatrixX<T>> jacobian,
+      Ref<Eigen::VectorX<T>> residual,
+      int& usedRows) final;
+  [[nodiscard]] size_t getJacobianSize() const final;
+
+  void setTargetParameters(const ModelParametersT<T>& params, const Eigen::VectorX<T>& weights) {
+    this->targetParameters_ = params;
+    this->targetWeights_ = weights;
+  }
+
+  [[nodiscard]] const ModelParametersT<T>& getTargetParameters() const {
+    return this->targetParameters_;
+  }
+
+  [[nodiscard]] const Eigen::VectorX<T>& getTargetWeights() const {
+    return this->targetWeights_;
+  }
+
+ private:
+  ModelParametersT<T> targetParameters_;
+  Eigen::VectorX<T> targetWeights_;
+
+ public:
+  static constexpr T kMotionWeight = 1e-1;
+};
+
+} // namespace momentum

include/momentum/character_solver/normal_error_function.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_solver/constraint_error_function.h>
+
+namespace momentum {
+
+/// Normal constraint data: data needed to compute the point-to-plane distance for
+/// "data-to-template" in ICP.
+template <typename T>
+struct NormalDataT : ConstraintData {
+  /// The template point defined in the local space
+  Vector3<T> localPoint;
+  /// The template normal defined in the local space
+  Vector3<T> localNormal;
+  /// The data point in world space
+  Vector3<T> globalPoint;
+
+  explicit NormalDataT(
+      const Vector3<T>& inLocalPt,
+      const Vector3<T>& inLocalNorm,
+      const Vector3<T>& inTargetPt,
+      size_t pIndex,
+      float w,
+      const std::string& n = {})
+      : ConstraintData(pIndex, w, n),
+        localPoint(inLocalPt),
+        localNormal(inLocalNorm.normalized()),
+        globalPoint(inTargetPt) {}
+};
+
+/// The NormalErrorFunction computes a "point-to-plane" (signed) distance from a target point to the
+/// plane defined by a local point and a local normal vector. The name is not informative but chosen
+/// for backwards compatibility purpose.
+template <typename T>
+class NormalErrorFunctionT : public ConstraintErrorFunctionT<T, NormalDataT<T>, 1, 2, 1> {
+ public:
+  explicit NormalErrorFunctionT(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : ConstraintErrorFunctionT<T, NormalDataT<T>, 1, 2, 1>(skel, pt, lossAlpha, lossC) {}
+
+  explicit NormalErrorFunctionT(
+      const Character& character,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : NormalErrorFunctionT<T>(
+            character.skeleton,
+            character.parameterTransform,
+            lossAlpha,
+            lossC) {}
+
+  /// Default constant weight in MarkerErrorFunction. This can be used for backwards compatibility
+  /// in setWeight().
+  static constexpr T kLegacyWeight = 1e-4f;
+
+ protected:
+  void evalFunction(
+      size_t constrIndex,
+      const JointStateT<T>& state,
+      Vector<T, 1>& f,
+      optional_ref<std::array<Vector3<T>, 2>> v = {},
+      optional_ref<std::array<Eigen::Matrix<T, 1, 3>, 2>> dfdv = {}) const final;
+};
+} // namespace momentum

include/momentum/character_solver/orientation_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_solver/constraint_error_function.h>
+
+namespace momentum {
+
+/// Constraint data on 3x3 rotation represented in quaternions
+template <typename T>
+struct OrientationDataT : ConstraintData {
+  /// Rotation offset in the parent space
+  Eigen::Quaternion<T> offset;
+  /// Target rotation in global space
+  Eigen::Quaternion<T> target;
+
+  /// param[in] inOffset: input rotation offset in the parent space
+  /// param[in] inTarget: input target rotation in global space
+  /// param[in] pIndex: joint index of the parent joint
+  /// param[in] w: weight
+  /// param[in] n: name of this constraint
+  explicit OrientationDataT(
+      const Eigen::Quaternion<T>& inOffset,
+      const Eigen::Quaternion<T>& inTarget,
+      size_t pIndex,
+      float w,
+      const std::string& n = "")
+      : ConstraintData(pIndex, w, n),
+        offset(inOffset.normalized()),
+        target(inTarget.normalized()) {}
+};
+
+/// The OrientationErrorFunction minimizes the F-norm of the element-wise difference of a 3x3
+/// rotation matrix to a target rotation. Therefore, function f has 9 numbers (FuncDim), and
+/// NumVec=3 for the 3 rotation axis.
+template <typename T>
+class OrientationErrorFunctionT : public ConstraintErrorFunctionT<T, OrientationDataT<T>, 9, 3, 0> {
+ public:
+  explicit OrientationErrorFunctionT(
+      const Skeleton& skel,
+      const ParameterTransform& pt,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : ConstraintErrorFunctionT<T, OrientationDataT<T>, 9, 3, 0>(skel, pt, lossAlpha, lossC) {}
+
+  explicit OrientationErrorFunctionT(
+      const Character& character,
+      const T& lossAlpha = GeneralizedLossT<T>::kL2,
+      const T& lossC = T(1))
+      : OrientationErrorFunctionT(
+            character.skeleton,
+            character.parameterTransform,
+            lossAlpha,
+            lossC) {}
+
+  /// Default constant weight in MarkerErrorFunction. This can be used for backwards compatibility
+  /// in setWeight().
+  static constexpr T kLegacyWeight = 1e-1f;
+
+ protected:
+  void evalFunction(
+      size_t constrIndex,
+      const JointStateT<T>& state,
+      Vector<T, 9>& f,
+      optional_ref<std::array<Vector3<T>, 3>> v = {},
+      optional_ref<std::array<Eigen::Matrix<T, 9, 3>, 3>> dfdv = {}) const final;
+};
+
+} // namespace momentum