Skip to content

[DRAFT] Add private FeatherPGS solver implementation#2331

Closed
dylanturpin wants to merge 1 commit intonewton-physics:mainfrom
dylanturpin:fpgs-private-api
Closed

[DRAFT] Add private FeatherPGS solver implementation#2331
dylanturpin wants to merge 1 commit intonewton-physics:mainfrom
dylanturpin:fpgs-private-api

Conversation

@dylanturpin
Copy link
Copy Markdown
Member

@dylanturpin dylanturpin commented Apr 6, 2026

Description

Work in progress. Draft PR for feather_pgs API in newton.

Checklist

  • New or existing tests cover these changes
  • The documentation is up to date with these changes
  • CHANGELOG.md has been updated (if user-facing change)

Test plan

Bug fix

Steps to reproduce:

Minimal reproduction:

import newton

# Code that demonstrates the bug

New feature / API change

import newton

# Code that demonstrates the new capability

Summary by CodeRabbit

  • Refactor

    • Improved velocity computation in the Featherstone solver for articulated body dynamics, consolidating multiple steps into a more efficient pipeline.
    • Enhanced handling of specific joint types with better velocity conventions during forward kinematics calculations.
  • Chores

    • Streamlined solver module structure for better code organization.

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai Bot commented Apr 6, 2026

📝 Walkthrough

Walkthrough

Two changes: (1) New feather_pgs/__init__.py module exports SolverFeatherPGS, (2) Refactored featherstone/kernels.py to consolidate FREE/DISTANCE velocity convention conversion into FK propagation, removing the standalone convert_articulation_free_distance_body_qd kernel and computing spatial velocities directly via rigid transforms.

Changes

Cohort / File(s) Summary
feather_pgs package initialization
newton/_src/solvers/feather_pgs/__init__.py
New module re-exports SolverFeatherPGS from solver_feather_pgs and defines __all__ to expose it as the public API symbol.
featherstone FK velocity refactoring
newton/_src/solvers/featherstone/kernels.py
Consolidated velocity convention conversion logic into eval_single_articulation_fk_with_velocity_conversion: removed velocity_at_point dependency; compute parent anchor-origin spatial velocity directly via rigid transforms; reworked child-body twist propagation using explicit twist transforms; integrated FREE/DISTANCE velocity conversion into FK path; adjusted COM writeback to convert origin-referenced twist to COM velocity; deleted standalone convert_articulation_free_distance_body_qd kernel.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

Suggested reviewers

  • mzamoramora-nvidia
  • eric-heiden
🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (2 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title accurately describes the main change: adding a private FeatherPGS solver implementation, which aligns with the new feather_pgs module and related updates.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@codecov
Copy link
Copy Markdown

codecov Bot commented Apr 6, 2026

❌ 6 Tests Failed:

Tests completed Failed Passed Skipped
1984 6 1978 392
View the top 3 failed test(s) by shortest run time
TestSimKinematics::test_solver_fk_prismatic_descendant_linear_velocity_matches_finite_difference_cpu
Stack Traces | 0.012s run time
Traceback (most recent call last):
  File ".../newton/tests/unittest_utils.py", line 282, in test_func
    func(self, device, **kwargs)
  File ".../newton/tests/test_kinematics.py", line 409, in test_solver_fk_prismatic_descendant_linear_velocity_matches_finite_difference
    assert_np_equal(origin_vel_fd, origin_vel_from_body_qd, tol=5.0e-3)
  File ".../newton/tests/unittest_utils.py", line 231, in assert_np_equal
    np.testing.assert_allclose(result.flatten(), expect.flatten(), atol=tol, equal_nan=True)
  File ".../newton/newton/.venv/lib/python3.12.../testing/_private/utils.py", line 1768, in assert_allclose
    assert_array_compare(compare, actual, desired, err_msg=str(err_msg),
  File ".../newton/newton/.venv/lib/python3.12.../testing/_private/utils.py", line 983, in assert_array_compare
    raise AssertionError(msg)
AssertionError: 
Not equal to tolerance rtol=1e-07, atol=0.005

Mismatched elements: 2 / 3 (66.7%)
Mismatch at indices:
 [0]: -1.214742660522461 (ACTUAL), -0.758348286151886 (DESIRED)
 [1]: 1.3178586959838867 (ACTUAL), 0.5672810077667236 (DESIRED)
Max absolute difference among violations: 0.7505777
Max relative difference among violations: 1.3231144
 ACTUAL: array([-1.214743e+00,  1.317859e+00,  5.960464e-04], dtype=float32)
 DESIRED: array([-0.758348,  0.567281,  0.      ], dtype=float32)
TestBodyForce::test_floating_body_angular_up_axis_Y_featherstone_cpu
Stack Traces | 0.014s run time
Traceback (most recent call last):
  File ".../newton/tests/unittest_utils.py", line 282, in test_func
    func(self, device, **kwargs)
  File ".../newton/tests/test_body_force.py", line 98, in test_floating_body
    test.assertAlmostEqual(body_qd[i], 0.0, delta=zero_velocity_tolerance)
AssertionError: np.float32(1.2000087) != 0.0 within 0.001 delta (np.float32(1.2000087) difference)
TestBodyForce::test_floating_body_angular_up_axis_Z_featherstone_cpu
Stack Traces | 0.018s run time
Traceback (most recent call last):
  File ".../newton/tests/unittest_utils.py", line 282, in test_func
    func(self, device, **kwargs)
  File ".../newton/tests/test_body_force.py", line 98, in test_floating_body
    test.assertAlmostEqual(body_qd[i], 0.0, delta=zero_velocity_tolerance)
AssertionError: np.float32(1.2000087) != 0.0 within 0.001 delta (np.float32(1.2000087) difference)
TestControlForce::test_floating_body_angular_featherstone_cpu
Stack Traces | 0.025s run time
Traceback (most recent call last):
  File ".../newton/tests/unittest_utils.py", line 282, in test_func
    func(self, device, **kwargs)
  File ".../newton/tests/test_control_force.py", line 57, in test_floating_body
    test.assertAlmostEqual(body_qd[i], 0.0, delta=1e-6)
AssertionError: np.float32(0.47999954) != 0.0 within 1e-06 delta (np.float32(0.47999954) difference)
TestSimKinematics::test_featherstone_fk_floating_base_descendant_linear_velocity_matches_finite_difference_cpu
Stack Traces | 0.048s run time
Traceback (most recent call last):
  File ".../newton/tests/unittest_utils.py", line 282, in test_func
    func(self, device, **kwargs)
  File ".../newton/tests/test_kinematics.py", line 466, in test_featherstone_fk_floating_base_descendant_linear_velocity_matches_finite_difference
    assert_np_equal(origin_vel_fd, origin_vel_from_body_qd, tol=5.0e-3)
  File ".../newton/tests/unittest_utils.py", line 231, in assert_np_equal
    np.testing.assert_allclose(result.flatten(), expect.flatten(), atol=tol, equal_nan=True)
  File ".../newton/newton/.venv/lib/python3.12.../testing/_private/utils.py", line 1768, in assert_allclose
    assert_array_compare(compare, actual, desired, err_msg=str(err_msg),
  File ".../newton/newton/.venv/lib/python3.12.../testing/_private/utils.py", line 983, in assert_array_compare
    raise AssertionError(msg)
AssertionError: 
Not equal to tolerance rtol=1e-07, atol=0.005

Mismatched elements: 1 / 3 (33.3%)
Mismatch at index:
 [1]: 0.8600000143051147 (ACTUAL), 0.8400000333786011 (DESIRED)
Max absolute difference among violations: 0.01999998
Max relative difference among violations: 0.0238095
 ACTUAL: array([0.  , 0.86, 0.  ], dtype=float32)
 DESIRED: array([0.  , 0.84, 0.  ], dtype=float32)
TestPhysicsValidation::test_momentum_conservation_featherstone_cpu
Stack Traces | 2.03s run time
Traceback (most recent call last):
  File ".../newton/tests/unittest_utils.py", line 282, in test_func
    func(self, device, **kwargs)
  File ".../newton/tests/test_physics_validation.py", line 491, in test_momentum_conservation
    test.assertLess(p_rel, 5e-4, f"Linear momentum drift: {p_rel:.6e}")
AssertionError: np.float64(0.5978395062896064) not less than 0.0005 : Linear momentum drift: 5.978395e-01

To view more test analytics, go to the Test Analytics Dashboard
📋 Got 3 mins? Take this short survey to help us improve Test Analytics.

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@newton/_src/solvers/featherstone/kernels.py`:
- Around line 1567-1607: The FK pass is mixing origin- and CoM-referenced
twists; keep body_qd as an origin-referenced spatial twist throughout the
recursion and remove the mid-recursion CoM conversions: when reading parent
twist (body_qd[parent]) do NOT add/subtract wp.cross terms using
body_com[parent] (remove the v_p = wp.spatial_top(...) + wp.cross(w_p, r_p)
adjustment that assumes COM-referenced storage), and when writing child twist
always assign the origin-referenced v_wc to body_qd[child] (remove the
special-case block that computes v_com and writes body_qd[child] only for
JointType.FREE/DISTANCE). Leave composition of X_wcj/X_wc and the spatial
transforms as-is, and perform any FREE/DISTANCE COM <-> origin conversions only
in the separate post-pass (e.g.,
eval_single_articulation_fk_with_velocity_conversion or a new
velocity-conversion kernel) so body_qd remains consistently origin-referenced
during recursion.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

Run ID: 5ca88532-3a71-4563-90db-053d216721c7

📥 Commits

Reviewing files that changed from the base of the PR and between 68fc4df and 106df2d.

📒 Files selected for processing (4)
  • newton/_src/solvers/feather_pgs/__init__.py
  • newton/_src/solvers/feather_pgs/kernels.py
  • newton/_src/solvers/feather_pgs/solver_feather_pgs.py
  • newton/_src/solvers/featherstone/kernels.py

Comment on lines +1567 to +1607
v_wpj = wp.spatial_vector()
if parent >= 0:
X_wp = body_q[parent]
X_wpj = X_wp * X_wpj
r_p = wp.transform_get_translation(X_wpj) - wp.transform_point(X_wp, body_com[parent])

v_wp = body_qd[parent]
w_p = wp.spatial_bottom(v_wp)
v_p = wp.spatial_top(v_wp) + wp.cross(w_p, r_p)
v_wpj = wp.spatial_vector(v_p, w_p)

# transform from world to joint anchor frame at child body
X_wcj = X_wpj * X_j
# transform from world to child body frame
X_wc = X_wcj * wp.transform_inverse(X_cj)

v_parent_origin = wp.vec3()
w_parent = wp.vec3()
if parent >= 0:
v_wp = body_qd[parent]
w_parent = wp.spatial_bottom(v_wp)
v_parent_origin = velocity_at_point(
v_wp, wp.transform_get_translation(X_wc) - wp.transform_get_translation(X_wp)
)
linear_vel = wp.transform_vector(X_wpj, wp.spatial_top(v_j))
angular_vel = wp.transform_vector(X_wpj, wp.spatial_bottom(v_j))

linear_joint_anchor = wp.transform_vector(X_wpj, wp.spatial_top(v_j))
angular_joint_world = wp.transform_vector(X_wpj, wp.spatial_bottom(v_j))
child_origin_offset_world = wp.transform_get_translation(X_wc) - wp.transform_get_translation(X_wcj)
linear_joint_origin = linear_joint_anchor + wp.cross(angular_joint_world, child_origin_offset_world)
if type == JointType.FREE or type == JointType.DISTANCE:
# Public FREE/DISTANCE joint_qd stores the linear term at the child COM.
# Convert back to the child-body origin convention before composing the
# world-space twist, then convert back to COM on body_qd writeback below.
r_com = wp.quat_rotate(wp.transform_get_rotation(X_wc), body_com[child])
linear_vel = linear_vel - wp.cross(angular_vel, r_com)

v_wc = wp.spatial_vector(v_parent_origin + linear_joint_origin, w_parent + angular_joint_world)
v_wc = v_wpj + wp.spatial_vector(linear_vel, angular_vel)

body_q[child] = X_wc
body_qd[child] = v_wc


@wp.kernel
def convert_articulation_free_distance_body_qd(
articulation_start: wp.array[int],
articulation_count: int,
articulation_mask: wp.array[bool],
articulation_indices: wp.array[int],
joint_type: wp.array[int],
joint_child: wp.array[int],
body_com: wp.array[wp.vec3],
body_q: wp.array[wp.transform],
body_qd: wp.array[wp.spatial_vector],
):
tid = wp.tid()

if articulation_indices:
articulation_id = articulation_indices[tid]
else:
articulation_id = tid

if articulation_id < 0 or articulation_id >= articulation_count:
return

if articulation_mask:
if not articulation_mask[articulation_id]:
return

joint_start = articulation_start[articulation_id]
joint_end = articulation_start[articulation_id + 1]

for i in range(joint_start, joint_end):
type = joint_type[i]
if type != JointType.FREE and type != JointType.DISTANCE:
continue

child = joint_child[i]
X_wc = body_q[child]
v_wc = body_qd[child]

v_origin = wp.spatial_top(v_wc)
omega = wp.spatial_bottom(v_wc)
r_com = wp.transform_point(X_wc, body_com[child])
v_com = v_origin + wp.cross(omega, r_com)
body_qd[child] = wp.spatial_vector(v_com, omega)
# Velocity conversion for FREE and DISTANCE joints:
# v_wc is a spatial twist at the origin, but body_qd should store COM velocity
# Transform: v_com = v_origin + ω x r_com
if type == JointType.FREE or type == JointType.DISTANCE:
v_origin = wp.spatial_top(v_wc)
omega = wp.spatial_bottom(v_wc)
r_com = wp.quat_rotate(wp.transform_get_rotation(X_wc), body_com[child])
v_com = v_origin + wp.cross(omega, r_com)
body_qd[child] = wp.spatial_vector(v_com, omega)
else:
body_qd[child] = v_wc
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

Don’t switch body_qd into CoM space mid-recursion.

This block now mixes reference points in one FK pass: Lines 1571-1576 read the parent as if body_qd[parent] were CoM-referenced, Lines 1583-1593 compose the child motion at the joint anchor, and Lines 1600-1607 only rewrite FREE/DISTANCE children back to CoM form. That makes the meaning of body_qd depend on the parent joint type, so descendants under nontrivial body_com / joint-offset configurations will accumulate the wrong ω × r terms. Please keep a single origin-referenced twist representation for the whole articulation traversal and do the FREE/DISTANCE CoM conversion only after FK completes (or into a separate output buffer).

Based on learnings: during eval_single_articulation_fk_with_velocity_conversion, body_qd[parent] held an origin-referenced spatial twist throughout the FK recursion, and FREE/DISTANCE COM conversion ran only in the separate post-pass kernel.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@newton/_src/solvers/featherstone/kernels.py` around lines 1567 - 1607, The FK
pass is mixing origin- and CoM-referenced twists; keep body_qd as an
origin-referenced spatial twist throughout the recursion and remove the
mid-recursion CoM conversions: when reading parent twist (body_qd[parent]) do
NOT add/subtract wp.cross terms using body_com[parent] (remove the v_p =
wp.spatial_top(...) + wp.cross(w_p, r_p) adjustment that assumes COM-referenced
storage), and when writing child twist always assign the origin-referenced v_wc
to body_qd[child] (remove the special-case block that computes v_com and writes
body_qd[child] only for JointType.FREE/DISTANCE). Leave composition of
X_wcj/X_wc and the spatial transforms as-is, and perform any FREE/DISTANCE COM
<-> origin conversions only in the separate post-pass (e.g.,
eval_single_articulation_fk_with_velocity_conversion or a new
velocity-conversion kernel) so body_qd remains consistently origin-referenced
during recursion.

@preist-nvidia preist-nvidia added this to the 1.3 Release milestone Apr 7, 2026
@preist-nvidia
Copy link
Copy Markdown
Member

@dylanturpin - @eric-heiden and I assigned 1.3 to this for now, if it makes it in earlier, also great but now there's no rush.

0.0, 0.0, 0.0, I[0, 0], I[0, 1], I[0, 2],
0.0, 0.0, 0.0, I[1, 0], I[1, 1], I[1, 2],
0.0, 0.0, 0.0, I[2, 0], I[2, 1], I[2, 2],
)
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we gain anything by combining this as a spatial matrix? naively I would think that keeping mass+inertia separate would be more efficient both in terms of memory loads and compute

r2 = wp.quat_rotate(q, wp.vec3(0.0, 1.0, 0.0))
r3 = wp.quat_rotate(q, wp.vec3(0.0, 0.0, 1.0))

R = wp.matrix_from_cols(r1, r2, r3)
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wouldn't that be R = wp.quat_from_matrix(q) ?

sib_art_b = body_to_articulation[sib_bb]
sib_ds_b = art_dof_start[sib_art_b]
for k in range(6):
v_out[sib_ds_b + k] = v_out[sib_ds_b + k] + mf_MiJt_b[world, sib, k] * sib_delta
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If I am reading correctly, each friction constraint row (2 per contacts) updates the constraint impulse for both friction directions; using once the diagonal Delassus coefficient for the first row, once for the second row. Would it make sense to use only one friction constraint, which would always handle both directions? The corresponding eff_mass_inv should be the max of the two to ensure convergence.
Alternatively, we could solve for both at once with distinct eff_mass_inv, but this requires either a polynomial or a root-finding solve -- should convergence in fewer iterations in theory but maybe not worth the additional cost.

self._mf_pgs_setup(state_aug, dt)
self.v_mf_accum.zero_()

for _pgs_iter in range(self.pgs_iterations):
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How feasible would it be to fuse all of this into a single (serial per world) kernel?

@dylanturpin
Copy link
Copy Markdown
Member Author

Superseded by #2432, which carries the rebased matrix-free-only private API cleanup on . Closing this older draft to keep review on the current branch.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants