[gym/common] Add average odometry velocity quantity.

python/gym_jiminy/common/gym_jiminy/common/bases/quantity.py

@@ -224,7 +224,7 @@ def __init__(self,
                  env: InterfaceJiminyEnv,
                  parent: Optional["AbstractQuantity"],
                  requirements: Dict[str, "QuantityCreator"],
-                 auto_refresh: bool) -> None:
+                 auto_refresh: bool = False) -> None:
         """
         :param env: Base or wrapped jiminy environment.
         :param parent: Higher-level quantity from which this quantity is a

python/gym_jiminy/common/gym_jiminy/common/bases/reward.py

@@ -128,6 +128,9 @@ def __init__(self,
         # Add quantity to the set of quantities managed by the environment
         self.env.quantities[self.name] = quantity
 
+        # Keep track of the underlying quantity
+        self.quantity = self.env.quantities.registry[self.name]
+
     @property
     def name(self) -> str:
         """Name uniquely identifying every reward. It will be used to add the

python/gym_jiminy/common/gym_jiminy/common/quantities/__init__.py

@@ -6,7 +6,9 @@
                       StackedQuantity,
                       AverageFrameSpatialVelocity,
                       MaskedQuantity)
-from .locomotion import CenterOfMass, ZeroMomentPoint
+from .locomotion import (AverageOdometryVelocity,
+                         CenterOfMass,
+                         ZeroMomentPoint)
 
 
 __all__ = [
@@ -16,6 +18,7 @@
     'StackedQuantity',
     'AverageFrameSpatialVelocity',
     'MaskedQuantity',
+    'AverageOdometryVelocity',
     'CenterOfMass',
     'ZeroMomentPoint',
 ]

python/gym_jiminy/common/gym_jiminy/common/quantities/generic.py

@@ -7,8 +7,9 @@
 from functools import partial
 from dataclasses import dataclass
 from typing import (
-    List, Dict, Set, Optional, Protocol, Sequence, Tuple, TypeVar, Union,
+    List, Dict, Set, Any, Optional, Protocol, Sequence, Tuple, TypeVar, Union,
     runtime_checkable)
+from typing_extensions import TypeAlias
 
 import numpy as np
 
@@ -17,9 +18,13 @@
 import pinocchio as pin
 
 from ..bases import InterfaceJiminyEnv, AbstractQuantity, QuantityCreator
-from ..utils import fill, matrix_to_rpy, matrix_to_quat
+from ..utils import (
+    fill, matrix_to_rpy, matrix_to_quat, quat_to_matrix,
+    quat_interpolate_middle)
 
 
+EllipsisType: TypeAlias = Any  # TODO: `EllipsisType` introduced in Python 3.10
+
 ValueT = TypeVar('ValueT')
 
 
@@ -515,12 +520,19 @@ class AverageFrameSpatialVelocity(AbstractQuantity[np.ndarray]):
 
     The average spatial velocity is obtained by finite difference. More
     precisely, it is defined here as the ratio of the geodesic distance in SE3
-    Lie  group between the pose of the frame at the end of previous and current
+    Lie Group between the pose of the frame at the end of previous and current
     step over the time difference between them. Notably, under this definition,
     the linear average velocity jointly depends on rate of change of the
     translation and rotation of the frame, which may be undesirable in some
     cases. Alternatively, the double geodesic distance could be used instead to
     completely decouple the translation from the rotation.
