Skip to content

Commit 4725c51

Browse files
cdtwiggmeta-codesync[bot]
authored andcommitted
Bind JointToJointSequenceErrorFunction to Python (#1528)
Summary: Pull Request resolved: #1528 Split out from the HOI relative-transform work (was bundled in D105769542). Adds the pybind11 binding for `JointToJointSequenceErrorFunction` in `solver2_sequence_error_functions.cpp`, with `add_constraint(source_joint, reference_joint, ...)`, regenerated `solver2.pyi` stubs, and a `test_solver2` case. Reviewed By: yutingye Differential Revision: D108807345 fbshipit-source-id: 1dac00257faae366d69a737a094b1c5c4a44f79a
1 parent 85f3e01 commit 4725c51

2 files changed

Lines changed: 170 additions & 0 deletions

File tree

‎pymomentum/solver2/solver2_sequence_error_functions.cpp‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
#include <momentum/character/skeleton_state.h>
1111
#include <momentum/character_sequence_solver/acceleration_sequence_error_function.h>
1212
#include <momentum/character_sequence_solver/jerk_sequence_error_function.h>
13+
#include <momentum/character_sequence_solver/joint_to_joint_sequence_error_function.h>
1314
#include <momentum/character_sequence_solver/model_parameters_sequence_error_function.h>
1415
#include <momentum/character_sequence_solver/sequence_solver.h>
1516
#include <momentum/character_sequence_solver/state_sequence_error_function.h>
@@ -21,6 +22,7 @@
2122
#include <pybind11/stl.h>
2223
#include <Eigen/Core>
2324

25+
#include <pymomentum/python_utility/eigen_quaternion.h>
2426
#include <pymomentum/solver2/solver2_sequence_error_functions.h>
2527
#include <pymomentum/solver2/solver2_utility.h>
2628

@@ -272,6 +274,92 @@ the gradient contribution is set to zero.
272274
&mm::VelocityMagnitudeSequenceErrorFunction::getTargetSpeeds,
273275
"Returns the per-joint target speeds.");
274276

277+
py::class_<
278+
mm::JointToJointSequenceErrorFunction,
279+
mm::SequenceErrorFunction,
280+
std::shared_ptr<mm::JointToJointSequenceErrorFunction>>(
281+
m, "JointToJointSequenceErrorFunction")
282+
.def(
283+
py::init<>([](const mm::Character& character,
284+
float weight,
285+
float positionWeight,
286+
float orientationWeight) {
287+
validateWeight(weight, "weight");
288+
validateWeight(positionWeight, "position_weight");
289+
validateWeight(orientationWeight, "orientation_weight");
290+
291+
auto result = std::make_shared<mm::JointToJointSequenceErrorFunction>(character);
292+
result->setWeight(weight);
293+
result->setPositionWeight(positionWeight);
294+
result->setOrientationWeight(orientationWeight);
295+
296+
return result;
297+
}),
298+
R"(A sequence error function that penalizes frame-to-frame changes in the relative
299+
transform between two joints.
300+
301+
For each constraint with source joint A and reference joint B, this function computes
302+
the relative transform at two consecutive frames and penalizes any change between them:
303+
304+
- Position error: ``relPos_{i+1} - relPos_i``
305+
- Orientation error: ``relRot_{i+1} - relRot_i``
306+
307+
This keeps the relative transform (e.g. a hand's grip on an object) stable from frame
308+
to frame.
309+
310+
:param character: The character to use.
311+
:param weight: The weight of the error function. Defaults to 1.0.
312+
:param position_weight: The weight of the position error. Defaults to 1.0.
313+
:param orientation_weight: The weight of the orientation error. Defaults to 1.0.)",
314+
py::arg("character"),
315+
py::kw_only(),
316+
py::arg("weight") = 1.0f,
317+
py::arg("position_weight") = 1.0f,
318+
py::arg("orientation_weight") = 1.0f)
319+
.def(
320+
"add_constraint",
321+
[](mm::JointToJointSequenceErrorFunction& self,
322+
size_t sourceJoint,
323+
size_t referenceJoint,
324+
float constraintWeight,
325+
const std::string& name) {
326+
self.addConstraint(sourceJoint, referenceJoint, constraintWeight, name);
327+
},
328+
R"(Add a constraint between two joints.
329+
330+
:param source_joint: Index of the source joint (the joint whose relative pose is tracked).
331+
:param reference_joint: Index of the reference joint (the frame of reference).
332+
:param weight: Weight for this constraint. Defaults to 1.0.
333+
:param name: Name for debugging. Defaults to empty string.)",
334+
py::arg("source_joint"),
335+
py::arg("reference_joint"),
336+
py::kw_only(),
337+
py::arg("weight") = 1.0f,
338+
py::arg("name") = "")
339+
.def(
340+
"clear_constraints",
341+
&mm::JointToJointSequenceErrorFunction::clearConstraints,
342+
"Clears all constraints from the error function.")
343+
.def_property_readonly(
344+
"num_constraints",
345+
&mm::JointToJointSequenceErrorFunction::getNumConstraints,
346+
"Returns the number of constraints.")
347+
.def_property(
348+
"weight",
349+
&mm::JointToJointSequenceErrorFunction::getWeight,
350+
&mm::JointToJointSequenceErrorFunction::setWeight,
351+
"The weight applied to the error function.")
352+
.def_property(
353+
"position_weight",
354+
&mm::JointToJointSequenceErrorFunction::getPositionWeight,
355+
&mm::JointToJointSequenceErrorFunction::setPositionWeight,
356+
"The weight of the position error.")
357+
.def_property(
358+
"orientation_weight",
359+
&mm::JointToJointSequenceErrorFunction::getOrientationWeight,
360+
&mm::JointToJointSequenceErrorFunction::setOrientationWeight,
361+
"The weight of the orientation error.");
362+
275363
py::class_<
276364
mm::StateSequenceErrorFunction,
277365
mm::SequenceErrorFunction,

