PyTorch 机器人运动学

• 并行且可微的正向运动学（FK）、雅可比矩阵计算，以及阻尼最小二乘逆向运动学（IK）
• 从 URDF、SDF 和 MJCF 格式加载机器人描述
• 通过 pytorch-volumetric 实现对不同配置和点的 SDF 查询批量处理

安装

pip install pytorch-kinematics

如需开发版本，请将仓库克隆至任意位置，然后运行 pip3 install -e . 以进入可编辑模式进行安装。

参考

如果您在研究中使用本包，请考虑引用以下文献：

@software{Zhong_PyTorch_Kinematics_2024,
author = {Zhong, Sheng and Power, Thomas and Gupta, Ashwin and Mitrano, Peter},
doi = {10.5281/zenodo.7700587},
month = feb,
title = {{PyTorch 运动学}},
version = {v0.7.1},
year = {2024}
}

使用方法

请参阅 tests 中的代码示例；部分示例也在此处展示。

加载机器人

import pytorch_kinematics as pk

urdf = "widowx/wx250s.urdf"
# 由于存在多个自然末端执行器链接，因此并非串联链
chain = pk.build_chain_from_urdf(open(urdf, mode="rb").read())
# 可视化各帧结构（字符串同样会返回）
chain.print_tree()
"""
base_link
└── shoulder_link
    └── upper_arm_link
        └── upper_forearm_link
            └── lower_forearm_link
                └── wrist_link
                    └── gripper_link
                        └── ee_arm_link
                            ├── gripper_prop_link
                            └── gripper_bar_link
                                └── fingers_link
                                    ├── left_finger_link
                                    ├── right_finger_link
                                    └── ee_gripper_link
"""

# 提取特定的串联链，例如用于逆向运动学
serial_chain = pk.SerialChain(chain, "ee_gripper_link", "base_link")
serial_chain.print_tree()
"""
base_link
└── shoulder_link
    └── upper_arm_link
        └── upper_forearm_link
            └── lower_forearm_link
                └── wrist_link
                    └── gripper_link
                        └── ee_arm_link
                            └── gripper_bar_link
                                └── fingers_link
                                    └── ee_gripper_link
"""

# 您也可以提取与原始链根节点不同的串联链
serial_chain = pk.SerialChain(chain, "ee_gripper_link", "gripper_link")
serial_chain.print_tree()
"""
gripper_link
└── ee_arm_link
    └── gripper_bar_link
        └── fingers_link
            └── ee_gripper_link
"""

正向运动学（FK）

import math
import pytorch_kinematics as pk

# 从 URDF 加载机器人描述，并指定末端执行器链接
chain = pk.build_serial_chain_from_urdf(open("kuka_iiwa.urdf").read(), "lbr_iiwa_link_7")
# 打印出链接的嵌套树结构
print(chain)
# 打印出关节名称列表
print(chain.get_joint_parameter_names())

# 指定关节值（可采用多种形式）
th = [0.0, -math.pi / 4.0, 0.0, math.pi / 2.0, 0.0, math.pi / 4.0, 0.0]
# 进行正向运动学计算并获取变换对象；若设置 end_only=False，则会返回所有链接的变换矩阵
ret = chain.forward_kinematics(th, end_only=False)
# 查找特定链接的变换矩阵
tg = ret['lbr_iiwa_link_7']
# 获取变换矩阵（1,4,4），随后将其转换为独立的位置和单位四元数
m = tg.get_matrix()
pos = m[:, :3, 3]
rot = pk.matrix_to_quaternion(m[:, :3, :3])

我们可以通过传入 2D 关节值来实现 FK 的并行化，同时在可用时也可使用 CUDA。

import torch
import pytorch_kinematics as pk

d = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float64

chain = pk.build_serial_chain_from_urdf(open("kuka_iiwa.urdf").read(), "lbr_iiwa_link_7")
chain = chain.to(dtype=dtype, device=d)

N = 1000
th_batch = torch.rand(N, len(chain.get_joint_parameter_names()), dtype=dtype, device=d)

