Add differentiable forward kinematics #420

fantaosha · 2023-01-03T02:15:20Z

Motivation and Context

How Has This Been Tested

Types of changes

Docs change / refactoring / dependency upgrade
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to change)

Checklist

My code follows the code style of this project.
My change requires a change to the documentation.
I have updated the documentation accordingly.
I have read the CONTRIBUTING document.
I have completed my CLA (see CONTRIBUTING)
I have added tests to cover my changes.
All new and existing tests passed.

cpaxton

lgtm in general, although i am a bit confused about some things

cpaxton · 2023-01-06T22:16:07Z

theseus/embodied/robot/forward_kinematics.py

+# TODO: Add support for joints with DOF>1
+
+
+def ForwardKinematicsFactory(robot: Robot, link_names: Optional[List[str]] = None):


Why set this up as a factory?

The outputs are forward kinematics for the given links of a specific robot. In some problems, we don't have to run the full forward kinematics if only interested in a subset of links.

cpaxton · 2023-01-06T23:12:50Z

theseus/embodied/robot/forward_kinematics.py

+
+    class ForwardKinematics(torch.autograd.Function):
+        @classmethod
+        def forward(cls, ctx, angles):


What is ctx? maybe a comment?

ctx is used in PyTorch for analytic backward propagation. Here we present analytical backward function for forward kinematics, which yields 10x speedup.

cpaxton · 2023-01-06T23:13:28Z

theseus/embodied/robot/link.py

@@ -60,6 +65,9 @@ def set_children(self, children: List[Any]):
    def set_ancestors(self, ancesotrs: List["Link"]):
        self._ancestors = ancesotrs

+    def set_angle_ids(self, angle_ids: List[int]):


what is this supposed to do?

This is the ids of non-fixed joints that a link depends on. For example, a 7-link robot arm, the second link only depends on the first and second joints. This is used to improve the computational efficiency of jacobians, in particular, for robot with branched chain topology.

We should have a more explicit/descriptive name for this than angle_ids.

Compared to tracing manually back to the root, the computational complexity isn't really affected by caching the dependent joints, since you still need linear time to go through them.
Also, performance shouldn't be an issue since you have a factory pattern to generate an object that has all of the computational paths traced.

We should have a more explicit/descriptive name for this than angle_ids.

To add to this, not all joint states are angles.

We should have a more explicit/descriptive name for this than angle_ids.

Now angle_ids -> ancestor_active_joint_ids.

theseus/embodied/robot/link.py

theseus/embodied/robot/forward_kinematics.py

luisenp

Great job adding all this stuff! I left some comments to address.

luisenp · 2023-01-04T15:35:08Z

theseus/embodied/robot/forward_kinematics.py

+# TODO: Add support for joints with DOF>1
+
+
+def ForwardKinematicsFactory(robot: Robot, link_names: Optional[List[str]] = None):


No return type annotations for this.

Is ->Any OK?

Why Any? What possible options can this return?

tests/embodied/robot/test_forward_kinematics.py

luisenp · 2023-01-24T14:44:59Z

theseus/embodied/robot/link.py

@@ -60,6 +65,9 @@ def set_children(self, children: List[Any]):
    def set_ancestors(self, ancesotrs: List["Link"]):
        self._ancestors = ancesotrs

+    def set_angle_ids(self, angle_ids: List[int]):


We should have a more explicit/descriptive name for this than angle_ids.

luisenp · 2023-01-24T14:52:23Z

theseus/embodied/robot/robot.py

