UM-ARM-Lab / pytorch_kinematics

Robot kinematics implemented in pytorch
MIT License
394 stars 34 forks source link

Feature request: coupled joints #6

Open johannespitz opened 2 years ago

johannespitz commented 2 years ago

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.

LemonPi commented 2 years ago

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

johannespitz commented 2 years ago

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.

PeterMitrano commented 10 months ago

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).

PeterMitrano commented 10 months ago

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

zoctipus commented 5 months ago

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