TorchControl

TorchControl is a library included within Polymetis that contains various tools and templates for writing your own controllers, including:

  • Common manipulation policies that can be used off-the-shelf or serve as a template for writing custom controllers.

  • Modules that contain frequently used control logic that can be used to piece together a controller.

  • A robot model that performs kinematic and dynamic computations of a given urdf.

  • A torch-scriptable rotation and transformation library.

Policies

All policies (a.k.a controllers) inherit from torchcontrol.PolicyModule, which allows policies to be scripted and run on the Polymetis server. TorchControl provides a set of commonly-used policies directly accessible at torchcontrol.policies, including impedance controllers, min-jerk trajectory controllers, and iLQR.

When starting to write your own policies, a good place to start is to go through the built-in policies as examples.

See autogenerated docs for exact API.

Modules

All modules inherit from torchcontrol.ControlModule and contain pieces frequently used control logic such as feedback control, feedforward gravity compensation, and min-jerk planning. Parameters of the module can also be updated while running through the update method, which is accessed through RobotInterface.update_current_policy when running on a robot. Just like torch.nn.Module, torchcontrol modules can also contain other modules, and operations such as update on a module can be propagated through all the submodules.

See autogenerated docs for exact API.

Models

We leverage Pinocchio, a fast, rigid body dynamics algorithms library for real-time computations of robot kinematics and dynamics.

See autogenerated docs for exact API.

Cheat sheet:

import torchcontrol as toco

# Creation
robot_model = toco.models.RobotModelPinocchio(urdf_path, end_effector_joint_name)

# Methods
pos_limits = robot_model.get_joint_angle_limits()
vel_limits = robot_model.get_joint_velocity_limits()
ee_pos, ee_quat = robot_model.forward_kinematics(joint_positions)
J = robot_model.compute_jacobian(joint_positions)
torques = robot_model.inverse_dynamics(joint_positions, joint_velocities, desired_joint_accelerations)

Transform

torchcontrol.transform provides a clean way of storing and manipulating transformations that is torch-scriptable. The API for transformation and rotation objects are modeled after scipy.spatial.transform.Rotation. The Eigen geometry module is used as a backend, enabling highly optimized geometry computations capable of running in a real-time controller.

Example code is provided below to demonstrate API usage.

Rotations

Rotations are stored as quaternions due to computational and memory efficiency. (Note that quaternions are expressed in the XYZW format, as with the other third-party libraries used in Polymetis, including Pinocchio, Pybullet, and Eigen.)

Creating rotation objects:

from torchcontrol.transform import Rotation as R

# Quaternion
R.from_quat(torch.Tensor([0., 0., 0., 1.]))
'''
RotationObj(quaternion=tensor([0., 0., 0., 1.]))
'''

# Rotation matrix (SO3)
R.from_matrix(torch.eye(3))
'''
RotationObj(quaternion=tensor([0., 0., 0., 1.]))
'''

# Rotation vector (axis-angle) (so3)
R.from_rotvec(torch.zeros(3))
'''
RotationObj(quaternion=tensor([0., 0., 0., 1.]))
'''

# Identity
R.identity()
'''
RotationObj(quaternion=tensor([0., 0., 0., 1.]))
'''

Operations:

import numpy as np
from torchcontrol.transform import Rotation as R

r1 = R.from_rotvec(torch.Tensor([0., 0., np.pi/2]))
r2 = R.from_rotvec(torch.Tensor([0., 0., -np.pi/2]))

# Inverse
r1.inv()
'''
RotationObj(quaternion=tensor([0., 0., -0.707, 0.707]))
'''

# Multiplication
r1 * r2
'''
RotationObj(quaternion=tensor([0., 0., 0., 1.]))
'''

# Applying rotation to a vector
r1.apply(torch.ones(3))
'''
tensor([-1., 1., 1.])
'''

Accessing components:

from torchcontrol.transform import Rotation as R

r = R.from_rotvec(torch.Tensor([1., 1., 0.]))

# Axis of rotation
r.axis()
'''
tensor([0.707, 0.707, 0.])
'''

# Magnitude of rotation (radians)
r.magnitude()
'''
tensor([1.4142])
'''

Convert to rotation representations:

from torchcontrol.transform import Rotation as R

r = R.identity()

# Quaternion
r.as_quat()
'''
tensor([0., 0., 0., 1.])
'''

# Rotation matrix (SO3)
r.as_matrix()
'''
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]])
'''

# Rotation vector (axis-angle) (so3)
r.as_rotvec()
'''
tensor([0., 0., 0.])
'''

Transformations

Transformations are represented as a pure rotation followed by a pure translation.

Creating transformation objects:

from torchcontrol.transform import Rotation as R
from torchcontrol.transform import Transformation as T

# Rotation & translation
T.from_rot_xyz(
    rotation=R.from_matrix(torch.eye(3)),
    translation=torch.zeros(3)
)
'''
TransformationObj(
    rotation=RotationObj(quaternion=tensor([0., 0., 0., 1.])),
    translation=tensor([0., 0., 0.])
)
'''

# Transformation matrix (SE3)
T.from_matrix(torch.eye(4))
'''
TransformationObj(
    rotation=RotationObj(quaternion=tensor([0., 0., 0., 1.])),
    translation=tensor([0., 0., 0.])
)
'''

# Identity
T.identity()
'''
TransformationObj(
    rotation=RotationObj(quaternion=tensor([0., 0., 0., 1.])),
    translation=tensor([0., 0., 0.])
)
'''

Operations:

from torchcontrol.transform import Transformation as T

t1 = T.from_rot_xyz(
    rotation=R.identity(),
    translation=torch.Tensor([1., 0., 0.])
)
t2 = T.from_rot_xyz(
    rotation=R.identity(),
    translation=torch.Tensor([0., 1., 0.])
)

# Inverse
t1.inv()
'''
TransformationObj(
    rotation=RotationObj(quaternion=tensor([0., 0., 0., 1.])),
    translation=tensor([-1., 0., 0.])
)
'''

# Multiplication
t1 * t2
'''
TransformationObj(
    rotation=RotationObj(quaternion=tensor([0., 0., 0., 1.])),
    translation=tensor([1., 1., 0.])
)
'''

# Applying rotation to a vector
t.apply(torch.ones(3))
'''
tensor([2., 1., 1.])
'''

Accessing components:

from torchcontrol.transform import Transformation as T

T.from_rot_xyz(
    rotation=R.from_quat(torch.Tensor([0., 0., 1., 0.])),
    translation=torch.ones(3)
)

# Translation
t.translation()
'''
tensor([1., 1., 1.])
'''

# Rotation
t.rotation()
'''
RotationObj(quaternion=tensor([0., 0., 1., 0.])),
'''

Convert to transformation representations:

from torchcontrol.transform import Transformation as T

t = T.identity()

# Transformation matrix (SE3)
matrix = t.as_matrix()
'''
tensor([[1., 0., 0., 0.],
        [0., 1., 0., 0.],
        [0., 0., 1., 0.],
        [0., 0., 0., 1.]])
'''

# Twist (translation + rotvec) (se3)
twist = R.as_twist()
'''
tensor([0., 0., 0., 0., 0., 0.])
'''