+                angle_ids = (
+                    link.parent.parent.angle_ids
+                    if isinstance(link.parent, FixedJoint)
+                    else link.parent.parent.angle_ids + [link.parent.id]


Let's add a property to Link class called something like parent_link, which returns parent.parent. It will be easier to read.

Yes I agree that parent.parent is rather misleading as it looks like you're traversing two links up the kinematic tree. Having something like parent_link, child_links will increase the readability of the code.
Or, refer to #417 (comment) in the previous PR for an alternative way to connect Links and Joints in the tree data structure.

Added Link.parent_link, and renamed some properties of a link, i.e., Link.parent->Link.parent_joint, Link.children->Link.child_joints, Link.angle_ids -> Link.ancestor_active_joint_ids, etc.

luisenp · 2023-01-24T15:44:32Z

theseus/embodied/robot/forward_kinematics.py

+from .joint import Joint
+
+
+# TODO: Add support for joints with DOF>1


Make a Github issue?

luisenp · 2023-01-24T21:13:14Z

theseus/embodied/robot/forward_kinematics.py

+    def _forward_kinematics_helper(angles: torch.Tensor):
+        if angles.ndim != 2 or angles.shape[1] != robot.dof:
+            raise ValueError(
+                f"Joint angles for {robot.name} should be {robot.dof}-D vectors"


Maybe add "{robot.dof}-D vectors with a batch dimension."?

Not sure. If I remember correctly, we only consider the second dimension in our Lie group code.

I'm just suggesting a more descriptive error message for what your code is already doing (you are checking ndim == 2 and dof being in dim 1). Your current message implies that ndim should be ==1.

luisenp · 2023-01-24T21:14:02Z

theseus/embodied/robot/forward_kinematics.py

+            )
+
+        poses: List[Optional[torch.Tensor]] = [None] * robot.num_links
+        poses[0] = angles.new_zeros(angles.shape[0], 3, 4)


Replace with se3.new_identity(angles).

Maybe we need more discussions about this the behavior of se3.new_identity(). For example, how to specify the dimensions.

Is there any reasonable option other than the code below?

g = se3.new_identity(batch_size) # g.shape = (batch_size, 3, 4)

The liegroup has no new_identity() method. This is definitely in our TODO list and I will go back to this once new_identity() is added to Lie group.

theseus/embodied/robot/forward_kinematics.py

luisenp · 2023-01-24T21:19:32Z

theseus/embodied/robot/forward_kinematics.py

+
+
+def get_forward_kinematics(robot: Robot, link_names: Optional[List[str]] = None):
+    ForwardKinematics, _, jforward_kinematics, _, _ = ForwardKinematicsFactory(


Rename output as ForwardKinematics --> fk_autograd_fn_cls?

Is there anything problematic about this suggestion? CamelCase is usually used for class definitions.

I'm OK for either. I will address this comment after all the PRs have been merged since lots of functions use ForwardKinematics.

theseus/embodied/robot/robot.py

theseus/embodied/robot/joint.py

exhaustin · 2023-01-31T19:18:12Z

theseus/embodied/robot/robot.py

+                angle_ids = (
+                    link.parent.parent.angle_ids
+                    if isinstance(link.parent, FixedJoint)
+                    else link.parent.parent.angle_ids + [link.parent.id]


Yes I agree that parent.parent is rather misleading as it looks like you're traversing two links up the kinematic tree. Having something like parent_link, child_links will increase the readability of the code.
Or, refer to #417 (comment) in the previous PR for an alternative way to connect Links and Joints in the tree data structure.

exhaustin · 2023-01-31T19:22:00Z

theseus/embodied/robot/link.py

@@ -60,6 +65,9 @@ def set_children(self, children: List[Any]):
    def set_ancestors(self, ancesotrs: List["Link"]):
        self._ancestors = ancesotrs

+    def set_angle_ids(self, angle_ids: List[int]):


Compared to tracing manually back to the root, the computational complexity isn't really affected by caching the dependent joints, since you still need linear time to go through them.
Also, performance shouldn't be an issue since you have a factory pattern to generate an object that has all of the computational paths traced.

exhaustin · 2023-01-31T19:24:40Z

theseus/embodied/robot/link.py

@@ -60,6 +65,9 @@ def set_children(self, children: List[Any]):
    def set_ancestors(self, ancesotrs: List["Link"]):
        self._ancestors = ancesotrs

+    def set_angle_ids(self, angle_ids: List[int]):


We should have a more explicit/descriptive name for this than angle_ids.

To add to this, not all joint states are angles.

exhaustin · 2023-01-31T19:32:18Z

theseus/embodied/robot/forward_kinematics.py

+        ancestors += [anc for anc in link.ancestors]
+        joint_ids += link.angle_ids
+    pose_ids = sorted(list(set([anc.id for anc in ancestors] + link_ids)))
+    joint_ids = sorted(list(set(joint_ids)))


Same comment as #417 (comment)

Yeah, I will refactor the code with link.joint and link.parent.

I think keeping the ancestor ids is important in lots of applications. This makes it possible to get the sparsity structure when the robot model is used in collision.

We have robot.joints and the users get the joint names by robot.joints[joint_id].name.

exhaustin · 2023-01-31T19:35:08Z

theseus/embodied/robot/forward_kinematics.py

+            jposes[:, :, joint.id : joint.id + 1] = (
+                se3.adjoint(poses[link.id]) @ joint.axis
+            )


Correct me if I'm wrong, but this line will only be correct if the joints are revolute joints?
I propose we implement joint jacobians within the Joint object, as mentioned in #415 (comment).
Thus, this Jacobian function will not only be general across different joint types, but will also be more readable.
Happy to discuss offline.

This code applies all types of joints, including both revolute and prismatic.

* add device to robot model * Refactor Joint.dof as a property (#423) * refactor joint dof * add virtual end-effector for tests

* add link * add id to joint * add id to link * add robot model * add name to robot * backup * refactor joint * refactor link * add robot model * robot model works * a minior refactor of joint * add ancestor to links * update dof computation * add forward kinematics function * Add differentiable forward kinematics (#420) * fixed a bug in FixedJoint * add angle_ids * add jacobians to FK * forward kinematics works * rename FK * add tests for differentiable kinematics * rename tests * simplify the computation of origin * add jacobian tests for forward kinematics * fixed a bug for vectorize=True in SE3 * improve the efficiency of forward kinematics * add tests for backward propagation * update lie group load paths * Add device to robot model (#422) * add device to robot model * Refactor Joint.dof as a property (#423) * refactor joint dof * add virtual end-effector for tests * seems to work after name changes * rename parent->parent_link and child->child_link * rename angle_ids -> ancestor_active_joint_ids * refactor tests with torch.testing.assert_close * small refactoring of the code * add Robot.get_links() * remove some setters for link and joint * made some functions private for encapsulation * consolidate if block and classmethod->staticmethod * refactor set_robot_from_urdf_file * rename ancestor_active_joint_ids -> ancestor_non_fixed_joint_ids * id=-1 -> id=None * refactor the code for simplicity * move link into joint

* add prismatic joint * refactor prismatic joint * Add robot model (#417) * add link * add id to joint * add id to link * add robot model * add name to robot * backup * refactor joint * refactor link * add robot model * robot model works * a minior refactor of joint * add ancestor to links * update dof computation * add forward kinematics function * Add differentiable forward kinematics (#420) * fixed a bug in FixedJoint * add angle_ids * add jacobians to FK * forward kinematics works * rename FK * add tests for differentiable kinematics * rename tests * simplify the computation of origin * add jacobian tests for forward kinematics * fixed a bug for vectorize=True in SE3 * improve the efficiency of forward kinematics * add tests for backward propagation * update lie group load paths * Add device to robot model (#422) * add device to robot model * Refactor Joint.dof as a property (#423) * refactor joint dof * add virtual end-effector for tests * seems to work after name changes * rename parent->parent_link and child->child_link * rename angle_ids -> ancestor_active_joint_ids * refactor tests with torch.testing.assert_close * small refactoring of the code * add Robot.get_links() * remove some setters for link and joint * made some functions private for encapsulation * consolidate if block and classmethod->staticmethod * refactor set_robot_from_urdf_file * rename ancestor_active_joint_ids -> ancestor_non_fixed_joint_ids * id=-1 -> id=None * refactor the code for simplicity * move link into joint

* add revolute joint * refactor revolute joint * add parent and child to joint * update the path to load so3 and se3 * Add prismatic joint (#416) * add prismatic joint * refactor prismatic joint * Add robot model (#417) * add link * add id to joint * add id to link * add robot model * add name to robot * backup * refactor joint * refactor link * add robot model * robot model works * a minior refactor of joint * add ancestor to links * update dof computation * add forward kinematics function * Add differentiable forward kinematics (#420) * fixed a bug in FixedJoint * add angle_ids * add jacobians to FK * forward kinematics works * rename FK * add tests for differentiable kinematics * rename tests * simplify the computation of origin * add jacobian tests for forward kinematics * fixed a bug for vectorize=True in SE3 * improve the efficiency of forward kinematics * add tests for backward propagation * update lie group load paths * Add device to robot model (#422) * Refactor Joint.dof as a property (#423) * refactor joint dof * add virtual end-effector for tests * seems to work after name changes * rename parent->parent_link and child->child_link * rename angle_ids -> ancestor_active_joint_ids * refactor tests with torch.testing.assert_close * small refactoring of the code * add Robot.get_links() * remove some setters for link and joint * made some functions private for encapsulation * consolidate if block and classmethod->staticmethod * refactor set_robot_from_urdf_file * rename ancestor_active_joint_ids -> ancestor_non_fixed_joint_ids * id=-1 -> id=None * refactor the code for simplicity * move link into joint * simplify the code * minor revision to be consistent with Lie group refatcoring * fixed a CI error * move files to labs

fantaosha added the enhancement New feature or request label Jan 3, 2023

fantaosha requested review from luisenp and mhmukadam January 3, 2023 02:15

fantaosha self-assigned this Jan 3, 2023

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jan 3, 2023

fantaosha changed the base branch from taoshaf.robot.prismatic_joint to taoshaf.robot.robot January 3, 2023 02:15

fantaosha requested review from cpaxton and exhaustin January 4, 2023 18:24

cpaxton reviewed Jan 6, 2023

View reviewed changes

fantaosha force-pushed the taoshaf.robot.robot branch from 6ca688b to 19a5326 Compare January 23, 2023 19:16

mhmukadam mentioned this pull request Jan 23, 2023

Add robot model #417

Merged

11 tasks

fantaosha added 13 commits January 23, 2023 14:06

fixed a bug in FixedJoint

b8ae954

add angle_ids

4e2b7cb

add jacobians to FK

f888c3d

forward kinematics works

101e212

rename FK

d3d517d

add tests for differentiable kinematics

de5ed34

rename tests

acb715a

simplify the computation of origin

36105aa

add jacobian tests for forward kinematics

b2ea6f3

fixed a bug for vectorize=True in SE3

b1dacc2

improve the efficiency of forward kinematics

3861b62

add tests for backward propagation

836c00b

update lie group load paths

1be2940

fantaosha force-pushed the taoshaf.robot.forward_kinematics branch from 928fb81 to 1be2940 Compare January 23, 2023 22:11

luisenp requested changes Jan 24, 2023

View reviewed changes

exhaustin reviewed Jan 31, 2023

View reviewed changes

fantaosha and others added 3 commits February 28, 2023 14:10

Add device to robot model (#422)

a2b9790

* add device to robot model * Refactor Joint.dof as a property (#423) * refactor joint dof * add virtual end-effector for tests

seems to work after name changes

6f0d22a

rename parent->parent_link and child->child_link

f5f312f

fantaosha added 4 commits February 28, 2023 13:10

rename angle_ids -> ancestor_active_joint_ids

04d0e1e

refactor tests with torch.testing.assert_close

804218f

small refactoring of the code

02d946e

add Robot.get_links()

1af136b

fantaosha requested a review from luisenp February 28, 2023 22:14

fantaosha added 2 commits February 28, 2023 15:58

remove some setters for link and joint

ff7a6d7

made some functions private for encapsulation

8fb70e2

fantaosha merged commit 765fcb5 into taoshaf.robot.robot Mar 1, 2023

fantaosha deleted the taoshaf.robot.forward_kinematics branch March 1, 2023 18:49

		# TODO: Add support for joints with DOF>1


		def ForwardKinematicsFactory(robot: Robot, link_names: Optional[List[str]] = None):

		from .joint import Joint


		# TODO: Add support for joints with DOF>1



		def get_forward_kinematics(robot: Robot, link_names: Optional[List[str]] = None):
		ForwardKinematics, _, jforward_kinematics, _, _ = ForwardKinematicsFactory(

Add differentiable forward kinematics #420

Add differentiable forward kinematics #420

Conversation

fantaosha commented Jan 3, 2023

Motivation and Context

How Has This Been Tested

Types of changes

Checklist

cpaxton left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fantaosha Jan 6, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

luisenp left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

exhaustin Jan 31, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

exhaustin Jan 31, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fantaosha Jan 6, 2023 •

edited

exhaustin Jan 31, 2023 •

edited

exhaustin Jan 31, 2023 •

edited