# 在并行环境下，按数量级排序能显著提升运算速度
# 对于 N=1000，使用并行方式运行时耗时 0.008678913116455078 秒
# （N,4,4）变换矩阵；默认情况下，仅返回末端执行器的变换矩阵，因为 end_only=True
tg_batch = chain.forward_kinematics(th_batch)

# 对于 N=1000，串行方式运行时耗时 8.44686508178711 秒
for i in range(N):
    tg = chain.forward_kinematics(th_batch[i])

我们还可以通过 FK 计算梯度。

import torch
import math
import pytorch_kinematics as pk

chain = pk.build_serial_chain_from_urdf(open("kuka_iiwa.urdf").read(), "lbr_iiwa_link_7")

# 需要对输入关节值进行梯度计算
th = torch.tensor([0.0, -math.pi / 4.0, 0.0, math.pi / 2.0, 0.0, math.pi / 4.0, 0.0], requires_grad=True)
tg = chain.forward_kinematics(th)
m = tg.get_matrix()
pos = m[:, :3, 3]
pos.norm().backward()
# 现在 th.grad 已被填充

我们同样可以加载 SDF 和 MJCF 描述，并通过字典形式传递关节值（未指定的关节默认为 th=0），适用于非串联链。

import math
import torch
import pytorch_kinematics as pk

chain = pk.build_chain_from_sdf(open("simple_arm.sdf").read())
ret = chain.forward_kinematics({'arm_elbow_pan_joint': math.pi / 2.0, 'arm_wrist_lift_joint': -0.5})
# 请记住，我们需指定关节值并获取链接变换
tg = ret['arm_wrist_roll']

# 也可以采用并行方式完成此操作
N = 100
ret = chain.forward_kinematics({'arm_elbow_pan_joint': torch.rand(N, 1), 'arm_wrist_lift_joint': torch.rand(N, 1)})
# （N, 4, 4）变换对象
tg = ret['arm_wrist_roll']

# 从 MJCF 文件构建机器人
chain = pk.build_chain_from_mjcf(open("ant.xml").read())
print(chain)
print(chain.get_joint_parameter_names())
th = {'hip_1': 1.0, 'ankle_1': 1}
ret = chain.forward_kinematics(th)

chain = pk.build_chain_from_mjcf(open("humanoid.xml").read())
print(chain)
print(chain.get_joint_parameter_names())
th = {'left_knee': 0.0, 'right_knee': 0.0}
ret = chain.forward_kinematics(th)

torch.compile 支持

通过 torch.compile(fullgraph=True) 可对 FK 计算进行编译，从而大幅提升性能，尤其适用于需要多次调用 FK 的应用（例如逆向运动学、轨迹优化等）。

使用 forward_kinematics_tensor，这是一种针对编译友好的变体，可接受并返回原始张量，而非字典或 Transform3d 对象：

import torch
import pytorch_kinematics as pk

chain = pk.build_serial_chain_from_urdf(open("kuka_iiwa.urdf").read(), "lbr_iiwa_link_7")

# 编译 FK 内核（一次性成本）
compiled_fk = torch.compile(chain.forward_kinematics_tensor, fullgraph=True, dynamic=True)

# 输入：(B, n_joints) 张量
th = torch.randn(1000, 7)

# 输出：(num_frames, B, 4, 4) 张量，包含所有帧的齐次变换矩阵
all_transforms = compiled_fk(th)

简化后的中文内容：

通过帧索引进行查找：使用 chain.frame_to_idx 来查询某个帧名称对应的索引

ee_idx = chain.frame_to_idx['lbr_iiwa_link_7'] ee_transform = all_transforms[ee_idx] # (B, 4, 4)

在 CPU 上的典型加速比（Kuka IIWA 7 自由度机器人）：

| 批次大小 | `forward_kinematics` | 编译后的 `forward_kinematics_tensor` | 加速比 |
|---:|---:|---:|---:|
| 1 | 0.21 ms | 0.04 ms | **4.7倍** |
| 64 | 0.26 ms | 0.08 ms | **3.5倍** |
| 1024 | 1.13 ms | 0.51 ms | **2.2倍** |

对于具有更多帧数的复杂机器人，加速比会更高（例如，在批次为 1 时，49 帧机器人可实现 6 倍以上的加速）。
借助 CUDA 图捕获的 GPU 加速，可在较大的批量规模下进一步提升性能。

标准的 `forward_kinematics` 方法（返回 `Dict[str, Transform3d]`）同样受益于经过重构的内部逻辑，无需任何代码变更。与 v0.7 相比，该方法在小批量规模下速度提升了 2–6 倍。

旋转转换函数（如 `matrix_to_quaternion`、`quaternion_to_axis_angle`、`axis_angle_to_quaternion` 等）也兼容 `torch.compile(fullgraph=True)`，可以单独编译，也可以作为更大编译图的一部分进行编译：

```python
import torch
import pytorch_kinematics as pk

# 编译一个旋转转换函数
compiled_m2q = torch.compile(pk.matrix_to_quaternion, fullgraph=True)
R = torch.randn(100, 3, 3)
q = compiled_m2q(R)

柱坐标计算

在运动学上下文中，柱坐标矩阵描述了末端执行器随关节值变化而发生的变化情况（其中 表示扭转量，或堆叠的线速度和角速度）：

对于 SerialChain，我们提供了一种可微且可并行化的方法，用于计算基于基座框架的柱坐标矩阵。

import math
import torch
import pytorch_kinematics as pk

# 可以通过选择末端执行器框架，将 Chain 转换为 SerialChain
chain = pk.build_chain_from_sdf(open("simple_arm.sdf").read())
# 打印链结构，查看可用的链接，用作末端执行器
# 请注意，任意链接都可以被选中；不必一定是没有子级的链接
chain = pk.SerialChain(chain, "arm_wrist_roll_frame")

chain = pk.build_serial_chain_from_urdf(open("kuka_iiwa.urdf").read(), "lbr_iiwa_link_7")
th = torch.tensor([0.0, -math.pi / 4.0, 0.0, math.pi / 2.0, 0.0, math.pi / 4.0, 0.0])
# (1,6,7) 张量，其中 7 代表机器人的自由度
J = chain.jacobian(th)

# 并行计算柱坐标矩阵，并在可用时利用 CUDA
N = 1000
d = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float64

chain = chain.to(dtype=dtype, device=d)
# 柱坐标计算是可微的
th = torch.rand(N, 7, dtype=dtype, device=d, requires_grad=True)
# (N,6,7)
J = chain.jacobian(th)

# 可以在距离末端执行器一定偏移的位置上获取柱坐标矩阵（位置指定在 EE 框架中）
# 默认情况下，位置位于 EE 框架的原点
loc = torch.rand(N, 3, dtype=dtype, device=d)
J = chain.jacobian(th, locations=loc)

与 FK 类似，柱坐标矩阵也有一个兼容 torch.compile 的变体——jacobian_tensor：

chain = pk.build_serial_chain_from_urdf(open("kuka_iiwa.urdf").read(), "lbr_iiwa_link_7")

# 编译柱坐标矩阵核函数（一次性成本）
compiled_jac = torch.compile(chain.jacobian_tensor, fullgraph=True)

# 输入：(B, n_joints) 张量；输出：(B, 6, n_joints) 柱坐标矩阵
th = torch.randn(100, 7)
J = compiled_jac(th)

为了在紧密循环中获得最佳性能（例如在 IK 中），可向 torch.compile 传递 mode='max-autotune'，以在编译成本较高的前提下，对核函数进行额外的优化调整。

在 CPU 上的典型加速比（Kuka IIWA 7 自由度机器人，对比旧版 v0.7 的 calc_jacobian）：

批次大小	旧版 calc_jacobian	新版 jacobian（即时计算）	编译后的 jacobian_tensor	编译版与旧版的加速比
1	0.63 ms	0.12 ms	0.03 ms	21倍
10	0.79 ms	0.15 ms	0.05 ms	15倍
100	0.75 ms	0.30 ms	0.09 ms	8倍
1,000	1.67 ms	1.23 ms	0.53 ms	3倍
10,000	7.93 ms	7.01 ms	4.91 ms	1.6倍
100,000	59.42 ms	64.58 ms	46.46 ms	1.3倍

在小到中等批量规模下（典型的 IK 应用场景），由于向量化计算取代了每帧的 Python 循环，新版的即时计算实现速度提升了 2–5 倍。 torch.compile 在此基础上又带来了 2–4 倍的性能提升，且在各个批量规模下都优于旧版实现。

柱坐标矩阵可用于逆运动学计算。有关如何进行逆运动学计算的详细信息，请参阅 IK 综述。需要注意的是，逆运动学可以通过其他方式更好地完成（不过，通过柱坐标矩阵进行逆运动学计算，可以实现端到端的可微化方法）。

逆运动学（IK）

逆运动学可通过阻尼最小二乘法实现（通过迭代步骤结合柱坐标伪逆，以避免在奇异点附近出现振荡）。与其他逆运动学库相比，其典型优势如下：

• 不依赖 ROS（许多逆运动学库需要将机器人模型发布至 ROS 参数服务器）
• 支持目标姿态的批量处理，以及从不同起始配置中进行重试
• 除了目标位置外，还支持目标方向

有关使用方法，请参阅 tests/test_inverse_kinematics.py，但通常您需要以下内容：

full_urdf = os.path.join(search_path, urdf)
chain = pk.build_serial_chain_from_urdf(open(full_urdf).read(), "lbr_iiwa_link_7")

# 目标以 **机器人框架中的 Transform3d 姿态** 进行指定
# 因此，如果您已将目标以世界坐标系中的姿态进行指定，则还需要将机器人坐标系中的姿态纳入世界坐标系中
pos = torch.tensor([0.0, 0.0, 0.0], device=device)
rot = torch.tensor([0.0, 0.0, 0.0], device=device)
rob_tf = pk.Transform3d(pos=pos, rot=rot, device=device)

# 将目标以世界坐标系中的 Transform3d 姿态进行指定
goal_in_world_frame_tf = ...
# 转换为机器人坐标系（如果目标已以机器人坐标系中指定，或若世界坐标系与机器人坐标系一致，则跳过此步骤）
goal_in_rob_frame_tf = rob_tf.inverse().compose(goal_tf)

# 获取机器人关节限位
lim = torch.tensor(chain.get_joint_limits(), device=device)

# 创建 IK 对象
# 请参阅构造函数，了解更多选项及其说明，例如收敛容差等
ik = pk.PseudoInverseIK(chain, max_iterations=30, num_retries=10,
                        joint_limits=lim.T,
                        early_stopping_any_converged=True,
                        early_stopping_no_improvement="all",
                        debug=False,
                        lr=0.2)
# 解决 IK
sol = ik.solve(goal_in_rob_frame_tf)
# 目标数量 × 重试次数 × 关节角度张量；若未收敛，仅显示迄今为止找到的最佳解
print(sol.solutions)
# 目标数量 × 重试次数，可检查每次运行的收敛情况
print(sol.converged)
# 目标数量 × 重试次数，可直接查看误差
print(sol.err_pos)
print(sol.err_rot)

编译 IK

对于那些会反复求解 IK 的工作负载（例如在规划循环中），请在 IK 循环内部，通过 use_compile=True 来编译 FK、雅可比矩阵以及阻尼最小二乘核函数。这样可以利用 PyTorch 的编译功能实现高效运行。而外层的 IK 循环及其基于数据的收敛性检查则仍保持在 eager 模式下。

ik = pk.PseudoInverseIK(chain, max_iterations=30, num_retries=10,
                         joint_limits=lim.T,
                         early_stopping_any_converged=True,
                         early_stopping_no_improvement="all",
                         lr=0.2,
                         use_compile=True)  # 编译内部核函数

# 第一次调用 solve() 会花费一次编译成本；后续调用则会更快
sol = ik.solve(goal_in_rob_frame_tf)

在 CPU 上的典型加速效果（Kuka IIWA 7 自由度，30 次迭代）：

重试次数	Eager (v0.7)	Eager (新版本)	编译后	编译与 v0.7 相比的加速比
10	25.6 ms	23.3 ms	11.8 ms	2.2x
50	31.5 ms	29.0 ms	15.8 ms	2.0x

注意：use_compile=True 需要 PyTorch 2.0+。由于外层循环中存在依赖于数据的控制流，您无法直接使用 torch.compile 来包装 ik.solve()——请改用此标志。

SDF 查询

最新详情请参见 pytorch-volumetric，部分说明已粘贴在此处：

对于许多应用，例如碰撞检测，掌握多连杆机器人在特定配置下的 SDF 是非常有用的。 首先，我们通过 PyTorch 运动学 创建机器人模型，并从 URDF、SDF、MJCF 等格式中加载模型。 例如，我们将使用 pybullet 数据中的 KUKA 7 自由度机械臂模型。

import os
import torch
import pybullet_data
import pytorch_kinematics as pk
import pytorch_volumetric as pv

urdf = "kuka_iiwa/model.urdf"
search_path = pybullet_data.getDataPath()
full_urdf = os.path.join(search_path, urdf)
chain = pk.build_serial_chain_from_urdf(open(full_urdf).read(), "lbr_iiwa_link_7")
d = "cuda" if torch.cuda.is_available() else "cpu"

chain = chain.to(device=d)
# 连杆网格文件的路径是根据其在 URDF 中的相对路径来指定的。为了正确加载这些文件，我们需要为它们提供路径前缀。
s = pv.RobotSDF(chain, path_prefix=os.path.join(search_path, "kuka_iiwa"))

默认情况下，每个连杆都会拥有一个 MeshSDF。若希望使用 CachedSDF 以提升查询速度，则可以这样做：

s = pv.RobotSDF(chain, path_prefix=os.path.join(search_path, "kuka_iiwa"),
                link_sdf_cls=pv.cache_link_sdf_factory(resolution=0.02, padding=1.0, device=d))

当可视化 y=0.02 的 SDF 切片时：

其中，表面点对应如下：

对这一 SDF 的查询会根据关节配置进行调整（默认情况下，所有关节配置均为零）。 查询会针对不同配置和查询点进行批量处理。例如，我们有一组需要查询的关节配置：

th = torch.tensor([0.0, -math.pi / 4.0, 0.0, math.pi / 2.0, 0.0, math.pi / 4.0, 0.0], device=d)
N = 200
th_perturbation = torch.randn(N - 1, 7, device=d) * 0.1
# N x 7 个关节值
th = torch.cat((th.view(1, -1), th_perturbation + th))

此外，我们还有一组需要查询的点（每组配置对应的点相同）：

y = 0.02
query_range = np.array([
    [-1, 0.5],
    [y, y],
    [-0.2, 0.8],
])
# M x 3 个点
coords, pts = pv.get_coordinates_and_points_in_grid(0.01, query_range, device=s.device)

我们设置了关节配置和查询的批量参数：

s.set_joint_configuration(th)
# N x M 个 SDF 值
# N x M x 3 个 SDF 偏导数
sdf_val, sdf_grad = s(pts)

资助来源

• pytorch_kinematics/transforms 从 pytorch3d 中提取而来，并进行了少量扩展。 之所以选择不将 pytorch3d 作为依赖项，是因为它安装起来较为困难，且大部分代码与我们的需求并无直接关联。 一个重要的区别在于，我们在机器人学中沿用左乘变换的方式（T * pt），而非右乘变换。
• pytorch_kinematics/urdf_parser_py 和 pytorch_kinematics/mjcf_parser 分别源自 kinpy，并结合了 FK 逻辑。该仓库将这些逻辑移植到 PyTorch 中，实现了并行化，并提供了一些扩展功能。

pytorch_kinematics 快速上手指南

pytorch_kinematics 是一个基于 PyTorch 的机器人运动学库，支持并行化、可微分的前向运动学（FK）、雅可比矩阵计算以及阻尼最小二乘法逆运动学（IK）。它可以直接从 URDF、SDF 和 MJCF 文件加载机器人描述，并充分利用 GPU 加速。

环境准备

• 操作系统：Linux, macOS, Windows
• Python 版本：建议 Python 3.8+
• 核心依赖：
  • torch (PyTorch)
  • numpy
• 可选依赖：若需使用 torch.compile 加速功能，建议使用较新版本的 PyTorch (2.0+)。

安装步骤

1. 基础安装

通过 pip 直接安装稳定版：

pip install pytorch-kinematics

提示：国内用户如遇下载缓慢，可使用清华或阿里镜像源加速：

pip install pytorch-kinematics -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 开发模式安装

如需修改源码或贡献代码，可克隆仓库后以可编辑模式安装：

git clone https://github.com/UM-ARM-Lab/pytorch_kinematics.git
cd pytorch_kinematics
pip3 install -e .

基本使用

以下示例展示如何加载 URDF 文件并执行前向运动学（FK）计算。

1. 加载机器人与前向运动学 (FK)

import math
import torch
import pytorch_kinematics as pk

# 1. 从 URDF 文件构建串联链 (SerialChain)
# 指定末端执行器链接名称，例如 "lbr_iiwa_link_7"
chain = pk.build_serial_chain_from_urdf(open("kuka_iiwa.urdf").read(), "lbr_iiwa_link_7")

# 查看关节参数名称
print(chain.get_joint_parameter_names())

# 2. 定义关节角度 (弧度制)
th = [0.0, -math.pi / 4.0, 0.0, math.pi / 2.0, 0.0, math.pi / 4.0, 0.0]

# 3. 执行前向运动学计算
# end_only=False 会返回所有链接的变换字典，默认为 True 仅返回末端
ret = chain.forward_kinematics(th, end_only=False)

# 4. 获取末端变换矩阵并解析位置和姿态
tg = ret['lbr_iiwa_link_7']
m = tg.get_matrix()  # 形状: (1, 4, 4)

# 提取位置 (x, y, z)
pos = m[:, :3, 3]

# 提取旋转矩阵并转换为四元数
rot = pk.matrix_to_quaternion(m[:, :3, :3])

print(f"Position: {pos}")
print(f"Quaternion: {rot}")

2. 批量计算与 GPU 加速

该库支持批量处理关节数据以利用并行计算，并可自动使用 CUDA。

import torch
import pytorch_kinematics as pk

# 选择设备
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float64

# 构建链并移至指定设备和精度
chain = pk.build_serial_chain_from_urdf(open("kuka_iiwa.urdf").read(), "lbr_iiwa_link_7")
chain = chain.to(dtype=dtype, device=device)

# 生成批量关节数据 (Batch Size = 1000)
N = 1000
th_batch = torch.rand(N, len(chain.get_joint_parameter_names()), dtype=dtype, device=device)

# 并行执行 FK (速度远快于循环调用)
tg_batch = chain.forward_kinematics(th_batch)

print(f"Batch result shape: {tg_batch.get_matrix().shape}")

3. 可微分特性 (用于梯度下降/优化)

由于基于 PyTorch，所有运动学计算均可自动求导，适用于轨迹优化或学习型逆运动学。

import torch
import math
import pytorch_kinematics as pk

chain = pk.build_serial_chain_from_urdf(open("kuka_iiwa.urdf").read(), "lbr_iiwa_link_7")

# 设置 requires_grad=True 以追踪梯度
th = torch.tensor([0.0, -math.pi / 4.0, 0.0, math.pi / 2.0, 0.0, math.pi / 4.0, 0.0], requires_grad=True)

tg = chain.forward_kinematics(th)
m = tg.get_matrix()
pos = m[:, :3, 3]

# 计算位置范数的梯度
loss = pos.norm()
loss.backward()

# th.grad 现在包含关节角度对位置损失的梯度
print(f"Gradients: {th.grad}")