[DRAFT] Add private FeatherPGS solver implementation#2331
[DRAFT] Add private FeatherPGS solver implementation#2331dylanturpin wants to merge 1 commit intonewton-physics:mainfrom
Conversation
📝 WalkthroughWalkthroughTwo changes: (1) New Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
❌ 6 Tests Failed:
View the top 3 failed test(s) by shortest run time
To view more test analytics, go to the Test Analytics Dashboard |
There was a problem hiding this comment.
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
📒 Files selected for processing (4)
newton/_src/solvers/feather_pgs/__init__.pynewton/_src/solvers/feather_pgs/kernels.pynewton/_src/solvers/feather_pgs/solver_feather_pgs.pynewton/_src/solvers/featherstone/kernels.py
| 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 |
There was a problem hiding this comment.
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.
|
@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], | ||
| ) |
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
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 |
There was a problem hiding this comment.
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): |
There was a problem hiding this comment.
How feasible would it be to fuse all of this into a single (serial per world) kernel?
|
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. |
Description
Work in progress. Draft PR for
feather_pgsAPI in newton.Checklist
CHANGELOG.mdhas been updated (if user-facing change)Test plan
Bug fix
Steps to reproduce:
Minimal reproduction:
New feature / API change
Summary by CodeRabbit
Refactor
Chores