# 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](https://github.com/facebookresearch/fairo/tree/main/polymetis/polymetis/python/torchcontrol/policies) as examples.

See [autogenerated docs](torchcontrol.policies) 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`](polymetis.robot_interface) 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](torchcontrol.modules) for exact API.

## Models

We leverage [Pinocchio](https://github.com/stack-of-tasks/pinocchio), a fast, rigid body dynamics algorithms library for real-time computations of robot kinematics and dynamics. 

See [autogenerated docs](torchcontrol.models) for exact API.

Cheat sheet:
```py
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](https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.transform.Rotation.html#scipy.spatial.transform.Rotation).
The [Eigen geometry module](https://eigen.tuxfamily.org/dox/group__Geometry__Module.html) 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:
```py
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:
```py
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:
```py
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:
```py
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:
```py
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:
```py
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:
```py
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:
```py
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.])
'''
```