Skip to content

[multibody] Adds fast BlockSystemJacobian calculation & Cassie benchmark for it #23090

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jun 12, 2025

Conversation

sherm1
Copy link
Member

@sherm1 sherm1 commented Jun 10, 2025

Adds a fast one-block-per-tree system Jacobian calculation. This is just Jv_V_WB, the spatial velocity Jacobian of each body frame w.r.t. the generalized velocity coordinates v. This is much less flexible than our current Jacobian functions, but considerably faster (I measured 8X faster for Cassie on my old Puget; that's a small system with just one tree).

MultibodyPlant gains three "(Internal use only)" public methods:

const std::vector<MatrixX<T>>& EvalBlockSystemJacobian(context);
MatrixX<T> CalcFullSystemJacobian(context);
const internal::LinkJointGraph& graph() const;

The block system Jacobian is cached, with the "Eval" method doing the usual update-if-necessary action.
The second method just calls the first and makes a full 6n x m Jacobian out of it (mostly useful for testing).
The graph() method gives access to the LinkJointGraph and SpanningForest.

Unit tests are provided showing that we get the same result as with our existing slow method. Also the Cassie benchmark is updated to show the relevant timings.

All changes are internal so there are no Release Notes needed. The design details are likely to change as we get some experience with these.


This change is Reviewable

@sherm1 sherm1 force-pushed the fast_system_Jacobian branch 3 times, most recently from cbc2873 to b1d49fd Compare June 10, 2025 23:44
@sherm1 sherm1 marked this pull request as ready for review June 10, 2025 23:49
Copy link
Member Author

@sherm1 sherm1 left a comment

Choose a reason for hiding this comment

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

+@amcastro-tri for design & feature review, please

Reviewable status: LGTM missing from assignee amcastro-tri, needs platform reviewer assigned, needs at least two assigned reviewers

Copy link
Contributor

@amcastro-tri amcastro-tri left a comment

Choose a reason for hiding this comment

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

quick first pass. A few questions mostly on documented (or not) definitions in the public API

Reviewed 1 of 11 files at r1, 4 of 9 files at r2.
Reviewable status: 7 unresolved discussions, LGTM missing from assignee amcastro-tri, needs platform reviewer assigned, needs at least two assigned reviewers


multibody/tree/block_system_jacobian_cache.cc line 13 at r2 (raw file):

      total_cols_(forest.num_velocities()),
      block_system_jacobian_(forest.trees().size()) {
  int row_count = 6, col_count = 0;  // Leave room for World.

I don't understand why the "room for the world"? per SpanningTree docs, world does not belong to any tree. So... we should not include it?


multibody/tree/block_system_jacobian_cache.cc line 18 at r2 (raw file):

    const int m = tree.nv();
    Eigen::MatrixX<T>& J = block_system_jacobian_[tree.index()];
    J.resize(n, m);

can we have m=0? I remember Eigen hating me in the past for attempting to resize with one of these being zero.
(though maybe it's only with maps, but definitely worth unit testing for it).


multibody/tree/block_system_jacobian_cache.h line 23 at r2 (raw file):

Note that locking and unlocking mobilizers does not affect the contents here;
the Jacobian reflects what would happen if a velocity variable changed
regardless of whether it can currently do so.

nit, consider adding this very important piece of information in the MbP's public API, where users can see it.

Code quote:

Note that locking and unlocking mobilizers does not affect the contents here;
the Jacobian reflects what would happen if a velocity variable changed
regardless of whether it can currently do so.

multibody/plant/multibody_plant.h line 4452 at r2 (raw file):

  ///@{

  /// Returns the System Jacobian Jv_V_WB_W(q) in block form. Each block is

minor, Consider our monogram shorthand Jv_V_WB. Per user feedback, people tend to hate overly decorated notation.

Code quote:

 Jv_V_WB_W(q)

multibody/plant/multibody_plant.h line 4452 at r2 (raw file):

  ///@{

  /// Returns the System Jacobian Jv_V_WB_W(q) in block form. Each block is

minor, Consider noting that V_WB = Jv_V_WB⋅v


multibody/plant/multibody_plant.h line 4455 at r2 (raw file):

  /// dense and corresponds to one Tree of the as-built SpanningForest. Each
  /// mobilized body generates a row and each mobility a column.
  const std::vector<Eigen::MatrixX<T>>& EvalBlockSystemJacobian(

thus far the best way I know to get tree structure is internal, with GetInternalTree(plant()).get\_topology(). Is there a public way to retrieve the forest? and have we defined "Tree" for users already?

BTW, I 'd be just as happy (if not happier) with a simple Calc version of this method. The less it goes through systems:: the easier it is to get a clean profile without a long chain of indirections.


multibody/plant/multibody_plant.h line 4464 at r2 (raw file):

  /// Evaluates the block system Jacobian, then uses it to fill in an equivalent
  /// full matrix of size 6n x m where n is the number of mobilized bodies and
  /// m the number of mobilities.

nit, in MbP's APIs we talk about num_velocities(), probably worth saying "number of velocities" or at least to reference num_velocities(). I believe we did not define "mobiblities" anywhere in public API (even if I agree with you that's the right term of art).

Code quote:

 number of mobilities.

Copy link
Member Author

@sherm1 sherm1 left a comment

Choose a reason for hiding this comment

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

All comments addressed, PTAL @amcastro-tri

Reviewable status: 3 unresolved discussions, LGTM missing from assignee amcastro-tri, needs platform reviewer assigned, needs at least two assigned reviewers


multibody/plant/multibody_plant.h line 4452 at r2 (raw file):

Previously, amcastro-tri (Alejandro Castro) wrote…

minor, Consider our monogram shorthand Jv_V_WB. Per user feedback, people tend to hate overly decorated notation.

Done


multibody/plant/multibody_plant.h line 4452 at r2 (raw file):

Previously, amcastro-tri (Alejandro Castro) wrote…

minor, Consider noting that V_WB = Jv_V_WB⋅v

Done, PTAL.


multibody/plant/multibody_plant.h line 4455 at r2 (raw file):

Previously, amcastro-tri (Alejandro Castro) wrote…

thus far the best way I know to get tree structure is internal, with GetInternalTree(plant()).get\_topology(). Is there a public way to retrieve the forest? and have we defined "Tree" for users already?

BTW, I 'd be just as happy (if not happier) with a simple Calc version of this method. The less it goes through systems:: the easier it is to get a clean profile without a long chain of indirections.

Since these things are internal:: currently, there is no official way to get them from the Plant. Rather than address that now, I marked these methods as "Advanced" and noted that the forest is internal. Anyone who doesn't know what that means shouldn't be using these functions!

Caching the Jacobian means that the memory gets allocated once at Finalize and then reused every time it is calculated. One of the reasons the old method (a Calc API) is so slow is that it requires many heap allocations. The API here is consistent with other useful things we cache in MbP. There is of course an underlying Calc method that is used to update the cache -- you could use that if you want but I think the Eval method is better, at least for now.


multibody/plant/multibody_plant.h line 4464 at r2 (raw file):

Previously, amcastro-tri (Alejandro Castro) wrote…

nit, in MbP's APIs we talk about num_velocities(), probably worth saying "number of velocities" or at least to reference num_velocities(). I believe we did not define "mobiblities" anywhere in public API (even if I agree with you that's the right term of art).

Done


multibody/tree/block_system_jacobian_cache.h line 23 at r2 (raw file):

Previously, amcastro-tri (Alejandro Castro) wrote…

nit, consider adding this very important piece of information in the MbP's public API, where users can see it.

Done. There was a lot of missing documentation on the public APIs. I added it -- PTAL there.


multibody/tree/block_system_jacobian_cache.cc line 13 at r2 (raw file):

Previously, amcastro-tri (Alejandro Castro) wrote…

I don't understand why the "room for the world"? per SpanningTree docs, world does not belong to any tree. So... we should not include it?

Done. Clarified what's going on here. As you expect, these blocks are for trees only -- there is no block corresponding to World.


multibody/tree/block_system_jacobian_cache.cc line 18 at r2 (raw file):

Previously, amcastro-tri (Alejandro Castro) wrote…

can we have m=0? I remember Eigen hating me in the past for attempting to resize with one of these being zero.
(though maybe it's only with maps, but definitely worth unit testing for it).

Done. I added a 0-dof tree to the last test in multibody_plant_jacobians_test.cc. PTAL, it does work.

@sherm1 sherm1 force-pushed the fast_system_Jacobian branch from e04af95 to 835b1e2 Compare June 11, 2025 17:40
Copy link
Contributor

@amcastro-tri amcastro-tri left a comment

Choose a reason for hiding this comment

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

Reviewed 1 of 11 files at r1, 4 of 9 files at r2, 3 of 3 files at r3, all commit messages.
Reviewable status: 4 unresolved discussions, LGTM missing from assignee amcastro-tri, needs platform reviewer assigned, needs at least two assigned reviewers


multibody/plant/multibody_plant.h line 4455 at r2 (raw file):

Since these things are internal:: currently, there is no official way to get them from the Plant

Fair. Is there a plan to expose the SpanningForest. It's just too beautiful to let it rot in internal::. I think we should expose it, it's ready for a typical robotics user that does need this information. I for sure would use it.
Actually, as an intermediate step. I still use the old topology code in my prototypes. Is there a way to retrieve the spanning forest witha smilar internal:: call?

Caching the Jacobian means that the memory gets allocated once at Finalize

that's okay in theory. Consider my use case, the new SAP (pressumably even GPU friendly). That software pre-allocates a pool of memory, including memory for these Jacobians. Therefore the ideal API would allow me to write into that pool.

Non-blocking though, but in general I've found it difficult to reuse memory that lives in the cache. That's why a C-like kernel could be useful to have in the future.

Obviusly this is non-blocking, it's just an fyi. Right now copying from this return into my data structs gets the job done and the cost is negligible.


multibody/tree/block_system_jacobian_cache.cc line 18 at r2 (raw file):

Previously, sherm1 (Michael Sherman) wrote…

Done. I added a 0-dof tree to the last test in multibody_plant_jacobians_test.cc. PTAL, it does work.

ah, then Eigen barks with maps and blocks, apparently actual MatrixX is okay. Good to know, thanks.


multibody/tree/multibody_tree.cc line 1353 at r3 (raw file):

      sjc->mutable_block_system_jacobian();

  // Each treei generates an ni x mi block.

nit, not sure if a typo, but it looks like one. Maybe consider "each i-th tree generates..."

Code quote:

// Each treei g

multibody/tree/multibody_tree.cc line 1357 at r3 (raw file):

    const int n = 6 * tree.num_mobods();
    const int m = tree.nv();
    MatrixX<T>& J = tree_jacobians[tree.index()];  // J_V_WB_W(tree)

who is initializing J to zero? Are you relying on the intialization to zero the very first time the BlcokSystemJacobianCache is allocated? Would that be too dangerous. Or at least document here that assumption

If you rely on this first initialization, you might even consider intializing to identity blocks corresponding to free-bodies and skiping them here (provided the Cassie benchamarks shows improvement)


multibody/plant/test/multibody_plant_jacobians_test.cc line 171 at r3 (raw file):

  // Calculate the System Jacobian Jv_V_WB_W two ways and compare.
  // TODO(sherm1) Should be num_mobods.

I agree. My preference: time to expose SpanningForest

Code quote:

// TODO(sherm1) Should be num_mobods.

multibody/plant/multibody_plant.h line 4456 at r3 (raw file):

  /// internal::SpanningForest. The blocks follow the Tree ordering defined by
  /// the SpanningForest, so are in TreeIndex order. The block for Treeᵢ is a
  /// MatrixX of size 6nᵢ x mᵢ, where nᵢ is the number of mobilized bodies in

For my use case, I'll actually need to extract the block corresponding to a particular body.
Unless there is a performance reason, could we instead return separate blocks for each body in the model? with entries in the std::vector<MatriX<T>> ordered by BodyIndex

Code quote:

  /// internal::SpanningForest. The blocks follow the Tree ordering defined by
  /// the SpanningForest, so are in TreeIndex order. The block for Treeᵢ is a
  /// MatrixX of size 6nᵢ x mᵢ, where nᵢ is the number of mobilized bodies in

Copy link
Contributor

@amcastro-tri amcastro-tri left a comment

Choose a reason for hiding this comment

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

Reviewable status: 5 unresolved discussions, LGTM missing from assignee amcastro-tri, needs platform reviewer assigned, needs at least two assigned reviewers


multibody/plant/multibody_plant.h line 4452 at r3 (raw file):

  ///@{

  /// (Advanced) Returns the System Jacobian Jv_V_WB(q) in block form. Each

minor, Consider (Internal) given this makes reference to internal:: APIs.

Code quote:

(Advanced)

@sherm1 sherm1 force-pushed the fast_system_Jacobian branch from 835b1e2 to 63b4ac6 Compare June 11, 2025 21:44
Copy link
Member Author

@sherm1 sherm1 left a comment

Choose a reason for hiding this comment

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

Comments addressed, PTAL. Everything is "internal use only" now and there is access to the internal:: graph. Due to my impending vacation I'm suggesting we leave this mostly as-is for now so you can have access to a faster Jacobian, then revise on my return (or you and Vince can revise it if you feel so inclined). There is still time to get this platform reviewed and merged if we don't make major changes now.

Reviewable status: 2 unresolved discussions, LGTM missing from assignee amcastro-tri, needs platform reviewer assigned, needs at least two assigned reviewers


multibody/plant/multibody_plant.h line 4455 at r2 (raw file):

Previously, amcastro-tri (Alejandro Castro) wrote…

Since these things are internal:: currently, there is no official way to get them from the Plant

Fair. Is there a plan to expose the SpanningForest. It's just too beautiful to let it rot in internal::. I think we should expose it, it's ready for a typical robotics user that does need this information. I for sure would use it.
Actually, as an intermediate step. I still use the old topology code in my prototypes. Is there a way to retrieve the spanning forest witha smilar internal:: call?

Caching the Jacobian means that the memory gets allocated once at Finalize

that's okay in theory. Consider my use case, the new SAP (pressumably even GPU friendly). That software pre-allocates a pool of memory, including memory for these Jacobians. Therefore the ideal API would allow me to write into that pool.

Non-blocking though, but in general I've found it difficult to reuse memory that lives in the cache. That's why a C-like kernel could be useful to have in the future.

Obviusly this is non-blocking, it's just an fyi. Right now copying from this return into my data structs gets the job done and the cost is negligible.

I've added an "internal use only" method that returns the LinkJointGraph. (And graph has a method for getting access to the SpanningForest.) PTAL

Now that everything is marked "internal" we can revise the design when I get back from vacation, including adding an API you can use to write into your own memory. I think if you and Vince want to use this while I'm gone we ought to merge as close to as-is now and revise later. It's fast, unit tested, and all-internal.


multibody/plant/multibody_plant.h line 4452 at r3 (raw file):

Previously, amcastro-tri (Alejandro Castro) wrote…

minor, Consider (Internal) given this makes reference to internal:: APIs.

Done


multibody/plant/multibody_plant.h line 4456 at r3 (raw file):

Previously, amcastro-tri (Alejandro Castro) wrote…

For my use case, I'll actually need to extract the block corresponding to a particular body.
Unless there is a performance reason, could we instead return separate blocks for each body in the model? with entries in the std::vector<MatriX<T>> ordered by BodyIndex

Per Slack discussion, IMO if we're going to pad out to the max number of columns mᵢ for a Tree, it's better to leave this organized by tree. It's very easy to access the individual body blocks via Eigen's m.block(6*i, 0, 6, mᵢ).

(BTW these are by MobodIndex rather than BodyIndex)


multibody/plant/test/multibody_plant_jacobians_test.cc line 171 at r3 (raw file):

Previously, amcastro-tri (Alejandro Castro) wrote…

I agree. My preference: time to expose SpanningForest

Done. Use MbP's new graph() method. graph().forest() gives you the SpanningForest.


multibody/tree/multibody_tree.cc line 1353 at r3 (raw file):

Previously, amcastro-tri (Alejandro Castro) wrote…

nit, not sure if a typo, but it looks like one. Maybe consider "each i-th tree generates..."

Hopefully clearer now.


multibody/tree/multibody_tree.cc line 1357 at r3 (raw file):

Previously, amcastro-tri (Alejandro Castro) wrote…

who is initializing J to zero? Are you relying on the intialization to zero the very first time the BlcokSystemJacobianCache is allocated? Would that be too dangerous. Or at least document here that assumption

If you rely on this first initialization, you might even consider intializing to identity blocks corresponding to free-bodies and skiping them here (provided the Cassie benchamarks shows improvement)

Done. Documented now. Pre-allocating floating body identities would require some thought. That has to be delegated to the bodies because we don't actually know they have identity H matrices -- that's only true because we happen to know we're using today's quaternion_floating or rpy_floating joints in the F frame. If (for example) we every had M-frame versions of those joints (god forbid!) they would still be floating bodies but wouldn't have identity Jacobians.

Currently I believe letting the floating bodies fill in their own identities is safe and probably not expensive. Worth investigating though if Jacobian calculation is still a bottleneck after this speedup.

Also, per discussion with Vince it appears the primary use of the Jacobians is to calculate JᵀGJ -- if so it would likely pay off more to target that calculation directly.

@sherm1 sherm1 added release notes: none This pull request should not be mentioned in the release notes and removed release notes: feature This pull request contains a new feature labels Jun 11, 2025
Copy link
Contributor

@amcastro-tri amcastro-tri left a comment

Choose a reason for hiding this comment

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

Thanks Sherm! there is no immediate rush on this though, Vince and I are still working on dev branches so we can alway grab this one commit if we don't get the reviews on time.

Reviewed 4 of 4 files at r4, all commit messages.
Reviewable status: 1 unresolved discussion, needs platform reviewer assigned, needs at least two assigned reviewers


multibody/plant/multibody_plant.h line 4456 at r3 (raw file):

Previously, sherm1 (Michael Sherman) wrote…

Per Slack discussion, IMO if we're going to pad out to the max number of columns mᵢ for a Tree, it's better to leave this organized by tree. It's very easy to access the individual body blocks via Eigen's m.block(6*i, 0, 6, mᵢ).

(BTW these are by MobodIndex rather than BodyIndex)

fair, this is internal API only, so I can give you better feedback after I integrate this into our new code.


multibody/tree/multibody_tree.cc line 1357 at r3 (raw file):

Jacobians is to calculate JᵀGJ

Correct, but when I measured the cost of this term, it was negligible already.
This level of optimization is good for now, thanks!


a discussion (no related file):
Thank you so much for pushing this Sherm! :lgtm_strong:

@amcastro-tri
Copy link
Contributor

+@rpoyner-tri for platform review please.

Copy link
Contributor

@rpoyner-tri rpoyner-tri left a comment

Choose a reason for hiding this comment

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

:lgtm:

Reviewed 1 of 11 files at r1, 5 of 9 files at r2, 1 of 3 files at r3, 4 of 4 files at r4, all commit messages.
Reviewable status: 1 unresolved discussion

Copy link
Contributor

@rpoyner-tri rpoyner-tri left a comment

Choose a reason for hiding this comment

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

Dismissed @amcastro-tri from a discussion.
Reviewable status: :shipit: complete! all discussions resolved, LGTM from assignees rpoyner-tri(platform),amcastro-tri

@rpoyner-tri rpoyner-tri merged commit d255b8f into RobotLocomotion:master Jun 12, 2025
10 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
priority: high release notes: none This pull request should not be mentioned in the release notes
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants