torchkin

A library for differentiable kinematics

Getting Started

Prerequisites

• We strongly recommend you install torchkin in a venv or conda environment with Python 3.8-3.10.
• torchkin requires torch installation. To install for your particular CPU/CUDA configuration, follow the instructions in the PyTorch website.

Installing

• pypi

  pip install torchkin
  

• From source

  The simplest way to install torchkin from source is by running the following

  git clone https://github.com/facebookresearch/theseus.git && cd theseus/torchkin
  pip install -e .
  

  If you are interested in contributing to torchkin, also install

  pip install -r ../requirements/dev.txt
  pre-commit install
  

  and follow the more detailed instructions in CONTRIBUTING.

Example

An inverse kinematics example is available in script.

import torch

import torchkin as kin

# We can load a robot model from a URDF file
dtype = torch.float64
device = "cuda"
robot = kin.Robot.from_urdf_file(YOUR_URDF_FILE, dtype=dtype, device=device)

# Print robot name, number of links and degrees of freedom
print(f"{robot.name} has {len(robot.get_links())} links and {robot.dof} degrees of freedom.\n")

# Print joint id and name
for id, name in enumerate(robot.joint_map):
    # A joint is not fixed if and only if id < robot.dof
    print(f"joint {id}: {name} is {'not fixed' if id < robot.dof else 'fixed'}")
print("\n")

# Print link id and name
for link in robot.get_links():
    print(f"link {link.id}: {link.name}")

# We can get differentiable forward kinematics functions for specific links
# by using `get_forward_kinematics_fns`. This function creates three differentiable
# functions for evaluating forward kinematics, body jacobian and spatial jacobian of
# the selected links, in that order. The return types of these functions are as
# follows:
#
# - fk: return a tuple of link poses in the order of link names
# - jfk_b: returns a tuple where the first is a list of link body jacobians, and the
#          second is a tuple of link poses---both are in the order of link names
# - jfk_s: same as jfk_b except returning the spatial jacobians
link_names = [LINK1, LINK2, LINK3]
fk, jfk_b, jfk_s = kin.get_forward_kinematics_fns(
    robot=robot, link_names=link_names)

batch_size = 10
# The joint states are in the order of the joint ids
joint_states = torch.rand(batch_size, robot.dof, dtype=dtype, device=device)

# Get link poses
link_poses = fk(joint_states)

# Get body jacobians and link poses
jacs_b, link_poses = jfk_b(joint_states) 

# Get spatial jacobians and link poses
jacs_s, link_poses = jfk_s(joint_states) 

Citing torchkin

If you use torchkin in your work, please cite the paper with the BibTeX below.

@article{pineda2022theseus,
  title   = {{Theseus: A Library for Differentiable Nonlinear Optimization}},
  author  = {Luis Pineda and Taosha Fan and Maurizio Monge and Shobha Venkataraman and Paloma Sodhi and Ricky TQ Chen and Joseph Ortiz and Daniel DeTone and Austin Wang and Stuart Anderson and Jing Dong and Brandon Amos and Mustafa Mukadam},
  journal = {Advances in Neural Information Processing Systems},
  year    = {2022}
}

License

torchkin is MIT licensed. See the LICENSE for details.

Additional Information

• Join the community on Github Discussions for questions and sugesstions.
• Use Github Issues for bugs and features.
• See CONTRIBUTING if interested in helping out.