UM-ARM-Lab/pytorch_kinematics

Feature request: coupled joints

Opened this issue · 5 comments

Hi,
I think it would be great if one could specify that certain joints are mechanically coupled together. It is really quite common to run into this problem. I had to implement it (for a simple 1:1 coupling) myself. Let me know if you are interested in adding this feature. I might be able to send a PR if I find some free time, but of course I wouldn't mind if you add it before I get to it :D

After loading the urdf one could specify with a matrix which joints are linked together, and what gear ratios are used.
And then the forward_kinematics and jacobian function would simply expect/return fewer inputs/outputs.

Hey it'd be great if you could submit a PR since you've already made it for yourself!

I mean this is all did so far:

class ForwardKinematics():
    def __init__(self):
        self.num_in_joints = 3
        self.num_fk_joints = 4
        self.chain = pk.build_serial_chain_from_urdf(
            open("finger.urdf").read(), "virtual_distal_link"
        )

    def _joints_from_x(self, x):
        return torch.stack(
            [
                x[..., 0],
                x[..., 1],
                x[..., 2],
                x[..., 2],
            ],
            dim=-1,
        )

    def forward(self, x):
        joints = self._joints_from_x(x)
        fk_object = self.chain.forward_kinematics(joints)
        matrix = fk_object.get_matrix()
        position = matrix[:, :3, 3]
        return position

    def jacobian(self, x):
        joints = self._joints_from_x(x)
        jacobian = self.chain.jacobian(joints)
        cartesian_jacobian = jacobian[:, : self.num_in_joints, :]
        coupled_joint = torch.stack(
            [
                torch.zeros((x.shape[0], self.num_in_joints)),
                torch.zeros((x.shape[0], self.num_in_joints)),
                cartesian_jacobian[:, :, self.num_fk_joints - 1],
            ],
            dim=-1,
        )
        coupled_jacobian = (
            cartesian_jacobian[:, :, : self.num_in_joints] + coupled_joint
        )
        return coupled_jacobian

Doing the general case will be a different story.
I might work on it next week, but I can't promise anything.

If we wanted to implement this more generally, one approach might be to get them from the robot specification file ("mimic" for urdf or "equality/joint" for mujoco XML).

This feature could be useful for Val, since it has this for its grippers.

wonder if this feature is planned, it will be super useful!