From 9b5a80555e55a4d4d44e018d251184049b5ff7d7 Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Mon, 25 Nov 2024 20:25:37 +0100 Subject: [PATCH 1/4] [gym_jiminy/common] Add 'compute_rpy' optional argument to mahony observer. --- .../gym_jiminy/common/blocks/mahony_filter.py | 83 +++++++++++++++---- .../unit_py/test_pipeline_control.py | 17 ++-- 2 files changed, 75 insertions(+), 25 deletions(-) diff --git a/python/gym_jiminy/common/gym_jiminy/common/blocks/mahony_filter.py b/python/gym_jiminy/common/gym_jiminy/common/blocks/mahony_filter.py index d37ae23cb..fde3a55ef 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/blocks/mahony_filter.py +++ b/python/gym_jiminy/common/gym_jiminy/common/blocks/mahony_filter.py @@ -13,6 +13,7 @@ from ..bases import BaseObs, BaseAct, BaseObserverBlock, InterfaceJiminyEnv from ..utils import (fill, + quat_to_rpy, matrices_to_quat, swing_from_vector, compute_tilt_from_quat, @@ -144,11 +145,14 @@ def update_twist(q: np.ndarray, ) -class MahonyFilter( - BaseObserverBlock[np.ndarray, Dict[str, np.ndarray], BaseObs, BaseAct] - ): +class MahonyFilter(BaseObserverBlock[ + Dict[str, np.ndarray], Dict[str, np.ndarray], BaseObs, BaseAct]): """Mahony's Nonlinear Complementary Filter on SO(3). + This observer estimate the 3D orientation of all the IMU of the robot from + the Gyrometer and Accelerometer measurements, i.e. the angular velocity in + local frame and the linear acceleration minus gravity in local frame. + .. seealso:: Robert Mahony, Tarek Hamel, and Jean-Michel Pflimlin "Nonlinear Complementary Filters on the Special Orthogonal Group" IEEE @@ -156,6 +160,18 @@ class MahonyFilter( Electronics Engineers, 2008, 53 (5), pp.1203-1217: https://hal.archives-ouvertes.fr/hal-00488376/document + .. note:: + The feature space of this observer is a dictionary storing quaternion + estimates under key 'quat', and optionally, their corresponding + Yaw-Pitch-Roll Euler angles representations under key 'key' if + `compute_rpy=True`. Both leaves are 2D array of shape (N, M), where + N is the number of components of the representation while M is the + number of IMUs of the robot. More specifically, `N=4` for quaternions + (Qx, Qy, Qz, Qw) and 'N=3' for Euler angles (Roll-Pitch-Yaw). Indeed, + the Yaw angle of the Yaw-Pitch-Roll Euler angles representations is + systematically included in the feature space, even if its value is + meaningless, i.e. `ignore_twist=True`. + .. warning:: This filter works best for 'observe_dt' smaller or equal to 5ms. Its performance drops rapidly beyond this point. Having 'observe_dt' equal @@ -170,6 +186,7 @@ def __init__(self, exact_init: bool = True, kp: Union[np.ndarray, float] = 1.0, ki: Union[np.ndarray, float] = 0.1, + compute_rpy: bool = True, update_ratio: int = 1) -> None: """ :param name: Name of the block. @@ -196,6 +213,10 @@ def __init__(self, Optional: `1.0` by default. :param mahony_ki: Integral gain used for gyro bias estimate. Optional: `0.1` by default. + :param compute_rpy: Whether to compute the Yaw-Pitch-Roll Euler angles + representation for the 3D orientation of the IMU, + in addition to the quaternion representation. + Optional: False by default. :param update_ratio: Ratio between the update period of the observer and the one of the subsequent observer. -1 to match the simulation timestep of the environment. @@ -209,6 +230,7 @@ def __init__(self, ki = np.full((num_imu_sensors,), ki) # Backup some of the user arguments + self.compute_rpy = compute_rpy self.exact_init = exact_init self.kp = kp self.ki = ki @@ -263,6 +285,13 @@ def __init__(self, # Initialize the observer super().__init__(name, env, update_ratio) + # Define some proxies for fast access + self._quat = self.observation["quat"] + if self.compute_rpy: + self._rpy = self.observation["rpy"] + else: + self._rpy = np.array([]) + def _initialize_state_space(self) -> None: """Configure the internal state space of the observer. @@ -291,22 +320,34 @@ def _initialize_state_space(self) -> None: def _initialize_observation_space(self) -> None: """Configure the observation space of the observer. + .. note:: + if `compute_rpy` is enabled, then the Roll, Pitch and Yaw angles + are guaranteed to be within range [-pi,pi], [-pi/2,pi/2], [-pi,pi]. + It corresponds to the current orientation estimate for all the IMUs of the robot at once, with special treatment for their twist part. See `__init__` documentation for details. """ + observation_space: Dict[str, gym.Space] = OrderedDict() num_imu_sensors = len(self.env.robot.sensors[ImuSensor.type]) - self.observation_space = gym.spaces.Box( + observation_space["quat"] = gym.spaces.Box( low=np.full((4, num_imu_sensors), -1.0 - 1e-9), high=np.full((4, num_imu_sensors), 1.0 + 1e-9), dtype=np.float64) + if self.compute_rpy: + high = np.array([np.pi, np.pi/2, np.pi]) + 1e-9 + observation_space["rpy"] = gym.spaces.Box( + low=-high[:, np.newaxis].repeat(num_imu_sensors, axis=1), + high=high[:, np.newaxis].repeat(num_imu_sensors, axis=1), + dtype=np.float64) + self.observation_space = gym.spaces.Dict(observation_space) def _setup(self) -> None: # Call base implementation super()._setup() # Fix initialization of the observation to be valid quaternions - self.observation[-1] = 1.0 + self._quat[-1] = 1.0 # Make sure observe update is discrete-time if self.env.observe_dt <= 0.0: @@ -317,7 +358,7 @@ def _setup(self) -> None: self.gyro, self.acc = np.split( self.env.sensor_measurements[ImuSensor.type], 2) - # Reset the sensor bias + # Reset the sensor bias estimate fill(self._bias, 0) # Reset the twist estimate @@ -345,10 +386,17 @@ def get_state(self) -> Dict[str, np.ndarray]: return self._state @property - def fieldnames(self) -> List[List[str]]: + def fieldnames(self) -> Dict[str, List[List[str]]]: imu_sensors = self.env.robot.sensors[ImuSensor.type] - return [[f"{sensor.name}.Quat{e}" for sensor in imu_sensors] - for e in ("x", "y", "z", "w")] + fieldnames: Dict[str, List[List[str]]] = {} + fieldnames["quat"] = [ + [f"{sensor.name}.Quat{e}" for sensor in imu_sensors] + for e in ("x", "y", "z", "w")] + if self.compute_rpy: + fieldnames["rpy"] = [ + [".".join((sensor.name, e)) for sensor in imu_sensors] + for e in ("Roll", "Pitch", "Yaw")] + return fieldnames def refresh_observation(self, measurement: BaseObs) -> None: # Re-initialize the quaternion estimate if no simulation running. @@ -365,7 +413,7 @@ def refresh_observation(self, measurement: BaseObs) -> None: # not very accurate because the initial acceleration is # often jerky, plus the tilt is not observable at all. acc = self.acc / np.linalg.norm(self.acc, axis=0) - swing_from_vector(acc, self.observation) + swing_from_vector(acc, self._quat) self._is_initialized = True if not self._is_initialized: @@ -378,18 +426,17 @@ def refresh_observation(self, measurement: BaseObs) -> None: imu_rots.append(rot) # Convert it to quaternions - matrices_to_quat(tuple(imu_rots), self.observation) + matrices_to_quat(tuple(imu_rots), self._quat) # Keep track of tilt if necessary if self._update_twist: - self._twist[:] = np.arctan2( - self.observation[2], self.observation[3]) + self._twist[:] = np.arctan2(self._quat[2], self._quat[3]) self._is_initialized = True return # Run an iteration of the filter, computing the next state estimate - mahony_filter(self.observation, + mahony_filter(self._quat, self._omega, self.gyro, self.acc, @@ -400,11 +447,15 @@ def refresh_observation(self, measurement: BaseObs) -> None: # Remove twist if requested if self._remove_twist: - remove_twist_from_quat(self.observation) + remove_twist_from_quat(self._quat) if self._update_twist: - update_twist(self.observation, + update_twist(self._quat, self._twist, self._omega, self.twist_time_constant_inv, self.observe_dt) + + # Compute the RPY representation if requested + if self.compute_rpy: + quat_to_rpy(self._quat, self._rpy) diff --git a/python/gym_jiminy/unit_py/test_pipeline_control.py b/python/gym_jiminy/unit_py/test_pipeline_control.py index 2616e9b88..0f466aae1 100644 --- a/python/gym_jiminy/unit_py/test_pipeline_control.py +++ b/python/gym_jiminy/unit_py/test_pipeline_control.py @@ -134,7 +134,8 @@ def test_mahony_filter(self): """ # Instantiate and reset the environment env = AtlasPDControlJiminyEnv() - assert isinstance(env.observer, MahonyFilter) + observer = env.observer + assert isinstance(observer, MahonyFilter) # Define a constant action that move the upper-body in all directions robot = env.robot @@ -150,16 +151,18 @@ def test_mahony_filter(self): # Check that the estimate IMU orientation is accurate over the episode for twist_time_constant in (None, float("inf"), 0.0): # Reinitialize the observer - env.observer = MahonyFilter( - env.observer.name, - env.observer.env, + env.observer = observer = MahonyFilter( + observer.name, + observer.env, kp=0.0, ki=0.0, twist_time_constant=twist_time_constant, - exact_init=True) + exact_init=True, + compute_rpy=True) # Reset the environment env.reset(seed=0) + rpy_est = observer.observation["rpy"][:, 0] # Run of few simulation steps for i in range(200): @@ -170,13 +173,9 @@ def test_mahony_filter(self): obs_true = matrix_to_quat(imu_rot) remove_twist_from_quat(obs_true) rpy_true = quat_to_rpy(obs_true) - obs_est = env.observer.observation[:, 0].copy() - remove_twist_from_quat(obs_est) - rpy_est = quat_to_rpy(obs_est) else: # The twist is either measured or estimated rpy_true = matrix_to_rpy(imu_rot) - rpy_est = quat_to_rpy(env.observer.observation[:, 0]) np.testing.assert_allclose(rpy_true, rpy_est, atol=5e-3) From 5ba42658d5fe585b7448a69bd840e2d76282743d Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Tue, 26 Nov 2024 07:54:15 +0100 Subject: [PATCH 2/4] [gym_jiminy/common] Add 'compute_rpy' optional argument to deformation estimator. --- .../common/blocks/deformation_estimator.py | 59 +++++++++++++++---- .../unit_py/test_deformation_estimator.py | 10 ++-- python/jiminy_py/unit_py/test_flexible_arm.py | 6 +- 3 files changed, 56 insertions(+), 19 deletions(-) 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 index 45ed04878..f82d2e7ef 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/blocks/deformation_estimator.py +++ b/python/gym_jiminy/common/gym_jiminy/common/blocks/deformation_estimator.py @@ -1,6 +1,7 @@ """Implementation of a stateless kinematic deformation estimator block compatible with gym_jiminy reinforcement learning pipeline environment design. """ +from collections import OrderedDict from collections.abc import Mapping from typing import List, Dict, Sequence, Tuple, Optional @@ -16,6 +17,7 @@ from ..bases import BaseAct, BaseObs, BaseObserverBlock, InterfaceJiminyEnv from ..utils import (DataNested, + quat_to_rpy, matrices_to_quat, quat_multiply, compute_tilt_from_quat, @@ -414,8 +416,8 @@ def get_flexibility_imu_frame_chains( return flex_imu_name_chains -class DeformationEstimator( - BaseObserverBlock[np.ndarray, np.ndarray, BaseObs, BaseAct]): +class DeformationEstimator(BaseObserverBlock[ + Dict[str, np.ndarray], np.ndarray, BaseObs, BaseAct]): """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. @@ -481,13 +483,16 @@ class DeformationEstimator( 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 theoretical configuration of the robot including its freeflyer, along with the name of the reference frame, ie the one having known orientation. + .. note:: + The feature space of this observer is the same as `MahonyFilter`. See + documentation for details. + .. seealso:: Matthieu Vigne, Antonio El Khoury, Marine Pétriaux, Florent Di Meglio, and Nicolas Petit "MOVIE: a Velocity-aided IMU Attitude Estimator for @@ -503,8 +508,9 @@ def __init__(self, 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", + nested_imu_key: Sequence[str] = ( + "features", "mahony_filter", "quat"), + compute_rpy: bool = True, update_ratio: int = 1) -> None: """ .. warning:: @@ -534,6 +540,10 @@ def __init__(self, to the IMU quaternion estimates. Their ordering must be consistent with the true IMU sensors of the robot. + :param compute_rpy: Whether to compute the Yaw-Pitch-Roll Euler angles + representation for the 3D orientation of the IMU, + in addition to the quaternion representation. + Optional: False by default. :param update_ratio: Ratio between the update period of the observer and the one of the subsequent observer. -1 to match the simulation timestep of the environment. @@ -549,6 +559,7 @@ def __init__(self, list, (imu_frame_names, flex_frame_names)) # Backup some of the user-argument(s) + self.compute_rpy = compute_rpy self.ignore_twist = ignore_twist # Create flexible dynamic model. @@ -687,6 +698,13 @@ def __init__(self, # Initialize the observer super().__init__(name, env, update_ratio) + # Define some proxies for fast access + self._quat = self.observation["quat"] + if self.compute_rpy: + self._rpy = self.observation["rpy"] + else: + self._rpy = np.array([]) + # Define chunks associated with each independent flexibility-imu chain self._deformation_flex_quats, self._obs_imu_indices = [], [] flex_start_index = 0 @@ -694,7 +712,7 @@ def __init__(self, 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]) + self._quat[:, 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( @@ -705,10 +723,18 @@ 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( + observation_space: Dict[str, gym.Space] = OrderedDict() + observation_space["quat"] = gym.spaces.Box( low=np.full((4, nflex), -1.0 - 1e-9), high=np.full((4, nflex), 1.0 + 1e-9), dtype=np.float64) + if self.compute_rpy: + high = np.array([np.pi, np.pi/2, np.pi]) + 1e-9 + observation_space["rpy"] = gym.spaces.Box( + low=-high[:, np.newaxis].repeat(nflex, axis=1), + high=high[:, np.newaxis].repeat(nflex, axis=1), + dtype=np.float64) + self.observation_space = gym.spaces.Dict(observation_space) def _setup(self) -> None: # Call base implementation @@ -725,7 +751,7 @@ def _setup(self) -> None: self.pinocchio_data_th = self.pinocchio_model_th.createData() # Fix initialization of the observation to be valid quaternions - self.observation[-1] = 1.0 + self._quat[-1] = 1.0 # Refresh flexibility and IMU frame orientation proxies self._kin_flex_rots.clear() @@ -762,9 +788,16 @@ def _setup(self) -> None: self._is_compiled = True @property - def fieldnames(self) -> List[List[str]]: - return [[f"{name}.Quat{e}" for name in self.flexibility_frame_names] - for e in ("x", "y", "z", "w")] + def fieldnames(self) -> Dict[str, List[List[str]]]: + fieldnames: Dict[str, List[List[str]]] = {} + fieldnames["quat"] = [ + [f"{name}.Quat{e}" for name in self.flexibility_frame_names] + for e in ("x", "y", "z", "w")] + if self.compute_rpy: + fieldnames["rpy"] = [ + [".".join((name, e)) for name in self.flexibility_frame_names] + for e in ("Roll", "Pitch", "Yaw")] + return fieldnames def refresh_observation(self, measurement: BaseObs) -> None: # Translate encoder data at joint level @@ -795,3 +828,7 @@ def refresh_observation(self, measurement: BaseObs) -> None: self._deformation_flex_quats): flexibility_estimator( self._obs_imu_quats, *args, self.ignore_twist) + + # Compute the RPY representation if requested + if self.compute_rpy: + quat_to_rpy(self._quat, self._rpy) diff --git a/python/gym_jiminy/unit_py/test_deformation_estimator.py b/python/gym_jiminy/unit_py/test_deformation_estimator.py index 627e6db63..649f27237 100644 --- a/python/gym_jiminy/unit_py/test_deformation_estimator.py +++ b/python/gym_jiminy/unit_py/test_deformation_estimator.py @@ -33,7 +33,8 @@ def _test_deformation_estimate(self, env, imu_atol, flex_atol): true_imu_rots.append(frame_rot) true_imu_quats = matrices_to_quat(tuple(true_imu_rots)) - est_imu_quats = env.observation['features']['mahony_filter'] + features = env.observation['features'] + est_imu_quats = features['mahony_filter']['quat'] np.testing.assert_allclose( true_imu_quats, est_imu_quats, atol=imu_atol) @@ -41,9 +42,8 @@ def _test_deformation_estimate(self, env, imu_atol, flex_atol): model_options = env.robot.get_model_options() flexibility_frame_names = [ flex_options["frameName"] - for flex_options in model_options["dynamics"]["flexibilityConfig"] - ] - est_flex_quats = env.observation['features']['deformation_estimator'] + for flex_options in model_options["dynamics"]["flexibilityConfig"]] + est_flex_quats = features['deformation_estimator']['quat'] est_flex_quats[:] *= np.sign(est_flex_quats[-1]) for frame_name, joint_index in zip( flexibility_frame_names, env.robot.flexibility_joint_indices): @@ -213,7 +213,7 @@ def _setup(self) -> None: model_options["dynamics"]["enableFlexibility"] = True self.robot.set_model_options(model_options) - # Create pipeline with Mahony filter and DeformationObserver blocks + # Create pipeline with Mahony filter and DeformationEstimator blocks PipelineEnv = build_pipeline( env_config=dict( cls=FlexAntJiminyEnv diff --git a/python/jiminy_py/unit_py/test_flexible_arm.py b/python/jiminy_py/unit_py/test_flexible_arm.py index f3c42c064..f7ac50f4b 100644 --- a/python/jiminy_py/unit_py/test_flexible_arm.py +++ b/python/jiminy_py/unit_py/test_flexible_arm.py @@ -1,6 +1,6 @@ -""" -@brief This file aims at verifying the sanity of the physics and the - integration method of jiminy on simple models. +# mypy: disable-error-code="no-untyped-def, var-annotated" +"""This file aims at verifying the sanity of the physics and the integration +method of jiminy on simple models. """ import os import io From e54394404d00c55b7fbc3714169ceed02c47323a Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Tue, 26 Nov 2024 07:56:26 +0100 Subject: [PATCH 3/4] [python|gym_jiminy] More natural figure layout to improve readability. --- .../common/gym_jiminy/common/envs/generic.py | 54 ++++--- .../gym_jiminy/common/quantities/generic.py | 2 +- .../common/gym_jiminy/common/utils/math.py | 14 +- python/jiminy_py/setup.py | 3 +- python/jiminy_py/src/jiminy_py/plot.py | 150 ++++++++++++------ 5 files changed, 145 insertions(+), 78 deletions(-) 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 794e882e9..fafa406c1 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py +++ b/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py @@ -13,7 +13,7 @@ from functools import partial from typing import ( Dict, Any, List, cast, no_type_check, Optional, Tuple, Callable, Union, - SupportsFloat, Iterator, Generic, Sequence as SequenceT, + SupportsFloat, Iterator, Generic, Sequence as SequenceT, Mapping as MappingT, MutableMapping as MutableMappingT) import numpy as np @@ -1047,32 +1047,44 @@ def plot(self, if not enable_block_states and key.endswith(".state"): continue - # Extract action hierarchical time series. + # Store fieldnames in dict systematically to avoid code duplication + if not isinstance(fieldnames, dict): + fieldnames = {"": fieldnames} + + # Extract hierarchical time series. # Fieldnames stored in a dictionary cannot be nested. In such a # case, keys corresponds to subplots, and values are individual # scalar data over time to be displayed to the same subplot. t = log_vars["Global.Time"] - tab_data: Dict[str, Union[np.ndarray, Dict[str, np.ndarray]]] = {} - if isinstance(fieldnames, dict): - for group, subfieldnames in fieldnames.items(): - if not isinstance(subfieldnames, list): - LOGGER.error( - "Action space not supported by this method.") - return figure - value_map = extract_variables_from_log( - log_vars, subfieldnames, "controller", as_dict=True) - tab_data[group] = { - key.split(".", 2)[2]: value - for key, value in value_map.items()} - elif isinstance(fieldnames, list): + base_name = key.replace(".", " ") + for group, subfieldnames in fieldnames.items(): + if not isinstance(subfieldnames, (list, tuple)): + LOGGER.error( + "Action space not supported by this method.") + return figure + + tab_name = " ".join(filter(None, (base_name, group))) value_map = extract_variables_from_log( - log_vars, fieldnames, "controller", as_dict=True) - tab_data.update({ - key.split(".", 2)[2]: value - for key, value in value_map.items()}) + log_vars, subfieldnames, "controller", as_dict=True) + tab_data = {key.split(".", 2)[2]: value + for key, value in value_map.items()} + + grid_spec: Tuple[Optional[int], Optional[int]] = (None, None) + nrows = len(subfieldnames) + if nrows and isinstance(subfieldnames[0], (list, tuple)): + ncols_all = set(map(len, subfieldnames)) + if len(ncols_all) == 1: + grid_spec = (nrows, next(iter(ncols_all))) - # Add action tab - figure.add_tab(key.replace(".", " "), t, tab_data) + try: + figure.add_tab(tab_name, + t, + tab_data, # type: ignore[arg-type] + nrows=grid_spec[0], + ncols=grid_spec[1]) + except ValueError: + LOGGER.error("Invalid plot spec for variable %s. Moving " + "to the next one", key) # Return figure for convenience and consistency with Matplotlib return figure diff --git a/python/gym_jiminy/common/gym_jiminy/common/quantities/generic.py b/python/gym_jiminy/common/gym_jiminy/common/quantities/generic.py index 9def585df..bd305869c 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/quantities/generic.py +++ b/python/gym_jiminy/common/gym_jiminy/common/quantities/generic.py @@ -1363,7 +1363,7 @@ def refresh(self) -> np.ndarray: @dataclass(unsafe_hash=True) class AverageFrameRollPitch(InterfaceQuantity[np.ndarray]): """Quaternion representation of the average Yaw-free orientation from the - Roll-Pitch_yaw decomposition of a given frame over the whole agent step. + Roll-Pitch-Yaw decomposition of a given frame over the whole agent step. .. seealso:: See `remove_yaw_from_quat` and `AverageFrameXYZQuat` for details about 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 c68855fac..721f4db58 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/utils/math.py +++ b/python/gym_jiminy/common/gym_jiminy/common/utils/math.py @@ -121,6 +121,9 @@ def quat_to_rpy(quat: np.ndarray, """Compute the Yaw-Pitch-Roll Euler angles representation of a single or a batch of quaternions. + The Roll, Pitch and Yaw angles are guaranteed to be within range [-pi,pi], + [-pi/2,pi/2], [-pi,pi], respectively. + :param quat: N-dimensional array whose first dimension gathers the 4 quaternion coordinates (qx, qy, qz, qw). :param out: Pre-allocated array in which to store the result. If not @@ -1111,7 +1114,9 @@ def remove_yaw_from_quat(quat: np.ndarray, Note that this decomposition is rarely used in practice, mainly because of singularity issues related to the Roll-Pitch-Yaw decomposition. It is usually preferable to remove the twist part of the Twist-after-Swing - decomposition. See `remove_twist_from_quat` documentation for details. + decomposition. Note that in both cases, the Roll and Pitch angles from + their corresponding Yaw-Pitch-Roll Euler angles representation matches + exactly. See `remove_twist_from_quat` documentation for details. :param quat: N-dimensional array whose first dimension gathers the 4 quaternion coordinates (qx, qy, qz, qw). @@ -1165,9 +1170,10 @@ def remove_twist_from_quat(quat: np.ndarray, R = R_z * R_s where R_z (the twist) is a rotation around e_z and R_s (the swing) is - the "smallest" rotation matrix such that t(R_s) = t(R). Note that although - the swing is not free of rotation around z-axis, the latter only depends on - the rotation around e_x, e_y, which is the main motivation for using this + the "smallest" rotation matrix (in terms of angle of its corresponding + Axis-Angle representation) such that s(R_s) = s(R). Note that although the + swing is not free of rotation around z-axis, the latter only depends on the + rotation around e_x, e_y, which is the main motivation for using this decomposition. One must use `remove_yaw_from_quat` to completely cancel you the rotation around z-axis. diff --git a/python/jiminy_py/setup.py b/python/jiminy_py/setup.py index 1527f0f96..b69220479 100644 --- a/python/jiminy_py/setup.py +++ b/python/jiminy_py/setup.py @@ -123,7 +123,8 @@ def finalize_options(self) -> None: extras_require={ "plot": [ # Standard library to generate figures. - # - 3.7.0: introduces 'outside' keyword for legend location + # - 3.7.0: introduces Button blitting and 'outside' keyword for + # legend location "matplotlib>=3.7.0" ], "meshcat": [ diff --git a/python/jiminy_py/src/jiminy_py/plot.py b/python/jiminy_py/src/jiminy_py/plot.py index 948837780..27511c271 100644 --- a/python/jiminy_py/src/jiminy_py/plot.py +++ b/python/jiminy_py/src/jiminy_py/plot.py @@ -31,7 +31,6 @@ # You can get a runtime error if Matplotlib is installed but cannot be # imported because of some conflicts with jupyter event loop for instance. raise ImportError("Matplotlib cannot be imported.") from e -from matplotlib import colors from matplotlib.axes import Axes from matplotlib.lines import Line2D from matplotlib.artist import Artist @@ -67,21 +66,47 @@ } -class _ButtonBlit(Button): - def _motion(self, event: Event) -> None: - if self.ignore(event): - return - c = self.hovercolor if event.inaxes == self.ax else self.color - if not colors.same_color(c, self.ax.get_facecolor()): - self.ax.set_facecolor(c) - if self.drawon: - # It is necessary to flush events beforehand to make sure - # figure refresh cannot get interrupted by button blitting. - # Otherwise the figure would be blank. - assert self.ax.figure is not None - self.ax.figure.canvas.flush_events() - self.ax.draw_artist(self.ax) - self.ax.figure.canvas.blit(self.ax.bbox) +def _get_grid_spec(width: float, + height: float, + nplots: int, + nrows: Optional[int] = None, + ncols: Optional[int] = None) -> Tuple[int, int]: + """Compute the optimal grid spec (i.e. number of rows and columns) for a + subplot grid, so that the width of each subplot is as close as possible to + their height. + + :param width: Width of the plotting area for the subplot grid. + :param height: Width of the plotting area for the subplot grid. + :param nplots: Expected number of subplots. + :param nrows: Desired number of rows for the subplot grid if any. + If not specified, it will be determined automatically, + so that the actual width of each subplot is as close as + possible to their height. + Optional: `None` by default. + :param ncols: Desired number of columns for the subplot grid if any. + If not specified, it will be determined automatically, + so that the actual width of each subplot is as close as + possible to their height. + Optional: `None` by default. + """ + if nrows is None and ncols is None: + ratio = width / height + nrows_1 = max(1, floor(sqrt(nplots / ratio))) + ncols_1 = ceil(nplots / nrows_1) + ncols_2 = ceil(sqrt(nplots * ratio)) + nrows_2 = ceil(nplots / ncols_2) + if nrows_1 * ncols_1 < nrows_2 * ncols_2: + nrows, ncols = map(int, (nrows_1, ncols_1)) + else: + nrows, ncols = map(int, (nrows_2, ncols_2)) + elif nrows is None and ncols is not None: + nrows = int(ceil(nplots / ncols)) + elif nrows is not None and ncols is None: + ncols = int(ceil(nplots / nrows)) + elif nrows is not None and ncols is not None: + assert nrows * ncols >= nplots + assert nrows is not None and ncols is not None + return nrows, ncols @dataclass @@ -93,6 +118,10 @@ class TabData: """Matpolib `Axes` handles of every subplots. """ + gridspec: Tuple[Optional[int], Optional[int]] + """The requested shape (nrows, ncols) of the subplot grid. + """ + legend_data: Tuple[List[Artist], List[str]] """Matpolib Legend data as returned by `get_legend_handles_labels` method. @@ -101,7 +130,7 @@ class TabData: them. """ - button: _ButtonBlit + button: Button """Button on which to click to switch to the tab at hand. """ @@ -193,7 +222,7 @@ def __init__(self, # pylint: disable=unused-argument # Customize figure subplot layout and reserve space for buttons # self.figure.get_layout_engine().set(w_pad=0.1, h_pad=0.1) self.subfigs = self.figure.subfigures( - 2, 1, wspace=0.1, height_ratios=[0.94, 0.06]) + 2, 1, wspace=0.1, height_ratios=[0.95, 0.06]) # Set window size if self.offscreen: @@ -241,18 +270,12 @@ def adjust_layout(self, # Re-arrange subplots in case figure aspect ratio has changed axes = self.tab_active.axes - num_subplots = len(axes) figure_extent = self.figure.get_window_extent() - figure_ratio = figure_extent.width / figure_extent.height - num_rows_1 = max(1, floor(sqrt(num_subplots / figure_ratio))) - num_cols_1 = ceil(num_subplots / num_rows_1) - num_cols_2 = ceil(sqrt(num_subplots * figure_ratio)) - num_rows_2 = ceil(num_subplots / num_cols_2) - if num_rows_1 * num_cols_1 < num_rows_2 * num_cols_2: - num_rows, num_cols = map(int, (num_rows_1, num_cols_1)) - else: - num_rows, num_cols = map(int, (num_rows_2, num_cols_2)) - grid_spec = self.subfigs[0].add_gridspec(num_rows, num_cols) + nrows, ncols = _get_grid_spec(figure_extent.width, + figure_extent.height, + len(axes), + *self.tab_active.gridspec) + grid_spec = self.subfigs[0].add_gridspec(nrows, ncols) for i, ax in enumerate(axes, 1): ax.set_subplotspec(grid_spec[i - 1]) @@ -340,7 +363,10 @@ def add_tab(self, # pylint: disable=unused-argument data: Union[np.ndarray, Dict[str, Union[ Dict[str, np.ndarray], np.ndarray]]], plot_method: Optional[ - Union[Callable[..., Any], str]] = None, *, + Union[Callable[..., Any], str]] = None, + *, + nrows: Optional[int] = None, + ncols: Optional[int] = None, refresh_canvas: bool = True, **kwargs: Any) -> None: """Create a new tab holding the provided data. @@ -372,6 +398,16 @@ def add_tab(self, # pylint: disable=unused-argument array in argument, or string instance method of `matplotlib.axes.Axes`. Optional: `step(..., where='post')` by default. + :param nrows: Desired number of rows for the subplot grid if any. + If not specified, it will be determined automatically, + so that the actual width of each subplot is as close as + possible to their height. + Optional: `None` by default. + :param ncols: Desired number of columns for the subplot grid if any. + If not specified, it will be determined automatically, + so that the actual width of each subplot is as close as + possible to their height. + Optional: `None` by default. :param refresh_canvas: Whether to refresh the figure. This step can be skipped if other tabs are going to be added or deleted soon, to avoid useless computation and @@ -395,12 +431,20 @@ def add_tab(self, # pylint: disable=unused-argument assert callable(plot_method) if isinstance(data, dict): - # Compute plot grid arrangement - n_cols = len(data) - n_rows = 1 - while n_cols > n_rows + 2: - n_rows = n_rows + 1 - n_cols = int(np.ceil(len(data) / n_rows)) + # Make sure that data are valid + if data and len(set( + tuple(sorted(e.keys())) + for e in data.values() if isinstance(e, dict))) > 1: + raise ValueError( + "Line names must be the same for all subplots.") + + # Compute (temporary) plot grid arrangement. + # It will be adjusted automatically at the end. + ncols = nplots = len(data) + nrows = 1 + while ncols > nrows + 2: + nrows = nrows + 1 + ncols = int(np.ceil(nplots / nrows)) # Initialize axes, and early return if none axes: List[plt.Axes] = [] @@ -408,7 +452,7 @@ def add_tab(self, # pylint: disable=unused-argument for i, plot_name in enumerate(data.keys()): uniq_label = '_'.join((tab_name, plot_name)) ax = self.subfigs[0].add_subplot( - n_rows, n_cols, i+1, label=uniq_label) + nrows, ncols, i+1, label=uniq_label) ax.autoscale(True, axis='x', tight=True) ax.autoscale(True, axis='y', tight=False) ax.ticklabel_format(axis='x', style='plain', useOffset=True) @@ -453,16 +497,22 @@ def add_tab(self, # pylint: disable=unused-argument uniq_label = '_'.join((tab_name, "button")) button_axcut = self.subfigs[1].add_axes( [0.0, 0.0, 0.0, 0.0], label=uniq_label) - button = _ButtonBlit(button_axcut, - tab_name.replace(' ', '\n'), - color='white') + button = Button(button_axcut, + tab_name.replace(' ', '\n'), + color='white') + button.label.set_fontsize(9) # Register buttons events button.on_clicked(self.__click) # Create new tab data container - self.tabs_data[tab_name] = TabData( - axes, legend_data, button, button_axcut, nav_stack=[], nav_pos=-1) + self.tabs_data[tab_name] = TabData(axes, + (nrows, ncols), + legend_data, + button, + button_axcut, + nav_stack=[], + nav_pos=-1) # Check if it is the first tab to be added if self.tab_active is None: @@ -842,8 +892,8 @@ def plot_log_interactive() -> None: [header[i] for header in matching_fieldnames]) # Create figure - n_plot = len(plotted_elements) - if not n_plot: + nplots = len(plotted_elements) + if not nplots: print("Nothing to plot. Exiting...") return fig = plt.figure(layout="constrained") @@ -856,13 +906,11 @@ def plot_log_interactive() -> None: fig.set_size_inches(14, 8) # Create subplots, arranging them in a rectangular fashion. - # Do not allow for n_cols to be more than n_rows + 2. - n_cols = n_plot - n_rows = 1 - while n_cols > n_rows + 2: - n_rows = int(n_rows + 1) - n_cols = int(np.ceil(n_plot / (1.0 * n_rows))) - axes = fig.subplots(n_rows, n_cols, sharex=True, squeeze=False).flat[:] + # Do not allow for ncols to be more than nrows + 2. + figure_extent = fig.get_window_extent() + nrows, ncols = _get_grid_spec( + figure_extent.width, figure_extent.height, nplots) + axes = fig.subplots(nrows, ncols, sharex=True, squeeze=False).flat[:] # Store lines in dictionary {file_name: plotted lines}, to enable to # toggle individually the visibility the data related to each of them. From c7ed1982a5524a7b9ba01a733fe21c3a8e076a14 Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Tue, 26 Nov 2024 08:11:27 +0100 Subject: [PATCH 4/4] [misc] Fix acrobot example eval max horizon. --- .github/workflows/macos.yml | 2 +- python/gym_jiminy/examples/rllib/acrobot_ppo.py | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/.github/workflows/macos.yml b/.github/workflows/macos.yml index 4b01b6159..b43ad8a09 100644 --- a/.github/workflows/macos.yml +++ b/.github/workflows/macos.yml @@ -10,7 +10,7 @@ on: jobs: build-and-test-osx: name: >- - (${{ matrix.PYTHON_VERSION }}) (${{ matrix.BUILD_TYPE }}) + (${{ matrix.OS }}) (${{ matrix.PYTHON_VERSION }}) (${{ matrix.BUILD_TYPE }}) Build and run the unit tests. Then generate and publish the wheels on PyPi. strategy: diff --git a/python/gym_jiminy/examples/rllib/acrobot_ppo.py b/python/gym_jiminy/examples/rllib/acrobot_ppo.py index 01b6f1274..34da9ec84 100644 --- a/python/gym_jiminy/examples/rllib/acrobot_ppo.py +++ b/python/gym_jiminy/examples/rllib/acrobot_ppo.py @@ -285,4 +285,7 @@ rl_module = build_module_from_checkpoint(checkpoint_path) policy_fn = build_module_wrapper(rl_module) for seed in (1, 1, 2): - env.evaluate(policy_fn, seed=seed) # type: ignore[attr-defined] + env.evaluate( # type: ignore[attr-defined] + policy_fn, + seed=seed, + horizon=env.spec.max_episode_steps) # type: ignore[union-attr]