From f39e6c73428c5add60a5749edd4cbfe93aa7628c Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Fri, 1 Mar 2024 14:04:59 +0100 Subject: [PATCH 1/2] [core] Fix broken 'getRigidVelocityFromFlexible'. --- core/src/robot/model.cc | 50 ++++++++++++++++++++--------------------- 1 file changed, 24 insertions(+), 26 deletions(-) diff --git a/core/src/robot/model.cc b/core/src/robot/model.cc index cc94125ba..2b37d33e2 100644 --- a/core/src/robot/model.cc +++ b/core/src/robot/model.cc @@ -1672,22 +1672,25 @@ namespace jiminy } } - void Model::getRigidPositionFromFlexible(const Eigen::VectorXd & qFlex, - Eigen::VectorXd & qRigid) const + + void Model::getFlexibleVelocityFromRigid(const Eigen::VectorXd & vRigid, + Eigen::VectorXd & vFlex) const { // Define some proxies - uint32_t nqFlex = pncModelFlexibleOrig_.nq; + uint32_t nvRigid = pinocchioModelOrig_.nv; + uint32_t nvFlex = pncModelFlexibleOrig_.nv; // Check the size of the input state - if (qFlex.size() != nqFlex) + if (vRigid.size() != nvRigid) { - THROW_ERROR(std::invalid_argument, "Size of qFlex inconsistent with flexible model."); + THROW_ERROR(std::invalid_argument, + "Size of vRigid inconsistent with theoretical model."); } - // Initialize the rigid state - qRigid = pinocchio::neutral(pinocchioModelOrig_); + // Initialize the flexible state + vFlex.setZero(nvFlex); - // Compute the rigid state based on the flexible state + // Compute the flexible state based on the rigid state int32_t idxRigid = 0; int32_t idxFlex = 0; for (; idxRigid < pinocchioModelOrig_.njoints; ++idxFlex) @@ -1700,32 +1703,30 @@ namespace jiminy const auto & jointFlex = pncModelFlexibleOrig_.joints[idxFlex]; if (jointRigid.idx_q() >= 0) { - qRigid.segment(jointRigid.idx_q(), jointRigid.nq()) = - qFlex.segment(jointFlex.idx_q(), jointFlex.nq()); + vFlex.segment(jointFlex.idx_v(), jointFlex.nv()) = + vRigid.segment(jointRigid.idx_v(), jointRigid.nv()); } ++idxRigid; } } } - void Model::getFlexibleVelocityFromRigid(const Eigen::VectorXd & vRigid, - Eigen::VectorXd & vFlex) const + void Model::getRigidPositionFromFlexible(const Eigen::VectorXd & qFlex, + Eigen::VectorXd & qRigid) const { // Define some proxies - uint32_t nvRigid = pinocchioModelOrig_.nv; - uint32_t nvFlex = pncModelFlexibleOrig_.nv; + uint32_t nqFlex = pncModelFlexibleOrig_.nq; // Check the size of the input state - if (vRigid.size() != nvRigid) + if (qFlex.size() != nqFlex) { - THROW_ERROR(std::invalid_argument, - "Size of vRigid inconsistent with theoretical model."); + THROW_ERROR(std::invalid_argument, "Size of qFlex inconsistent with flexible model."); } - // Initialize the flexible state - vFlex.setZero(nvFlex); + // Initialize the rigid state + qRigid = pinocchio::neutral(pinocchioModelOrig_); - // Compute the flexible state based on the rigid state + // Compute the rigid state based on the flexible state int32_t idxRigid = 0; int32_t idxFlex = 0; for (; idxRigid < pinocchioModelOrig_.njoints; ++idxFlex) @@ -1738,8 +1739,8 @@ namespace jiminy const auto & jointFlex = pncModelFlexibleOrig_.joints[idxFlex]; if (jointRigid.idx_q() >= 0) { - vFlex.segment(jointFlex.idx_v(), jointFlex.nv()) = - vRigid.segment(jointRigid.idx_v(), jointRigid.nv()); + qRigid.segment(jointRigid.idx_q(), jointRigid.nq()) = + qFlex.segment(jointFlex.idx_q(), jointFlex.nq()); } ++idxRigid; } @@ -1778,10 +1779,7 @@ namespace jiminy vRigid.segment(jointRigid.idx_v(), jointRigid.nv()) = vFlex.segment(jointFlex.idx_v(), jointFlex.nv()); } - } - else - { - ++idxFlex; + ++idxRigid; } } } From a0f99dc16935124265f8e99d53a2c615a7bab067 Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Sat, 2 Mar 2024 15:48:55 +0100 Subject: [PATCH 2/2] [gym/common] Add 'DeformationEstimator' block. --- core/src/control/abstract_controller.cc | 2 +- data/toys_models/ant/ant_hardware.toml | 28 +- .../common/gym_jiminy/common/bases/block.py | 5 +- .../common/gym_jiminy/common/bases/generic.py | 14 +- .../gym_jiminy/common/blocks/__init__.py | 2 + .../common/blocks/deformation_estimator.py | 757 ++++++++++++++++++ .../common/gym_jiminy/common/envs/generic.py | 8 +- .../common/gym_jiminy/common/utils/math.py | 5 +- .../envs/gym_jiminy/envs/__init__.py | 6 +- python/gym_jiminy/envs/gym_jiminy/envs/ant.py | 48 +- .../gym_jiminy/unit_py/data/flexible_arm.urdf | 154 ++++ .../unit_py/test_deformation_estimator.py | 254 ++++++ .../unit_py/test_pipeline_control.py | 12 +- .../unit_py/test_pipeline_design.py | 12 +- python/jiminy_py/src/jiminy_py/plot.py | 2 +- python/jiminy_py/src/jiminy_py/simulator.py | 2 +- 16 files changed, 1244 insertions(+), 67 deletions(-) create mode 100644 python/gym_jiminy/common/gym_jiminy/common/blocks/deformation_estimator.py create mode 100644 python/gym_jiminy/unit_py/data/flexible_arm.urdf create mode 100644 python/gym_jiminy/unit_py/test_deformation_estimator.py diff --git a/core/src/control/abstract_controller.cc b/core/src/control/abstract_controller.cc index d74abf630..363c6be23 100644 --- a/core/src/control/abstract_controller.cc +++ b/core/src/control/abstract_controller.cc @@ -203,7 +203,7 @@ namespace jiminy { return element.first == *fieldIt; }); if (variableIt != registeredVariables.end()) { - THROW_ERROR(lookup_error, "Variable already registered."); + THROW_ERROR(lookup_error, "Variable '", *fieldIt, "' already registered."); } registeredVariables.emplace_back(*fieldIt, &values[i]); } diff --git a/data/toys_models/ant/ant_hardware.toml b/data/toys_models/ant/ant_hardware.toml index e88bcbff9..6141d2115 100644 --- a/data/toys_models/ant/ant_hardware.toml +++ b/data/toys_models/ant/ant_hardware.toml @@ -7,20 +7,6 @@ contactFrameNames = [] #['tip_1', 'tip_2', 'tip_3', 'tip_4'] # ============== Motors =============== -[Motor.SimpleMotor.hip_4] -joint_name = "hip_4" -mechanicalReduction = 150.0 -armature = 1.0 -frictionViscousPositive = -1.0 -frictionViscousNegative = -1.0 - -[Motor.SimpleMotor.ankle_4] -joint_name = "ankle_4" -mechanicalReduction = 150.0 -armature = 1.0 -frictionViscousPositive = -1.0 -frictionViscousNegative = -1.0 - [Motor.SimpleMotor.hip_1] joint_name = "hip_1" mechanicalReduction = 150.0 @@ -63,6 +49,20 @@ armature = 1.0 frictionViscousPositive = -1.0 frictionViscousNegative = -1.0 +[Motor.SimpleMotor.hip_4] +joint_name = "hip_4" +mechanicalReduction = 150.0 +armature = 1.0 +frictionViscousPositive = -1.0 +frictionViscousNegative = -1.0 + +[Motor.SimpleMotor.ankle_4] +joint_name = "ankle_4" +mechanicalReduction = 150.0 +armature = 1.0 +frictionViscousPositive = -1.0 +frictionViscousNegative = -1.0 + # ============= Encoder sensors =============== [Sensor.EncoderSensor.hip_1] diff --git a/python/gym_jiminy/common/gym_jiminy/common/bases/block.py b/python/gym_jiminy/common/gym_jiminy/common/bases/block.py index 94dc5ec8d..22bedd361 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/bases/block.py +++ b/python/gym_jiminy/common/gym_jiminy/common/bases/block.py @@ -12,7 +12,7 @@ import gymnasium as gym -from ..utils import FieldNested, DataNested, get_fieldnames, fill, zeros +from ..utils import FieldNested, DataNested, get_fieldnames, zeros from .generic import (ObsT, ActT, @@ -159,9 +159,6 @@ def _setup(self) -> None: # Compute the update period self.observe_dt = self.env.observe_dt * self.update_ratio - # Set default observation - fill(self.observation, 0) - # Make sure the controller period is lower than environment timestep assert self.observe_dt <= self.env.step_dt, ( "The observer update period must be lower than or equal to the " diff --git a/python/gym_jiminy/common/gym_jiminy/common/bases/generic.py b/python/gym_jiminy/common/gym_jiminy/common/bases/generic.py index 0e2d94837..baa29c1bb 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/bases/generic.py +++ b/python/gym_jiminy/common/gym_jiminy/common/bases/generic.py @@ -16,7 +16,7 @@ from jiminy_py.simulator import Simulator from jiminy_py.viewer.viewer import is_display_available -from ..utils import DataNested, fill +from ..utils import DataNested # Temporal resolution of simulator steps @@ -155,7 +155,7 @@ def compute_reward(self, :returns: Aggregated reward for the current step. """ - raise NotImplementedError + return 0.0 # Note that `InterfaceJiminyEnv` must inherit from `InterfaceObserver` @@ -220,10 +220,6 @@ def _setup(self) -> None: # The observation must always be refreshed after setup self.__is_observation_refreshed = False - # Reset observation and action buffers - fill(self.observation, 0) - fill(self.action, 0) - @no_type_check def _observer_handle(self, t: float, @@ -259,7 +255,11 @@ def _observer_handle(self, sensor_measurements_it = iter(sensor_measurements.values()) for sensor_type in self._sensors_types: measurement_sensors[sensor_type] = next(sensor_measurements_it) - self.refresh_observation(measurement) + try: + self.refresh_observation(measurement) + except RuntimeError as e: + raise RuntimeError( + "The observation space must be constant.") from e self.__is_observation_refreshed = True def _controller_handle(self, diff --git a/python/gym_jiminy/common/gym_jiminy/common/blocks/__init__.py b/python/gym_jiminy/common/gym_jiminy/common/blocks/__init__.py index 3b3d69732..f272bd154 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/blocks/__init__.py +++ b/python/gym_jiminy/common/gym_jiminy/common/blocks/__init__.py @@ -3,10 +3,12 @@ from .mahony_filter import MahonyFilter from .motor_safety_limit import MotorSafetyLimit from .proportional_derivative_controller import PDController +from .deformation_estimator import DeformationEstimator __all__ = [ 'MahonyFilter', 'MotorSafetyLimit', 'PDController', + 'DeformationEstimator' ] diff --git a/python/gym_jiminy/common/gym_jiminy/common/blocks/deformation_estimator.py b/python/gym_jiminy/common/gym_jiminy/common/blocks/deformation_estimator.py new file mode 100644 index 000000000..60dc3334f --- /dev/null +++ b/python/gym_jiminy/common/gym_jiminy/common/blocks/deformation_estimator.py @@ -0,0 +1,757 @@ +from collections.abc import Mapping +from typing import List, Dict, Sequence, Tuple, Optional + +import numpy as np +import numba as nb +import gymnasium as gym + +import jiminy_py.core as jiminy +from jiminy_py.core import ( # pylint: disable=no-name-in-module + EncoderSensor as encoder, ImuSensor as imu, get_frame_indices) + +import pinocchio as pin + +from ..bases import BaseActT, BaseObsT, BaseObserverBlock, InterfaceJiminyEnv +from ..utils import (DataNested, + matrices_to_quat, + quat_multiply, + compute_tilt_from_quat, + swing_from_vector) + + +# FIXME: Enabling cache causes segfault on Apple Silicon +@nb.jit(nopython=True, cache=False, inline='always') +def _compute_orientation_error(obs_imu_quats: np.ndarray, + obs_imu_indices: Tuple[int, ...], + inv_obs_imu_quats: np.ndarray, + kin_imu_rots: Tuple[np.ndarray, ...], + kin_imu_quats: np.ndarray, + dev_imu_quats: np.ndarray, + ignore_twist: bool) -> None: + """Compute total deviation of observed IMU data wrt to rigid model in world + frame. + """ + # Re-order IMU data + for i, imu_index in enumerate(obs_imu_indices): + inv_obs_imu_quats[:, i] = obs_imu_quats[:, imu_index] + + # Compute orientation error + if ignore_twist: + # Compute swing of deviation between observed and kinematic IMU data: + # 1. Extract observed IMU tilt + # 2. Apply kinematic IMU rotation on it + # 3. Reconstruct swing from tilt of deviation + # This approach is equivalent to removing the twist from the exact + # deviation directly, but much faster as it does not require computing + # the observed IMU quaternions nor to compose rotations. + + # Extract observed tilt: w_R_obs.T @ e_z + obs_imu_tilts = np.stack(compute_tilt_from_quat(inv_obs_imu_quats), 1) + + # Apply rigid kinematic IMU rotation on observed tilt. + # The result can be interpreted as the tilt error between observed + # and rigid kinematic IMU orientation in world frame, ie: + # w_tilt_err = (w_R_kin @ w_R_obs.T) @ e_z + for i, kin_imu_rot in enumerate(kin_imu_rots): + obs_imu_tilts[i] = kin_imu_rot @ obs_imu_tilts[i] + + # Compute "smallest" rotation that can explain the tilt error. + swing_from_vector(obs_imu_tilts.T, dev_imu_quats) + + # Conjugate quaternion of IMU deviation + dev_imu_quats[-1] *= -1 + else: + # Convert kinematic IMU rotation matrices to quaternions + matrices_to_quat(kin_imu_rots, kin_imu_quats) + + # Conjugate quaternion of IMU orientation + inv_obs_imu_quats[-1] *= -1 + + # Compute one-by-one difference between observed and rigid + # kinematic IMU orientations. + quat_multiply(kin_imu_quats, + inv_obs_imu_quats, + dev_imu_quats) + + +@nb.jit(nopython=True, cache=True, inline='always') +def _compute_deformation_from_deviation(kin_flex_rots: Tuple[np.ndarray, ...], + kin_flex_quats: np.ndarray, + is_flex_flipped: np.ndarray, + is_chain_orphan: Tuple[bool, bool], + dev_imu_quats: np.ndarray, + inv_child_dev_imu_quats: np.ndarray, + parent_dev_imu_quats: np.ndarray, + deformation_flex_quats: np.ndarray + ) -> None: + """Cast observed IMU deviations in world frame into deformations at their + corresponding flexibility frame. + """ + # Convert kinematic flexibility rotation matrices to quaternions + matrices_to_quat(kin_flex_rots, kin_flex_quats) + + # Translate observed IMU deviation in world frame to their respective + # parent and child flex frames. + is_flex_parent_orphan, is_flex_child_orphan = is_chain_orphan + if is_flex_parent_orphan: + quat_multiply( + dev_imu_quats, kin_flex_quats[:, 1:], parent_dev_imu_quats[:, 1:]) + if is_flex_child_orphan: + quat_multiply(dev_imu_quats, + kin_flex_quats[:, :-1], + inv_child_dev_imu_quats[:, :-1]) + else: + quat_multiply( + dev_imu_quats, kin_flex_quats, inv_child_dev_imu_quats) + else: + nflexs = len(kin_flex_rots) + quat_multiply( + dev_imu_quats[:, :nflexs], kin_flex_quats, parent_dev_imu_quats) + if is_flex_child_orphan: + quat_multiply(dev_imu_quats[:, 1:], + kin_flex_quats[:, :-1], + inv_child_dev_imu_quats[:, :-1]) + else: + quat_multiply( + dev_imu_quats[:, 1:], kin_flex_quats, inv_child_dev_imu_quats) + + if is_flex_parent_orphan: + parent_dev_imu_quats[:, 0] = kin_flex_quats[:, 0] + if is_flex_child_orphan: + inv_child_dev_imu_quats[:, -1] = kin_flex_quats[:, -1] + + inv_child_dev_imu_quats[-1] *= -1 + + # Project observed total deformation on flexibility frames + quat_multiply(inv_child_dev_imu_quats, + parent_dev_imu_quats, + deformation_flex_quats) + + # Conjugate deformation quaternion if triplet (parent IMU, flex, child IMU) + # is reversed wrt to the standard joint ordering from URDF. + deformation_flex_quats[-1][is_flex_flipped] *= -1 + + +# FIXME: Enabling cache causes segfault on Apple Silicon +@nb.jit(nopython=True, cache=False) +def flexibility_estimator(obs_imu_quats: np.ndarray, + obs_imu_indices: Tuple[int, ...], + inv_obs_imu_quats: np.ndarray, + kin_imu_rots: Tuple[np.ndarray, ...], + kin_imu_quats: np.ndarray, + kin_flex_rots: Tuple[np.ndarray, ...], + kin_flex_quats: np.ndarray, + is_flex_flipped: np.ndarray, + is_chain_orphan: Tuple[bool, bool], + dev_imu_quats: np.ndarray, + inv_child_dev_imu_quats: np.ndarray, + parent_dev_imu_quats: np.ndarray, + deformation_flex_quats: np.ndarray, + ignore_twist: bool) -> None: + """Compute the local deformation at an arbitrary set of flexibility points + that are presumably responsible for most of the whole deformation of the + mechanical structure. + + .. warning:: + The local deformation at each flexibility frame must be observable, ie + the flexibility and IMU frames interleave with each others in a unique + and contiguous sub-chain in the rigid kinematic tree of the robot. + + :param obs_imu_quats: Orientation estimates of an unordered arbitrary set + of IMUs as a 2D array whose first dimension gathers + the 4 quaternion coordinates [qx, qy, qz, qw] while + the second corresponds to N independent IMU data. + :param obs_imu_indices: M-tuple of ordered IMU indices of interest for + which the total deviation will be computed. + :param inv_obs_imu_quats: Pre-allocated memory for storing the inverse of + the orientation estimates for an ordered subset + of the IMU data `obs_imu_quats` according to + `obs_imu_indices`. + :param kin_imu_rots: Tuple of M kinematic frame orientations corresponding + to the ordered subset of IMUs `obs_imu_indices`, for + the rigid configuration of the theoretical model of + the robot. + :param kin_imu_quats: Pre-allocated memory for storing the kinematic frame + orientations of the ordered subset of IMUs of + interest as a 2D array whose first dimension gathers + the 4 quaternion coordinates while the second + corresponds to the M independent IMUs. + :param kin_flex_rots: Tuple of K kinematic frame orientations for all the + flexibility points that interleaves with the ordered + subset of IMUs of interest in the kinematic tree. + :param kin_flex_quats: Pre-allocated memory for storing the kinematic frame + orientations of the flexibility points that + interleaves with the ordered subset of IMUs of + interest as a 2D array whose first dimension gathers + the 4 quaternion coordinates while the second + corresponds to the K independent flexibility points. + :param is_flex_flipped: Whether local deformation estimates for each + flexibility point must be inverted to be consistent + with standard URDF convention as 1D boolean array. + :param is_chain_orphan: 2-Tuple stating whether first and last flexibility + point is orphan respectively, ie only a single IMU + is available for estimating its local deformation. + :param dev_imu_quats: Total deviation observed IMU data wrt to rigid model + in world frame for the ordered subset of IMUs of + interest as a 2D array whose first dimension gathers + the 4 quaternion coordinates while the second + corresponds to the M independent IMUs. + :param inv_child_dev_imu_quats: + Total deviation observed IMU data in child flexibility frame as a 2D + array whose first dimension gathers the 4 quaternion coordinates while + the second corresponds to the K independent flexibility frames. + :param parent_dev_imu_quats: + Total deviation observed IMU data in parent flexibility frame as a 2D + array whose first dimension gathers the 4 quaternion coordinates while + the second corresponds to the K independent flexibility frames. + :param deformation_flex_quats: + Pre-allocated memory for storing the local deformation estimates for + each flexibility point flexibility points as a 2D array whose first + dimension gathers the 4 quaternion coordinates while the second + corresponds to the K independent flexibility points. + :param ignore_twist: Whether to ignore the twist of the orientation + estimates of the ordered subset of IMUs of interest, + and incidentally the twist of deformation at the + flexibility points. + """ + # Compute error between observed and rigid kinematic IMU orientation + _compute_orientation_error(obs_imu_quats, + obs_imu_indices, + inv_obs_imu_quats, + kin_imu_rots, + kin_imu_quats, + dev_imu_quats, + ignore_twist) + + # Project observed total deformation on flexibility frames + _compute_deformation_from_deviation(kin_flex_rots, + kin_flex_quats, + is_flex_flipped, + is_chain_orphan, + dev_imu_quats, + inv_child_dev_imu_quats, + parent_dev_imu_quats, + deformation_flex_quats) + + +def get_flexibility_imu_frame_chains( + pinocchio_model: pin.Model, + flex_joint_names: Sequence[str], + imu_frame_names: Sequence[str]) -> Sequence[Tuple[ + Sequence[str], Sequence[Optional[str]], Sequence[bool]]]: + """Extract the minimal set of contiguous sub-chains in the kinematic tree + of a given flexible model that goes through all the specified flexibility + and IMU frames. + + :param pinocchio_model: Model for which to extract sub-chains. + :param flex_joint_names: Unordered sequence of joint names that must be + considered as associated with flexibility frames. + :param imu_frame_names: Unordered sequence of frame names that must be + considered as associated with IMU frames. + """ + # Determine the leaf joints of the kinematic tree, ie all the joints that + # are not parent of any other joint. + parents = pinocchio_model.parents + leaf_joint_indices = set(range(len(parents))) - set(parents) + + # Compute the support of each leaf joint, ie the sub-chain going from each + # leaf to the root joint. + supports = [] + for joint_index in leaf_joint_indices: + support = [] + while joint_index > 0: + support.append(joint_index) + joint_index = parents[joint_index] + supports.append(support) + + # Deduce all the kinematic chains. + # For each support, check if there is a match in any other chain. + # The first match (in order) must be the only one to be considered. + # It always exists, as a root joint is shared by all supports. + chains = [] + for i, support_1 in enumerate(supports): + for support_2 in supports[(i + 1):]: + for joint_index in support_1: + if joint_index in support_2: + break + support_left = support_1[:(support_1.index(joint_index) + 1)] + support_right = support_2[:support_2.index(joint_index)] + chains.append([*support_left, *support_right[::-1]]) + + # Special case when there is a chain in the kinematic tree + if not chains: + chains.append(supports[0]) + + # Determine the parent joint of all flexibility and IMU frames + flex_joint_map, imu_joint_map = ({ + pinocchio_model.frames[index].parent: name + for name, index in zip( + frame_names, get_frame_indices(pinocchio_model, frame_names))} + for frame_names in (flex_joint_names, imu_frame_names)) + flex_joint_indices, imu_joint_indices = ( + set(joint_map.keys()) + for joint_map in (flex_joint_map, imu_joint_map)) + + # Make sure that the robot has no IMU on its root joint if fixed-based + root_type = jiminy.get_joint_type(pinocchio_model, 1) + if root_type != jiminy.JointModelType.FREE: + if 1 in imu_joint_indices: + raise ValueError( + "There must not be an IMU frame attached to the root joint of " + "the robot if it has a fixed based (no freeflyer).") + if not imu_joint_indices.issuperset(leaf_joint_indices): + raise ValueError( + "There must be an IMU frame attached to all the leaf joints " + "of the robot if it has a fixed based (no freeflyer).") + + # Remove all joints not being flex or having an IMU attached + flex_imu_joint_chains_all = [] + flex_or_imu_joint_indices = flex_joint_indices | imu_joint_indices + for chain in chains: + flex_imu_chain = list( + i for i in chain if i in flex_or_imu_joint_indices) + if len(flex_imu_chain) > 1: + flex_imu_joint_chains_all.append(flex_imu_chain) + + # Remove redundant chains, ie the subchains of some other + flex_imu_joint_chains = [] + for i, chain_i in enumerate(flex_imu_joint_chains_all): + for chain_j in flex_imu_joint_chains_all: + is_subchain = False + for i in range(0, len(chain_j) - len(chain_i)): + if chain_i == chain_j[:len(chain_i)]: + is_subchain = True + break + if is_subchain: + break + else: + flex_imu_joint_chains.append(chain_i) + + # Duplicate flexibility and IMU frames sharing the same joint indices. + # Go through each chain and check that IMU and flex joints interleaves. + duplicated_flex_imu_joint_chains = [] + flex_and_imu_joint_indices = flex_joint_indices & imu_joint_indices + for chain in flex_imu_joint_chains_all: + is_imu_joint = False + duplicated_chain = [] + for joint_index in chain: + duplicated_chain.append(joint_index) + if joint_index in flex_and_imu_joint_indices: + duplicated_chain.append(joint_index) + else: + if (is_imu_joint and joint_index in flex_joint_indices) or ( + not is_imu_joint and joint_index in imu_joint_indices): + is_imu_joint = not is_imu_joint + continue + raise ValueError("Flexibility and IMU frames must interleave.") + duplicated_flex_imu_joint_chains.append(duplicated_chain) + + # Extract triplets (parent IMU, flex, child IMU), where either the parent + # or child IMU is None for orphan flexibility frames. + flex_imu_joint_triplets_all: List[ + Tuple[Optional[int], int, Optional[int]]] = [] + imu_and_not_flex_joint_indices = imu_joint_indices - flex_joint_indices + for chain in duplicated_flex_imu_joint_chains: + start_index = 0 + if chain[0] not in imu_and_not_flex_joint_indices: + flex_imu_joint_triplets_all.append( + (None, *chain[:2])) # type: ignore[arg-type] + i = 2 + for j in range(start_index, len(chain) - 2, 2): + flex_imu_joint_triplets_all.append( + chain[j:(j + 3)]) # type: ignore[arg-type] + if chain[-1] not in imu_and_not_flex_joint_indices: + flex_imu_joint_triplets_all.append( + (*chain[-2:], None)) # type: ignore[arg-type] + + # Remove redundant triplet, ie associated with the same flexibility frame. + # Complete triplet must be preferred over orphan triplet. + flex_imu_joint_triplets = [] + flex_joint_triplets = [] + is_orphan_triplets = [] + for triplet in flex_imu_joint_triplets_all: + parent_imu_joint, flex_joint, child_imu_joint = triplet + is_orphan = parent_imu_joint is None or child_imu_joint is None + if flex_joint not in flex_joint_triplets: + flex_imu_joint_triplets.append(triplet) + flex_joint_triplets.append(flex_joint) + is_orphan_triplets.append(is_orphan) + elif not is_orphan: + triplet_index = flex_joint_triplets.index(flex_joint) + if is_orphan_triplets[triplet_index]: + flex_imu_joint_triplets[triplet_index] = triplet + flex_joint_triplets[triplet_index] = flex_joint + + # Concatenate triplets back in frame chains. + # Note that computations can be vectorized for each independent chain. + flex_imu_name_chains = [] # [(flexs, imus, flipped), ...] + imu_grp: List[Optional[str]] = [] + child_joint_imu_prev: Optional[int] = -1 + for triplet in flex_imu_joint_triplets: + parent_imu_joint, flex_joint, child_imu_joint = triplet + if parent_imu_joint is None: + assert child_imu_joint is not None + is_flipped = flex_joint >= child_imu_joint + elif child_imu_joint is None: + assert parent_imu_joint is not None + is_flipped = parent_imu_joint >= flex_joint + else: + is_flipped = triplet != sorted( + triplet) # type: ignore[type-var] + if child_joint_imu_prev != parent_imu_joint: + flex_grp: List[str] = [] + imu_grp = [imu_joint_map.get(parent_imu_joint, None)] + is_flipped_grp: List[bool] = [] + flex_imu_name_chains.append((flex_grp, imu_grp, is_flipped_grp)) + child_joint_imu_prev = child_imu_joint + flex_grp.append(flex_joint_map[flex_joint]) + imu_grp.append(imu_joint_map.get(child_imu_joint, None)) + is_flipped_grp.append(is_flipped) + + return flex_imu_name_chains + + +class DeformationEstimator( + BaseObserverBlock[np.ndarray, np.ndarray, BaseObsT, BaseActT]): + """Compute the local deformation at an arbitrary set of flexibility points + that are presumably responsible for most of the whole deformation of the + mechanical structure. + + .. details:: + The number of IMU sensors and flexibility frames must be consistent: + * If the robot has no freeflyer, there must be as many IMU sensors + as flexibility frames (0), ie + + *---o---x---o---x---o---x + | + | + x---o---x + + * Otherwise, it can either have one more IMU than flexibility + frames (1), the same number (2), or up to one less IMU per + branch in the kinematic tree (3). + + (1) x---o---x---o---x---o---x + | + | + x---o---x + + (2) +---o---x---o---x---o---x + | + | + x---o---x + + (3) +---o---x---o---x---o---+ + | + | + x---o---+ + + *: Fixed base, +: leaf frame, x: IMU frame, o: flexibility frame + + (1): The pose of the freeflyer is ignored when estimating the + deformation at the flexibility frames in local frame. Mathematically, + it is the same as (0) when considered a virtual IMU with fixed + orientation to identity for the root joint. + + (2): One has to compensate for the missing IMU by providing the + configuration of the freeflyer. More precisely, one should ensure that + the orientation of the parent frame of the orphan flexibility frame + matches the reality when considering the rigid configuration. This + usually requires making some assumptions to guess to pose of the frame + that is not directly observable. Any discrepancy will be aggregated to + the estimated deformation for the orphan flexibility frame specifically + since both cannot be distinguished. This issue typically happens when + there is no IMUs in the feet of a legged robot. In such a case, there + is no better option than assuming one of the active contact bodies + remains flat on the ground. If the twist of the IMUs are ignored, then + the twist of the contact body does not matter either, otherwise it must + be set appropriately. If it cannot be observed by some exteroceptive + sensor such as vision, then the most reasonable assumption is to + suppose that it matches the twist of the IMU coming right after in the + kinematic tree. This way, they will cancel out each other without + adding bias to the twist of the orphan flexibility frame. + + (3): This case is basically the same as (2), with the addition that + only the deformation of one of the orphan flexibility frames can + be estimated at once, namely the one whose parent frame is declared as + having known orientation. The other ones will be set to identity. For a + legged robot, this corresponds to one of the contact bodies, usually + the one holding most of the total weight. + + .. warning:: + (2) and (3) are not supported for now, as it requires using one + additional observation layer responsible for estimating the rigid + configuration of the robot including its freeflyer, along with the + name of the reference frame, ie the one having known orientation. + + .. seealso:: + Matthieu Vigne, Antonio El Khoury, Marine Pétriaux, Florent Di Meglio, + and Nicolas Petit "MOVIE: a Velocity-aided IMU Attitude Estimator for + Observing and Controlling Multiple Deformations on Legged Robots" IEEE + Robotics and Automation Letters, Institute of Electrical and + Electronics Engineers, 2022, 7 (2): + https://hal.science/hal-03511198/document + """ + def __init__(self, + name: str, + env: InterfaceJiminyEnv[BaseObsT, BaseActT], + *, + imu_frame_names: Sequence[str], + flex_frame_names: Sequence[str], + ignore_twist: bool = True, + nested_imu_key: Sequence[str] = ("features", "mahony_filter"), + # root_frame: str = "root_joint", + update_ratio: int = 1) -> None: + """ + .. warning:: + The user-specified ordering of the flexibility frames will not be + honored. Indeed, they are reordered to be consistent with the + kinematic tree according to URDF standard. + + :param name: Name of the block. + :param env: Environment to connect with. + :param imu_frame_names: Unordered sequence of IMU frame names that must + be used to estimate the local deformation at + all flexibility frames. + :param flex_frame_names: Unordered sequence of flexibility frame names. + It does not have to match the flexibility + frames of the true dynamic model, which is + only known in simulation. It is up to the user + to choose them appropriately. Ideally, they + must be chosen so as to explain as much as + possible the effect of the deformation on + kinematic quantities, eg the vertical position + of the end effectors, considering the number + and placement of the IMUs at hand. + :param ignore_twist: Whether to ignore the twist of the IMU quaternion + estimate. If so, then the twist of the deformation + will not be estimated either. + :param nested_imu_key: Nested key from environment observation mapping + to the IMU quaternion estimates. Their ordering + must be consistent with the true IMU sensors of + the robot. + :param update_ratio: Ratio between the update period of the controller + and the one of the subsequent controller. + Optional: `1` by default. + """ + # Sanitize user argument(s) + imu_frame_names, flex_frame_names = map( + list, (imu_frame_names, flex_frame_names)) + + # Backup some of the user-argument(s) + self.ignore_twist = ignore_twist + + # Initialize jiminy model + self.model = jiminy.Model() + self.model.initialize(env.robot.pinocchio_model_th) + + # Define proxies for fast access + self.pinocchio_model_th = self.model.pinocchio_model_th + self.pinocchio_data_th = self.model.pinocchio_data_th + + # Create flexible dynamical model. + # Dummy physical parameters are specified as they have no effect on + # kinematic computations. It is only used for computing the flexible + # configuration if requested. + options = self.model.get_options() + for frame_name in flex_frame_names: + options["dynamics"]["flexibilityConfig"].append( + { + "frameName": frame_name, + "stiffness": np.ones(3), + "damping": np.ones(3), + "inertia": np.ones(3), + } + ) + self.model.set_options(options) + + # Extract contiguous chains of flexibility and IMU frames for which + # computations can be vectorized. It also stores the information of + # whether or not the sign of the deformation must be reversed to be + # consistent with standard convention. + flexible_joint_names = self.model.flexible_joint_names + flex_imu_frame_names_chains = get_flexibility_imu_frame_chains( + self.model.pinocchio_model, flexible_joint_names, imu_frame_names) + + # Replace actual flex joint name by corresponding rigid frame + self.flex_imu_frame_names_chains = [] + for flex_frame_names_, imu_frame_names_, is_flipped in ( + flex_imu_frame_names_chains): + flex_frame_names_ = [ + flex_frame_names[flexible_joint_names.index(name)] + for name in flex_frame_names_] + self.flex_imu_frame_names_chains.append( + (flex_frame_names_, imu_frame_names_, is_flipped)) + + # Check if a freeflyer estimator is required + if self.model.has_freeflyer: + for _, imu_frame_names_, _ in self.flex_imu_frame_names_chains: + if None in imu_frame_names_: + raise NotImplementedError( + "Freeflyer estimator is not supported for now.") + + # Backup flexibility frame names + self.flexible_frame_names = [ + name for flex_frame_names, _, _ in self.flex_imu_frame_names_chains + for name in flex_frame_names] + + # Define flexibility and IMU frame orientation proxies for fast access. + # Note that they will be initialized in `_setup` method. + self._kin_flex_rots: List[Tuple[np.ndarray, ...]] = [] + self._kin_imu_rots: List[Tuple[np.ndarray, ...]] = [] + + # Pre-allocate memory for intermediary computations + self._is_chain_orphan = [] + self._is_flex_flipped = [] + self._kin_flex_quats = [] + self._inv_obs_imu_quats = [] + self._kin_imu_quats = [] + self._dev_imu_quats = [] + self._dev_parent_imu_quats = [] + self._dev_child_imu_quats = [] + for flex_frame_names_, imu_frame_names_, is_flipped in ( + self.flex_imu_frame_names_chains): + num_flexs = len(flex_frame_names_) + num_imus = len(tuple(filter(None, imu_frame_names_))) + self._kin_flex_quats.append(np.empty((4, num_flexs))) + self._kin_imu_quats.append(np.empty((4, num_imus))) + self._inv_obs_imu_quats.append(np.empty((4, num_imus))) + self._dev_imu_quats.append(np.empty((4, num_imus))) + dev_parent_imu_quats = np.empty((4, num_flexs)) + is_parent_orphan = imu_frame_names_[0] is None + is_child_orphan = imu_frame_names_[-1] is None + self._dev_parent_imu_quats.append(dev_parent_imu_quats) + dev_child_imu_quats = np.empty((4, num_flexs)) + self._dev_child_imu_quats.append(dev_child_imu_quats) + self._is_flex_flipped.append(np.array(is_flipped)) + self._is_chain_orphan.append((is_parent_orphan, is_child_orphan)) + + # Define IMU and encoder measurement proxy for fast access + obs_imu_quats: DataNested = env.observation + for key in nested_imu_key: + assert isinstance(obs_imu_quats, Mapping) + obs_imu_quats = obs_imu_quats[key] + assert isinstance(obs_imu_quats, np.ndarray) + self._obs_imu_quats = obs_imu_quats + + # Get mapping from IMU frame to index + imu_frame_map: Dict[str, int] = {} + for sensor_name in env.robot.sensor_names[imu.type]: + sensor = env.robot.get_sensor(imu.type, sensor_name) + assert isinstance(sensor, imu) + imu_frame_map[sensor.frame_name] = sensor.index + + # Make sure that the robot has one encoder per rigid joint + encoder_sensor_names = env.robot.sensor_names[encoder.type] + if len(encoder_sensor_names) < len(self.model.rigid_joint_indices): + raise ValueError( + "The robot must have one encoder per joint in theoretical " + "model (excluding floating base if any).") + + # Extract mapping from encoders to rigid configuration. + # Note that revolute unbounded joints are not supported for now. + self.encoder_to_config = [-1 for _ in range(env.robot.nmotors)] + for i, sensor_name in enumerate(encoder_sensor_names): + sensor = env.robot.get_sensor(encoder.type, sensor_name) + assert isinstance(sensor, encoder) + if sensor.joint_type == jiminy.JointModelType.ROTARY_UNBOUNDED: + raise ValueError( + "Revolute unbounded joints are not supported for now.") + encoder_joint = self.pinocchio_model_th.joints[sensor.joint_index] + self.encoder_to_config[i] = encoder_joint.idx_q + + # Extract measured joint positions for fast access. + # Note that they will be initialized in `_setup` method. + self.encoder_data = np.array([]) + + # Buffer storing the rigid configuration. + self._q_rigid = pin.neutral(self.pinocchio_model_th) + + # Whether the observer has been compiled already + self._is_compiled = False + + # Initialize the observer + super().__init__(name, env, update_ratio) + + # Define chunks associated with each independent flexibility-imu chain + self._deformation_flex_quats, self._obs_imu_indices = [], [] + flex_start_index = 0 + for flex_frame_names_, imu_frame_names_, _ in ( + self.flex_imu_frame_names_chains): + flex_last_index = flex_start_index + len(flex_frame_names_) + self._deformation_flex_quats.append( + self.observation[:, flex_start_index:flex_last_index]) + flex_start_index = flex_last_index + imu_frame_names_filtered = tuple(filter(None, imu_frame_names_)) + imu_indices = tuple( + imu_frame_map[name] for name in imu_frame_names_filtered) + self._obs_imu_indices.append(imu_indices) + + def _initialize_observation_space(self) -> None: + nflex = sum( + len(flex_frame_names) + for flex_frame_names, _, _ in self.flex_imu_frame_names_chains) + self.observation_space = gym.spaces.Box( + low=np.full((4, nflex), -1.0 - 1e-9), + high=np.full((4, nflex), 1.0 + 1e-9), + dtype=np.float64) + + def _setup(self) -> None: + # Call base implementation + super()._setup() + + # Fix initialization of the observation to be valid quaternions + self.observation[-1] = 1.0 + + # Refresh flexibility and IMU frame orientation proxies + self._kin_flex_rots.clear() + self._kin_imu_rots.clear() + for flex_frame_names, imu_frame_names_, _ in ( + self.flex_imu_frame_names_chains): + imu_frame_names = list(filter(None, imu_frame_names_)) + kin_flex_rots, kin_imu_rots = (tuple( + self.pinocchio_data_th.oMf[frame_index].rotation + for frame_index in get_frame_indices( + self.pinocchio_model_th, frame_names)) + for frame_names in (flex_frame_names, imu_frame_names)) + self._kin_flex_rots.append(kin_flex_rots) + self._kin_imu_rots.append(kin_imu_rots) + + # Refresh measured motor positions and velocities proxies + self.encoder_data, _ = self.env.sensor_measurements[encoder.type] + + # Call `refresh_observation` manually to pre-compile it if necessary + if not self._is_compiled: + self.refresh_observation(self.env.observation) + self._is_compiled = True + + @property + def fieldnames(self) -> List[List[str]]: + return [[f"{name}.Quat{e}" for name in self.flexible_frame_names] + for e in ("x", "y", "z", "w")] + + def refresh_observation(self, measurement: BaseObsT) -> None: + # Estimate the rigid configuration of the robot from encoder data + self._q_rigid[self.encoder_to_config] = self.encoder_data + + # Update kinematic quantities according to estimate rigid configuration + pin.framesForwardKinematics( + self.pinocchio_model_th, self.pinocchio_data_th, self._q_rigid) + + # Estimate all the deformations in their local frame. + # It loops over each flexibility-imu chain independently. + for args in zip( + self._obs_imu_indices, + self._inv_obs_imu_quats, + self._kin_imu_rots, + self._kin_imu_quats, + self._kin_flex_rots, + self._kin_flex_quats, + self._is_flex_flipped, + self._is_chain_orphan, + self._dev_imu_quats, + self._dev_child_imu_quats, + self._dev_parent_imu_quats, + self._deformation_flex_quats): + flexibility_estimator( + self._obs_imu_quats, *args, self.ignore_twist) diff --git a/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py b/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py index 20adbfbb8..7ac438946 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py +++ b/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py @@ -167,13 +167,7 @@ def __init__(self, else: render_mode = 'rgb_array' - # Make sure the rendering mode is valid. - if render_mode == 'human' and { - **kwargs, **simulator.viewer_kwargs - }.get('backend') == 'panda3d-sync': - raise ValueError( - "render_mode='human' is incompatible with " - "backend='panda3d-sync'.") + # Make sure that rendering mode is valid assert render_mode in self.metadata['render_modes'] # Backup some user arguments diff --git a/python/gym_jiminy/common/gym_jiminy/common/utils/math.py b/python/gym_jiminy/common/gym_jiminy/common/utils/math.py index 1aeb70499..bc9e5aea4 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/utils/math.py +++ b/python/gym_jiminy/common/gym_jiminy/common/utils/math.py @@ -160,11 +160,14 @@ def matrix_to_quat(mat: np.ndarray, # TODO: Merge this method with `matrix_to_quat` by leverage compile-time # implementation dispatching via `nb.generated_jit` or `nb.overload`. @nb.jit(nopython=True, cache=True) -def matrices_to_quat(mat_list: Tuple[np.ndarray], +def matrices_to_quat(mat_list: Tuple[np.ndarray, ...], out: Optional[np.ndarray] = None) -> np.ndarray: """Compute the [qx, qy, qz, qw] Quaternion representation of multiple rotation matrices. + .. seealso:: + See https://math.stackexchange.com/a/3183435. + :param mat: Tuple of N arrays corresponding to independent 3D rotation matrices. :param out: A pre-allocated array into which the result is stored. If not diff --git a/python/gym_jiminy/envs/gym_jiminy/envs/__init__.py b/python/gym_jiminy/envs/gym_jiminy/envs/__init__.py index 5b865c9d2..a26d9d074 100644 --- a/python/gym_jiminy/envs/gym_jiminy/envs/__init__.py +++ b/python/gym_jiminy/envs/gym_jiminy/envs/__init__.py @@ -2,7 +2,7 @@ from .cartpole import CartPoleJiminyEnv from .acrobot import AcrobotJiminyEnv -from .ant import AntEnv +from .ant import AntJiminyEnv from .cassie import CassieJiminyEnv, CassiePDControlJiminyEnv from .digit import DigitJiminyEnv, DigitPDControlJiminyEnv from .anymal import ANYmalJiminyEnv, ANYmalPDControlJiminyEnv @@ -15,7 +15,7 @@ __all__ = [ 'CartPoleJiminyEnv', 'AcrobotJiminyEnv', - 'AntEnv', + 'AntJiminyEnv', 'CassieJiminyEnv', 'CassiePDControlJiminyEnv', 'DigitJiminyEnv', @@ -46,7 +46,7 @@ ) register( id='ant', - entry_point='gym_jiminy.envs:AntEnv', + entry_point='gym_jiminy.envs:AntJiminyEnv', max_episode_steps=1000, reward_threshold=6000.0, order_enforce=False, diff --git a/python/gym_jiminy/envs/gym_jiminy/envs/ant.py b/python/gym_jiminy/envs/gym_jiminy/envs/ant.py index 9f0c0d4bf..c8b6108da 100644 --- a/python/gym_jiminy/envs/gym_jiminy/envs/ant.py +++ b/python/gym_jiminy/envs/gym_jiminy/envs/ant.py @@ -13,7 +13,7 @@ from jiminy_py.simulator import Simulator from gym_jiminy.common.bases import InfoType, EngineObsType from gym_jiminy.common.envs import BaseJiminyEnv -from gym_jiminy.common.utils import sample, copyto, squared_norm_2, clip +from gym_jiminy.common.utils import sample, copyto, squared_norm_2 if sys.version_info < (3, 9): from importlib_resources import files @@ -25,7 +25,7 @@ STEP_DT = 0.05 -class AntEnv(BaseJiminyEnv[np.ndarray, np.ndarray]): +class AntJiminyEnv(BaseJiminyEnv[np.ndarray, np.ndarray]): """ TODO: Write documentation. """ @@ -65,10 +65,13 @@ def __init__(self, debug: bool = False, **kwargs: Any) -> None: self._obs_slices: Tuple[np.ndarray, ...] = () # Define base orientation and external forces proxies for fast access. - # Note that they will be initialized in `_initialize_buffers`. self._base_rot = np.array([]) self._f_external: Tuple[np.ndarray, ...] = () + # Rigid configuration and velocity of the robot. + self._q_rigid = np.array([]) + self._v_rigid = np.array([]) + # Initialize base class super().__init__( simulator=simulator, @@ -120,10 +123,10 @@ def _initialize_observation_space(self) -> None: The observation space comprises: - - robot configuration vector (absolute position (x, y) excluded), - - robot velocity vector (with base linear velocity in world frame), - - flatten external forces applied on each body in local frame, ie - centered at their respective center of mass. + - rigid configuration (absolute position (x, y) excluded), + - rigid velocity (with base linear velocity in world frame), + - flattened external forces applied on each body in local frame, + ie centered at their respective center of mass. """ # http://www.mujoco.org/book/APIreference.html#mjData @@ -157,7 +160,12 @@ def _initialize_buffers(self) -> None: self._base_rot = self.robot.pinocchio_data.oMf[1].rotation # Initialize vector of external forces - self._f_external = tuple(f.vector for f in self.robot_state.f_external) + self._f_external = tuple( + self.robot_state.f_external[joint_index].vector + for joint_index in self.robot.rigid_joint_indices) + + # Refresh buffers manually to initialize them early + self._refresh_buffers() # Re-initialize observation slices. # Note that the base linear velocity is isolated as it will be computed @@ -165,9 +173,9 @@ def _initialize_buffers(self) -> None: obs_slices = [] obs_index_first = 0 for data in ( - self._robot_state_q[2:], - self._robot_state_v[:3], - self._robot_state_v[3:], + self._q_rigid[2:], + self._v_rigid[:3], + self._v_rigid[3:], *self._f_external): obs_index_last = obs_index_first + len(data) obs_slices.append(self.observation[obs_index_first:obs_index_last]) @@ -177,18 +185,26 @@ def _initialize_buffers(self) -> None: # Initialize previous torso position along x-axis self._xpos_prev = self._robot_state_q[0] + def _refresh_buffers(self) -> None: + self._q_rigid = self.robot.get_rigid_position_from_flexible( + self._robot_state_q) + self._v_rigid = self.robot.get_rigid_velocity_from_flexible( + self._robot_state_v) + def refresh_observation(self, measurement: EngineObsType) -> None: # Update observation copyto(self._obs_slices[:-1], ( - self._robot_state_q[2:], - self._robot_state_v[3:], - *self._f_external)) + self._q_rigid[2:], self._v_rigid[3:], *self._f_external)) # Transform observed linear velocity to be in world frame self._obs_slices[-1][:] = self._base_rot @ self._robot_state_v[:3] - # Clip observation to make sure it is not out of bounds - clip(self.observation, self.observation_space) + # Clip observation in-place to make sure it is not out of bounds + assert isinstance(self.observation_space, gym.spaces.Box) + np.clip(self.observation, + self.observation_space.low, + self.observation_space.high, + out=self.observation) def has_terminated(self) -> Tuple[bool, bool]: """ TODO: Write documentation. diff --git a/python/gym_jiminy/unit_py/data/flexible_arm.urdf b/python/gym_jiminy/unit_py/data/flexible_arm.urdf new file mode 100644 index 000000000..071a7a3d5 --- /dev/null +++ b/python/gym_jiminy/unit_py/data/flexible_arm.urdf @@ -0,0 +1,154 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/python/gym_jiminy/unit_py/test_deformation_estimator.py b/python/gym_jiminy/unit_py/test_deformation_estimator.py new file mode 100644 index 000000000..448615dd1 --- /dev/null +++ b/python/gym_jiminy/unit_py/test_deformation_estimator.py @@ -0,0 +1,254 @@ +""" TODO: Write documentation +""" +import os +import unittest +from typing import Any + +import numpy as np + +import jiminy_py.core as jiminy +from jiminy_py.simulator import Simulator + +from gym_jiminy.common.envs import BaseJiminyEnv +from gym_jiminy.common.blocks import PDController, MahonyFilter, DeformationEstimator +from gym_jiminy.common.bases import ObservedJiminyEnv, ControlledJiminyEnv +from gym_jiminy.common.utils import matrices_to_quat, build_pipeline +from gym_jiminy.envs import AntJiminyEnv + + +class DeformationEstimatorBlock(unittest.TestCase): + """ TODO: Write documentation + """ + def _test_deformation_estimate(self, env, imu_atol, flex_atol): + # Check that quaternion estimates from MahonyFilter are valid + true_imu_rots = [] + for frame_name in env.robot.sensor_names['ImuSensor']: + frame_index = env.robot.pinocchio_model.getFrameId(frame_name) + frame_rot = env.robot.pinocchio_data.oMf[frame_index].rotation + true_imu_rots.append(frame_rot) + true_imu_quats = matrices_to_quat(tuple(true_imu_rots)) + + est_imu_quats = env.observation['features']['mahony_filter'] + np.testing.assert_allclose( + true_imu_quats, est_imu_quats, atol=imu_atol) + + # Check that deformation estimates from DeformationEstimator are valid + model_options = env.robot.get_model_options() + flexible_frame_names = [ + flex_options["frameName"] + for flex_options in model_options["dynamics"]["flexibilityConfig"] + ] + est_flex_quats = env.observation['features']['deformation_estimator'] + est_flex_quats[:] *= np.sign(est_flex_quats[-1]) + for frame_name, joint_index in zip( + flexible_frame_names, env.robot.flexible_joint_indices): + idx_q = env.robot.pinocchio_model.joints[joint_index].idx_q + true_flex_quat = env.robot_state.q[idx_q:(idx_q + 4)] + flex_index = env.observer.flexible_frame_names.index(frame_name) + est_flex_quat = est_flex_quats[:, flex_index] + np.testing.assert_allclose( + true_flex_quat, est_flex_quat, atol=flex_atol) + + def test_arm(self): + """ TODO: Write documentation + """ + import numpy as np + + import jiminy_py.core as jiminy + + # First mount the drive + data_dir = os.path.join(os.path.dirname(__file__), "data") + urdf_path = f"{data_dir}/flexible_arm.urdf" + robot = jiminy.Robot() + robot.initialize(urdf_path, has_freeflyer=False) + + # Add motor + motor_joint_name = 'base_to_link1' + motor = jiminy.SimpleMotor(motor_joint_name) + robot.attach_motor(motor) + motor.initialize(motor_joint_name) + + # Add sensors + encoder = jiminy.EncoderSensor(motor_joint_name) + robot.attach_sensor(encoder) + encoder.initialize(motor_joint_name) + + for i in range(1, 5): + imu = jiminy.ImuSensor(f"link{i + 1}") + robot.attach_sensor(imu) + imu.initialize(f"link{i + 1}") + + # We set inertia along non-moving axis to 1.0 for numerical stability + k_j, d_j = 50.0, 5.0 + model_options = robot.get_model_options() + model_options['dynamics']['flexibilityConfig'] = [{ + 'frameName': f"link{i}_to_link{i+1}", + 'stiffness': k_j * np.ones(3), + 'damping': d_j * np.ones(3), + 'inertia': np.array([1.0, 1.0, 0.0]) + } for i in range(1, 5) + ] + model_options['joints']['enablePositionLimit'] = False + model_options['joints']['enableVelocityLimit'] = False + robot.set_model_options(model_options) + + # Create a simulator using this robot and controller + simulator = Simulator(robot) + + # Configure the controller and sensor update periods + engine_options = simulator.engine.get_options() + engine_options['stepper']['controllerUpdatePeriod'] = 0.0001 + engine_options['stepper']['sensorsUpdatePeriod'] = 0.0001 + simulator.engine.set_options(engine_options) + + # Instantiate the environment + env = BaseJiminyEnv(simulator, step_dt=0.0001, debug=True) + + # Add controller and observer blocks + pd_controller = PDController( + "pd_controller", + env, + kp=150.0, + kd=0.03, + order=1, + update_ratio=1) + env = ControlledJiminyEnv(env, pd_controller) + + mahony_filter = MahonyFilter( + "mahony_filter", + env, + kp=0.0, + ki=0.0, + twist_time_constant=None, + exact_init=True, + update_ratio=1) + env = ObservedJiminyEnv(env, mahony_filter) + + deformation_estimator = DeformationEstimator( + "deformation_estimator", + env, + imu_frame_names=robot.sensor_names['ImuSensor'], + flex_frame_names=robot.flexible_joint_names, + ignore_twist=True, + update_ratio=1) + env = ObservedJiminyEnv(env, deformation_estimator) + + # Simulate for a while + env.reset(seed=0) + for _ in range(10000): + env.step(env.action) + + # Check that deformation estimates matches ground truth + self._test_deformation_estimate(env, imu_atol=1e-6, flex_atol=1e-5) + + def test_ant(self): + """ TODO: Write documentation + """ + FLEXIBLE_FRAME_NAMES = ("leg_1", + "leg_2", + "leg_3", + "leg_4", + "ankle_1", + "ankle_2", + "ankle_3", + "ankle_4") + + IMU_FRAME_NAMES = ("torso", + "aux_1", + "foot_1", + "aux_2", + "foot_2", + "aux_3", + "foot_3", + "aux_4", + "foot_4") + + # Overload the environment to add deformation points and IMU frames + class FlexAntJiminyEnv(AntJiminyEnv): + def __init__(self, debug: bool = False, **kwargs: Any) -> None: + # Call base implementation + super().__init__(debug, **kwargs) + + # Add IMU frames + for frame_name in IMU_FRAME_NAMES: + sensor = jiminy.ImuSensor(frame_name) + self.robot.attach_sensor(sensor) + sensor.initialize(frame_name) + + def _setup(self) -> None: + # Call base implementation + super()._setup() + + # Configure the controller and sensor update periods + engine_options = self.engine.get_options() + engine_options['stepper']['controllerUpdatePeriod'] = 0.0001 + engine_options['stepper']['sensorsUpdatePeriod'] = 0.0001 + self.engine.set_options(engine_options) + + # Add flexibility frames + model_options = self.robot.get_model_options() + model_options["dynamics"]["flexibilityConfig"] = [] + for frame_name in FLEXIBLE_FRAME_NAMES: + model_options["dynamics"]["flexibilityConfig"].append( + { + "frameName": frame_name, + "stiffness": np.array([10.0, 10.0, 10.0]), + "damping": np.array([0.1, 0.1, 0.1]), + "inertia": np.array([0.01, 0.01, 0.01]), + } + ) + model_options["dynamics"]["enableFlexibleModel"] = True + self.robot.set_model_options(model_options) + + # Create pipeline with Mahony filter and DeformationObserver blocks + PipelineEnv = build_pipeline( + env_config=dict( + cls=FlexAntJiminyEnv + ), + layers_config=[ + dict( + block=dict( + cls=PDController, + kwargs=dict( + kp=150.0, + kd=0.03, + order=1, + update_ratio=1, + ) + ) + ), + dict( + block=dict( + cls=MahonyFilter, + kwargs=dict( + kp=0.0, + ki=0.0, + twist_time_constant=None, + exact_init=True, + update_ratio=1, + ) + ) + ), dict( + block=dict( + cls=DeformationEstimator, + kwargs=dict( + imu_frame_names=IMU_FRAME_NAMES, + flex_frame_names=FLEXIBLE_FRAME_NAMES, + ignore_twist=False, + update_ratio=1, + ) + ) + ) + ] + ) + + # Instantiate the environment + env = PipelineEnv() + + # Run a simulation + env.reset(seed=0) + for _ in range(150): + env.step(env.action) + + # Check that deformation estimates matches ground truth + self._test_deformation_estimate(env, imu_atol=1e-4, flex_atol=5e-3) diff --git a/python/gym_jiminy/unit_py/test_pipeline_control.py b/python/gym_jiminy/unit_py/test_pipeline_control.py index e96a31c81..ab214639d 100644 --- a/python/gym_jiminy/unit_py/test_pipeline_control.py +++ b/python/gym_jiminy/unit_py/test_pipeline_control.py @@ -37,7 +37,7 @@ def _test_pid_standing(self): """ TODO: Write documentation """ # Reset the environment - self.env.reset() + self.env.reset(seed=0) # Zero target motors velocities, so that the robot stands still action = np.zeros(self.env.robot.nmotors) @@ -146,13 +146,13 @@ def test_mahony_filter(self): env.observer = MahonyFilter( env.observer.name, env.observer.env, + kp=0.0, + ki=0.0, twist_time_constant=twist_time_constant, - exact_init=True, - kp=env.observer.kp, - ki=env.observer.ki) + exact_init=True) # Reset the environment - env.reset() + env.reset(seed=0) # Run of few simulation steps for i in range(200): @@ -179,7 +179,7 @@ def test_pid_controller(self): """ # Instantiate the environment and run a simulation with random action env = AtlasPDControlJiminyEnv() - env.reset() + env.reset(seed=0) env.unwrapped._height_neutral = float("-inf") while env.stepper_state.t < 2.0: env.step(0.2 * env.action_space.sample()) diff --git a/python/gym_jiminy/unit_py/test_pipeline_design.py b/python/gym_jiminy/unit_py/test_pipeline_design.py index 289a9f01e..6e5163f91 100644 --- a/python/gym_jiminy/unit_py/test_pipeline_design.py +++ b/python/gym_jiminy/unit_py/test_pipeline_design.py @@ -89,14 +89,14 @@ def test_load_files(self): toml_file = os.path.join(data_dir, "anymal_pipeline.toml") ANYmalPipelineEnv = load_pipeline(toml_file) env = ANYmalPipelineEnv(debug=True) - env.reset() + env.reset(seed=0) env.step(env.action) # Load JSON pipeline description, create env and perform a step json_file = os.path.join(data_dir, "anymal_pipeline.json") ANYmalPipelineEnv = load_pipeline(json_file) env = ANYmalPipelineEnv(debug=True) - env.reset() + env.reset(seed=0) env.step(env.action) def test_override_default(self): @@ -119,7 +119,7 @@ def test_memory_leak(self): impossible for the garbage collector to release memory. """ env = self.ANYmalPipelineEnv() - env.reset() + env.reset(seed=0) proxy = weakref.proxy(env) env = None gc.collect() @@ -130,7 +130,7 @@ def test_initial_state(self): """ # Get initial observation env = self.ANYmalPipelineEnv() - obs, _ = env.reset() + obs, _ = env.reset(seed=0) # Controller target is observed, and has right name self.assertTrue('actions' in obs and 'pd_controller' in obs['actions']) @@ -166,7 +166,7 @@ def test_step_state(self): """ # Perform a single step env = self.ANYmalPipelineEnv() - env.reset() + env.reset(seed=0) action = env.env.observation['actions']['pd_controller'].copy() action += 1.0e-3 obs, *_ = env.step(action) @@ -211,7 +211,7 @@ def configure_telemetry() -> InterfaceJiminyEnv: env.simulator.engine.set_options(engine_options) return env - env.reset(options=dict(reset_hook=configure_telemetry)) + env.reset(seed=0, options=dict(reset_hook=configure_telemetry)) env.step(env.action) # Check that the command is updated 1/2 low-level controller update diff --git a/python/jiminy_py/src/jiminy_py/plot.py b/python/jiminy_py/src/jiminy_py/plot.py index f6cd267d3..ce173d4eb 100644 --- a/python/jiminy_py/src/jiminy_py/plot.py +++ b/python/jiminy_py/src/jiminy_py/plot.py @@ -155,7 +155,7 @@ def __init__(self, # pylint: disable=unused-argument if self.offscreen: self.figure.set_size_inches(18, 12) else: - self.figure.set_size_inches(14, 8) + self.figure.set_size_inches(12, 8) # Register 'on resize' event callback to adjust layout self.figure.canvas.mpl_connect('resize_event', self.adjust_layout) diff --git a/python/jiminy_py/src/jiminy_py/simulator.py b/python/jiminy_py/src/jiminy_py/simulator.py index c6db5d2f1..127ecc631 100644 --- a/python/jiminy_py/src/jiminy_py/simulator.py +++ b/python/jiminy_py/src/jiminy_py/simulator.py @@ -607,7 +607,7 @@ def render(self, self.viewer.set_camera_transform(*camera_pose) # Make sure the graphical window is open if required - if not return_rgb_array: + if not return_rgb_array and self.viewer.backend != "panda3d-sync": Viewer.open_gui() # Try refreshing the viewer