# mypy: disable-error-code="attr-defined, name-defined"
"""Helpers to ease computation of kinematic and dynamic quantities.
"""
# pylint: disable=invalid-name,no-member
import logging
from copy import deepcopy
from typing import Optional, Tuple, Sequence, Callable, TypedDict, Any
import numpy as np
import eigenpy
import hppfcl
import pinocchio as pin
from pinocchio.rpy import (rpyToMatrix, # pylint: disable=import-error
matrixToRpy,
computeRpyJacobian,
computeRpyJacobianInverse)
from . import core as jiminy
LOGGER = logging.getLogger(__name__)
# #####################################################################
# ######################### Generic math ##############################
# #####################################################################
[docs]
def SE3ToXYZRPY(M: pin.SE3) -> np.ndarray:
"""Convert Pinocchio SE3 object to [X,Y,Z,Roll,Pitch,Yaw] vector.
"""
return np.concatenate((M.translation, matrixToRpy(M.rotation)))
[docs]
def XYZRPYToSE3(xyzrpy: np.ndarray) -> np.ndarray:
"""Convert [X,Y,Z,Roll,Pitch,Yaw] vector to Pinocchio SE3 object.
"""
return pin.SE3(rpyToMatrix(xyzrpy[3:]), xyzrpy[:3])
[docs]
def XYZRPYToXYZQuat(xyzrpy: np.ndarray) -> np.ndarray:
"""Convert [X,Y,Z,Roll,Pitch,Yaw] to [X,Y,Z,Qx,Qy,Qz,Qw].
"""
xyz, rpy = xyzrpy[:3], xyzrpy[3:]
return np.concatenate((xyz, pin.Quaternion(rpyToMatrix(rpy)).coeffs()))
[docs]
def XYZQuatToXYZRPY(xyzquat: np.ndarray) -> np.ndarray:
"""Convert [X,Y,Z,Qx,Qy,Qz,Qw] to [X,Y,Z,Roll,Pitch,Yaw].
"""
xyz, quat = xyzquat[:3], xyzquat[3:]
return np.concatenate((xyz, matrixToRpy(pin.Quaternion(quat).matrix())))
[docs]
def velocityXYZRPYToXYZQuat(xyzrpy: np.ndarray,
dxyzrpy: np.ndarray) -> np.ndarray:
"""Convert the derivative of [X,Y,Z,Roll,Pitch,Yaw] to [X,Y,Z,Qx,Qy,Qz,Qw].
.. warning::
Linear velocity in XYZRPY must be local-world-aligned frame, while
returned linear velocity in XYZQuat is in local frame.
"""
rpy = xyzrpy[3:]
R, J_rpy = rpyToMatrix(rpy), computeRpyJacobian(rpy)
return np.concatenate((R.T @ dxyzrpy[:3], J_rpy @ dxyzrpy[3:]))
[docs]
def velocityXYZQuatToXYZRPY(xyzquat: np.ndarray,
v: np.ndarray) -> np.ndarray:
"""Convert the derivative of [X,Y,Z,Qx,Qy,Qz,Qw] to [X,Y,Z,Roll,Pitch,Yaw].
.. note:
No need to estimate yaw angle to get RPY velocity, since
`computeRpyJacobianInverse` only depends on Roll and Pitch angles.
However, it is not the case for the linear velocity.
.. warning::
Linear velocity in XYZQuat must be local frame, while returned linear
velocity in XYZRPY is in local-world-aligned frame.
"""
quat = pin.Quaternion(xyzquat[3:])
rpy = matrixToRpy(quat.matrix())
J_rpy_inv = computeRpyJacobianInverse(rpy)
return np.concatenate((quat * v[:3], J_rpy_inv @ v[3:]))
# #####################################################################
# #################### State and Trajectory ###########################
# #####################################################################
[docs]
class State:
"""Store the kinematics and dynamics data of the robot at a given time.
"""
def __init__(self, # pylint: disable=unused-argument
t: float,
q: np.ndarray,
v: Optional[np.ndarray] = None,
a: Optional[np.ndarray] = None,
tau: Optional[np.ndarray] = None,
contact_frames: Optional[Sequence[str]] = None,
f_ext: Optional[Sequence[np.ndarray]] = None,
copy: bool = False,
**kwargs: Any) -> None:
"""
:param t: Time.
:param q: Configuration vector.
:param v: Velocity vector.
:param a: Acceleration vector.
:param tau: Joint efforts.
:param contact_frames: Name of the contact frames.
:param f_ext: Joint external forces.
:param copy: Force to copy the arguments.
"""
# Time
self.t = t
# Configuration vector
self.q = deepcopy(q) if copy else q
# Velocity vector
self.v = deepcopy(v) if copy else v
# Acceleration vector
self.a = deepcopy(a) if copy else a
# Effort vector
self.tau = deepcopy(tau) if copy else tau
# Frame names of the contact points
if copy:
self.contact_frames = deepcopy(contact_frames)
else:
self.contact_frames = contact_frames
# External forces
self.f_ext = deepcopy(f_ext) if copy else f_ext
def __repr__(self) -> str:
"""Convert the kinematics and dynamics data into string.
:returns: The kinematics and dynamics data as a string.
"""
msg = ""
for key, val in self.__dict__.items():
if val is not None:
msg += f"{key} : {val}\n"
return msg
[docs]
class TrajectoryDataType(TypedDict, total=False):
"""Basic data structure storing the required information about a trajectory
to later replay it using `jiminy_py.viewer.play_trajectories`.
"""
# List of State objects of increasing time
evolution_robot: Sequence[State]
# Optional Jiminy robot
robot: Optional[jiminy.Robot]
# Whether to use the theoretical model or the extended simulation model
use_theoretical_model: bool
# #####################################################################
# ################### Kinematic and dynamics ##########################
# #####################################################################
[docs]
def update_quantities(robot: jiminy.Model,
position: np.ndarray,
velocity: Optional[np.ndarray] = None,
acceleration: Optional[np.ndarray] = None,
update_physics: bool = True,
update_com: bool = True,
update_energy: bool = True,
update_jacobian: bool = False,
update_collisions: bool = True,
use_theoretical_model: bool = True) -> None:
"""Compute all quantities using position, velocity and acceleration
configurations.
Run multiple algorithms to compute all quantities which can be known with
the model position, velocity and acceleration.
This includes:
- body spatial transforms,
- body spatial velocities,
- body spatial drifts,
- body transform acceleration,
- body transform jacobians,
- center-of-mass position,
- center-of-mass velocity,
- center-of-mass drift,
- center-of-mass acceleration,
- center-of-mass jacobian,
- articular inertia matrix,
- non-linear effects (Coriolis + gravity)
- collisions and distances
.. note::
Computation results are stored internally in the robot, and can
be retrieved with associated getters.
.. warning::
This function modifies the internal robot data.
.. warning::
It does not called overloaded pinocchio methods provided by
`jiminy_py.core` but the original pinocchio methods instead. As a
result, it does not take into account the rotor inertias / armatures.
One is responsible of calling the appropriate methods manually instead
of this one if necessary. This behavior is expected to change in the
future.
:param robot: Jiminy robot.
:param position: Robot position vector.
:param velocity: Robot velocity vector.
:param acceleration: Robot acceleration vector.
:param update_physics: Whether to compute the non-linear effects and
internal/external forces.
Optional: True by default.
:param update_com: Whether to compute the COM of the robot AND each link
individually. The global COM is the first index.
Optional: False by default.
:param update_energy: Whether to compute the energy of the robot.
Optional: False by default
:param update_jacobian: Whether to compute the jacobians.
Optional: False by default.
:param use_theoretical_model: Whether the state corresponds to the
theoretical model when updating and fetching
the state of the robot.
Optional: True by default.
"""
if use_theoretical_model:
model = robot.pinocchio_model_th
data = robot.pinocchio_data_th
else:
model = robot.pinocchio_model
data = robot.pinocchio_data
if (update_physics and update_com and
update_energy and update_jacobian and
velocity is not None and acceleration is None):
pin.computeAllTerms(model, data, position, velocity)
else:
if velocity is None:
pin.forwardKinematics(model, data, position)
elif acceleration is None:
pin.forwardKinematics(model, data, position, velocity)
else:
pin.forwardKinematics(
model, data, position, velocity, acceleration)
if update_com:
if velocity is None:
kinematic_level = pin.POSITION
elif acceleration is None:
kinematic_level = pin.VELOCITY
else:
kinematic_level = pin.ACCELERATION
pin.centerOfMass(model, data, kinematic_level, False)
if update_jacobian:
if update_com:
pin.jacobianCenterOfMass(model, data)
pin.computeJointJacobians(model, data)
if update_physics:
if velocity is not None:
pin.nonLinearEffects(model, data, position, velocity)
pin.crba(model, data, position)
if update_energy:
if velocity is not None:
pin.computeKineticEnergy(model, data)
pin.computePotentialEnergy(model, data)
pin.updateFramePlacements(model, data)
if update_collisions:
pin.updateGeometryPlacements(
model, data, robot.collision_model, robot.collision_data)
pin.computeCollisions(
robot.collision_model, robot.collision_data,
stop_at_first_collision=False)
pin.computeDistances(robot.collision_model, robot.collision_data)
for dist_req in robot.collision_data.distanceResults:
if np.linalg.norm(dist_req.normal) < 1e-6:
pin.computeDistances(
robot.collision_model, robot.collision_data)
break
[docs]
def get_body_world_transform(robot: jiminy.Model,
body_name: str,
use_theoretical_model: bool = True,
copy: bool = True) -> pin.SE3:
"""Get the transform from world frame to body frame for a given body.
.. warning::
It is assumed that `update_quantities` has been called beforehand.
:param robot: Jiminy robot.
:param body_name: Name of the body.
:param use_theoretical_model: Whether the state corresponds to the
theoretical model when updating and fetching
the state of the robot.
:param copy: Whether to return the internal buffers (which could be
altered) or copy them.
Optional: True by default. It is less efficient but safer.
:returns: Body transform.
"""
# Pick the right pinocchio model and data
if use_theoretical_model:
model = robot.pinocchio_model_th
data = robot.pinocchio_data_th
else:
model = robot.pinocchio_model
data = robot.pinocchio_data
# Get frame index and make sure it exists
body_id = model.getFrameId(body_name)
assert body_id < model.nframes, f"Frame '{body_name}' does not exits."
# Get body transform in world frame
transform = data.oMf[body_id]
if copy:
transform = transform.copy()
return transform
[docs]
def get_body_world_velocity(robot: jiminy.Model,
body_name: str,
use_theoretical_model: bool = True) -> pin.SE3:
"""Get the spatial velocity in world frame.
.. warning::
It is assumed that `update_quantities` has been called beforehand.
:param robot: Jiminy robot.
:param body_name: Name of the body.
:param use_theoretical_model: Whether the state corresponds to the
theoretical model when updating and fetching
the state of the robot.
Optional: True by default.
:returns: Spatial velocity.
"""
# Pick the right pinocchio model and data
if use_theoretical_model:
model = robot.pinocchio_model_th
data = robot.pinocchio_data_th
else:
model = robot.pinocchio_model
data = robot.pinocchio_data
# Get frame index and make sure it exists
body_id = model.getFrameId(body_name)
assert body_id < model.nframes, f"Frame '{body_name}' does not exits."
return pin.getFrameVelocity(model, data, body_id, pin.WORLD)
[docs]
def get_body_world_acceleration(robot: jiminy.Model,
body_name: str,
use_theoretical_model: bool = True) -> pin.SE3:
"""Get the body spatial acceleration in world frame.
The moment of this tensor (i.e linear part) is NOT the linear acceleration
of the center of the body frame, expressed in the world frame.
.. warning::
It is assumed that `update_quantities` has been called.
:param robot: Jiminy robot.
:param body_name: Name of the body.
:param use_theoretical_model: Whether the state corresponds to the
theoretical model when updating and fetching
the state of the robot.
Optional: True by default.
:returns: Spatial acceleration.
"""
# Pick the right pinocchio model and data
if use_theoretical_model:
model = robot.pinocchio_model_th
data = robot.pinocchio_data_th
else:
model = robot.pinocchio_model
data = robot.pinocchio_data
# Get frame index and make sure it exists
body_id = model.getFrameId(body_name)
assert body_id < model.nframes, f"Frame '{body_name}' does not exits."
return pin.getFrameAcceleration(model, data, body_id, pin.WORLD)
[docs]
def compute_freeflyer_state_from_fixed_body(
robot: jiminy.Model,
position: np.ndarray,
velocity: Optional[np.ndarray] = None,
acceleration: Optional[np.ndarray] = None,
fixed_body_name: Optional[str] = None,
ground_profile: Optional[Callable[
[np.ndarray], Tuple[float, np.ndarray]]] = None,
use_theoretical_model: Optional[bool] = None) -> None:
"""Fill rootjoint data from articular data when a body is fixed and
aligned with world frame.
This method computes the position of freeflyer in the fixed body frame.
If **fixed_body_name** is omitted, it will default to either:
- the set of three contact points
- a single collision body
In such a way that no contact points nor collision bodies are going through
the ground and at least contact points or a collision body are touching it.
`None` is returned if their is no contact frame or if the robot does not
have a freeflyer.
.. warning::
This function modifies the internal robot data.
.. warning::
The method fills in freeflyer data instead of returning an updated copy
for efficiency.
:param robot: Jiminy robot.
:param position: Must contain current articular data. The freeflyer data
can contain any value, it will be ignored and replaced.
:param velocity: See position.
:param acceleration: See position.
:param fixed_body_name: Name of the body frame that is considered fixed
parallel to world frame.
Optional: It will be infered from the set of
contact points and collision bodies.
:param ground_profile: Ground profile callback.
:param use_theoretical_model:
Whether the state corresponds to the theoretical model when updating
and fetching the state of the robot. Must be False if `fixed_body_name`
is not specified.
Optional: True by default if `fixed_body_name` is specified, False
otherwise.
:returns: Name of the contact frame, if any.
"""
# Early return if no freeflyer
if not robot.has_freeflyer:
return
# Handling of default arguments
if use_theoretical_model is None:
use_theoretical_model = fixed_body_name is not None
# Clear freeflyer position, velocity and acceleration
position[:6].fill(0.0)
position[6] = 1.0
if velocity is not None:
velocity[:6].fill(0.0)
if acceleration is not None:
acceleration[:6].fill(0.0)
# Update kinematics, frame placements and collision information
update_quantities(robot,
position,
velocity,
acceleration,
update_physics=False,
update_com=False,
update_energy=False,
use_theoretical_model=use_theoretical_model)
if fixed_body_name is None:
if use_theoretical_model:
raise RuntimeError(
"Cannot infer contact transform for the theoretical model.")
w_M_ff = compute_transform_contact(robot, ground_profile)
else:
ff_M_fixed_body = get_body_world_transform(
robot, fixed_body_name, use_theoretical_model, copy=False)
if ground_profile is not None:
ground_translation = np.zeros(3)
ground_translation[2], normal = ground_profile(
ff_M_fixed_body.translation[:2])
ground_rotation = pin.Quaternion.FromTwoVectors(
np.array([0.0, 0.0, 1.0]), normal).matrix()
w_M_ground = pin.SE3(ground_rotation, ground_translation)
else:
w_M_ground = pin.SE3.Identity()
w_M_ff = w_M_ground.act(ff_M_fixed_body.inverse())
position[:7] = pin.SE3ToXYZQUAT(w_M_ff)
if fixed_body_name is not None:
if velocity is not None:
ff_v_fixed_body = get_body_world_velocity(
robot, fixed_body_name, use_theoretical_model)
base_link_velocity = - ff_v_fixed_body
velocity[:6] = base_link_velocity.vector
if acceleration is not None:
ff_a_fixedBody = get_body_world_acceleration(
robot, fixed_body_name, use_theoretical_model)
base_link_acceleration = - ff_a_fixedBody
acceleration[:6] = base_link_acceleration.vector
[docs]
def compute_efforts_from_fixed_body(
robot: jiminy.Model,
position: np.ndarray,
velocity: np.ndarray,
acceleration: np.ndarray,
fixed_body_name: str,
use_theoretical_model: bool = True) -> Tuple[np.ndarray, pin.Force]:
"""Compute the efforts using RNEA method.
.. warning::
This function modifies the internal robot data.
:param robot: Jiminy robot.
:param position: Robot configuration vector.
:param velocity: Robot velocity vector.
:param acceleration: Robot acceleration vector.
:param fixed_body_name: Name of the body frame.
:param use_theoretical_model: Whether the state corresponds to the
theoretical model when updating and fetching
the state of the robot.
Optional: True by default.
:returns: articular efforts and external forces.
"""
# Pick the right pinocchio model and data
if use_theoretical_model:
model = robot.pinocchio_model_th
data = robot.pinocchio_data_th
else:
model = robot.pinocchio_model
data = robot.pinocchio_data
# Apply a first run of rnea without explicit external forces
jiminy.rnea(model, data, position, velocity, acceleration)
# Initialize vector of exterior forces to zero
f_ext = pin.StdVec_Force()
f_ext.extend(len(model.names) * (pin.Force.Zero(),))
# Compute the force at the contact frame
pin.forwardKinematics(model, data, position)
support_joint_index = model.frames[
model.getBodyId(fixed_body_name)].parent
f_ext[support_joint_index] = data.oMi[support_joint_index].actInv(
data.oMi[1]).act(data.f[1])
# Recompute the efforts with RNEA and the correct external forces
tau = jiminy.rnea(model, data, position, velocity, acceleration, f_ext)
f_ext = f_ext[support_joint_index]
return tau, f_ext
[docs]
def compute_inverse_dynamics(robot: jiminy.Model,
position: np.ndarray,
velocity: np.ndarray,
acceleration: np.ndarray,
use_theoretical_model: bool = False
) -> np.ndarray:
"""Compute the motor torques through inverse dynamics, assuming to external
forces except the one resulting from the anyaltical constraints applied on
the model.
.. warning::
This function modifies the internal robot data.
:param robot: Jiminy robot.
:param position: Robot configuration vector.
:param velocity: Robot velocity vector.
:param acceleration: Robot acceleration vector.
:param use_theoretical_model:
Whether the position, velocity and acceleration are associated with the
theoretical model instead of the extended one.
Optional: False by default.
:returns: motor torques.
"""
if not robot.has_constraints:
raise NotImplementedError(
"Robot without active constraints is not supported for now.")
# Convert theoretical position, velocity and acceleration if necessary
if use_theoretical_model and robot.is_flexibility_enabled:
position = robot.get_extended_position_from_theoretical(position)
velocity = robot.get_extended_velocity_from_theoretical(velocity)
acceleration = (
robot.get_extended_velocity_from_theoretical(acceleration))
# Define some proxies for convenience
model = robot.pinocchio_model
data = robot.pinocchio_data
motor_velocity_indices = robot.motor_velocity_indices
# Updating kinematics quantities
pin.forwardKinematics(model, data, position, velocity, acceleration)
pin.updateFramePlacements(model, data)
# Compute constraint jacobian and drift
robot.compute_constraints(position, velocity)
J, drift = robot.get_constraints_jacobian_and_drift()
# No need to compute the internal matrix using `crba` nor to perform the
# cholesky decomposition since it is already done by `compute_constraints`
# internally.
M_inv = pin.cholesky.computeMinv(model, data)
M_inv_mcol = M_inv[:, motor_velocity_indices]
# Compute non-linear effects
pin.nonLinearEffects(model, data, position, velocity)
nle = data.nle
# Compute constraint forces
jiminy.computeJMinvJt(model, data, J)
a_f = jiminy.solveJMinvJtv(data, + J @ M_inv @ nle - drift)
B_f = jiminy.solveJMinvJtv(data, - J @ M_inv_mcol, False)
# compute feedforward term
a_ydd = (M_inv @ (J.T @ a_f - nle) - acceleration)[motor_velocity_indices]
B_ydd = (M_inv_mcol + M_inv @ J.T @ B_f)[motor_velocity_indices]
# Compute motor torques
u = eigenpy.LDLT(B_ydd).solve(- a_ydd)
return u
# #####################################################################
# ################### State sequence wrappers #########################
# #####################################################################
[docs]
def compute_freeflyer(trajectory_data: TrajectoryDataType,
freeflyer_continuity: bool = True) -> None:
"""Compute the freeflyer positions and velocities.
.. warning::
This function modifies the internal robot data.
:param trajectory_data: Sequence of States for which to retrieve the
freeflyer.
:param freeflyer_continuity: Whether to enforce the continuity in position
of the freeflyer.
Optional: True by default.
"""
robot = trajectory_data['robot']
contact_frame_prev: Optional[str] = None
w_M_ff_offset = pin.SE3.Identity()
w_M_ff_prev = None
for s in trajectory_data['evolution_robot']:
# Compute freeflyer using contact frame as reference frame
compute_freeflyer_state_from_fixed_body(
robot, s.q, s.v, s.a, s.contact_frame, None)
# Move freeflyer to ensure continuity over time, if requested
if freeflyer_continuity:
# Extract the current freeflyer transform
w_M_ff = pin.XYZQUATToSE3(s.q[:7])
# Update the internal buffer of the freeflyer transform
if (contact_frame_prev is not None and
contact_frame_prev != s.contact_frame):
w_M_ff_offset = w_M_ff_offset * w_M_ff_prev * w_M_ff.inverse()
contact_frame_prev = s.contact_frame
w_M_ff_prev = w_M_ff
# Add the appropriate offset to the freeflyer
w_M_ff = w_M_ff_offset * w_M_ff
s.q[:7] = pin.SE3ToXYZQUAT(w_M_ff)
[docs]
def compute_efforts(trajectory_data: TrajectoryDataType) -> None:
"""Compute the efforts in the trajectory using RNEA method.
:param trajectory_data: Sequence of States for which to compute the
efforts.
"""
robot = trajectory_data['robot']
use_theoretical_model = trajectory_data['use_theoretical_model']
for s in trajectory_data['evolution_robot']:
assert s.q is not None and s.v is not None and s.a is not None
assert s.contact_frames is not None
s.tau, s.f_ext = compute_efforts_from_fixed_body(
robot, s.q, s.v, s.a, s.contact_frame, use_theoretical_model)