PyTorch Batched interface

This module implements the batched version of the Rigid Body Dynamics algorithms using PyTorch. This module uses Jax under the hood, then the functions are jax.vmap-ed and jax.jit-ed to run on batches of inputs, which are ultimately converted to PyTorch using the jax2torch conversion functions.

Note

The first time you run a function from this module, it will take a bit longer to execute as they are being compiled by Jax.

Note

If the GPU support for JAX is needed, follow the instructions in the Jax documentation.

---

class adam.pytorch.computation_batch.KinDynComputationsBatch(urdfstring: str, joints_name_list: list = None, root_link: str = None, gravity: array = Array([0., 0., -9.80665, 0., 0., 0.], dtype=float32))

Bases: object

This is a small class that retrieves robot quantities using Jax for Floating Base systems. These functions are vmapped and jit compiled and passed to jax2torch to convert them to PyTorch functions.

set_frame_velocity_representation(representation: Representations) → None

Sets the representation of the velocity of the frames

Parameters:

representation (Representations) – The representation of the velocity

mass_matrix(base_transform: Tensor, joint_positions: Tensor) → Tensor

Returns the Mass Matrix functions computed the CRBA

Parameters:

• base_transform (torch.Tensor) – The batch of homogenous transforms from base to world frame

• joint_positions (torch.Tensor) – The batch of joints position

Returns:

The batch Mass Matrix

Return type:

M (torch.Tensor)

mass_matrix_fun()

Returns the Mass Matrix functions computed the CRBA as a pytorch function

Returns:

Mass Matrix

Return type:

M (pytorch function)

centroidal_momentum_matrix(base_transform: Tensor, joint_positions: Tensor) → Tensor

Returns the Centroidal Momentum Matrix functions computed the CRBA

Parameters:

• base_transform (torch.Tensor) – The homogenous transform from base to world frame

• joint_positions (torch.Tensor) – The joints position

Returns:

Centroidal Momentum matrix

Return type:

Jcc (torch.Tensor)

centroidal_momentum_matrix_fun()

Returns the Centroidal Momentum Matrix functions computed the CRBA as a pytorch function

Returns:

Centroidal Momentum matrix

Return type:

Jcc (pytorch function)

relative_jacobian(frame: str, joint_positions: Tensor) → Tensor

relative_jacobian_fun(frame: str)

Returns the Jacobian between the root link and a specified frame frames as a pytorch function

Parameters:

frame (str) – The tip of the chain

Returns:

The Jacobian between the root and the frame

Return type:

J (pytorch function)

jacobian_dot(frame: str, base_transform: Tensor, joint_positions: Tensor, base_velocity: Tensor, joint_velocities: Tensor) → Tensor

Returns the Jacobian derivative relative to the specified frame

Parameters:

• frame (str) – The frame to which the jacobian will be computed

• base_transform (torch.Tensor) – The homogenous transform from base to world frame

• joint_positions (torch.Tensor) – The joints position

• base_velocity (torch.Tensor) – The base velocity

• joint_velocities (torch.Tensor) – The joint velocities

Returns:

The Jacobian derivative relative to the frame

Return type:

Jdot (torch.Tensor)

jacobian_dot_fun(frame: str)

Returns the Jacobian derivative between the root and the specified frame as a pytorch function

Parameters:

frame (str) – The frame to which the jacobian will be computed

Returns:

The Jacobian derivative between the root and the frame

Return type:

Jdot (pytorch function)

forward_kinematics(frame: str, base_transform: Tensor, joint_positions: Tensor) → Tensor

Computes the forward kinematics between the root and the specified frame

Parameters:

frame (str) – The frame to which the fk will be computed

Returns:

The fk represented as Homogenous transformation matrix

Return type:

H (torch.Tensor)

forward_kinematics_fun(frame: str)

Computes the forward kinematics between the root and the specified frame as a pytorch function

Parameters:

