[gym/common] Add 'compute_tilt', 'swing_from_vector', and 'remove_twi…

…st_from_quat' utilities.

python/gym_jiminy/common/gym_jiminy/common/blocks/mahony_filter.py

@@ -2,7 +2,7 @@
 reinforcement learning pipeline environment design.
 """
 import logging
-from typing import List, Union, Optional, Tuple, no_type_check
+from typing import List, Union, Optional
 
 import numpy as np
 import numba as nb
@@ -12,32 +12,18 @@
     ImuSensor as imu)
 
 from ..bases import BaseObsT, BaseActT, BaseObserverBlock, InterfaceJiminyEnv
-from ..utils import ArrayOrScalar, fill, matrices_to_quat
+from ..utils import (fill,
+                     matrices_to_quat,
+                     swing_from_vector,
+                     compute_tilt,
+                     remove_twist_from_quat)
 
 
 EARTH_SURFACE_GRAVITY = 9.81
-TWIST_SWING_SINGULARITY_THR = 1e-6
 
 LOGGER = logging.getLogger(__name__)
 
 
-@nb.jit(nopython=True, cache=True)
-def compute_tilt(q: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-    """Compute e_z in R(q) frame (Euler-Rodrigues Formula): R(q).T @ e_z.
-
-    :param q: Array whose rows are the 4 components of quaternions (x, y, z, w)
-              and columns are N independent orientations.
-
-    :returns: Tuple of arrays corresponding to the 3 individual components
-              (a_x, a_y, a_z) of N independent tilt axes.
-    """
-    q_x, q_y, q_z, q_w = q
-    v_x = 2 * (q_x * q_z - q_y * q_w)
-    v_y = 2 * (q_y * q_z + q_w * q_x)
-    v_z = 1 - 2 * (q_x * q_x + q_y * q_y)
-    return (v_x, v_y, v_z)
-
-
 @nb.jit(nopython=True, nogil=True, cache=True)
 def mahony_filter(q: np.ndarray,
                   omega: np.ndarray,
@@ -108,83 +94,6 @@ def mahony_filter(q: np.ndarray,
     bias_hat -= dt * ki * omega_mes
 
 
-# FIXME: Enabling cache causes segfault on Apple Silicon
-@no_type_check
-@nb.jit(nopython=True, cache=False)
-def swing_from_vector(
-        v_a: Tuple[ArrayOrScalar, ArrayOrScalar, ArrayOrScalar],
-        q: np.ndarray) -> None:
-    """Compute the "smallest" rotation transforming vector 'v_a' in 'e_z'.
-
-    :param v_a: Tuple of arrays corresponding to the 3 individual components
-                (a_x, a_y, a_z) of N independent tilt axes.
-    :param q: Array where the result will be stored. The rows are the 4
-              components of quaternions (x, y, z, w) and columns are the N
-              independent orientations.
-    """
-    # Extract individual tilt components
-    v_x, v_y, v_z = v_a
-
-    # There is a singularity when the rotation axis of orientation estimate
-    # and z-axis are nearly opposites, i.e. v_z ~= -1. One solution that
-    # ensure continuity of q_w is picked arbitrarily using SVD decomposition.
-    # See `Eigen::Quaternion::FromTwoVectors` implementation for details.
-    if q.ndim > 1:
-        is_singular = np.any(v_z < -1.0 + TWIST_SWING_SINGULARITY_THR)
-    else:
-        is_singular = v_z < -1.0 + TWIST_SWING_SINGULARITY_THR
-    if is_singular:
-        if q.ndim > 1:
-            for i, q_i in enumerate(q.T):
-                swing_from_vector((v_x[i], v_y[i], v_z[i]), q_i)
-        else:
-            _, _, v_h = np.linalg.svd(np.array((
-                (v_x, v_y, v_z),
-                (0.0, 0.0, 1.0))
-            ), full_matrices=True)
-            w_2 = (1 + max(v_z, -1)) / 2
-            q[:3], q[3] = v_h[-1] * np.sqrt(1 - w_2), np.sqrt(w_2)
-    else:
-        s = np.sqrt(2 * (1 + v_z))
-        q[0], q[1], q[2], q[3] = v_y / s, - v_x / s, 0.0, s / 2
-
-    # First order quaternion normalization to prevent compounding of errors.
-    # If not done, shit may happen with removing twist again and again on the
-    # same quaternion, which is typically the case when the IMU is steady, so
-    # that the mahony filter update is actually skipped internally.
-    q *= (3.0 - np.sum(np.square(q), 0)) / 2
-
-
-# FIXME: Enabling cache causes segfault on Apple Silicon
-@nb.jit(nopython=True, cache=False)
-def remove_twist(q: np.ndarray) -> None:
-    """Remove the twist part of the Twist-after-Swing decomposition of given
-    orientations in quaternion representation.
-
-    Any rotation R can be decomposed as:
-
-        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).
-
-    .. seealso::
-        * See "Estimation and control of the deformations of an exoskeleton
-          using inertial sensors", PhD Thesis, M. Vigne, 2021, p. 130.
-        * See "Swing-twist decomposition in Clifford algebra", P. Dobrowolski,
-          2015 (https://arxiv.org/abs/1506.05481)
-
-    :param q: Array whose rows are the 4 components of quaternions (x, y, z, w)
-              and columns are N independent orientations from which to remove
-              the swing part. It will be updated in-place.
-    """
-    # Compute e_z in R(q) frame (Euler-Rodrigues Formula): R(q).T @ e_z
-    v_a = compute_tilt(q)
-
-    # Compute the "smallest" rotation transforming vector 'v_a' in 'e_z'
-    swing_from_vector(v_a, q)
-
-
 @nb.jit(nopython=True, nogil=True, cache=True)
 def update_twist(q: np.ndarray,
                  twist: np.ndarray,
@@ -269,8 +178,8 @@ def __init__(self,
             integrator used to estimate the twist part of twist-after-swing
             decomposition of the estimated orientation in place of the Mahony
             Filter. If `0.0`, then its is kept constant equal to zero. `None`
-            to kept the original estimate provided by Mahony Filter.
-            See `remove_twist` and `update_twist` documentation for details.
+            to kept the original estimate provided by Mahony Filter. See
+            `remove_twist_from_quat` and `update_twist` doc for details.
             Optional: `0.0` by default.
         :param exact_init: Whether to initialize orientation estimate using
                            accelerometer measurements or ground truth. `False`
@@ -470,7 +379,7 @@ def refresh_observation(self, measurement: BaseObsT) -> None:
 
         # Remove twist if requested
         if self._remove_twist:
-            remove_twist(self.observation)
+            remove_twist_from_quat(self.observation)
 
         if self._update_twist:
             update_twist(self.observation,

python/gym_jiminy/common/gym_jiminy/common/utils/__init__.py

@@ -12,7 +12,10 @@
                    quat_multiply,
                    quat_average,
                    rpy_to_matrix,
-                   rpy_to_quat)
+                   rpy_to_quat,
+                   compute_tilt,
+                   swing_from_vector,
+                   remove_twist_from_quat)
 from .spaces import (DataNested,
                      FieldNested,
                      ArrayOrScalar,
@@ -39,19 +42,22 @@
 
 
 __all__ = [
-    "squared_norm_2",
-    "matrix_to_quat",
-    "matrices_to_quat",
-    "matrix_to_rpy",
-    "matrix_to_yaw",
-    "quat_to_matrix",
-    "quat_to_rpy",
-    "quat_to_yaw",
-    "quat_to_yaw_cos_sin",
-    "quat_multiply",
-    "quat_average",
-    "rpy_to_matrix",
-    "rpy_to_quat",
+    'squared_norm_2',
+    'matrix_to_quat',
+    'matrices_to_quat',
+    'matrix_to_rpy',
+    'matrix_to_yaw',
+    'quat_to_matrix',
+    'quat_to_rpy',
+    'quat_to_yaw',
+    'quat_to_yaw_cos_sin',
+    'quat_multiply',
+    'quat_average',
+    'rpy_to_matrix',
+    'rpy_to_quat',
+    'compute_tilt',
+    'swing_from_vector',
+    'remove_twist_from_quat',
     'DataNested',
     'FieldNested',
     'ArrayOrScalar',

python/gym_jiminy/common/gym_jiminy/common/utils/math.py

@@ -1,10 +1,15 @@
 """ TODO: Write documentation.
 """
-from typing import Union, Tuple, Optional
+from typing import Union, Tuple, Optional, no_type_check
 
 import numpy as np
 import numba as nb
 
+from .spaces import ArrayOrScalar
+
+
+TWIST_SWING_SINGULARITY_THR = 1e-6
+
 
 @nb.jit(nopython=True, cache=True)
 def squared_norm_2(array: np.ndarray) -> float:
@@ -332,6 +337,100 @@ def quat_multiply(quat_left: np.ndarray,
     return out_
 
 
+@nb.jit(nopython=True, cache=True)
+def compute_tilt(q: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Compute e_z in R(q) frame (Euler-Rodrigues Formula): R(q).T @ e_z.
+
+    :param q: Array whose rows are the 4 components of quaternions (x, y, z, w)
+              and columns are N independent orientations.
+
+    :returns: Tuple of arrays corresponding to the 3 individual components
+              (a_x, a_y, a_z) of N independent tilt axes.
+    """
+    q_x, q_y, q_z, q_w = q
+    v_x = 2 * (q_x * q_z - q_y * q_w)
+    v_y = 2 * (q_y * q_z + q_w * q_x)
+    v_z = 1 - 2 * (q_x * q_x + q_y * q_y)
+    return (v_x, v_y, v_z)
+
+
+# FIXME: Enabling cache causes segfault on Apple Silicon
+@no_type_check
+@nb.jit(nopython=True, cache=False)
+def swing_from_vector(
+        v_a: Tuple[ArrayOrScalar, ArrayOrScalar, ArrayOrScalar],
+        q: np.ndarray) -> None:
+    """Compute the "smallest" rotation transforming vector 'v_a' in 'e_z'.
+
+    :param v_a: Tuple of arrays corresponding to the 3 individual components
+                (a_x, a_y, a_z) of N independent tilt axes.
+    :param q: Array where the result will be stored. The rows are the 4
+              components of quaternions (x, y, z, w) and columns are the N
+              independent orientations.
+    """
+    # Extract individual tilt components
+    v_x, v_y, v_z = v_a
+
+    # There is a singularity when the rotation axis of orientation estimate
+    # and z-axis are nearly opposites, i.e. v_z ~= -1. One solution that
+    # ensure continuity of q_w is picked arbitrarily using SVD decomposition.
+    # See `Eigen::Quaternion::FromTwoVectors` implementation for details.
+    if q.ndim > 1:
+        is_singular = np.any(v_z < -1.0 + TWIST_SWING_SINGULARITY_THR)
+    else:
+        is_singular = v_z < -1.0 + TWIST_SWING_SINGULARITY_THR
+    if is_singular:
+        if q.ndim > 1:
+            for i, q_i in enumerate(q.T):
+                swing_from_vector((v_x[i], v_y[i], v_z[i]), q_i)
+        else:
+            _, _, v_h = np.linalg.svd(np.array((
+                (v_x, v_y, v_z),
+                (0.0, 0.0, 1.0))
+            ), full_matrices=True)
+            w_2 = (1 + max(v_z, -1)) / 2
+            q[:3], q[3] = v_h[-1] * np.sqrt(1 - w_2), np.sqrt(w_2)
+    else:
+        s = np.sqrt(2 * (1 + v_z))
+        q[0], q[1], q[2], q[3] = v_y / s, - v_x / s, 0.0, s / 2
+
+    # First order quaternion normalization to prevent compounding of errors.
+    # If not done, shit may happen with removing twist again and again on the
+    # same quaternion, which is typically the case when the IMU is steady, so
+    # that the mahony filter update is actually skipped internally.
+    q *= (3.0 - np.sum(np.square(q), 0)) / 2
+
+
+# FIXME: Enabling cache causes segfault on Apple Silicon
+@nb.jit(nopython=True, cache=False)
+def remove_twist_from_quat(q: np.ndarray) -> None:
+    """Remove the twist part of the Twist-after-Swing decomposition of given
+    orientations in quaternion representation.
+
+    Any rotation R can be decomposed as:
+
+        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).
+
+    .. seealso::
+        * See "Estimation and control of the deformations of an exoskeleton
+          using inertial sensors", PhD Thesis, M. Vigne, 2021, p. 130.
+        * See "Swing-twist decomposition in Clifford algebra", P. Dobrowolski,
+          2015 (https://arxiv.org/abs/1506.05481)
+
+    :param q: Array whose rows are the 4 components of quaternions (x, y, z, w)
+              and columns are N independent orientations from which to remove
+              the swing part. It will be updated in-place.
+    """
+    # Compute e_z in R(q) frame (Euler-Rodrigues Formula): R(q).T @ e_z
+    v_a = compute_tilt(q)
+
+    # Compute the "smallest" rotation transforming vector 'v_a' in 'e_z'
+    swing_from_vector(v_a, q)
+
+
 def quat_average(quat: np.ndarray,
                  axis: Optional[Union[Tuple[int, ...], int]] = None
                  ) -> np.ndarray:

python/gym_jiminy/toolbox/gym_jiminy/toolbox/math/__init__.py

@@ -1,12 +1,10 @@
 # pylint: disable=missing-module-docstring
 
 from .qhull import ConvexHull, compute_distance_convex_to_point
-from .generic import quat_average
 
 __all__ = [
     "ConvexHull",
-    "compute_distance_convex_to_point",
-    "quat_average"
+    "compute_distance_convex_to_point"
 ]
 
 try:

python/gym_jiminy/unit_py/test_pipeline_control.py

@@ -20,8 +20,8 @@
 from gym_jiminy.envs import (
     AtlasPDControlJiminyEnv, CassiePDControlJiminyEnv, DigitPDControlJiminyEnv)
 from gym_jiminy.common.blocks import PDController, MahonyFilter
-from gym_jiminy.common.blocks.mahony_filter import remove_twist
-from gym_jiminy.common.utils import quat_to_rpy, matrix_to_rpy, matrix_to_quat
+from gym_jiminy.common.utils import (
+    quat_to_rpy, matrix_to_rpy, matrix_to_quat, remove_twist_from_quat)
 
 
 IMAGE_DIFF_THRESHOLD = 5.0
@@ -161,10 +161,10 @@ def test_mahony_filter(self):
                 if twist_time_constant == 0.0:
                     # The twist must be ignored as it is not observable
                     obs_true = matrix_to_quat(imu_rot)
-                    remove_twist(obs_true)
+                    remove_twist_from_quat(obs_true)
                     rpy_true = quat_to_rpy(obs_true)
                     obs_est = env.observer.observation[:, 0].copy()
-                    remove_twist(obs_est)
+                    remove_twist_from_quat(obs_est)
                     rpy_est = quat_to_rpy(obs_est)
                 else:
                     # The twist is either measured or estimated