Add forward kinematics package to `labs` #415

fantaosha · 2023-01-02T03:54:21Z

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.

luisenp

Looks good! Requesting some small changes that should be easy to address.

luisenp · 2023-01-12T16:31:59Z

theseus/embodied/robot/joint.py

+
+            if dtype is not None and origin.dtype != dtype:
+                warnings.warn(
+                    f"tensor.dtype {origin.dtype} does not match given dtype {dtype}, "


Maybe clearer to write "The origin's dtype ({origin.dtype}) does not match the desired dtype. Origin's dtype will take precendence".

Refactored.

luisenp · 2023-01-12T16:53:23Z

theseus/embodied/robot/joint.py

+            angle_axis = angle_axis.view(-1, 1)
+
+        if angle_axis.ndim != 2 or angle_axis.shape != (3, 1):
+            raise ValueError("The angle axis must be a 3-D vector.")


The ndim != 2 seems redundant since you are checking the exact angle_axis shape. Also, in the message "3-D vector" seems ambiguous (e.g., a user wouldn't know if (1, 3), (3,) are also valid shapes).

Maybe better to check that angle_axis.numel() == 3, and if that's the case you can just do self._axis[3:] = angle_axis.view(-1). In this case the current error message seems appropriate.

numel() is used now and the error message is "The revolute axis must have 3 elements.".

luisenp · 2023-01-12T16:55:53Z

theseus/embodied/robot/joint.py

+    def _rotation_impl(self, angle: torch.Tensor) -> torch.Tensor:
+        if angle.ndim == 2 and angle.shape[1] == 1:
+            angle = angle.view(-1)
+        if angle.ndim != 1:


It's possible that people could pass ndim > 2 by mistake. I think we can do the following:

angle = angle.squeeze() if angle.ndim >= 2: raise ValueError("The joint angle must be a vector.") angle = angle.view(-1) # in case the squeeze collapsed it to shape `[]`

luisenp · 2023-01-12T17:04:58Z

theseus/embodied/robot/joint.py

+    def axis(self) -> torch.Tensor:
+        return self._axis
+
+    def set_parent(self, parent: int):


Python properties provide a mechanism for setting properties so that one can do x.parent = the_parent instead of x.set_parent(the_parent). You can do:

@parent.setter def parent(self, parent: int): ...

All the setters have been removed for encapsulation.

luisenp · 2023-01-23T16:26:33Z

theseus/embodied/robot/joint.py

+        return 1
+
+    @abc.abstractmethod
+    def _rotation_impl(self, angle: torch.Tensor) -> torch.Tensor:


_rotation_impl implies that there is another rotation() method that's public and that works something like the following:

def rotation(self, angle: torch.Tensor) -> torch.Tensor: # common initialization such as checking angle shape pre_out = self._rotation_impl(angle) # post process pre_out return processed_out

If just a private/internal implementation is needed then def _rotation(...) is sufficient.

rotation() and translation() are defined.

luisenp · 2023-01-23T16:27:45Z

theseus/embodied/robot/joint.py

+        self._axis[4] = 1
+
+    def _rotation_impl(self, angle: torch.Tensor) -> torch.Tensor:
+        if angle.ndim == 2 and angle.shape[1] == 1:


If you implement rotation() in the super class as mentioned above, all this angle checking code can go there and you don't need to reimplement this check in all subclasses.

Refactored.

mhmukadam · 2023-01-23T17:45:16Z

theseus/embodied/robot/joint.py

+from typing import Optional
+import torch
+
+from theseus.geometry.functional import so3, se3


Based on offline discussion we'll move the files to /theseus/labs/kinematics/...

Also the import needs updating.

mhmukadam · 2023-01-23T17:46:43Z

theseus/embodied/robot/joint.py

+from theseus.geometry.functional import so3, se3
+
+
+class Joint(abc.ABC):


Should we have any unit tests accompanying these or will that come in later PRs that add other data structures.

theseus/embodied/robot/joint.py

mhmukadam · 2023-01-23T17:54:17Z

theseus/embodied/robot/joint.py

+            origin = torch.zeros(1, 3, 4, dtype=dtype)
+            origin[:, 0, 0] = 1
+            origin[:, 1, 1] = 1
+            origin[:, 2, 2] = 1


Could we add a identity SE3 helper function and then call that here?

I like this suggestion. Maybe origin = se3.identity()? (this looks like an identity element for SE3, is this correct?)

Will go back to this once identity() is added to the Lie group library.

mhmukadam · 2023-01-23T18:30:14Z

theseus/embodied/robot/joint.py

+        pass
+
+    @abc.abstractmethod
+    def relative_pose(self, angle: torch.Tensor) -> torch.Tensor:


Add a comment about this, for instance relative to what?

I am guessing that as joints are basically pose transforms between links, relative_pose refers to the pose of the child link relative to the parent link.

Maybe a function name like compute_joint_transform or parent_to_child_transform would be more descriptive?

Assuming your interpretation is correct, I like the parent_to_chid_transform name. A shorter version of it would be nice, but nothing comes to mind right now.

Maybe relative_transformation?

mhmukadam · 2023-01-23T18:49:10Z

theseus/embodied/robot/joint.py

+    def dof(self) -> int:
+        return 1
+
+    @abc.abstractmethod


If you have abstract methods then you want to derive from abc as well
class _RevoluteJointImpl(Joint, abc.ABC):

mhmukadam · 2023-01-23T18:56:24Z

theseus/embodied/robot/joint.py

+        return rot
+
+
+class RevoluteJointZ(_RevoluteJointImpl):


Wondering if the following alternative is better/worse in any way: only one RevoluteJoint class that has a full, x, y, or z type as a property (this would get rid of the Impl intermediate as well).

Yeah, this crossed my mind as well. I think I'm OK with either.

I am okay with either as well. We should consider tradeoffs like ease of use and potential overheads.

Wondering if the following alternative is better/worse in any way: only one RevoluteJoint class that has a full, x, y, or z type as a property (this would get rid of the Impl intermediate as well).

I agree with the simpler implementation as well. If the need to have an abstract interface and separate implementations for each class arises, we can always easily refactor.

Now there are two classes for RevoluteJoint, i.e., RevoluteJointFull and RevoluteJointXYZ

mhmukadam

Looks good so far. A few comments for us to hash out.

luisenp · 2023-01-24T21:08:43Z

theseus/embodied/robot/joint.py

+    def set_parent(self, parent: int):
+        self._parent = parent
+
+    def set_child(self, child: int):


On further thought, I think a lot of these properties (maybe all) can be removed and replaced with public attributes. Using a property is more appropriate if we want to hide some extra computation under the hood (e.g., checking for correct inputs).

All the setters have been removed.

exhaustin · 2023-01-27T02:52:58Z

theseus/embodied/robot/joint.py

+        return rot
+
+
+class RevoluteJointZ(_RevoluteJointImpl):


Wondering if the following alternative is better/worse in any way: only one RevoluteJoint class that has a full, x, y, or z type as a property (this would get rid of the Impl intermediate as well).

I agree with the simpler implementation as well. If the need to have an abstract interface and separate implementations for each class arises, we can always easily refactor.

exhaustin · 2023-01-31T00:19:07Z

theseus/embodied/robot/joint.py

+    @property
+    def axis(self) -> torch.Tensor:
+        return self._axis


Only revolute joints have axes, so I don't think a general Joint abstract class should have this as an attribute.

exhaustin · 2023-01-31T00:39:35Z

theseus/embodied/robot/joint.py

+        pass
+
+    @abc.abstractmethod
+    def relative_pose(self, angle: torch.Tensor) -> torch.Tensor:


I am guessing that as joints are basically pose transforms between links, relative_pose refers to the pose of the child link relative to the parent link.

exhaustin · 2023-01-31T00:43:38Z

theseus/embodied/robot/joint.py

+    @abc.abstractmethod
+    def relative_pose(self, angle: torch.Tensor) -> torch.Tensor:
+        pass


Why output an se3 vector instead of an SE3 object?

Should each joint also have its own Jacobian method?

Different types of joints have different parametrizations. A revolute joint requires a single angle value, but that is not the case for other types of joints. Since this is a general abstract Joint class, maybe this function should have a more general input, such as a parameter vector of length self.dof?

For the first question, it might be because this is based on the initial version of the new Lie groups library, which was all regular tensor to regular tensors. This will probably change with the upcoming class-based API (see #450 and #452, for a quick peek).

Agree on the third point, I think this would be preferable.

@fantaosha any thoughts on the second question? (or anything I missed above)

For any types of joints, the input tensor shape should be batch_size x dof.

exhaustin · 2023-01-31T00:43:40Z

theseus/embodied/robot/joint.py

+        pass
+
+    @abc.abstractmethod
+    def relative_pose(self, angle: torch.Tensor) -> torch.Tensor:


Maybe a function name like compute_joint_transform or parent_to_child_transform would be more descriptive?

* 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

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

fantaosha requested review from luisenp and mhmukadam January 2, 2023 03:54

fantaosha self-assigned this Jan 2, 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 2, 2023

Base automatically changed from taoshaf.functional.se3.log to taoshaf.functional.se3.left_project January 12, 2023 20:11

Base automatically changed from taoshaf.functional.se3.left_project to taoshaf.functional.se3.left_act January 12, 2023 20:11

Base automatically changed from taoshaf.functional.se3.left_act to taoshaf.functional.lift_and_project January 12, 2023 20:12

Base automatically changed from taoshaf.functional.lift_and_project to taoshaf.functional.se3.compose January 12, 2023 20:13

Base automatically changed from taoshaf.functional.se3.compose to taoshaf.functional.hat_and_vee January 12, 2023 20:13

Base automatically changed from taoshaf.functional.hat_and_vee to taoshaf.inverse January 12, 2023 20:14

Base automatically changed from taoshaf.inverse to taoshaf.functional.se3.adjoint January 12, 2023 20:14

Base automatically changed from taoshaf.functional.se3.adjoint to taoshaf.functional.se3.rand January 12, 2023 20:15

fantaosha force-pushed the taoshaf.functional.se3.rand branch from d2efac1 to 69762bc Compare January 12, 2023 20:43

Base automatically changed from taoshaf.functional.se3.rand to taoshaf.functional.so3.compose January 12, 2023 20:43

luisenp force-pushed the taoshaf.functional.so3.compose branch from f0f12de to 04fa10a Compare January 20, 2023 18:20

Base automatically changed from taoshaf.functional.so3.compose to main January 20, 2023 22:45

luisenp requested changes Jan 23, 2023

View reviewed changes

mhmukadam reviewed Jan 23, 2023

View reviewed changes

theseus/embodied/robot/joint.py Outdated Show resolved Hide resolved

mhmukadam reviewed Jan 23, 2023

View reviewed changes

theseus/embodied/robot/joint.py Outdated Show resolved Hide resolved

mhmukadam reviewed Jan 23, 2023

View reviewed changes

fantaosha force-pushed the taoshaf.robot.revolute_joint branch from 1a2ce88 to 51702ff Compare January 23, 2023 18:37

mhmukadam reviewed Jan 23, 2023

View reviewed changes

luisenp reviewed Jan 24, 2023

View reviewed changes

exhaustin reviewed Jan 31, 2023

View reviewed changes

exhaustin mentioned this pull request Jan 31, 2023

Add differentiable forward kinematics #420

Merged

11 tasks

luisenp changed the title ~~Add revolute joint~~ Add forward kinematics package to labs Mar 7, 2023

fantaosha and others added 10 commits March 8, 2023 10:22

add revolute joint

788e48c

refactor revolute joint

6a9731d

refactor revolute joint

edfb38c

refactor revolute joint

9ab009c

add parent and child to joint

5ffc17b

update the path to load so3 and se3

ab44d3b

simplify the code

bb3b347

simplify the code

d612e16

minor revision to be consistent with Lie group refacoring

93670ae

fantaosha force-pushed the taoshaf.robot.revolute_joint branch from 22ff507 to 93670ae Compare March 8, 2023 18:36

fantaosha added 2 commits March 8, 2023 14:16

fixed a CI error

8ae533f

move files to labs

992cd12

fantaosha merged commit cd2fc2c into main Mar 10, 2023

fantaosha deleted the taoshaf.robot.revolute_joint branch March 10, 2023 21:42

		from theseus.geometry.functional import so3, se3


		class Joint(abc.ABC):

Add forward kinematics package to labs #415

Add forward kinematics package to labs #415

Conversation

fantaosha commented Jan 2, 2023

Motivation and Context

How Has This Been Tested

Types of changes

Checklist

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

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

luisenp Jan 23, 2023 • edited Loading

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

mhmukadam left a comment

Choose a reason for hiding this comment

luisenp Jan 24, 2023 • edited Loading

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

Add forward kinematics package to `labs` #415

Add forward kinematics package to `labs` #415

luisenp Jan 23, 2023 •

edited

Loading

luisenp Jan 24, 2023 •

edited

Loading