Add forward kinematics package to labs

tests/labs/embodied/robot/test_forward_kinematics.py

@@ -0,0 +1,86 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import pytest
+import os
+import torch
+from tests.decorators import run_if_labs
+
+
+URDF_REL_PATH = "../../../embodied/kinematics/data/panda_no_gripper.urdf"
+urdf_path = os.path.join(os.path.dirname(__file__), URDF_REL_PATH)
+
+
+@run_if_labs()
+@pytest.mark.parametrize("batch_size", [1, 20, 40])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_backward(batch_size: int, dtype: torch.dtype):
+    from theseus.labs.lie.functional.constants import TEST_EPS
+    from theseus.labs.embodied.robot.forward_kinematics import Robot
+    from theseus.labs.embodied.robot.forward_kinematics import (
+        get_forward_kinematics,
+        ForwardKinematicsFactory,
+    )
+
+    robot = Robot.from_urdf_file(urdf_path, dtype)
+    selected_links = ["panda_link2", "panda_link5", "panda_virtual_ee_link"]
+    _, fkin_impl, _, _, _ = ForwardKinematicsFactory(robot, selected_links)
+    fkin, _ = get_forward_kinematics(robot, selected_links)
+
+    rng = torch.Generator()
+    rng.manual_seed(0)
+    angles = torch.rand(batch_size, robot.dof, generator=rng, dtype=dtype)
+
+    jacs_impl = torch.autograd.functional.jacobian(fkin_impl, angles, vectorize=True)
+    jacs = torch.autograd.functional.jacobian(fkin, angles, vectorize=True)
+
+    for jac, jac_impl in zip(jacs, jacs_impl):
+        torch.testing.assert_close(
+            actual=jac_impl, expected=jac, atol=TEST_EPS, rtol=1e-5
+        )
+
+    grads = []
+    for func in [fkin, fkin_impl]:
+        temp = angles.clone()
+        temp.requires_grad = True
+        loss = torch.tensor(0, dtype=dtype)
+        for pose in func(temp):
+            loss = loss + (pose**2).sum()
+        loss.backward()
+        grads.append(temp.grad)
+    torch.testing.assert_close(actual=grads[0], expected=grads[1], atol=1e-6, rtol=1e-5)
+
+
+@run_if_labs()
+@pytest.mark.parametrize("batch_size", [1, 20, 40])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_jacobian(batch_size: int, dtype: torch.dtype):
+    from theseus.labs.lie.functional import se3
+    from theseus.labs.embodied.robot.forward_kinematics import Robot
+    from theseus.labs.embodied.robot.forward_kinematics import get_forward_kinematics
+
+    robot = Robot.from_urdf_file(urdf_path, dtype)
+    selected_links = ["panda_link2", "panda_link5", "panda_virtual_ee_link"]
+    fkin, jfkin = get_forward_kinematics(robot, selected_links)
+
+    rng = torch.Generator()
+    rng.manual_seed(0)
+    angles = torch.rand(batch_size, robot.dof, generator=rng, dtype=dtype)
+
+    jacs_actual, poses = jfkin(angles)
+
+    sels = range(batch_size)
+    jacs_dense = torch.autograd.functional.jacobian(fkin, angles, vectorize=True)
+    jacs_expected = []
+    for pose, jac_dense in zip(poses, jacs_dense):
+        jac_dense = jac_dense[sels, :, :, sels].transpose(-1, 1).transpose(-1, -2)
+        jac_expected = se3.left_project(pose, jac_dense).transpose(-1, -2)
+        jac_expected[:, 3:] *= 0.5
+        jacs_expected.append(jac_expected)
+
+    for jac_actual, jac_expected in zip(jacs_actual, jacs_expected):
+        torch.testing.assert_close(
+            actual=jac_actual, expected=jac_expected, atol=1e-6, rtol=1e-5
+        )

theseus/labs/embodied/robot/forward_kinematics.py