‎pymomentum/test/test_solver2.py‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3253,3 +3253,85 @@ def test_transform_pose_single(self) -> None:
32533253
np.testing.assert_allclose(
32543254
skel_out[:, :3], skel_init[:, :3] + transform[:3], atol=1e-2
32553255
)
3256+
3257+
def test_joint_to_joint_sequence_error_function(self) -> None:
3258+
"""Test JointToJointSequenceErrorFunction creation and basic behavior."""
3259+
character = pym_test_utils.create_test_character(num_joints=4)
3260+
n_params = character.parameter_transform.size
3261+
3262+
# Create error function
3263+
ef = pym_solver2.JointToJointSequenceErrorFunction(
3264+
character,
3265+
weight=1.0,
3266+
position_weight=2.0,
3267+
orientation_weight=0.5,
3268+
)
3269+
self.assertEqual(ef.weight, 1.0)
3270+
self.assertEqual(ef.position_weight, 2.0)
3271+
self.assertEqual(ef.orientation_weight, 0.5)
3272+
3273+
# Add constraints
3274+
ef.add_constraint(source_joint=2, reference_joint=0)
3275+
ef.add_constraint(
3276+
source_joint=1,
3277+
reference_joint=3,
3278+
weight=0.5,
3279+
name="test_constraint",
3280+
)
3281+
self.assertEqual(ef.num_constraints, 2)
3282+
3283+
# Same parameters on both frames: zero error
3284+
model_params_same = np.zeros((2, n_params), dtype=np.float32)
3285+
error = ef.get_error(model_params_same)
3286+
self.assertAlmostEqual(error, 0.0, places=6)
3287+
3288+
# Different parameters on both frames: non-zero error
3289+
rng = np.random.RandomState(42)
3290+
model_params_diff = rng.randn(2, n_params).astype(np.float32) * 0.1
3291+
error = ef.get_error(model_params_diff)
3292+
self.assertGreater(error, 0.0)
3293+
3294+
# Clear constraints
3295+
ef.clear_constraints()
3296+
self.assertEqual(ef.num_constraints, 0)
3297+
3298+
# After clearing, error should be zero
3299+
error = ef.get_error(model_params_diff)
3300+
self.assertAlmostEqual(error, 0.0, places=6)
3301+
3302+
def test_joint_to_joint_sequence_error_function_solve_reduces_error(
3303+
self,
3304+
) -> None:
3305+
"""Test solve_sequence() can reduce joint-to-joint sequence error."""
3306+
character = pym_test_utils.create_test_character(num_joints=4)
3307+
n_params = character.parameter_transform.size
3308+
n_frames = 2
3309+
3310+
rng = np.random.RandomState(7)
3311+
model_params_init = np.zeros((n_frames, n_params), dtype=np.float32)
3312+
model_params_init[1] = (rng.randn(n_params) * 0.25).astype(np.float32)
3313+
3314+
rel_error = pym_solver2.JointToJointSequenceErrorFunction(
3315+
character,
3316+
weight=10.0,
3317+
position_weight=10.0,
3318+
orientation_weight=10.0,
3319+
)
3320+
rel_error.add_constraint(source_joint=2, reference_joint=0)
3321+
initial_error = rel_error.get_error(model_params_init)
3322+
self.assertGreater(initial_error, 1e-4)
3323+
3324+
solver_function = pym_solver2.SequenceSolverFunction(character, n_frames)
3325+
solver_function.set_enabled_parameters(np.ones(n_params, dtype=bool))
3326+
solver_function.add_sequence_error_function(0, rel_error)
3327+
3328+
solver_options = pym_solver2.SequenceSolverOptions()
3329+
solver_options.max_iterations = 25
3330+
solver_options.regularization = 1e-6
3331+
3332+
model_params_final = pym_solver2.solve_sequence(
3333+
solver_function, model_params_init, solver_options
3334+
)
3335+
3336+
final_error = rel_error.get_error(model_params_final)
3337+
self.assertLess(final_error, initial_error * 1e-3)

0 commit comments

Comments
 (0)