# Copyright (C) Istituto Italiano di Tecnologia (IIT). All rights reserved.
import warnings
import numpy as np
import torch
from adam.core.constants import Representations
from adam.core.rbd_algorithms import RBDAlgorithms
from adam.model import Model, build_model_factory
from adam.model.kindyn_mixin import KinDynFactoryMixin
from adam.pytorch.torch_like import SpatialMath
from adam.core.array_api_math import spec_from_reference
[docs]
class KinDynComputations(KinDynFactoryMixin):
"""This is a small class that retrieves robot quantities using Pytorch for Floating Base systems."""
def __init__(
self,
urdfstring: str,
joints_name_list: list = None,
device: torch.device = (
torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
),
dtype: torch.dtype = torch.float32,
root_link: str = None,
gravity: np.array = torch.tensor([0, 0, -9.80665, 0, 0, 0]),
) -> None:
"""
Args:
urdfstring (str): path/string of a URDF or a MuJoCo MjModel.
NOTE: The parameter name `urdfstring` is deprecated and will be renamed to `model` in a future release.
joints_name_list (list): list of the actuated joints
root_link (str, optional): Deprecated. The root link is automatically chosen as the link with no parent in the URDF. Defaults to None.
"""
ref = torch.tensor(0.0, dtype=dtype, device=device)
spec = spec_from_reference(ref)
math = SpatialMath(spec=spec)
factory = build_model_factory(description=urdfstring, math=math)
model = Model.build(factory=factory, joints_name_list=joints_name_list)
[docs]
self.rbdalgos = RBDAlgorithms(model=model, math=math)
[docs]
self.NDoF = self.rbdalgos.NDoF
[docs]
self.g = gravity.to(dtype=dtype, device=device)
if root_link is not None:
warnings.warn(
"The root_link argument is not used. The root link is automatically chosen as the link with no parent in the URDF",
DeprecationWarning,
stacklevel=2,
)
[docs]
def set_frame_velocity_representation(
self, representation: Representations
) -> None:
"""Sets the representation of the velocity of the frames
Args:
representation (Representations): The representation of the velocity
"""
self.rbdalgos.set_frame_velocity_representation(representation)
[docs]
def mass_matrix(
self, base_transform: torch.Tensor, joint_positions: torch.Tensor
) -> torch.Tensor:
"""Returns the Mass Matrix functions computed the CRBA
Args:
base_transform (torch.tensor): The homogenous transform from base to world frame
joint_positions (torch.tensor): The joints position
Returns:
M (torch.tensor): Mass Matrix
"""
[M, _] = self.rbdalgos.crba(base_transform, joint_positions)
return M.array
[docs]
def centroidal_momentum_matrix(
self, base_transform: torch.Tensor, joint_positions: torch.Tensor
) -> torch.Tensor:
"""Returns the Centroidal Momentum Matrix functions computed the CRBA
Args:
base_transform (torch.tensor): The homogenous transform from base to world frame
joint_positions (torch.tensor): The joints position
Returns:
Jcc (torch.tensor): Centroidal Momentum matrix
"""
[_, Jcm] = self.rbdalgos.crba(base_transform, joint_positions)
return Jcm.array
[docs]
def forward_kinematics(
self, frame, base_transform: torch.Tensor, joint_positions: torch.Tensor
) -> torch.Tensor:
"""Computes the forward kinematics relative to the specified frame
Args:
frame (str): The frame to which the fk will be computed
base_transform (torch.tensor): The homogenous transform from base to world frame
joint_positions (torch.tensor): The joints position
Returns:
H (torch.tensor): The fk represented as Homogenous transformation matrix
"""
return (
self.rbdalgos.forward_kinematics(
frame,
base_transform,
joint_positions,
)
).array
[docs]
def jacobian(
self,
frame: str,
base_transform: torch.Tensor,
joint_positions: torch.Tensor,
) -> torch.Tensor:
"""Returns the Jacobian relative to the specified frame
Args:
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:
J_tot (torch.tensor): The Jacobian relative to the frame
"""
return self.rbdalgos.jacobian(frame, base_transform, joint_positions).array
[docs]
def relative_jacobian(self, frame, joint_positions: torch.Tensor) -> torch.Tensor:
"""Returns the Jacobian between the root link and a specified frame frames
Args:
frame (str): The tip of the chain
joint_positions (torch.tensor): The joints position
Returns:
J (torch.tensor): The Jacobian between the root and the frame
"""
return self.rbdalgos.relative_jacobian(frame, joint_positions).array
[docs]
def jacobian_dot(
self,
frame: str,
base_transform: torch.Tensor,
joint_positions: torch.Tensor,
base_velocity: torch.Tensor,
joint_velocities: torch.Tensor,
) -> torch.Tensor:
"""Returns the Jacobian derivative relative to the specified frame
Args:
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:
Jdot (torch.Tensor): The Jacobian derivative relative to the frame
"""
return self.rbdalgos.jacobian_dot(
frame, base_transform, joint_positions, base_velocity, joint_velocities
).array
[docs]
def CoM_position(
self, base_transform: torch.Tensor, joint_positions: torch.Tensor
) -> torch.Tensor:
"""Returns the CoM position
Args:
base_transform (torch.tensor): The homogenous transform from base to world frame
joint_positions (torch.tensor): The joints position
Returns:
CoM (torch.tensor): The CoM position
"""
return self.rbdalgos.CoM_position(
base_transform, joint_positions
).array.squeeze()
[docs]
def CoM_jacobian(
self, base_transform: torch.Tensor, joint_positions: torch.Tensor
) -> torch.Tensor:
"""Returns the CoM Jacobian
Args:
base_transform (torch.tensor): The homogenous transform from base to world frame
joint_positions (torch.tensor): The joints position
Returns:
Jcom (torch.tensor): The CoM Jacobian
"""
return self.rbdalgos.CoM_jacobian(base_transform, joint_positions).array
[docs]
def bias_force(
self,
base_transform: torch.Tensor,
joint_positions: torch.Tensor,
base_velocity: torch.Tensor,
joint_velocities: torch.Tensor,
) -> torch.Tensor:
"""Returns the bias force of the floating-base dynamics equation,
using a reduced RNEA (no acceleration and external forces)
Args:
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:
h (torch.tensor): the bias force
"""
return self.rbdalgos.rnea(
base_transform,
joint_positions,
base_velocity,
joint_velocities,
self.g,
).array.squeeze()
[docs]
def coriolis_term(
self,
base_transform: torch.Tensor,
joint_positions: torch.Tensor,
base_velocity: torch.Tensor,
joint_velocities: torch.Tensor,
) -> torch.Tensor:
"""Returns the coriolis term of the floating-base dynamics equation,
using a reduced RNEA (no acceleration and external forces)
Args:
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:
C (torch.tensor): the Coriolis term
"""
# set in the bias force computation the gravity term to zero
return self.rbdalgos.rnea(
base_transform,
joint_positions,
base_velocity,
joint_velocities,
torch.zeros(6, dtype=base_transform.dtype, device=base_transform.device),
).array.squeeze()
[docs]
def gravity_term(
self, base_transform: torch.Tensor, joint_positions: torch.Tensor
) -> torch.Tensor:
"""Returns the gravity term of the floating-base dynamics equation,
using a reduced RNEA (no acceleration and external forces)
Args:
base_transform (torch.tensor): The homogenous transform from base to world frame
joint_positions (torch.tensor): The joints positions
Returns:
G (torch.tensor): the gravity term
"""
return self.rbdalgos.rnea(
base_transform,
joint_positions,
torch.zeros(6, dtype=base_transform.dtype, device=base_transform.device),
torch.zeros(
self.NDoF, dtype=base_transform.dtype, device=base_transform.device
),
self.g,
).array.squeeze()
[docs]
def aba(
self,
base_transform: torch.Tensor,
joint_positions: torch.Tensor,
base_velocity: torch.Tensor,
joint_velocities: torch.Tensor,
joint_torques: torch.Tensor,
external_wrenches: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
"""Featherstone Articulated-Body Algorithm (floating base, O(n)).
Args:
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
joint_torques (torch.Tensor): The joint torques
external_wrenches (dict[str, torch.Tensor], optional): External wrenches applied to the robot. Defaults to None.
Returns:
torch.Tensor: The base acceleration and the joint accelerations
"""
return self.rbdalgos.aba(
base_transform,
joint_positions,
base_velocity,
joint_velocities,
joint_torques,
self.g,
external_wrenches,
).array.squeeze()
[docs]
def get_total_mass(self) -> float:
"""Returns the total mass of the robot
Returns:
mass: The total mass
"""
return self.rbdalgos.get_total_mass()