-
Notifications
You must be signed in to change notification settings - Fork 121
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
Add forward kinematics package to labs
#415
Conversation
d2efac1
to
69762bc
Compare
f0f12de
to
04fa10a
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks good! Requesting some small changes that should be easy to address.
theseus/embodied/robot/joint.py
Outdated
|
||
if dtype is not None and origin.dtype != dtype: | ||
warnings.warn( | ||
f"tensor.dtype {origin.dtype} does not match given dtype {dtype}, " |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe clearer to write "The origin's dtype ({origin.dtype}) does not match the desired dtype. Origin's dtype will take precendence".
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Refactored.
theseus/embodied/robot/joint.py
Outdated
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.") |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
numel()
is used now and the error message is "The revolute axis must have 3 elements."
.
theseus/embodied/robot/joint.py
Outdated
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: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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 `[]`
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed.
theseus/embodied/robot/joint.py
Outdated
def axis(self) -> torch.Tensor: | ||
return self._axis | ||
|
||
def set_parent(self, parent: int): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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):
...
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
All the setters have been removed for encapsulation.
theseus/embodied/robot/joint.py
Outdated
return 1 | ||
|
||
@abc.abstractmethod | ||
def _rotation_impl(self, angle: torch.Tensor) -> torch.Tensor: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
_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
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If just a private/internal implementation is needed then def _rotation(...)
is sufficient.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
rotation()
and translation()
are defined.
theseus/embodied/robot/joint.py
Outdated
self._axis[4] = 1 | ||
|
||
def _rotation_impl(self, angle: torch.Tensor) -> torch.Tensor: | ||
if angle.ndim == 2 and angle.shape[1] == 1: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Refactored.
theseus/embodied/robot/joint.py
Outdated
from typing import Optional | ||
import torch | ||
|
||
from theseus.geometry.functional import so3, se3 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Based on offline discussion we'll move the files to /theseus/labs/kinematics/...
Also the import needs updating.
theseus/embodied/robot/joint.py
Outdated
from theseus.geometry.functional import so3, se3 | ||
|
||
|
||
class Joint(abc.ABC): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
Outdated
origin = torch.zeros(1, 3, 4, dtype=dtype) | ||
origin[:, 0, 0] = 1 | ||
origin[:, 1, 1] = 1 | ||
origin[:, 2, 2] = 1 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Could we add a identity SE3 helper function and then call that here?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I like this suggestion. Maybe origin = se3.identity()
? (this looks like an identity element for SE3, is this correct?)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Will go back to this once identity()
is added to the Lie group library.
theseus/embodied/robot/joint.py
Outdated
pass | ||
|
||
@abc.abstractmethod | ||
def relative_pose(self, angle: torch.Tensor) -> torch.Tensor: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Add a comment about this, for instance relative to what?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe a function name like compute_joint_transform
or parent_to_child_transform
would be more descriptive?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe relative_transformation
?
1a2ce88
to
51702ff
Compare
theseus/embodied/robot/joint.py
Outdated
def dof(self) -> int: | ||
return 1 | ||
|
||
@abc.abstractmethod |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If you have abstract methods then you want to derive from abc as well
class _RevoluteJointImpl(Joint, abc.ABC):
theseus/embodied/robot/joint.py
Outdated
return rot | ||
|
||
|
||
class RevoluteJointZ(_RevoluteJointImpl): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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).
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah, this crossed my mind as well. I think I'm OK with either.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I am okay with either as well. We should consider tradeoffs like ease of use and potential overheads.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Now there are two classes for RevoluteJoint
, i.e., RevoluteJointFull
and RevoluteJointXYZ
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks good so far. A few comments for us to hash out.
theseus/embodied/robot/joint.py
Outdated
def set_parent(self, parent: int): | ||
self._parent = parent | ||
|
||
def set_child(self, child: int): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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).
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
All the setters have been removed.
theseus/embodied/robot/joint.py
Outdated
return rot | ||
|
||
|
||
class RevoluteJointZ(_RevoluteJointImpl): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
theseus/embodied/robot/joint.py
Outdated
@property | ||
def axis(self) -> torch.Tensor: | ||
return self._axis |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Only revolute joints have axes, so I don't think a general Joint
abstract class should have this as an attribute.
theseus/embodied/robot/joint.py
Outdated
pass | ||
|
||
@abc.abstractmethod | ||
def relative_pose(self, angle: torch.Tensor) -> torch.Tensor: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
theseus/embodied/robot/joint.py
Outdated
@abc.abstractmethod | ||
def relative_pose(self, angle: torch.Tensor) -> torch.Tensor: | ||
pass |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
- 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 abstractJoint
class, maybe this function should have a more general input, such as a parameter vector of lengthself.dof
?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
For any types of joints, the input tensor
shape should be batch_size x dof
.
theseus/embodied/robot/joint.py
Outdated
pass | ||
|
||
@abc.abstractmethod | ||
def relative_pose(self, angle: torch.Tensor) -> torch.Tensor: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
22ff507
to
93670ae
Compare
Motivation and Context
How Has This Been Tested
Types of changes
Checklist