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 f82d2e7ef..abb4babf6 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 @@ -3,7 +3,7 @@ """ from collections import OrderedDict from collections.abc import Mapping -from typing import List, Dict, Sequence, Tuple, Optional +from typing import List, Dict, Sequence, Tuple, Union, Optional import numpy as np import numba as nb @@ -16,6 +16,7 @@ import pinocchio as pin from ..bases import BaseAct, BaseObs, BaseObserverBlock, InterfaceJiminyEnv +from ..wrappers.observation_layout import NestedData from ..utils import (DataNested, quat_to_rpy, matrices_to_quat, @@ -508,7 +509,7 @@ def __init__(self, imu_frame_names: Sequence[str], flex_frame_names: Sequence[str], ignore_twist: bool = True, - nested_imu_key: Sequence[str] = ( + nested_imu_key: NestedData = ( "features", "mahony_filter", "quat"), compute_rpy: bool = True, update_ratio: int = 1) -> None: @@ -649,11 +650,26 @@ def __init__(self, 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 + # Define IMU observation 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] + if isinstance(key, str): + assert isinstance(obs_imu_quats, Mapping) + obs_imu_quats = obs_imu_quats[key] + elif isinstance(key, int): + assert isinstance(obs_imu_quats, Sequence) + obs_imu_quats = obs_imu_quats[key] + else: + assert isinstance(obs_imu_quats, np.ndarray) + slices: List[Union[int, slice]] = [] + for start_end in key: + if isinstance(start_end, int): + slices.append(start_end) + elif not start_end: + slices.append(slice(None,)) + else: + slices.append(slice(*start_end)) + obs_imu_quats = obs_imu_quats[tuple(slices)] assert isinstance(obs_imu_quats, np.ndarray) self._obs_imu_quats = obs_imu_quats 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 fde3a55ef..6f9f77d96 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 @@ -162,15 +162,16 @@ class MahonyFilter(BaseObserverBlock[ .. 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`. + estimates under key 'quat', optionally, their corresponding + Yaw-Pitch-Roll Euler angles representations under key 'rpy' if + `compute_rpy=True`, and finally, the angular velocity in local frame + estimates under key 'omega'. 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. Specifically, `N=4` for quaternions + (Qx, Qy, Qz, Qw), 'N=3' for both the Euler angles (Roll-Pitch-Yaw) and + the angular velocity. 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 @@ -186,7 +187,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, + compute_rpy: bool = False, update_ratio: int = 1) -> None: """ :param name: Name of the block. @@ -205,14 +206,14 @@ def __init__(self, is not recommended because the robot is often free-falling at init, which is not realistic anyway. Optional: `True` by default. - :param mahony_kp: Proportional gain used for gyro-accel sensor fusion. - Set it to 0.0 to use only the gyroscope. In such a - case, the orientation estimate would be exact if the - sensor is bias- and noise-free, and the update period - matches the simulation integrator timestep. - Optional: `1.0` by default. - :param mahony_ki: Integral gain used for gyro bias estimate. - Optional: `0.1` by default. + :param kp: Proportional gain used for gyro-accel sensor fusion. Set it + to 0.0 to use only the gyroscope. In such a case, the + orientation estimate would be exact if the sensor is bias- + and noise-free, and the update period matches the + simulation integrator timestep. + Optional: `1.0` by default. + :param 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. @@ -279,9 +280,6 @@ def __init__(self, if self._update_twist: self._state["twist"] = self._twist - # Store the estimate angular velocity to avoid redundant computations - self._omega = np.zeros((3, num_imu_sensors)) - # Initialize the observer super().__init__(name, env, update_ratio) @@ -291,6 +289,7 @@ def __init__(self, self._rpy = self.observation["rpy"] else: self._rpy = np.array([]) + self._omega = self.observation["omega"] def _initialize_state_space(self) -> None: """Configure the internal state space of the observer. @@ -340,6 +339,11 @@ def _initialize_observation_space(self) -> None: low=-high[:, np.newaxis].repeat(num_imu_sensors, axis=1), high=high[:, np.newaxis].repeat(num_imu_sensors, axis=1), dtype=np.float64) + observation_space["omega"] = gym.spaces.Box( + low=float('-inf'), + high=float('inf'), + shape=(3, num_imu_sensors), + dtype=np.float64) self.observation_space = gym.spaces.Dict(observation_space) def _setup(self) -> None: @@ -390,8 +394,11 @@ def fieldnames(self) -> Dict[str, List[List[str]]]: imu_sensors = self.env.robot.sensors[ImuSensor.type] 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")] + [".".join((sensor.name, f"Quat{e}")) for sensor in imu_sensors] + for e in ("X", "Y", "Z", "W")] + fieldnames["omega"] = [ + [".".join((sensor.name, e)) for sensor in imu_sensors] + for e in ("X", "Y", "Z")] if self.compute_rpy: fieldnames["rpy"] = [ [".".join((sensor.name, e)) for sensor in imu_sensors] @@ -425,7 +432,7 @@ def refresh_observation(self, measurement: BaseObs) -> None: rot = robot.pinocchio_data.oMf[sensor.frame_index].rotation imu_rots.append(rot) - # Convert it to quaternions + # Convert the rotation matrices of the IMUs to quaternions matrices_to_quat(tuple(imu_rots), self._quat) # Keep track of tilt if necessary @@ -433,6 +440,11 @@ def refresh_observation(self, measurement: BaseObs) -> None: self._twist[:] = np.arctan2(self._quat[2], self._quat[3]) self._is_initialized = True + + # Compute the RPY representation if requested + if self.compute_rpy: + quat_to_rpy(self._quat, self._rpy) + return # Run an iteration of the filter, computing the next state estimate