+
+    .. note::
+        The local frame for which the velocity is expressed is defined as the
+        midpoint interpolation between the previous and current frame pose.
+        This definition is arbitrary, in a sense that any other point for an
+        interpolation ratio going from 0.0 (previous pose) to 1.0 (current
+        pose) would be equally valid.
     """
 
     frame_name: str
@@ -574,8 +586,9 @@ def __init__(self,
         # Inverse step size
         self._inv_step_dt = 0.0
 
-        # Define proxy to the current frame pose (translation, rotation matrix)
-        self._rot_mat = np.eye(3)
+        # Allocate memory for the average quaternion and rotation matrix
+        self._quat_mean = np.zeros(4)
+        self._rot_mat_mean = np.eye(3)
 
         # Pre-allocate memory for the spatial velocity
         self._v_spatial: np.ndarray = np.zeros(6)
@@ -590,11 +603,6 @@ def initialize(self) -> None:
         # Compute inverse step size
         self._inv_step_dt = 1.0 / self.env.step_dt
 
-        # Refresh proxy to current frame pose
-        frame_index = self.pinocchio_model.getFrameId(self.frame_name)
-        transform = self.pinocchio_data.oMf[frame_index]
-        self._rot_mat = transform.rotation
-
         # Re-initialize pre-allocated buffers
         fill(self._v_spatial, 0)
 
@@ -612,8 +620,17 @@ def refresh(self) -> np.ndarray:
 
         # Translate local velocity to world frame
         if self.reference_frame == pin.LOCAL_WORLD_ALIGNED:
+            # Define world frame as the "middle" between prev and next pose.
+            # The orientation difference has an effect on the translation
+            # difference, but not the other way around. Here, we only care
+            # about the middle rotation, so we can consider SO3 Lie Group
+            # algebra instead of SE3.
+            quat_interpolate_middle(
+                xyzquat_prev[-4:], xyzquat[-4:], self._quat_mean)
+            quat_to_matrix(self._quat_mean, self._rot_mat_mean)
+
             # TODO: x2 speedup can be expected using `np.dot` with  `nb.jit`
-            self._v_lin_ang[:] = self._rot_mat @ self._v_lin_ang
+            self._v_lin_ang[:] = self._rot_mat_mean @ self._v_lin_ang
 
         return self._v_spatial
 
@@ -624,14 +641,16 @@ class MaskedQuantity(AbstractQuantity[np.ndarray]):
     array along an axis.
 
     Elements will be extract by copy unless the indices of the elements to
-    extract to be written equivalently by a slice, ie they are evenly spaced.
+    extract to be written equivalently by a slice (ie they are evenly spaced),
+    and the array can be flattened while preserving memory contiguity if 'axis'
+    is `None`.
     """
 
     quantity: AbstractQuantity
     """Base quantity whose elements must be extracted.
     """
 
-    indices: Tuple[int]
+    indices: Tuple[int, ...]
     """Indices of the elements to extract.
     """
 
@@ -643,7 +662,7 @@ def __init__(self,
                  env: InterfaceJiminyEnv,
                  parent: Optional[AbstractQuantity],
                  quantity: QuantityCreator[np.ndarray],
-                 key: Union[Sequence[int], Sequence[bool], slice],
+                 key: Union[Sequence[int], Sequence[bool]],
                  axis: Optional[int] = None
                  ) -> None:
         """
@@ -653,17 +672,15 @@ def __init__(self,
         :param quantity: Tuple gathering the class of the quantity whose values
                          must be extracted, plus all its constructor keyword-
                          arguments except environment 'env' and parent 'parent.
-        :param key: Sequence of indices, boolean mask, or slice that will be
-                    used to extract elements from the quantity along one axis.
+        :param key: Sequence of indices or boolean mask that will be used to
+                    extract elements from the quantity along one axis.
         :param axis: Axis over which to extract elements. `None` to consider
                      flattened array.
                      Optional: `None` by default.
         """