frame (str) – The frame to which the fk will be computed

Returns:

The fk represented as Homogenous transformation matrix

Return type:

H (pytorch function)

jacobian(frame: str, base_transform: Tensor, joint_positions: Tensor) → Tensor

Returns the Jacobian relative to the specified frame

Parameters:

• frame (str) – The frame to which the jacobian will be computed

• base_transform (torch.Tensor) – The homogenous transform from base to world frame

• joint_positions (torch.Tensor) – The joints position

Returns:

The Jacobian between the root and the frame

Return type:

J (torch.Tensor)

jacobian_fun(frame: str)

Returns the Jacobian relative to the specified frame as a pytorch function

Parameters:

frame (str) – The frame to which the jacobian will be computed

Returns:

The Jacobian relative to the frame

Return type:

J (pytorch function)

bias_force(base_transform: Tensor, joint_positions: Tensor, base_velocity: Tensor, joint_velocities: Tensor) → array

Returns the bias force of the floating-base dynamics equation, using a reduced RNEA (no acceleration and external forces)

Parameters:

• base_transform (torch.Tensor) – The homogenous transform from base to world frame

• joint_positions (torch.Tensor) – The joints position

• base_velocity (torch.Tensor) – The base velocity

• joint_velocities (torch.Tensor) – The joints velocity

Returns:

the bias force

Return type:

h (torch.Tensor)

bias_force_fun()

Returns the bias force of the floating-base dynamics equation as a pytorch function

Returns:

the bias force

Return type:

h (pytorch function)

coriolis_term(base_transform: Tensor, joint_positions: Tensor, base_velocity: Tensor, joint_velocities: Tensor) → Tensor

Returns the coriolis term of the floating-base dynamics equation, using a reduced RNEA (no acceleration and external forces)

Parameters:

• base_transform (torch.Tensor) – The homogenous transform from base to world frame

• joint_positions (torch.Tensor) – The joints position

• base_velocity (torch.Tensor) – The base velocity

• joint_velocities (torch.Tensor) – The joints velocity

Returns:

the Coriolis term

Return type:

C (torch.Tensor)

coriolis_term_fun()

Returns the coriolis term of the floating-base dynamics equation as a pytorch function

Returns:

the Coriolis term

Return type:

C (pytorch function)

gravity_term(base_transform: Tensor, joint_positions: Tensor) → Tensor

Returns the gravity term of the floating-base dynamics equation, using a reduced RNEA (no acceleration and external forces)

Parameters:

• base_transform (torch.Tensor) – The homogenous transform from base to world frame

• joint_positions (torch.Tensor) – The joints position

Returns:

the gravity term

Return type:

G (jnp.array)

gravity_term_fun()

Returns the gravity term of the floating-base dynamics equation as a pytorch function

Returns:

the gravity term

Return type:

G (pytorch function)

CoM_position(base_transform: Tensor, joint_positions: Tensor) → Tensor

Returns the CoM position

Parameters:

• base_transform (torch.Tensor) – The homogenous transform from base to world frame

• joint_positions (torch.Tensor) – The joints position

Returns:

The CoM position

Return type:

CoM (torch.Tensor)

CoM_position_fun()

Returns the CoM position as a pytorch function

Returns:

The CoM position

Return type:

CoM (pytorch function)

CoM_jacobian(base_transform: Tensor, joint_positions: Tensor) → Tensor

Returns the CoM Jacobian

Parameters:

• base_transform (torch.Tensor) – The homogenous transform from base to world frame

• joint_positions (torch.Tensor) – The joints position

Returns:

The CoM Jacobian

Return type:

Jcom (torch.Tensor)

CoM_jacobian_fun()

Returns the CoM Jacobian as a pytorch function

Returns:

The CoM Jacobian

Return type:

Jcom (pytorch function)

get_total_mass() → float

Returns the total mass of the robot

Returns:

The total mass

Return type:

mass