@@ -0,0 +1,137 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+from typing import List, Optional
+
+from theseus.labs.lie.functional import se3
+from .robot import Robot, Joint, Link
+
+# TODO: Add support for joints with DOF>1
+
+
+def ForwardKinematicsFactory(robot: Robot, link_names: Optional[List[str]] = None):
+    links: List[Link] = robot.get_links(link_names)
+    link_ids: List[int] = [link.id for link in links]
+
+    ancestor_links: List[Link] = []
+    joint_ids: List[int] = []
+    for link in links:
+        ancestor_links += [anc for anc in link.ancestor_links]
+        joint_ids += link.ancestor_non_fixed_joint_ids
+    pose_ids = sorted(list(set([anc.id for anc in ancestor_links] + link_ids)))
+    joint_ids = sorted(list(set(joint_ids)))
+
+    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"
+            )
+
+        poses: List[Optional[torch.Tensor]] = [None] * robot.num_links
+        poses[0] = angles.new_zeros(angles.shape[0], 3, 4)
+        poses[0][:, 0, 0] = 1
+        poses[0][:, 1, 1] = 1
+        poses[0][:, 2, 2] = 1
+
+        for id in pose_ids[1:]:
+            curr: Link = robot.links[id]
+            joint: Joint = robot.links[id].parent_joint
+            prev: Link = joint.parent_link
+            relative_pose = (
+                joint.relative_pose(angles[:, joint.id])
+                if joint.id < robot.dof
+                else joint.relative_pose()
+            )
+            poses[curr.id] = se3.compose(poses[prev.id], relative_pose)
+
+        return poses
+
+    def _forward_kinematics_impl(angles: torch.Tensor):
+        poses = _forward_kinematics_helper(angles)
+        return tuple(poses[id] for id in link_ids)
+
+    def _jforward_kinematics_helper(
+        poses: List[Optional[torch.Tensor]],
+    ) -> torch.Tensor:
+        jposes = poses[0].new_zeros(poses[0].shape[0], 6, robot.dof)
+
+        for id in pose_ids[1:]:
+            link: Link = robot.links[id]
+            joint: Joint = link.parent_joint
+            if joint.id >= robot.dof:
+                break
+            jposes[:, :, joint.id : joint.id + 1] = se3.adj(poses[link.id]) @ joint.axis
+
+        return jposes
+
+    def _jforward_kinematics_impl(angles: torch.Tensor):
+        poses = _forward_kinematics_helper(angles)
+        jposes = _jforward_kinematics_helper(poses)
+
+        rets = tuple(poses[id] for id in link_ids)
+        jacs: List[torch.Tensor] = []
+
+        for link_id in link_ids:
+            pose = poses[link_id]
+            jac = jposes.new_zeros(angles.shape[0], 6, robot.dof)
+            sel = robot.links[link_id].ancestor_non_fixed_joint_ids
+            jac[:, :, sel] = se3.adj(se3.inv(pose)) @ jposes[:, :, sel]
+            jacs.append(jac)
+
+        return jacs, rets
+
+    class ForwardKinematics(torch.autograd.Function):
+        @classmethod
+        def forward(cls, ctx, angles):
+            poses = _forward_kinematics_helper(angles)
+            ctx.poses = poses
+            rets = tuple(poses[id] for id in link_ids)
+            ctx.rets = rets
+            return rets
+
+        @classmethod
+        def backward(cls, ctx, *grad_outputs):
+            if not hasattr(ctx, "jposes"):
+                ctx.jposes: torch.Tensor = _jforward_kinematics_helper(ctx.poses)
+            rets: tuple(torch.Tensor) = ctx.rets
+            grad_pose = grad_outputs[0].new_zeros(
+                grad_outputs[0].shape[0], 6, robot.dof
+            )
+            grad_input = grad_outputs[0].new_zeros(grad_outputs[0].shape[0], robot.dof)
+
+            for link_id, ret, grad_output in zip(link_ids, rets, grad_outputs):
+                ancestor_non_fixed_joint_ids = robot.links[
+                    link_id
+                ].ancestor_non_fixed_joint_ids
+                temp = se3.project(
+                    torch.cat(
+                        (grad_output @ ret.transpose(1, 2), grad_output[:, :, 3:]),
+                        dim=-1,
+                    )
+                ).unsqueeze(-1)
+                grad_pose[:, :, ancestor_non_fixed_joint_ids] += temp
+            grad_input[:, joint_ids] = (
+                ctx.jposes[:, :, joint_ids] * grad_pose[:, :, joint_ids]
+            ).sum(-2)
+
+            return grad_input
+
+    return (
+        ForwardKinematics,
+        _forward_kinematics_impl,
+        _jforward_kinematics_impl,
+        _forward_kinematics_helper,
+        _jforward_kinematics_helper,
+    )
+
+
+def get_forward_kinematics(robot: Robot, link_names: Optional[List[str]] = None):
+    ForwardKinematics, _, jforward_kinematics, _, _ = ForwardKinematicsFactory(
+        robot, link_names
+    )
+    forward_kinematics = ForwardKinematics.apply
+    return forward_kinematics, jforward_kinematics