-        # Check if a slice, indices or a mask has been provided
-        if key is slice:
-            pass
-        elif all(isinstance(e, bool) for e in key):
-            key, _ = np.nonzero(key)
+        # Check if indices or boolean mask has been provided
+        if all(isinstance(e, bool) for e in key):
+            key = tuple(np.flatnonzero(key))
         elif not all(isinstance(e, int) for e in key):
             raise ValueError(
                 "Argument 'key' invalid. It must either be a "
@@ -679,19 +696,23 @@ def __init__(self,
                 "No indices to extract from quantity. Data would be empty.")
 
         # Check if the indices are evenly spaced
-        self._slices: Optional[slice] = None
+        self._slices: Tuple[Union[slice, EllipsisType], ...] = ()
+        stride: Optional[int] = None
         if len(self.indices) == 1:
             stride = 1
-        if len(self.indices) > 1:
+        if len(self.indices) > 1 and all(e >= 0 for e in self.indices):
             spacing = np.unique(np.diff(self.indices))
-            stride = stride[0] if spacing.size == 1 else None
+            if spacing.size == 1:
+                stride = spacing[0]
         if stride is not None:
             slice_ = slice(self.indices[0], self.indices[-1] + 1, stride)
-            if axis > 0:
-                self._slices = (slice(None),) * axis + (slice_,)
+            if axis is None:
+                self._slices = (slice_,)
+            elif axis > 0:
+                self._slices = (*((slice(None),) * axis), slice_)
             else:
                 self._slices = (
-                    Ellipsis, slice_) + (slice(None),) * (- axis - 1)
+                    Ellipsis, slice_, *((slice(None),) * (- axis - 1)))
 
         # Call base implementation
         super(). __init__(env,
@@ -708,9 +729,13 @@ def initialize(self) -> None:
 
     def refresh(self) -> np.ndarray:
         # Extract elements from quantity
-        if self._slices is None:
+        if not self._slices:
             # Note that `take` is faster than classical advanced indexing via
             # `operator[]` (`__getitem__`) because the latter is more generic.
             # Notably, `operator[]` supports boolean mask but `take` does not.
             return self.data.take(self.indices, axis=self.axis)
+        if self.axis is None:
+            # `reshape` must be used instead of `flat` to get a view that can
+            # be sliced without copy.
+            return self.data.reshape((-1,))[self._slices]
         return self.data[self._slices]

python/gym_jiminy/common/gym_jiminy/common/quantities/locomotion.py

@@ -9,6 +9,44 @@
 from ..bases import InterfaceJiminyEnv, AbstractQuantity
 from ..utils import fill
 
+from ..quantities import MaskedQuantity, AverageFrameSpatialVelocity
+
+
+@dataclass(unsafe_hash=True)
+class AverageOdometryVelocity(AbstractQuantity[np.ndarray]):
+    """Average odometry velocity in local-world-aligned frame at the end of the
+    agent step.
+
+    The odometry pose fully specifies the position of the robot in 2D world
+    plane. As such, it comprises the linear translation (X, Y) and the angular
+    velocity around Z axis (namely rate of change of Yaw Euler angle). The
+    average spatial velocity is obtained by finite difference. See
+    `AverageFrameSpatialVelocity` documentation for details.
+    """
+
+    def __init__(self,
+                 env: InterfaceJiminyEnv,
+                 parent: Optional[AbstractQuantity]) -> None:
+        """
+        :param env: Base or wrapped jiminy environment.
+        :param parent: Higher-level quantity from which this quantity is a
+                       requirement if any, `None` otherwise.
+        """
+        # Call base implementation
+        super().__init__(
+            env,
+            parent,
+            requirements=dict(
+                data=(MaskedQuantity, dict(
+                    quantity=(AverageFrameSpatialVelocity, dict(
+                        frame_name="root_joint",
+                        reference_frame=pin.LOCAL_WORLD_ALIGNED)),
+                    key=(0, 1, 5)))),
+            auto_refresh=False)
+
+    def refresh(self) -> np.ndarray:
+        return self.data
+
 
 @dataclass(unsafe_hash=True)
 class CenterOfMass(AbstractQuantity[np.ndarray]):

python/gym_jiminy/common/gym_jiminy/common/rewards/locomotion.py

@@ -1,39 +1,65 @@
-
+"""Rewards mainly relevant for locomotion tasks on floating-base robots.
+"""
 from typing import Sequence
 
 import numpy as np
 
 from ..bases import InterfaceJiminyEnv, BaseQuantityReward
-from ..quantities import MaskedQuantity, AverageFrameSpatialVelocity
+from ..quantities import AverageOdometryVelocity
 
 from .generic import radial_basis_function
 
 
 class OdometryVelocityReward(BaseQuantityReward):
-    """ TODO: Write documentation.
+    """Reward the agent for tracking a non-stationary target odometry velocity.
+
+    The error transform in a normalized reward to maximize by applying RBF
+    kernel on the error. The reward will be 0.0 if the error cancels out
+    completely and less than 0.01 above the user-specified cutoff threshold.
     """
     def __init__(self,
                  env: InterfaceJiminyEnv,
                  target: Sequence[float],
                  cutoff: float) -> None:
-        """ TODO: Write documentation.
+        """
+        :param target: Initial target average odometry velocity (vX, vY, vYaw).
+                       The target can be updated in necessary by calling
+                       `set_target`.
+        :param cutoff: Cutoff threshold for the RBF kernel transform.
         """
         # Backup some user argument(s)
-        self.target = target
+        self._target = np.asarray(target)
         self.cutoff = cutoff
 
         # Call base implementation
         super().__init__(
             env,
             "reward_odometry_velocity",
-            (MaskedQuantity, dict(
-                quantity=(AverageFrameSpatialVelocity, dict(frame_name="root_joint")),
-                key=(0, 1, 5))),
+            (AverageOdometryVelocity, {}),
             self._transform,
             is_normalized=True,
             is_terminal=False)
 
-    def _transform(self, value: np.ndarray) -> np.ndarray:
-        """ TODO: Write documentation.
+    @property
+    def target(self) -> np.ndarray:
+        """Get current target odometry velocity.
+        """
+        return self._target
+
+    @target.setter
+    def target(self, target: Sequence[float]) -> None:
+        """Set current target odometry velocity.
+        """
+        self._target = np.asarray(target)
+
+    def _transform(self, value: np.ndarray) -> float:
+        """Apply Radial Base Function transform to the residual error between
+        the current and target average odometry velocity.
+
+        .. note::
+            The user must call `set_target` method before `compute_reward` to
+            update the target odometry velocity if non-stationary.
+
+        :param value: Current average odometry velocity.
         """
         return radial_basis_function(value - self.target, self.cutoff)

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

@@ -11,6 +11,7 @@
                    quat_to_yaw_cos_sin,
                    quat_multiply,
                    quat_average,
+                   quat_interpolate_middle,
                    rpy_to_matrix,
                    rpy_to_quat,
                    transforms_to_vector,
@@ -55,6 +56,7 @@
     'quat_to_yaw_cos_sin',
     'quat_multiply',
     'quat_average',
+    'quat_interpolate_middle',
     'rpy_to_matrix',
     'rpy_to_quat',
     'transforms_to_vector',

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

@@ -522,3 +522,37 @@ def quat_average(quat: np.ndarray,
     q_flat = q_perm.reshape((*q_perm.shape[:-len(axis)], -1))
     _, eigvec = np.linalg.eigh(q_flat @ np.swapaxes(q_flat, -1, -2))
     return np.moveaxis(eigvec[..., -1], -1, 0)
+
+
+@nb.jit(nopython=True, cache=True, fastmath=True)
+def quat_interpolate_middle(quat1: np.ndarray,
+                            quat2: np.ndarray,
+                            out: Optional[np.ndarray] = None) -> np.ndarray:
+    """Compute the midpoint interpolation between two batches of quaternions
+    [qx, qy, qz, qw].
+
+    The midpoint interpolation of two quaternion is defined as the integration
+    of half the difference between them, starting from the first one, ie
+    `q_mid = integrate(q1, 0.5 * difference(q1, d2))`, which is a special case
+    of the `slerp` method (spherical linear interpolation) for `alpha=0.5`.
+
+    For the midpoint in particular, one can show that the middle quaternion is
+    simply normalized sum of the previous and next quaternions.
+
+    :param quat1: First batch of quaternions as a N-dimensional array whose
+                  first dimension gathers the 4 quaternion coordinates.
+    :param quat2: Second batch of quaternions as a N-dimensional array.
+    :param out: A pre-allocated array into which the result is stored. If not
+                provided, a new array is freshly-allocated, which is slower.
+    """
+    assert quat1.ndim >= 1 and quat1.shape == quat2.shape
+    if out is None:
+        out_ = np.empty((4, *quat1.shape[1:]))
+    else:
+        assert out.shape == (4, *quat1.shape[1:])
+        out_ = out
+
+    dot = np.expand_dims(np.sum(quat1 * quat2, axis=0), axis=0)
+    out_[:] = (quat1 + np.sign(dot) * quat2) / np.sqrt(2 * (1 + np.abs(dot)))
+
+    return out_