[DRAFT] Refactor SolverFeatherstone

[etaoxing]

Description

(work in progress) Refactor SolverFeatherstone to follow similar code structure as SolverSemiImplicit.

• attempts to separate construction of custom State and Model attributes into StateFeatherstone, ModelFeatherstone, ModelBuilderFeatherstone
• custom _state, _model, _builder objects are prefixed with an underscore
• attempts to improves differentiability of SolverFeatherstone (based on changes made in Rewarped)

Newton Migration Guide

Please ensure the migration guide for warp.sim users is up-to-date with the changes made in this PR.

• The migration guide in docs/migration.rst is up-to date

Before your PR is "Ready for review"

• Necessary tests have been added and new examples are tested (see newton/tests/test_examples.py)
• Documentation is up-to-date
• Code passes formatting and linting checks with pre-commit run -a

Summary by CodeRabbit

• New Features
  • Introduced a Featherstone-based solver surface with dedicated model/state containers and a builder-driven initialization.
  • Added an option to enable triangle–triangle contact handling via a toggle.
• Refactor
  • Updated the step API to accept explicit input/output states and centralized data flow through solver-managed model/state objects.
  • Standardized public torque naming by renaming the exposed joint torque parameter/output to “joint_tau” for consistency across joints.
  • Reorganized kernel imports to streamline system-matrix operations.

[coderabbitai]

📝 Walkthrough

Walkthrough

Refactors kernel parameter names from tau to joint_tau and updates all references. Introduces Featherstone-specific State and Model classes, a ModelBuilder to finalize these, modifies SolverFeatherstone construction and step control flow, switches kernel import paths, and adds optional triangle-triangle contact evaluation gating.

Changes

Cohort / File(s)	Summary
Featherstone kernel torque rename newton/_src/solvers/featherstone/kernels_body.py	Renamed public parameter/output tau to joint_tau in jcalc_tau and eval_rigid_tau; updated all internal references and return flows across BALL, FREE, DISTANCE branches.
Featherstone solver architecture and flow newton/_src/solvers/featherstone/solver_featherstone.py	Added StateFeatherstone, ModelFeatherstone, and ModelBuilderFeatherstone; replaced legacy finalize with builder pipeline; updated SolverFeatherstone.__init__ (adds enable_tri_contact) and step(...) signature; reorganized state/model buffers; switched imports to .kernels_body; added optional triangle contact evaluation; imported Devicelike.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  actor User
  participant Solver as SolverFeatherstone
  participant Model as ModelFeatherstone
  participant StateIn as StateFeatherstone (in)
  participant StateOut as StateFeatherstone (out)
  participant Kernels as kernels_body
  participant Contact as eval_triangle_contact_forces

  User->>Solver: step(StateIn, StateOut, Control?, Contacts?, dt)
  Solver->>Model: Access matrices/buffers (M, J, P, H, L)
  Solver->>StateIn: Read joint/body state arrays
  Solver->>Kernels: assemble/evaluate dynamics (uses joint_tau, qdd, S_s, etc.)
  Kernels-->>Solver: Updated system matrices/vectors
  alt enable_tri_contact and Contacts provided
    Solver->>Contact: evaluate triangle-triangle contacts
    Contact-->>Solver: contact forces/impulses
  end
  Solver->>StateOut: Write updated matrices and per-state outputs
  Solver-->>User: return

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• Refactor SolverSemiImplicit #880 — Touches Featherstone joint/kernels surface; relocates joint_force to semi_implicit/kernels_body, aligning with this PR’s kernel renames and solver import updates.

Suggested reviewers

• mmacklin
• eric-heiden

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 10.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ 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 clearly summarizes the primary change of refactoring SolverFeatherstone and directly reflects the main objective of the pull request. Although it includes a draft status tag, the core message remains concise and specific.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

📜 Recent review details

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between f0b6dbf and af84ee1.

📒 Files selected for processing (2)

• newton/_src/solvers/featherstone/kernels_body.py (5 hunks)
• newton/_src/solvers/featherstone/solver_featherstone.py (19 hunks)

🧰 Additional context used

🧬 Code graph analysis (2)

newton/_src/solvers/featherstone/kernels_body.py (1)

newton/_src/sim/joints.py (1)

• JointType (20-44)

newton/_src/solvers/featherstone/solver_featherstone.py (7)

newton/_src/core/types.py (1)

• override (36-37)

newton/_src/sim/contacts.py (1)

• Contacts (23-93)

newton/_src/sim/control.py (1)

• Control (21-77)

newton/_src/sim/model.py (2)

• Model (29-649)
• control (476-509)

newton/_src/sim/state.py (3)

• State (21-112)
• requires_grad (78-84)
• joint_dof_count (108-112)

newton/_src/solvers/semi_implicit/kernels_contact.py (4)

• eval_body_contact (398-559)
• eval_particle_body_contact_forces (644-687)
• eval_particle_contact_forces (562-582)
• eval_triangle_contact_forces (585-601)

newton/_src/solvers/featherstone/kernels_body.py (6)

• compute_spatial_inertia (28-46)
• compute_com_transforms (50-57)
• create_inertia_matrix_cholesky_kernel (1082-1120)
• create_inertia_matrix_kernel (1032-1061)
• eval_rigid_jacobian (889-928)
• eval_dense_solve_batched (1300-1313)

---

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

• X
• Mastodon
• Reddit
• LinkedIn

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

[eric-heiden]

• attempts to separate construction of custom State and Model attributes into StateFeatherstone, ModelFeatherstone, ModelBuilderFeatherstone

@etaoxing I would hold off from making these changes - we want to keep the newton.Model(Builder) and newton.State the single source of truth. We are currently trying to come up with a way to add custom attributes to the Model(Builder) and State which should allow us to express these solver-specific attributes in a more elegant way: #710. Once that lands we will probably revisit the Style3D solver/model/state to adapt this mechanism.

Having your diffsim improvements would definitely be a great contribution though!

[etaoxing]

Sounds good @eric-heiden! After trying out the change, definitely agree solver-specific StateFeatherstone, ModelFeatherstone, ModelBuilderFeatherstone is not a good way to handle custom attributes.

Will keep this PR as a draft then (and hold off on messing with other solvers) until #710 is ready!