Gradients w.r.t. joint displacements

[JonathanKuelz]

Dear pytorch kinematics team, thank you for this amazing repo!

I am interested in computing the FK gradients, but not only w.r.t joint angles, but also with respect to the robot morphology (for my purpose, it is sufficient to say it is captured by the joint displacements). A naive approach is to just monkey-patch the current implementation and setting requires_grad=True for the displacement matrices.

However, usually, joint displacements are parameterized in robotics, most commonly using (modified) DH-parameters. Alternatively, a fully flexible 6D (displacement + axis angle) parameterization would be of interest to me to compute the gradient w.r.t. these parameters only.

Did you do any development in this direction yet? If not so, I could start working on a PR if this is a feature of interest for the repository -- would this be a feature you would be willing to support in the future given an initial implementation?

[JonathanKuelz]

Related to #8 that came with a similar request but was not followed up upon by the original poster

[LemonPi]

Hi Jonathan, thanks for your interest! We haven't worked towards anything like this, but it sounds interesting. Feel free to work on it and post your thoughts and questions here, and we can chime in. I suspect it'll take some work though.

Some initial thoughts: The offset matrices are currently converted from denser parameterizations inside the description files, but there is no easy to convert back to those parameterizations. You would need to implement an additional method on chain that lets you explicitly set the offset matrices via some parameterization (e.g. DH-parameters). Then you can define requires_grad=True on the parameters you input to this method and get gradients on it. One thing to think about is how batching will work when you have batch offset parameters in addition to batch joint angles.

[JonathanKuelz]

I will start working on it to provide a prototype based on M-DH parameters after Craig. To this purpose, I wrote a ParameterizedTransform class that inherits from the Transform3d and, given some parameters, provides the matrices.

Regarding the batching: My first attempt was to create a DiffChain class that comes with joint offsets that need to be of the abovementioned class. However, thinking about it, the easiest way to implement this feature would possible just be to allow passing not only joint angles, but also joint displacement parameters to the fk method (and jacobian etc, but I want to start small). Do you have any thoughts on that?

[LemonPi]

I see two approaches:

1. keep all the current method APIs the same, but add in an additional method on Chain that lets you replace the offsets with your parameterized versions (e.g. with ParameterizedTransform). Method implementations could probably then be adjusted to allow for batch offsets without breaking current behavior.
2. add optional joint offsets and link offsets parameters to the method arguments.

Either way, I don't think you need a separate DiffChain class. Option 2 is more direct, and probably the one I would try first. You would have to pass in the matrix link and joint offsets directly, which could be computed via some other method from your parameterization.

Currently, for efficiency reasons the link and joint offsets are computed at construction time and cached (assumed constant). They're used inside the methods in ways such as:

                    link_offset_i = self.link_offsets[chain_idx]
                    if link_offset_i is not None:
                        frame_transform = frame_transform @ link_offset_i

                    joint_offset_i = self.joint_offsets[chain_idx]
                    if joint_offset_i is not None:
                        frame_transform = frame_transform @ joint_offset_i

So it's pretty easy to insert in arguments link_offsets and joint_offsets, which if you don't provide any will default to the precomputed ones.

[JonathanKuelz]

I pushed a first draft to #32 and would be happy to discuss it