From 2f1a834af83b569b3c877b019d8ed778bc972f4b Mon Sep 17 00:00:00 2001 From: Alexis DUBURCQ Date: Tue, 11 Jun 2024 18:25:06 +0200 Subject: [PATCH] [gym/common] Add angular momentum, stability, and friction rewards. (#808) * [gym/common] Remove error-prone, confusing and slow '__getattr__' fallback in pipelines. * [gym/common] Add angular momentum minimisation reward. * [gym/common] Add projected support polygon zmp stability margin maximisation reward. * [gym/common] Add ground contact friction minimisation reward. * [misc] Fix segfault on windows due to vcruntime conflicts. --- .pylintrc | 18 +- CMakeLists.txt | 4 +- build_tools/build_install_deps_windows.ps1 | 7 +- .../gym_jiminy/common/bases/interfaces.py | 25 +- .../gym_jiminy/common/bases/pipeline.py | 49 +- .../common/compositions/__init__.py | 8 +- .../gym_jiminy/common/compositions/generic.py | 5 +- .../common/compositions/locomotion.py | 36 +- .../gym_jiminy/common/compositions/mixin.py | 23 +- .../common/gym_jiminy/common/envs/generic.py | 35 +- .../gym_jiminy/common/quantities/__init__.py | 6 +- .../gym_jiminy/common/quantities/generic.py | 11 +- .../common/quantities/locomotion.py | 94 ++- .../gym_jiminy/common/quantities/transform.py | 7 +- .../common/wrappers/observation_stack.py | 4 + .../gym_jiminy/envs/gym_jiminy/envs/atlas.py | 13 +- .../toolbox/compositions/__init__.py | 8 + .../toolbox/compositions/locomotion.py | 102 ++++ .../gym_jiminy/toolbox/math/__init__.py | 6 +- .../toolbox/gym_jiminy/toolbox/math/qhull.py | 569 +++++++++++++----- .../gym_jiminy/toolbox/quantities/__init__.py | 9 + .../toolbox/quantities/locomotion.py | 220 +++++++ .../toolbox/wrappers/frame_rate_limiter.py | 2 - python/gym_jiminy/toolbox/setup.py | 7 +- python/gym_jiminy/unit_py/test_misc.py | 3 +- .../unit_py/test_pipeline_control.py | 6 +- .../unit_py/test_pipeline_design.py | 1 + python/gym_jiminy/unit_py/test_quantities.py | 24 +- python/gym_jiminy/unit_py/test_rewards.py | 43 +- python/jiminy_py/unit_py/utilities.py | 2 +- 30 files changed, 1124 insertions(+), 223 deletions(-) create mode 100644 python/gym_jiminy/toolbox/gym_jiminy/toolbox/compositions/__init__.py create mode 100644 python/gym_jiminy/toolbox/gym_jiminy/toolbox/compositions/locomotion.py create mode 100644 python/gym_jiminy/toolbox/gym_jiminy/toolbox/quantities/__init__.py create mode 100644 python/gym_jiminy/toolbox/gym_jiminy/toolbox/quantities/locomotion.py diff --git a/.pylintrc b/.pylintrc index f6f9034f2..a8239b7ec 100644 --- a/.pylintrc +++ b/.pylintrc @@ -22,15 +22,15 @@ generated-members = torch, jiminy [tool.pylint.basic] # Good variable names which should always be accepted, separated by a comma good-names = - i, j, k, l, N, # Python: for-loop indices - tb, np, nb, mp, tp, # Python: classical modules - fd, _, # Python: contexte - t, q, v, x, u, s, qx, qy, qz, qw, # Physics: state, action - I, R, H, T, M, dt, # Physics: dynamics - a, b, c, y, z, n, e, # Maths / Algebra : variables - f, rg, lo, hi, op, fn, # Maths / Algebra : operators - kp, kd, ki, # Control: Gains - ax # Matplotlib + i, j, k, l, N, # Python: for-loop indices + tb, np, nb, mp, tp, # Python: classical modules + fd, _, # Python: contexte + t, q, v, x, u, s, qx, qy, qz, qw, # Physics: state, action + I, R, H, T, M, dt, # Physics: dynamics + A, a, b, c, y, z, n, e, # Maths / Algebra: variables + f, rg, lo, hi, op, fn, # Maths / Algebra: operators + kp, kd, ki, # Control: Gains + ax # Matplotlib [tool.pylint.format] # Regexp for a line that is allowed to be longer than the limit diff --git a/CMakeLists.txt b/CMakeLists.txt index 16d4ab678..8803b650b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -71,7 +71,9 @@ if(MSVC) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /fp:contract") endif() endif() - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /MP /EHsc /bigobj /Gy /Zc:inline /Zc:preprocessor /Zc:__cplusplus /permissive- /DWIN32 /D_USE_MATH_DEFINES /DNOMINMAX") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} \ + /MP /EHsc /bigobj /Gy /Zc:inline /Zc:preprocessor /Zc:__cplusplus /permissive- /DWIN32 \ + /D_USE_MATH_DEFINES /D_DISABLE_CONSTEXPR_MUTEX_CONSTRUCTOR /DNOMINMAX") set(CMAKE_CXX_FLAGS_DEBUG "/Zi /Od") set(CMAKE_CXX_FLAGS_RELEASE "/DNDEBUG /O2 /Ob3") set(CMAKE_CXX_FLAGS_RELWITHDEBINFO "${CMAKE_CXX_FLAGS_RELEASE} /Zi") diff --git a/build_tools/build_install_deps_windows.ps1 b/build_tools/build_install_deps_windows.ps1 index f774d980b..aee2380c7 100644 --- a/build_tools/build_install_deps_windows.ps1 +++ b/build_tools/build_install_deps_windows.ps1 @@ -44,7 +44,12 @@ if (-not (Test-Path env:GENERATOR)) { } ### Set common CMAKE_C/CXX_FLAGS -${CMAKE_CXX_FLAGS} = "${env:CMAKE_CXX_FLAGS} /MP2 /EHsc /bigobj /Gy /Zc:inline /Zc:preprocessor /Zc:__cplusplus /permissive- /DWIN32 /D_USE_MATH_DEFINES /DNOMINMAX" +# * Flag "DISABLE_CONSTEXPR_MUTEX_CONSTRUCTOR" is a dirty workaround to deal with VC runtime +# conflict related to different search path ordering at compile-time / run-time causing segfault: +# https://github.com/actions/runner-images/issues/10004 +${CMAKE_CXX_FLAGS} = "${env:CMAKE_CXX_FLAGS} $( +) /MP2 /EHsc /bigobj /Gy /Zc:inline /Zc:preprocessor /Zc:__cplusplus /permissive- $( +) /DWIN32 /D_USE_MATH_DEFINES /D_DISABLE_CONSTEXPR_MUTEX_CONSTRUCTOR /DNOMINMAX" if (${BUILD_TYPE} -eq "Debug") { ${CMAKE_CXX_FLAGS} = "${CMAKE_CXX_FLAGS} /Zi /Od" } else { diff --git a/python/gym_jiminy/common/gym_jiminy/common/bases/interfaces.py b/python/gym_jiminy/common/gym_jiminy/common/bases/interfaces.py index f42efc7f7..1d25bf655 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/bases/interfaces.py +++ b/python/gym_jiminy/common/gym_jiminy/common/bases/interfaces.py @@ -17,6 +17,7 @@ from ..utils import DataNested if TYPE_CHECKING: + from ..envs.generic import BaseJiminyEnv from ..quantities import QuantityManager @@ -220,6 +221,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # Call super to allow mixing interfaces through multiple inheritance super().__init__(*args, **kwargs) + # Define convenience proxy for quantity manager + self.quantities = self.unwrapped.quantities + def _setup(self) -> None: """Configure the observer-controller. @@ -335,11 +339,26 @@ def _controller_handle(self, # '_controller_handle' as it is never called more often than necessary. self.__is_observation_refreshed = False + def stop(self) -> None: + """Stop the episode immediately without waiting for a termination or + truncation condition to be satisfied. + + .. note:: + This method is mainly intended for data analysis and debugging. + Stopping the episode is necessary to log the final state, otherwise + it will be missing from plots and viewer replay. Moreover, sensor + data will not be available during replay using object-oriented + method `replay`. Helper method `play_logs_data` must be preferred + to replay an episode that cannot be stopped at the time being. + """ + self.simulator.stop() + @property - def unwrapped(self) -> "InterfaceJiminyEnv": - """Base environment of the pipeline. + @abstractmethod + def unwrapped(self) -> "BaseJiminyEnv": + """The "underlying environment at the basis of the pipeline from which + this environment is part of. """ - return self @property @abstractmethod diff --git a/python/gym_jiminy/common/gym_jiminy/common/bases/pipeline.py b/python/gym_jiminy/common/gym_jiminy/common/bases/pipeline.py index c1cc53caf..4d7164878 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/bases/pipeline.py +++ b/python/gym_jiminy/common/gym_jiminy/common/bases/pipeline.py @@ -9,14 +9,16 @@ * a wrapper to combine a controller block and a `BaseJiminyEnv` environment, eventually already wrapped, so that it appears as a black-box environment. """ +import sys import math +import logging from weakref import ref from copy import deepcopy from abc import abstractmethod from collections import OrderedDict from typing import ( Dict, Any, List, Optional, Tuple, Union, Generic, TypeVar, SupportsFloat, - Callable, cast) + Callable, cast, TYPE_CHECKING) import numpy as np @@ -37,6 +39,8 @@ from .blocks import BaseControllerBlock, BaseObserverBlock from ..utils import DataNested, is_breakpoint, zeros, build_copyto, copy +if TYPE_CHECKING: + from ..envs.generic import BaseJiminyEnv OtherObsT = TypeVar('OtherObsT', bound=DataNested) @@ -46,6 +50,9 @@ TransformedActT = TypeVar('TransformedActT', bound=DataNested) +LOGGER = logging.getLogger(__name__) + + class BasePipelineWrapper( InterfaceJiminyEnv[ObsT, ActT], Generic[ObsT, ActT, BaseObsT, BaseActT]): @@ -101,7 +108,20 @@ def __getattr__(self, name: str) -> Any: It enables to get access to the attribute and methods of the wrapped environment directly without having to do it through `env`. + + .. warning:: + This fallback incurs a significant runtime overhead. As such, it + must only be used for debug and manual analysis between episodes. + Calling this method in script mode while a simulation is already + running would trigger a warning to avoid relying on it by mistake. """ + if self.is_simulation_running and not hasattr(sys, 'ps1'): + # `hasattr(sys, 'ps1')` is used to detect whether the method was + # called from an interpreter or within a script. For details, see: + # https://stackoverflow.com/a/64523765/4820605 + LOGGER.warning( + "Relying on fallback attribute getter is inefficient and " + "strongly discouraged at runtime.") return getattr(self.__getattribute__('env'), name) def __dir__(self) -> List[str]: @@ -143,9 +163,7 @@ def np_random(self, value: np.random.Generator) -> None: self.env.np_random = value @property - def unwrapped(self) -> InterfaceJiminyEnv: - """Base environment of the pipeline. - """ + def unwrapped(self) -> "BaseJiminyEnv": return self.env.unwrapped @property @@ -236,8 +254,7 @@ def step(self, # type: ignore[override] self._copyto_action(action) # Make sure that the pipeline has not change since last reset - env_derived = ( - self.unwrapped.derived) # type: ignore[attr-defined] + env_derived = self.unwrapped.derived if env_derived is not self: raise RuntimeError( "Pipeline environment has changed. Please call 'reset' " @@ -532,14 +549,14 @@ def __init__(self, # Register the observer's internal state and feature to the telemetry if state is not None: try: - self.env.register_variable( # type: ignore[attr-defined] + self.unwrapped.register_variable( 'state', state, None, self.observer.name) except ValueError: pass - self.env.register_variable('feature', # type: ignore[attr-defined] - self.observer.observation, - self.observer.fieldnames, - self.observer.name) + self.unwrapped.register_variable('feature', + self.observer.observation, + self.observer.fieldnames, + self.observer.name) def _setup(self) -> None: """Configure the wrapper. @@ -750,14 +767,14 @@ def __init__(self, # Register the controller's internal state and target to the telemetry if state is not None: try: - self.env.register_variable( # type: ignore[attr-defined] + self.unwrapped.register_variable( 'state', state, None, self.controller.name) except ValueError: pass - self.env.register_variable('action', # type: ignore[attr-defined] - self.action, - self.controller.fieldnames, - self.controller.name) + self.unwrapped.register_variable('action', + self.action, + self.controller.fieldnames, + self.controller.name) def _setup(self) -> None: """Configure the wrapper. diff --git a/python/gym_jiminy/common/gym_jiminy/common/compositions/__init__.py b/python/gym_jiminy/common/gym_jiminy/common/compositions/__init__.py index 9dfd9b037..a0ac75db7 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/compositions/__init__.py +++ b/python/gym_jiminy/common/gym_jiminy/common/compositions/__init__.py @@ -1,6 +1,7 @@ # pylint: disable=missing-module-docstring -from .mixin import (radial_basis_function, +from .mixin import (CUTOFF_ESP, + radial_basis_function, AdditiveMixtureReward, MultiplicativeMixtureReward) from .generic import (BaseTrackingReward, @@ -11,9 +12,11 @@ TrackingCapturePointReward, TrackingFootPositionsReward, TrackingFootOrientationsReward, - MinimizeAngularMomentumReward) + MinimizeAngularMomentumReward, + MinimizeFrictionReward) __all__ = [ + "CUTOFF_ESP", "radial_basis_function", "AdditiveMixtureReward", "MultiplicativeMixtureReward", @@ -25,5 +28,6 @@ "TrackingFootPositionsReward", "TrackingFootOrientationsReward", "MinimizeAngularMomentumReward", + "MinimizeFrictionReward", "SurviveReward" ] diff --git a/python/gym_jiminy/common/gym_jiminy/common/compositions/generic.py b/python/gym_jiminy/common/gym_jiminy/common/compositions/generic.py index 5afcd804b..0de5a6b0d 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/compositions/generic.py +++ b/python/gym_jiminy/common/gym_jiminy/common/compositions/generic.py @@ -54,9 +54,10 @@ class BaseTrackingReward(BaseQuantityReward): otherwise an exception will be risen. See `DatasetTrajectoryQuantity` and `AbstractQuantity` documentations for details. - The error transform in a normalized reward to maximize by applying RBF + The error is transformed 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. + completely and less than 'CUTOFF_ESP' above the user-specified cutoff + threshold. """ def __init__(self, env: InterfaceJiminyEnv, diff --git a/python/gym_jiminy/common/gym_jiminy/common/compositions/locomotion.py b/python/gym_jiminy/common/gym_jiminy/common/compositions/locomotion.py index 68044d684..6932e895f 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/compositions/locomotion.py +++ b/python/gym_jiminy/common/gym_jiminy/common/compositions/locomotion.py @@ -9,8 +9,9 @@ from ..bases import ( InterfaceJiminyEnv, StateQuantity, QuantityEvalMode, BaseQuantityReward) from ..quantities import ( - MaskedQuantity, UnaryOpQuantity, AverageBaseOdometryVelocity, - AverageBaseMomentum, MultiFootRelativeXYZQuat, CapturePoint) + MaskedQuantity, UnaryOpQuantity, AverageBaseOdometryVelocity, CapturePoint, + MultiFootRelativeXYZQuat, MultiContactRelativeForceTangential, + AverageBaseMomentum) from ..quantities.locomotion import sanitize_foot_frame_names from ..utils import quat_difference @@ -221,3 +222,34 @@ def __init__(self, partial(radial_basis_function, cutoff=self.cutoff, order=2), is_normalized=True, is_terminal=False) + + +class MinimizeFrictionReward(BaseQuantityReward): + """Reward the agent for minimizing the tangential forces at all the contact + points and collision bodies, and to avoid jerky intermittent contact state. + + The L2-norm is used to aggregate all the local tangential forces. While the + L1-norm would be more natural in this specific cases, using the L2-norm is + preferable as it promotes space-time regularity, ie balancing the force + distribution evenly between all the candidate contact points and avoiding + jerky contact forces over time (high-frequency vibrations), phenomena to + which the L1-norm is completely insensitive. + """ + def __init__(self, + env: InterfaceJiminyEnv, + cutoff: float) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param cutoff: Cutoff threshold for the RBF kernel transform. + """ + # Backup some user argument(s) + self.cutoff = cutoff + + # Call base implementation + super().__init__( + env, + "reward_friction", + (MultiContactRelativeForceTangential, dict()), + partial(radial_basis_function, cutoff=self.cutoff, order=2), + is_normalized=True, + is_terminal=False) diff --git a/python/gym_jiminy/common/gym_jiminy/common/compositions/mixin.py b/python/gym_jiminy/common/gym_jiminy/common/compositions/mixin.py index e9689bb8a..80bb0dab8 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/compositions/mixin.py +++ b/python/gym_jiminy/common/gym_jiminy/common/compositions/mixin.py @@ -13,7 +13,7 @@ # Reward value at cutoff threshold -RBF_CUTOFF_ESP = 1.0e-2 +CUTOFF_ESP = 1.0e-2 ArrayOrScalar = Union[np.ndarray, float] @@ -46,12 +46,23 @@ def radial_basis_function(error: ArrayOrScalar, :param cutoff: Cut-off threshold to consider. :param order: Order of Lp-Norm that will be used as distance metric. """ - error_ = np.asarray(error).reshape((-1,)) - if order == 2: - squared_dist_rel = np.dot(error_, error_) / math.pow(cutoff, 2) + error_ = np.asarray(error) + is_contiguous = error_.flags.f_contiguous or error_.flags.c_contiguous + if is_contiguous or order != 2: + if error_.ndim > 1 and not is_contiguous: + error_ = np.ascontiguousarray(error_) + if error_.flags.c_contiguous: + error1d = np.asarray(error_).ravel() + else: + error1d = np.asarray(error_.T).ravel() + if order == 2: + squared_dist_rel = np.dot(error1d, error1d) / math.pow(cutoff, 2) + else: + squared_dist_rel = math.pow( + np.linalg.norm(error1d, order) / cutoff, 2) else: - squared_dist_rel = math.pow(np.linalg.norm(error_, order) / cutoff, 2) - return math.pow(RBF_CUTOFF_ESP, squared_dist_rel) + squared_dist_rel = np.sum(np.square(error_)) / math.pow(cutoff, 2) + return math.pow(CUTOFF_ESP, squared_dist_rel) class AdditiveMixtureReward(BaseMixtureReward): 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 bce0eb64e..03bf129b7 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py +++ b/python/gym_jiminy/common/gym_jiminy/common/envs/generic.py @@ -100,6 +100,12 @@ class BaseJiminyEnv(InterfaceJiminyEnv[ObsT, ActT], to implement one. It has been designed to be highly flexible and easy to customize by overloading it to fit the vast majority of users' needs. """ + + derived: "InterfaceJiminyEnv" + """Top-most block from which this environment is part of when leveraging + modular pipeline design capability. + """ + def __init__(self, simulator: Simulator, step_dt: float, @@ -186,8 +192,8 @@ def __init__(self, self.sensor_measurements: SensorMeasurementStackMap = OrderedDict( self.robot.sensor_measurements) - # Top-most block of the pipeline to which the environment is part of - self.derived: InterfaceJiminyEnv = self + # Top-most block of the pipeline is the environment itself by default + self.derived = self # Store references to the variables to register to the telemetry self._registered_variables: MutableMappingT[ @@ -215,6 +221,9 @@ def __init__(self, self.num_steps = np.array(-1, dtype=np.int64) self._num_steps_beyond_terminate: Optional[int] = None + # Initialize a quantity manager for later use + self.quantities = QuantityManager(self) + # Initialize the interfaces through multiple inheritance super().__init__() # Do not forward extra arguments, if any @@ -233,9 +242,6 @@ def __init__(self, "`BaseJiminyEnv.compute_command` must be overloaded in case " "of custom action spaces.") - # Initialize a quantity manager for later use - self.quantities = QuantityManager(self) - # Define specialized operators for efficiency. # Note that a partial view of observation corresponding to measurement # must be extracted since only this one must be updated during refresh. @@ -560,6 +566,10 @@ def step_dt(self) -> float: def is_training(self) -> bool: return self._is_training + @property + def unwrapped(self) -> "BaseJiminyEnv": + return self + def train(self) -> None: self._is_training = True @@ -599,8 +609,8 @@ def reset(self, # type: ignore[override] if seed is not None: self._initialize_seed(seed) - # Stop the simulator - self.simulator.stop() + # Stop the episode if one is still running + self.stop() # Remove external forces, if any self.simulator.remove_all_forces() @@ -854,7 +864,7 @@ def step(self, # type: ignore[override] self.simulator.step(self.step_dt) except Exception: # Stop the simulation before raising the exception - self.simulator.stop() + self.stop() raise # Make sure there is no 'nan' value in observation @@ -1023,8 +1033,8 @@ def replay(self, **kwargs: Any) -> None: kwargs['close_backend'] = not self.simulator.is_viewer_available # Stop any running simulation before replay if `has_terminated` is True - if self.is_simulation_running and any(self.has_terminated({})): - self.simulator.stop() + if any(self.has_terminated({})): + self.stop() with viewer_lock: # Call render before replay in order to take into account custom @@ -1135,8 +1145,7 @@ def _interact(key: Optional[str] = None) -> bool: # Stop the simulation to unlock the robot. # It will enable to display contact forces for replay. - if self.simulator.is_simulation_running: - self.simulator.stop() + self.stop() # Disable play interactive mode flag and restore training flag self._is_interactive = False @@ -1213,7 +1222,7 @@ def evaluate(self, action = policy_fn(obs, reward, terminated or truncated, info) obs, reward, terminated, truncated, info = env.step(action) info_episode.append(info) - self.simulator.stop() + self.stop() except KeyboardInterrupt: pass diff --git a/python/gym_jiminy/common/gym_jiminy/common/quantities/__init__.py b/python/gym_jiminy/common/gym_jiminy/common/quantities/__init__.py index 0d75ab7a6..d22168ed0 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/quantities/__init__.py +++ b/python/gym_jiminy/common/gym_jiminy/common/quantities/__init__.py @@ -26,7 +26,9 @@ MultiFootRelativeXYZQuat, CenterOfMass, CapturePoint, - ZeroMomentPoint) + ZeroMomentPoint, + MultiContactRelativeForceTangential, + translate_position_odom) __all__ = [ @@ -57,4 +59,6 @@ 'CenterOfMass', 'CapturePoint', 'ZeroMomentPoint', + 'MultiContactRelativeForceTangential', + 'translate_position_odom' ] 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 08eea09c2..73e94c2d9 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/quantities/generic.py +++ b/python/gym_jiminy/common/gym_jiminy/common/quantities/generic.py @@ -230,6 +230,9 @@ def initialize(self) -> None: # Re-allocate memory as the number of frames is not known in advance. # Note that Fortran memory layout (column-major) is used for speed up # because it preserves contiguity when copying frame data. + # Anyway, C memory layout (row-major) does not make sense in this case + # since chunks of columns are systematically extracted, which means + # that the returned array would NEVER be contiguous. nframes = len(self.frame_names) self._rot_mat_batch = np.zeros((3, 3, nframes), order='F') @@ -387,9 +390,9 @@ def initialize(self) -> None: # Re-allocate memory as the number of frames is not known in advance nframes = len(self.frame_names) if self.type in (OrientationType.EULER, OrientationType.ANGLE_AXIS): - self._data_batch = np.zeros((3, nframes), order='C') + self._data_batch = np.zeros((3, nframes), order='F') elif self.type == OrientationType.QUATERNION: - self._data_batch = np.zeros((4, nframes), order='C') + self._data_batch = np.zeros((4, nframes), order='F') # Re-assign mapping from chunks of frame names to corresponding data if self.type is not OrientationType.MATRIX: @@ -640,7 +643,7 @@ def initialize(self) -> None: # Re-allocate memory as the number of frames is not known in advance nframes = len(self.frame_names) - self._pos_batch = np.zeros((3, nframes), order='C') + self._pos_batch = np.zeros((3, nframes), order='F') # Refresh proxies self._pos_views.clear() @@ -908,7 +911,7 @@ def __init__(self, auto_refresh=False) # Pre-allocate memory for storing the pose XYZQuat of all frames - self._xyzquats = np.zeros((7, len(frame_names)), order='C') + self._xyzquats = np.zeros((7, len(frame_names)), order='F') def refresh(self) -> np.ndarray: # Copy the position of all frames at once in contiguous buffer diff --git a/python/gym_jiminy/common/gym_jiminy/common/quantities/locomotion.py b/python/gym_jiminy/common/gym_jiminy/common/quantities/locomotion.py index 8d9c911d8..493122016 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/quantities/locomotion.py +++ b/python/gym_jiminy/common/gym_jiminy/common/quantities/locomotion.py @@ -1,14 +1,15 @@ """Quantities mainly relevant for locomotion tasks on floating-base robots. """ import math -from typing import Optional, Tuple, Sequence, Literal, Union +from typing import List, Optional, Tuple, Sequence, Literal, Union from dataclasses import dataclass import numpy as np import numba as nb import jiminy_py.core as jiminy -from jiminy_py.core import array_copyto # pylint: disable=no-name-in-module +from jiminy_py.core import ( # pylint: disable=no-name-in-module + array_copyto, multi_array_copyto) from jiminy_py.dynamics import update_quantities import pinocchio as pin @@ -902,3 +903,92 @@ def refresh(self) -> np.ndarray: translate_position_odom(self._dcm, self.odom_pose, self._dcm) return self._dcm + + +@dataclass(unsafe_hash=True) +class MultiContactRelativeForceTangential(InterfaceQuantity[np.ndarray]): + """Standardized tangential forces apply on all contact points and collision + bodies. + + The tangential force is rescaled by the weight of the robot rather than the + actual vertical force. It has the advantage to guarantee that the resulting + quantity is never poorly conditioned, which would be the case otherwise. + Moreover, the effect of the vertical force is not canceled out, which is + interesting for deriving a reward, as it allows for indirectly penalize + jerky contact states and violent impacts. The side effect is not being able + to guarantee that this quantity is bounded. Indeed, only the ratio of the + norm of the tangential force at every contact point (or the resulting one) + is bounded by the product of the friction coefficient by the vertical + force, not the tangential force itself. This issue is a minor inconvenience + as all it requires is normalization using RBF kernel to make it finite. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + parent: Optional[InterfaceQuantity]) -> 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={}, + auto_refresh=False) + + # Weight of the robot + self._robot_weight: float = -1 + + # Stacked tangential forces on all contact points and collision bodies + self._force_tangential_batch = np.array([]) + + # Define proxy for views of individual tangential forces in the stack + self._force_tangential_views: Tuple[np.ndarray, ...] = () + + # Define proxy for the current tangential forces + self._force_tangential_refs: Tuple[np.ndarray, ...] = () + + def initialize(self) -> None: + # Call base implementation + super().initialize() + + # Compute the weight of the robot + gravity = abs(self.env.robot.pinocchio_model.gravity.linear[2]) + robot_mass = self.env.robot.pinocchio_data.mass[0] + self._robot_weight = robot_mass * gravity + + # Get all frame constraints associated with contacts and collisions + frame_constraints: List[jiminy.FrameConstraint] = [] + for constraint in self.env.robot.constraints.contact_frames.values(): + assert isinstance(constraint, jiminy.FrameConstraint) + frame_constraints.append(constraint) + for constraints_body in self.env.robot.constraints.collision_bodies: + for constraint in constraints_body: + assert isinstance(constraint, jiminy.FrameConstraint) + frame_constraints.append(constraint) + + # Extract references to all the tangential forces + force_tangential_refs = [] + for constraint in frame_constraints: + # f_lin, f_ang = lambda[:3], np.array([0.0, 0.0, lambda[3]]) + force_tangential_refs.append(constraint.lambda_c[:2]) + self._force_tangential_refs = tuple(force_tangential_refs) + + # Pre-allocated memory for stacked tangential forces + self._force_tangential_batch = np.zeros( + (2, len(self._force_tangential_refs)), order='F') + + # Refresh proxies + self._force_tangential_views = tuple(self._force_tangential_batch.T) + + def refresh(self) -> np.ndarray: + # Copy all tangential forces in contiguous buffer + multi_array_copyto(self._force_tangential_views, + self._force_tangential_refs) + + # Scale the tangential forces by the weight of the robot + self._force_tangential_batch /= self._robot_weight + + return self._force_tangential_batch diff --git a/python/gym_jiminy/common/gym_jiminy/common/quantities/transform.py b/python/gym_jiminy/common/gym_jiminy/common/quantities/transform.py index f879e7c87..f471af826 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/quantities/transform.py +++ b/python/gym_jiminy/common/gym_jiminy/common/quantities/transform.py @@ -116,7 +116,8 @@ class MaskedQuantity(InterfaceQuantity[np.ndarray]): 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), and the array can be flattened while preserving memory contiguity if 'axis' - is `None`. + is `None`, which means that the result will be different between C- and F- + contiguous arrays. """ quantity: InterfaceQuantity @@ -204,9 +205,9 @@ def refresh(self) -> np.ndarray: # Notably, `operator[]` supports boolean mask but `take` does not. return self.data.take(self.indices, self.axis) if self.axis is None: - # `reshape` must be used instead of `flat` to get a view that can + # `ravel` 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.ravel(order="K")[self._slices] return self.data[self._slices] diff --git a/python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_stack.py b/python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_stack.py index e700a63b1..6499d0c4b 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_stack.py +++ b/python/gym_jiminy/common/gym_jiminy/common/wrappers/observation_stack.py @@ -146,6 +146,10 @@ def __init__(self, # Whether the stack has been shifted to the left since last update self._was_stack_shifted = True + # Bind action of the base environment + assert self.action_space.contains(self.env.action) + self.action = self.env.action + def _initialize_action_space(self) -> None: """Configure the action space. """ diff --git a/python/gym_jiminy/envs/gym_jiminy/envs/atlas.py b/python/gym_jiminy/envs/gym_jiminy/envs/atlas.py index f11ee2b82..a49c89a2a 100644 --- a/python/gym_jiminy/envs/gym_jiminy/envs/atlas.py +++ b/python/gym_jiminy/envs/gym_jiminy/envs/atlas.py @@ -18,7 +18,7 @@ PDAdapter, MahonyFilter) from gym_jiminy.common.utils import build_pipeline -from gym_jiminy.toolbox.math import ConvexHull +from gym_jiminy.toolbox.math import ConvexHull2D if sys.version_info < (3, 9): from importlib_resources import files @@ -100,18 +100,19 @@ def _cleanup_contact_points(env: WalkerJiminyEnv) -> None: contact_frame_indices = env.robot.contact_frame_indices contact_frame_names = env.robot.contact_frame_names - num_contacts = int(len(env.robot.contact_frame_indices) // 2) + num_contacts = len(env.robot.contact_frame_indices) // 2 for contact_slice in (slice(num_contacts), slice(num_contacts, None)): contact_positions = np.stack([ env.robot.pinocchio_data.oMf[frame_index].translation - for frame_index in contact_frame_indices[contact_slice]], axis=0) + for frame_index in contact_frame_indices[contact_slice] + ], axis=0) contact_bottom_index = np.argsort( - contact_positions[:, 2])[:int(num_contacts//2)] - convex_hull = ConvexHull(contact_positions[contact_bottom_index, :2]) + contact_positions[:, 2])[:(num_contacts // 2)] + convex_hull = ConvexHull2D(contact_positions[contact_bottom_index, :2]) env.robot.remove_contact_points([ contact_frame_names[contact_slice][i] for i in set(range(num_contacts)).difference( - contact_bottom_index[convex_hull._vertex_indices])]) + contact_bottom_index[convex_hull.indices])]) class AtlasJiminyEnv(WalkerJiminyEnv): diff --git a/python/gym_jiminy/toolbox/gym_jiminy/toolbox/compositions/__init__.py b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/compositions/__init__.py new file mode 100644 index 000000000..9264558bd --- /dev/null +++ b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/compositions/__init__.py @@ -0,0 +1,8 @@ +# pylint: disable=missing-module-docstring + +from .locomotion import tanh_normalization, MaximizeStability + +__all__ = [ + "tanh_normalization", + "MaximizeStability" +] diff --git a/python/gym_jiminy/toolbox/gym_jiminy/toolbox/compositions/locomotion.py b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/compositions/locomotion.py new file mode 100644 index 000000000..1d814f6cc --- /dev/null +++ b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/compositions/locomotion.py @@ -0,0 +1,102 @@ +"""Rewards mainly relevant for locomotion tasks on floating-base robots. +""" +import math +from functools import partial + +import numba as nb + +from gym_jiminy.common.compositions import CUTOFF_ESP +from gym_jiminy.common.bases import ( + InterfaceJiminyEnv, QuantityEvalMode, BaseQuantityReward) + +from ..quantities import StabilityMarginProjectedSupportPolygon + + +@nb.jit(nopython=True, cache=True) +def tanh_normalization(value: float, + cutoff_low: float, + cutoff_high: float) -> float: + """Normalize a given quantity between 0.0 and 1.0. + + The extremum 0.0 and 1.0 correspond to the upper and lower cutoff + respectively, if the lower cutoff is smaller than the upper cutoff. The + other way around otherwise. These extremum are reached asymptotically, + which is that the gradient is never zero but rather vanishes exponentially. + The gradient will be steeper if the cutoff range is tighter and the other + way around. + + :param value: Value of the scalar floating-point quantity. The quantity may + be bounded or unbounded, and signed or not, without + restrictions. + :param cutoff: Cut-off threshold to consider. + :param order: Order of Lp-Norm that will be used as distance metric. + """ + value_rel = ( + cutoff_high + cutoff_low - 2 * value) / (cutoff_high - cutoff_low) + return 1.0 / (1.0 + math.pow(CUTOFF_ESP / (1.0 - CUTOFF_ESP), value_rel)) + + +class MaximizeStability(BaseQuantityReward): + """Encourage the agent to maintain itself in postures as robust as possible + to external disturbances. + + The signed distance is transformed in a normalized reward to maximize by + applying rescaled tanh. The reward is smaller than CUTOFF_ESP if the ZMP is + outside the projected support polygon and further away from the border than + the upper cutoff. Conversely, the reward is larger than 1.0 - CUTOFF_ESP if + the ZMP is inside the projected support polygon and further away from the + border than the lower cutoff. + + The agent may opt from one of the two very different strategies to maximize + this reward: + * Foot placement: reshaping the projected support polygon by moving the + feet (aka the candidate contact points in the direction of the ZMP + without actually moving the ZMP itself. + * Torso/Ankle control: Modulating the linear and angular momentum of its + upper-body to move the ZMP closer to the Chebyshev center of the + projected support polygon while holding the feet at the exact same + location. + + These two strategies are complementary rather than mutually exclusive. + Usually, ankle control is preferred for small disturbances. Foot placement + comes to place when ankle control is no longer sufficient to keep balance. + Indeed, the first strategy is only capable of recovering 0-step capturable + disturbances, while the second one is only limited to inf-step capturable + disturbances, which includes and dramatically extends 0-step capturability. + """ + def __init__(self, + env: InterfaceJiminyEnv, + cutoff_inner: float, + cutoff_outer: float) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param cutoff_inner: Cutoff threshold when the ZMP lies inside the + support polygon. The reward will be larger than + '1.0 - CUTOFF_ESP' if the distance from the border + is larger than 'cutoff_inner'. + :param cutoff_outer: Cutoff threshold when the ZMP lies outside the + support polygon. The reward will be smaller than + 'CUTOFF_ESP' if the ZMP is further away from the + border of the support polygon than 'cutoff_outer'. + """ + # Backup some user argument(s) + self.cutoff_inner = cutoff_inner + self.cutoff_outer = cutoff_outer + + # The cutoff thresholds must be positive + if self.cutoff_inner < 0.0 or self.cutoff_outer < 0.0: + raise ValueError( + "The inner and outer cutoff must both be positive.") + + # Call base implementation + super().__init__( + env, + "reward_momentum", + (StabilityMarginProjectedSupportPolygon, dict( + mode=QuantityEvalMode.TRUE + )), + partial(tanh_normalization, + cutoff_low=self.cutoff_inner, + cutoff_high=-self.cutoff_outer), + is_normalized=True, + is_terminal=False) diff --git a/python/gym_jiminy/toolbox/gym_jiminy/toolbox/math/__init__.py b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/math/__init__.py index 4f8da3806..8f11c59c8 100644 --- a/python/gym_jiminy/toolbox/gym_jiminy/toolbox/math/__init__.py +++ b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/math/__init__.py @@ -1,10 +1,10 @@ # pylint: disable=missing-module-docstring -from .qhull import ConvexHull, compute_distance_convex_to_point +from .qhull import ConvexHull2D, compute_convex_chebyshev_center __all__ = [ - "ConvexHull", - "compute_distance_convex_to_point" + "ConvexHull2D", + "compute_convex_chebyshev_center" ] try: diff --git a/python/gym_jiminy/toolbox/gym_jiminy/toolbox/math/qhull.py b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/math/qhull.py index 45944151b..e5467b9ac 100644 --- a/python/gym_jiminy/toolbox/gym_jiminy/toolbox/math/qhull.py +++ b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/math/qhull.py @@ -1,16 +1,32 @@ -""" TODO: Write documentation. +"""This module proposes a basic yet reasonably efficient representation of the +convex hull of a set of points in 2D Euclidean space. + +This comes with helper methods to compute various types of distances, including +for degenerated cases. """ -from typing import Optional +import math +from functools import cached_property +from typing import Tuple, List import numpy as np import numba as nb from numba.np.extensions import cross2d -from scipy.spatial import _qhull + +from scipy.optimize import linprog + +try: + from matplotlib.pyplot import Figure +except ImportError: + Figure = type(None) # type: ignore[misc,assignment,unused-ignore] + +from jiminy_py.viewer import interactive_mode @nb.jit(nopython=True, cache=True, inline='always') def _amin_last_axis(array: np.ndarray) -> np.ndarray: - """ TODO: Write documentation. + """Compute the minimum of a 2-dimensional array along the second axis. + + :param array: Input array. """ res = np.empty(array.shape[0]) for i in range(array.shape[0]): @@ -20,7 +36,10 @@ def _amin_last_axis(array: np.ndarray) -> np.ndarray: @nb.jit(nopython=True, cache=True, inline='always') def _all_last_axis(array: np.ndarray) -> np.ndarray: - """ TODO: Write documentation. + """Test whether all array elements along the second axis of a 2-dimensional + array evaluate to True. + + :param array: Input array. """ res = np.empty(array.shape[0], dtype=np.bool_) for i in range(array.shape[0]): @@ -28,184 +47,462 @@ def _all_last_axis(array: np.ndarray) -> np.ndarray: return res -@nb.jit(nopython=True, cache=True) -def compute_distance_convex_to_point(points: np.ndarray, - vertex_indices: np.ndarray, - queries: np.ndarray) -> np.ndarray: - """ TODO: Write documentation. +@nb.jit(nopython=True, cache=True, fastmath=True) +def compute_convex_hull_vertices(points: np.ndarray) -> np.ndarray: + """Determine the vertices of the convex hull defined by a set of points in + 2D Euclidean space. + + As a reminder, the convex hull of a set of points is defined as the + smallest convex polygon that can enclose all the points. + + .. seealso:: + Internally, it leverages using Andrew's monotone chain algorithm, which + as almost optimal complexity O(n*log(n)) but is only applicable in 2D + space. For an overview of all the existing algorithms to this day, see: + https://en.wikipedia.org/wiki/Convex_hull_algorithms + + :param points: Set of 2D points from which to compute the convex hull, + as a 2D array whose first dimension corresponds to the + number of points, and the second gathers the 2 position + coordinates (X, Y). """ - # Compute the equations of the edges - points_1 = points[np.roll(vertex_indices, 1)].T - points_0 = points[vertex_indices].T - vectors = points_1 - points_0 - normals = np.stack((-vectors[1], vectors[0]), 0) - normals /= np.sqrt(np.sum(np.square(normals), 0)) - offsets = - np.sum(normals * points_0, 0) - equations = np.concatenate((normals, np.expand_dims(offsets, 0))) + # Sorting points lexicographically (by x and then by y). + indices = np.argsort(points[:, 0], kind='mergesort') + indices = indices[np.argsort(points[indices, 1], kind='mergesort')] + + # If there is less than 3 points, the hull comprises all the points + npoints = len(points) + if npoints <= 3: + return indices + + # Combining the lower and upper hulls to get the full convex hull. + # The first point of each hull is omitted because it is the same as the + # last point of the other hull. + vertex_indices: List[int] = [] + + # Build the upper hull. + # The upper hull is similar to the lower hull but runs along the top of the + # set of points. It is constructed by starting with the right-most point + # and moving left. + upper: List[np.ndarray] = [] + for i in indices[::-1]: + point = points[i] + while len(upper) > 1: + # Check if the point is inside or outside of the hull at the time + # being by checking the sign of the 2D cross product. + if (upper[-1][0] - upper[-2][0]) * (point[1] - upper[-2][1]) > ( + point[0] - upper[-2][0]) * (upper[-1][1] - upper[-2][1]): + break + upper.pop() + vertex_indices.pop() + if upper: + vertex_indices.append(i) + upper.append(point) + + # Build the lower hull. + # The lower hull is the part of the convex hull that runs along the bottom + # of the set of points when they are sorted by their x-coordinates (from + # left to right). It is constructed by starting with the left-most point, + # points are added to the lower hull. If adding a new point creates a + # "right turn" (or non-left turn) with the last two points in the lower + # hull, the second-to-last point is removed. + lower: List[np.ndarray] = [] + for i in indices: + point = points[i] + while len(lower) > 1: + if (lower[-1][0] - lower[-2][0]) * (point[1] - lower[-2][1]) > ( + lower[-1][1] - lower[-2][1]) * (point[0] - lower[-2][0]): + break + lower.pop() + vertex_indices.pop() + if lower: + vertex_indices.append(i) + lower.append(point) + + return np.array(vertex_indices) + + +@nb.jit(nopython=True, cache=True, inline='always') +def compute_vectors_from_convex(vertices: np.ndarray) -> np.ndarray: + """Compute the un-normalized oriented direction vector of the edges. + + A point is inside the convex hull if it lies on the left side of all the + edges. + + :param vertices: Vertices of the convex hull with counter-clockwise + ordering, as a 2D array whose first dimension corresponds + to individual vertices while the second dimensions gathers + the 2 position coordinates (X, Y). + :returns: Direction of all the edges with the same ordering of the provided + vertices, as a 2D array whose first dimension corresponds to individual + edges while the second gathers the 2 components of the direction. + """ + vectors = np.empty((2, len(vertices))).T + vectors[0] = vertices[-1] - vertices[0] + vectors[1:] = vertices[:-1] - vertices[1:] + return vectors + + +@nb.jit(nopython=True, cache=True, inline='always') +def compute_equations_from_convex(vertices: np.ndarray, + vectors: np.ndarray) -> np.ndarray: + """Compute the (normalized) equations of the edges for a convex hull in 2D + Euclidean space. + + The equation of a edge is fully specified by its normal vector 'a' and a + scalar floating-point offset 'c'. A given point 'x' is on the line + defined by the edge of `np.dot(a, x) + d = 0.0`, inside if negative, + outside otherwise. + + :param vertices: Vertices of the convex hull with counter-clockwise + ordering, as a 2D array whose first dimension corresponds + to individual vertices while the second dimensions gathers + the 2 position coordinates (X, Y). + :param vectors: Direction of all the edges with the same ordering of the + provided vertices, as a 2D array whose first dimension + corresponds to individual edges while the second gathers + the 2 components of the direction. + + :returns: Equations of all the edges with the same ordering of the provided + vertices, as a 2D array whose first dimension corresponds to individual + edges while the second gathers the 2 components of the normal (ax, ay) plus + the offset d. The normal vector is normalized. + """ + equations = np.empty((3, len(vertices))) + normals, offsets = equations[:2], equations[-1] + normals[0] = - vectors[:, 1] + normals[1] = + vectors[:, 0] + normals /= np.sqrt(np.sum(np.square(normals), axis=0)) + offsets[:] = - np.sum(normals * vertices.T, axis=0) + return equations.T + + +@nb.jit(nopython=True, cache=True, inline='always') +def compute_distance_convex_to_point(vertices: np.ndarray, + vectors: np.ndarray, + queries: np.ndarray) -> np.ndarray: + """Compute the signed distance of query points from a convex hull in 2D + Euclidean space. + + Positive distance corresponds to a query point lying outside the convex + hull. + + .. warning: + The convex hull must be non-degenerated, ie having at least 3 points. + + :param vertices: Vertices of the convex hull with counter-clockwise + ordering, as a 2D array whose first dimension corresponds + to individual vertices while the second dimensions gathers + the 2 position coordinates (X, Y). + :param vectors: Direction of all the edges with the same ordering of the + provided vertices, as a 2D array whose first dimension + corresponds to individual edges while the second gathers + the 2 components of the direction. + :param queries: Set of 2D points for which to compute the distance from the + convex hull, as a 2D array whose first dimension + corresponds to the individual query points while the second + dimensions gathers the 2 position coordinates (X, Y). + """ # Determine for each query point if it lies inside or outside - queries = np.ascontiguousarray(queries) + queries_rel = np.expand_dims(queries, -1) - vertices.T sign_dist = 1.0 - 2.0 * _all_last_axis( - queries @ equations[:-1] + equations[-1] < 0.0) + queries_rel[:, 0] * vectors[:, 1] > queries_rel[:, 1] * vectors[:, 0]) # Compute the distance from the convex hull, as the min distance # from every segment of the convex hull. - ratios = np.sum( - (np.expand_dims(queries, -1) - points_0) * vectors, axis=1 - ) / np.sum(np.square(vectors), 0) - ratios = np.clip(ratios, 0.0, 1.0) - projs = np.expand_dims(ratios, 1) * vectors + points_0 + ratios = np.expand_dims(np.sum( + queries_rel * vectors.T, axis=1 + ), 1) / np.sum(np.square(vectors), axis=1) + ratios = np.minimum(np.maximum(ratios, 0.0), 1.0) + projs = ratios * vectors.T + vertices.T dist = np.sqrt(_amin_last_axis(np.sum(np.square( - np.expand_dims(queries, -1) - projs), 1))) - - # Compute the resulting signed distance (negative if inside) - signed_dist = sign_dist * dist + np.expand_dims(queries, -1) - projs), axis=1))) - return signed_dist + # Resulting the signed distance (negative if inside) + return sign_dist * dist -@nb.jit(nopython=True, cache=True) +@nb.jit(nopython=True, cache=True, inline='always') def compute_distance_convex_to_ray( - points: np.ndarray, - vertex_indices: np.ndarray, - query_vector: np.ndarray, + vertices: np.ndarray, + vectors: np.ndarray, + query_dir: np.ndarray, query_origin: np.ndarray) -> float: - """ TODO: Write documentation. + """Compute ray-casting signed distance (aka. time-of-flight) from a convex + hull to a oriented ray originating at a given position and pointing in a + specific direction, in 2D Euclidean space. + + The distance is negative if the origin of the ray lies inside the convex + hull. The distance is 'inf' is there is no intersection between the + oriented ray and the convex hull. It never happens if the origin lays + inside the convex hull, which means that the distance is negative, but + there is no guarantee otherwise. + + .. warning: + The convex hull must be non-degenerated, ie having at least 3 points. + + :param vertices: Vertices of the convex hull with counter-clockwise + ordering, as a 2D array whose first dimension corresponds + to individual vertices while the second dimensions gathers + the 2 position coordinates (X, Y). + :param vectors: Direction of all the edges with the same ordering of the + provided vertices, as a 2D array whose first dimension + corresponds to individual edges while the second gathers + the 2 components of the direction. + :param query_dir: Direction in which the ray is casting, as a 1D array. + It does not have to be normalized. + :param query_origin: Position from which the ray is casting, as a 1D array. """ - # Compute the direction vectors of the edges - points_1 = points[np.roll(vertex_indices, 1)] - points_0 = points[vertex_indices] - vectors = points_1 - points_0 - - # Compute the distance from the convex hull, as the only edge intersecting - # with the oriented line. - ratios = cross2d(query_origin - points_0, query_vector) / \ - cross2d(vectors, query_vector) - - if np.sum(np.logical_and(0.0 <= ratios, ratios < 1.0)) != 2: - raise ValueError("Query point origin not lying inside convex hull.") - + # Compute the distance from the convex hull. + # The distance only edge intersecting with the oriented line. + # The follow ratio corresponds to the relative position of the intersection + # point from each edge, 0.0 and 1.0 correspond to start vertex 'point_0' + # and end vertex 'point_1' respectively. This ratio is unbounded. Values + # outside range [0.0, 1.0] means that there is no intersection with the + # corresponding edge. + ratios = (cross2d(query_origin - vertices, query_dir) / + cross2d(vectors, query_dir)) + + # Compute the minimum casting distance and count how many edges are crossed + collide_num = 0 + casting_dist = math.inf for j, ratio in enumerate(ratios): if 0.0 <= ratio < 1.0: - proj = ratio * vectors[j] + points_0[j] - query_origin - if proj.dot(query_vector) > 0.0: - return np.linalg.norm(proj) # type: ignore[return-value] - - return 0.0 # This case cannot happens because for the explicit check + collision = ratio * vectors[j] + vertices[j] + oriented_ray = collision - query_origin + if oriented_ray.dot(query_dir) > 0.0: + casting_dist = min( # type: ignore[assignment] + np.linalg.norm(oriented_ray), casting_dist) + collide_num += 1 + + # If the ray is intersecting with two edges and the sign of the oriented + # casting ray from origin to collision point is positive for only one of + # them, then the origin is laying inside the convex hull, which means that + # the sign of the distance should be negative. On the contrary, if both + # are positive, then it is outside and the distance is the minimum between + # them. In all other cases, the distance is undefine, returning 'inf'. + if collide_num == 1: + casting_dist *= -1 + return casting_dist -class ConvexHull: - """ TODO: Write documentation. +@nb.jit(nopython=True, cache=True, inline='always') +def compute_distance_convex_to_convex(vertices_1: np.ndarray, + vectors_1: np.ndarray, + vertices_2: np.ndarray, + vectors_2: np.ndarray) -> float: + """Compute the distance between two convex hulls in 2D Euclidean space. + + .. warning: + Both convex hull must be non-degenerated, ie having at least 3 points + each. + + :param vertices_1: Vertices of the first convex hull with counter-clockwise + ordering, as a 2D array whose first dimension + corresponds to individual vertices while the second + dimensions gathers the 2 position coordinates (X, Y). + :param vectors_1: Direction of all the edges of the first convex hull with + the same ordering of the provided vertices, as a 2D array + whose first dimension corresponds to individual edges and + the second gathers the 2 components of the direction. + :param vertices_2: Vertices of the second convex hull with counter- + clockwise ordering, as a 2D array. See `vertices_1` for + details. + :param vectors_2: Direction of all the edges of the second convex hull with + the same ordering of the provided vertices, as a 2D + array. See `vertices_2` for details. + """ + distance_1 = np.min( + compute_distance_convex_to_point(vertices_1, vectors_1, vertices_2)) + distance_2 = np.min( + compute_distance_convex_to_point(vertices_2, vectors_2, vertices_1)) + return min(distance_1, distance_2) + + +def compute_convex_chebyshev_center( + equations: np.ndarray) -> Tuple[np.ndarray, float]: + r"""Compute the Chebyshev center of a convex polyhedron in N-dimensional + Euclidean space. + + The Chebyshev center is the point that is furthest inside a convex + polyhedron. Alternatively, it is the center of the largest hypersphere of + inscribed in the polyhedron. This can easily be computed using linear + programming. Considering halfspaces of the form :math:`Ax + b \leq 0`, + solving the linear program: + + .. math:: + max \: y + s.t. Ax + y ||A_i|| \leq -b + + With :math:`A_i` being the rows of A, i.e. the normals to each plane. The + equations outputted by Qhull are always normalized. For reference, see: + https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.HalfspaceIntersection.html + + .. warning: + The convex hull must be non-degenerated, ie having at least 3 points. + + :param equations: Equations of the edges as a 2D array whose first dimension + corresponds to individual edges while the second gathers + the 2 components of the normal plus the offset. + + :return: Pair (center, radius) where 'center' is a 1D array, and 'radius' + is a positive scalar floating point value. + """ # noqa: E501 # pylint: disable=line-too-long + # Compute the centroid of the polyhedron as the initial guess + num_dims = equations.shape[1] + A = np.concatenate(( + equations[:, :-1], np.ones((len(equations), 1))), axis=1) + b = - equations[:, -1:] + c = np.array([*(0.0,) * (num_dims - 1), -1.0]) + res = linprog(c, A_ub=A, b_ub=b, bounds=(None, None)) + return res.x[:-1], res.x[-1] + + +class ConvexHull2D: + """Class representing the convex hulls of a set of points in 2D Euclidean + space. """ + def __init__(self, points: np.ndarray) -> None: - """Compute the convex hull defined by a set of points. + """Compute the convex hull defined by a set of points in 2D Euclidean + space. - :param points: N-D points whose to computed the associated convex hull, + :param points: Set of 2D points from which to compute the convex hull, as a 2D array whose first dimension corresponds to the - number of points, and the second to the N-D coordinates. + number of points, and the second gathers the 2 position + coordinates (X, Y). """ assert len(points) > 0, "The length of 'points' must be at least 1." # Backup user argument(s) - self._points = np.ascontiguousarray(points) - - # Create convex full if possible - if len(self._points) > 2: - try: - self._hull = _qhull._Qhull(points=self._points, - options=b"", - mode_option=b"i", - required_options=b"Qt", - furthest_site=False, - incremental=False, - interior_point=None) - except _qhull.QhullError as e: - raise ValueError( - f"Impossible to compute convex hull ({self._points})." - ) from e - self._vertex_indices = self._hull.get_extremes_2d() - else: - self._hull = None - - # Buffer to cache center computation - self._center: Optional[np.ndarray] = None - - @property - def center(self) -> np.ndarray: - """Get the center of the convex hull. + self.points = points + self.npoints = len(self.points) - .. note:: - Degenerated convex hulls corresponding to len(points) == 1 or 2 are - handled separately. + # Compute the vertices of the convex hull + self.indices = compute_convex_hull_vertices(self.points) - :returns: 1D float vector with N-D coordinates of the center. + # Extract vertices of the convex hull with counter-clockwise ordering + self.vertices = self.points[self.indices] + + @cached_property + def vectors(self) -> np.ndarray: + """Un-normalized oriented direction vector of the edges. """ - if self._center is None: - if len(self._points) > 3: - vertices = self._points[self._vertex_indices] - else: - vertices = self._points - self._center = np.mean(vertices, axis=0) - return self._center + return compute_vectors_from_convex(self.vertices) - def get_distance_to_point(self, queries: np.ndarray) -> np.ndarray: - """Compute the signed distance of query points from the convex hull. + @cached_property + def equations(self) -> np.ndarray: + """Normalized equations of the edges. + """ + return compute_equations_from_convex(-self.vertices, self.vectors) - Positive distance corresponds to a query point lying outside the convex - hull. + @cached_property + def center(self) -> np.ndarray: + """Barycenter. .. note:: - Degenerated convex hulls corresponding to len(points) == 1 or 2 are - handled separately. The distance from a point and a segment is used - respectively. + The barycenter must be distinguished from the Chebyshev center, + which is defined as the center of the largest circle inscribed in + the polyhedron. Computing the latter involves solving a Linear + Program, which is known to have a unique solution that can always + be found in finite time. However, it is several order of magnitude + to compute than the barycenter. For details about the Chebyshev + center, see `compute_convex_chebyshev_center`. + """ + return np.mean(self.vertices, axis=0) - .. warning:: - This method only supports 2D space for the non-degenerated case. + def get_distance_to_point(self, points: np.ndarray) -> np.ndarray: + """Compute the signed distance of a single or a batch of query points + from the convex hull. - :param queries: N-D query points for which to compute distance from the - convex hull, as a 2D array. + Positive distance corresponds to a query point lying outside the convex + hull. - :returns: 1D float vector of the same length than `queries`. + .. note:: + Degenerated convex hulls are handled separately. The distance from + a point and a segment is used respectively. + + :param points: Set of 2D points for which to compute the distance from + the convex hull, as a 2D array whose first dimension + corresponds to the individual query points while the + second dimensions gathers the 2 position coordinates. + Note that the first dimension can be omitted if there is + a single query point. """ - if len(self._points) > 2: - # Compute the signed distance between query points and convex hull - if self._points.shape[1] != 2: - raise NotImplementedError + # Make sure that the input is at least 2D + if points.ndim < 2: + points = np.atleast_2d(points) + + # Compute the signed distance between query points and convex hull + if self.npoints > 2: return compute_distance_convex_to_point( - self._points, self._vertex_indices, queries) + self.vertices, self.vectors, points) - if len(self._points) == 2: - # Compute the distance between query points and segment - vec = self._points[1] - self._points[0] - ratio = (queries - self._points[0]) @ vec / np.dot(vec, vec) - proj = self._points[0] + np.outer(np.clip(ratio, 0.0, 1.0), vec) - return np.linalg.norm(queries - proj, 2, axis=1) + # Compute the distance between query points and segment + if self.npoints == 2: + vec = self.vertices[1] - self.vertices[0] + ratio = (points - self.vertices[0]) @ vec / np.dot(vec, vec) + proj = self.vertices[0] + np.outer(np.clip(ratio, 0.0, 1.0), vec) + return np.linalg.norm(points - proj, axis=1) # Compute the distance between query points and point - return np.linalg.norm(queries - self._points, 2, axis=1) + return np.linalg.norm(points - self.vertices, axis=1) def get_distance_to_ray(self, - query_vector: np.ndarray, - query_origin: np.ndarray) -> np.ndarray: - """Compute the distance of single ray from the convex hull. - - .. warning:: - It is assumed that the query origins are lying inside the convex - hull. + vector: np.ndarray, + origin: np.ndarray) -> np.ndarray: + """Compute the signed distance of single oriented ray from this convex + hull. - .. warning:: - This method only supports 2D space. + The distance is negative if the origin of the ray lies inside the + convex hull. The distance is 'inf' is there is no intersection between + the oriented ray and the convex hull. .. warning:: - Degenerated convex hulls corresponding to len(points) == 1 or 2 are - not supported. + Degenerated convex hulls are not supported. - :param query_vector: Direction of the ray. - :param query_origin: Origin of the ray. + :param vector: Direction in which the ray is casting, as a 1D array. + This vector does not have to be normalized. + :param origin: Position from which the ray is casting, as a 1D array. """ - if len(self._points) < 3: - raise NotImplementedError - if self._points.shape[1] != 2: + if self.npoints < 3: raise NotImplementedError return compute_distance_convex_to_ray( - self._points, self._vertex_indices, query_vector, query_origin) + self.vertices, self.vectors, vector, origin) + + def get_distance_to_convex(self, other: "ConvexHull2D") -> float: + """Compute the distance between two convex hulls in 2D Euclidean space. + + :param other: Convex hull from which to compute the distance wrt the + one characterized by this instance. + """ + return compute_distance_convex_to_convex( + self.vertices, self.vectors, other.vertices, other.vectors) + + def plot(self) -> Figure: + """Display the original points along their convex hull. + """ + # Make sure matplotlib is available + try: + # pylint: disable=import-outside-toplevel + import matplotlib.pyplot as plt + except ImportError as e: + raise ImportError("Matplotlib library not available. Please " + "install it before calling this method.") from e + + # Create new figure and return it to the user + fig = plt.figure() + ax = fig.add_subplot(111) + ax.set_aspect('equal') + plt.plot(*self.points.T, 'o') + plt.plot(*np.stack(( + self.vertices.T, + np.roll(self.vertices, 1, axis=0).T), axis=-1), 'k-') + plt.show(block=False) + + # Show figure, without blocking for interactive python sessions only + if interactive_mode() < 2: + fig.show() + + return fig diff --git a/python/gym_jiminy/toolbox/gym_jiminy/toolbox/quantities/__init__.py b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/quantities/__init__.py new file mode 100644 index 000000000..bd2e8c3e0 --- /dev/null +++ b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/quantities/__init__.py @@ -0,0 +1,9 @@ +# pylint: disable=missing-module-docstring + +from .locomotion import (ProjectedSupportPolygon, + StabilityMarginProjectedSupportPolygon) + +__all__ = [ + "ProjectedSupportPolygon", + "StabilityMarginProjectedSupportPolygon" +] diff --git a/python/gym_jiminy/toolbox/gym_jiminy/toolbox/quantities/locomotion.py b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/quantities/locomotion.py new file mode 100644 index 000000000..838917c88 --- /dev/null +++ b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/quantities/locomotion.py @@ -0,0 +1,220 @@ +"""Quantities mainly relevant for locomotion tasks on floating-base robots. +""" +from typing import List, Optional, Tuple +from dataclasses import dataclass + +import numpy as np +from scipy.spatial import ConvexHull + +from jiminy_py.core import ( # pylint: disable=no-name-in-module + multi_array_copyto) +import pinocchio as pin + +from gym_jiminy.common.bases import ( + InterfaceJiminyEnv, InterfaceQuantity, AbstractQuantity, QuantityEvalMode) +from gym_jiminy.common.quantities import ( + BaseOdometryPose, ZeroMomentPoint, translate_position_odom) + +from ..math import ConvexHull2D + + +@dataclass(unsafe_hash=True) +class ProjectedSupportPolygon(AbstractQuantity[ConvexHull2D]): + """Projected support polygon of the robot. + + The projected support polygon is defined as the 2D convex hull of all the + candidate contact points. It has the major advantage to be agnostic to the + contact state, unlike the true support polygon. This criterion can be + viewed as an anticipation of future impact. This decouples vertical foot + landing and timings from horizontal foot placement in world plane, which + makes it easier to derive additive reward components. + + .. note:: + This quantity is only supported for robots with specified contact + points but no collision bodies for now. + """ + + reference_frame: pin.ReferenceFrame + """Whether to compute the projected support polygon in local frame + specified by the odometry pose of floating base of the robot or the frame + located on the position of the floating base with axes kept aligned with + world frame. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + parent: Optional[InterfaceQuantity], + *, + reference_frame: pin.ReferenceFrame = pin.LOCAL, + mode: QuantityEvalMode = QuantityEvalMode.TRUE) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param parent: Higher-level quantity from which this quantity is a + requirement if any, `None` otherwise. + :param reference_frame: Whether to compute the support polygon in local + odometry frame, aka 'pin.LOCAL', or aligned + with world axes, aka 'pin.LOCAL_WORLD_ALIGNED'. + :param mode: Desired mode of evaluation for this quantity. + Optional: 'QuantityEvalMode.TRUE' by default. + """ + # Make sure at requested reference frame is supported + if reference_frame not in (pin.LOCAL, pin.LOCAL_WORLD_ALIGNED): + raise ValueError("Reference frame must be either 'pin.LOCAL' or " + "'pin.LOCAL_WORLD_ALIGNED'.") + + # Backup some user argument(s) + self.reference_frame = reference_frame + + # Call base implementation + super().__init__( + env, + parent, + requirements=dict( + odom_pose=(BaseOdometryPose, dict(mode=mode)) + ), + mode=mode, + auto_refresh=False) + + # Stacked position in world plane of all the candidate contact points + self._candidate_xy_batch = np.array([]) + + # Define proxy for views of individual position vectors in the stack + self._candidate_xy_views: Tuple[np.ndarray, ...] = () + + # Define proxy for positions in world plane of candidate contact points + self._candidate_xy_refs: Tuple[np.ndarray, ...] = () + + def initialize(self) -> None: + # Call base implementation + super().initialize() + + # Make sure that the robot has contact points but no collision bodies. + # This must be done at runtime as it may change dynamically. + if self.robot.collision_body_names: + raise RuntimeError( + "Robots having collision bodies are not supported for now.") + if not self.robot.contact_frame_names: + raise RuntimeError( + "Only robots having at least one contact are supported.") + + # Get all groups of contact points having the same parent body + contact_joint_indices: List[int] = [] + contact_positions: List[List[np.ndarray]] = [] + for frame_index in self.robot.contact_frame_indices: + frame = self.robot.pinocchio_model.frames[frame_index] + joint_index = frame.parent + try: + contact_index = contact_joint_indices.index(joint_index) + except ValueError: + contact_index = len(contact_joint_indices) + contact_joint_indices.append(joint_index) + contact_positions.append([]) + transform = self.robot.pinocchio_data.oMf[frame_index] + contact_positions[contact_index].append(transform.translation) + + # Filter out candidate contact points that will never be part of the + # projected support polygon no matter what, then gather their position + # in world plane in a single list. + # One can show that computing of 3D convex hull of the 3D volume before + # computing the 2D convex hull of its projection on a given plan yields + # to the exact same result. This 2-steps process is advantageous over, + # as the first step can be done only once, and computing the 2D convex + # hull will be faster if there are least points to consider. + # Note that this procedure is only applicable for fixed 3D volume. This + # implies that the convex hull of each contact group must be computed + # separately rather than all at once. + candidate_xy_refs: List[np.ndarray] = [] + for positions in contact_positions: + convhull = ConvexHull(np.stack(positions, axis=0)) + candidate_indices = set( + range(len(positions))).intersection(convhull.vertices) + candidate_xy_refs += ( + positions[j][:2] for j in candidate_indices) + self._candidate_xy_refs = tuple(candidate_xy_refs) + + # Allocate memory for stacked position of candidate contact points. + # Note that Fortran memory layout (row-major) is used for speed up + # because it preserves contiguity when copying frame data, and because + # the `ConvexHull2D` would perform one extra copy otherwise. + self._candidate_xy_batch = np.empty( + (len(self._candidate_xy_refs), 2), order="C") + + # Refresh proxies + self._candidate_xy_views = tuple(self._candidate_xy_batch) + + def refresh(self) -> ConvexHull2D: + # Copy all translations in contiguous buffer + multi_array_copyto(self._candidate_xy_views, self._candidate_xy_refs) + + # Translate candidate contact points from world to local odometry frame + if self.reference_frame == pin.LOCAL: + translate_position_odom(self._candidate_xy_batch, + self.odom_pose, + self._candidate_xy_batch) + + # Compute the 2D convex hull in world plane + return ConvexHull2D(self._candidate_xy_batch) + + +@dataclass(unsafe_hash=True) +class StabilityMarginProjectedSupportPolygon(InterfaceQuantity[float]): + """Signed distance of the Zero Moment Point (ZMP) from the borders of the + projected support polygon on world plane. + + The distance is positive if the ZMP lies inside the support polygon, + negative otherwise. + + To keep balance on flat ground, it is sufficient to maintain the ZMP inside + the projected support polygon at all time. In addition, the larger the + distance from the borders of the support polygon, the more "stable" the + robot. This means that it can withstand larger external impulse forces in + any direction before starting tiling around an edge of the convex hull. + + .. seealso:: + For an primer about the most common stability assessment criteria for + legged robots, see "Learning and Optimization of the Locomotion with an + Exoskeleton for Paraplegic People", A. Duburcq, 2022, Chap.2.2.2, p.33. + + .. note:: + This quantity is only supported for robots with specified contact + points but no collision bodies for now. + """ + + mode: QuantityEvalMode + """Specify on which state to evaluate this quantity. See `Mode` + documentation for details about each mode. + + .. warning:: + Mode `REFERENCE` requires a reference trajectory to be selected + manually prior to evaluating this quantity for the first time. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + parent: Optional[InterfaceQuantity], + *, + mode: QuantityEvalMode = QuantityEvalMode.TRUE) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param parent: Higher-level quantity from which this quantity is a + requirement if any, `None` otherwise. + :param mode: Desired mode of evaluation for this quantity. + """ + # Backup some user argument(s) + self.mode = mode + + # Call base implementation + super().__init__( + env, + parent, + requirements=dict( + support_polygon=(ProjectedSupportPolygon, dict( + reference_frame=pin.LOCAL_WORLD_ALIGNED, + mode=mode)), + zmp=(ZeroMomentPoint, dict( + reference_frame=pin.LOCAL_WORLD_ALIGNED, + mode=mode))), + auto_refresh=False) + + def refresh(self) -> float: + return - self.support_polygon.get_distance_to_point(self.zmp).item() diff --git a/python/gym_jiminy/toolbox/gym_jiminy/toolbox/wrappers/frame_rate_limiter.py b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/wrappers/frame_rate_limiter.py index 828c16a67..4dfa81737 100644 --- a/python/gym_jiminy/toolbox/gym_jiminy/toolbox/wrappers/frame_rate_limiter.py +++ b/python/gym_jiminy/toolbox/gym_jiminy/toolbox/wrappers/frame_rate_limiter.py @@ -10,7 +10,6 @@ from jiminy_py.viewer import sleep from gym_jiminy.common.bases import ObsT, ActT, InfoType, InterfaceJiminyEnv -from gym_jiminy.common.envs import BaseJiminyEnv class FrameRateLimiter(gym.Wrapper, # [ObsT, ActT, ObsT, ActT], @@ -44,7 +43,6 @@ def __init__(self, # pylint: disable=unused-argument self.human_only = human_only # Extract proxies for convenience - assert isinstance(env.unwrapped, BaseJiminyEnv) self._step_dt_rel = env.unwrapped.step_dt / speed_ratio # Buffer to keep track of the last time `step` method was called diff --git a/python/gym_jiminy/toolbox/setup.py b/python/gym_jiminy/toolbox/setup.py index db22695aa..8c44f04d1 100644 --- a/python/gym_jiminy/toolbox/setup.py +++ b/python/gym_jiminy/toolbox/setup.py @@ -32,12 +32,11 @@ package_data={"gym_jiminy.toolbox": ["py.typed"]}, install_requires=[ f"gym-jiminy~={gym_jiminy_version}", - # Used to compute convex hull. + # Used to compute 3D convex hulls (using custom implementation in 2D), + # and to solve the Linear Program for finding their Chebyshev center. # No wheel is distributed on pypi for PyPy, and pip requires to install # `libatlas-base-dev` system dependency to build it from source. - # 1.8.0: `scipy.spatial.qhull` low-level API changes. - # 1.9.2: First release to support Python 3.11 - "scipy>=1.9.2" + "scipy" ], zip_safe=False ) diff --git a/python/gym_jiminy/unit_py/test_misc.py b/python/gym_jiminy/unit_py/test_misc.py index 63a7a41be..8baac24de 100644 --- a/python/gym_jiminy/unit_py/test_misc.py +++ b/python/gym_jiminy/unit_py/test_misc.py @@ -190,4 +190,5 @@ def test_math(self): np.testing.assert_allclose( rpy_i, pin.rpy.matrixToRpy(rot_mat_i.reshape((3, 3)))) np.testing.assert_allclose(math.rpy_to_quat(rpy), quat) - np.testing.assert_allclose(math.matrix_to_quat(rot_mat), quat) + np.testing.assert_allclose( + math.matrix_to_quat(rot_mat), quat, atol=1e-9) diff --git a/python/gym_jiminy/unit_py/test_pipeline_control.py b/python/gym_jiminy/unit_py/test_pipeline_control.py index c34202b7d..103d0131f 100644 --- a/python/gym_jiminy/unit_py/test_pipeline_control.py +++ b/python/gym_jiminy/unit_py/test_pipeline_control.py @@ -49,6 +49,7 @@ def _test_pid_standing(self): # Run the simulation while self.env.stepper_state.t < 9.0: self.env.step(action) + self.env.stop() # Export figure fd, pdf_path = mkstemp(prefix="plot_", suffix=".pdf") @@ -263,6 +264,7 @@ def test_pd_controller(self): env.unwrapped._height_neutral = float("-inf") while env.stepper_state.t < 2.0: env.step(0.2 * env.action_space.sample()) + env.stop() # Extract the target position and velocity of a single motor adapter_name, controller_name = adapter.name, controller.name @@ -284,9 +286,9 @@ def test_pd_controller(self): command_vel[(update_ratio-1)::update_ratio], atol=TOLERANCE) np.testing.assert_allclose( - target_accel_diff, target_accel[1:], atol=TOLERANCE) + target_accel_diff[:-1], target_accel[1:-1], atol=TOLERANCE) np.testing.assert_allclose( - target_vel_diff, target_vel[1:], atol=TOLERANCE) + target_vel_diff[:-1], target_vel[1:-1], atol=TOLERANCE) # Make sure that the position and velocity targets are within bounds motor = env.robot.motors[-1] diff --git a/python/gym_jiminy/unit_py/test_pipeline_design.py b/python/gym_jiminy/unit_py/test_pipeline_design.py index 0c39e76aa..973a18112 100644 --- a/python/gym_jiminy/unit_py/test_pipeline_design.py +++ b/python/gym_jiminy/unit_py/test_pipeline_design.py @@ -246,6 +246,7 @@ def configure_telemetry() -> InterfaceJiminyEnv: env.reset(seed=0, options=dict(reset_hook=configure_telemetry)) env.step(env.action) + env.stop() controller = env.env.env.env.controller assert isinstance(controller, PDController) diff --git a/python/gym_jiminy/unit_py/test_quantities.py b/python/gym_jiminy/unit_py/test_quantities.py index 1290b1045..b57175047 100644 --- a/python/gym_jiminy/unit_py/test_quantities.py +++ b/python/gym_jiminy/unit_py/test_quantities.py @@ -30,7 +30,8 @@ ActuatedJointsPosition, CenterOfMass, CapturePoint, - ZeroMomentPoint) + ZeroMomentPoint, + MultiContactRelativeForceTangential) class Quantities(unittest.TestCase): @@ -523,3 +524,24 @@ def test_foot_relative_pose(self): np.testing.assert_allclose(value[:3], pos_rel[:, :-1]) np.testing.assert_allclose(value[-4:], quat_rel[:, :-1]) + + def test_tangential_forces(self): + """ TODO: Write documentation + """ + env = gym.make("gym_jiminy.envs:atlas") + + env.quantities["force_tangential_rel"] = ( + MultiContactRelativeForceTangential, {}) + + env.reset(seed=0) + for _ in range(10): + env.step(env.action) + + gravity = abs(env.robot.pinocchio_model.gravity.linear[2]) + robot_mass = env.robot.pinocchio_data.mass[0] + force_tangential_rel = np.stack(tuple( + constraint.lambda_c[:2] + for constraint in env.robot.constraints.contact_frames.values()), + axis=-1) / (robot_mass * gravity) + np.testing.assert_allclose( + force_tangential_rel, env.quantities["force_tangential_rel"]) diff --git a/python/gym_jiminy/unit_py/test_rewards.py b/python/gym_jiminy/unit_py/test_rewards.py index 0931e6dfc..096311b1f 100644 --- a/python/gym_jiminy/unit_py/test_rewards.py +++ b/python/gym_jiminy/unit_py/test_rewards.py @@ -9,15 +9,19 @@ from jiminy_py.log import extract_trajectory_from_log from gym_jiminy.common.compositions import ( + CUTOFF_ESP, TrackingActuatedJointPositionsReward, TrackingBaseOdometryVelocityReward, TrackingBaseHeightReward, TrackingCapturePointReward, TrackingFootPositionsReward, TrackingFootOrientationsReward, + MinimizeFrictionReward, SurviveReward, - MinimizeAngularMomentumReward, AdditiveMixtureReward) +from gym_jiminy.toolbox.compositions import ( + tanh_normalization, + MaximizeStability) class Rewards(unittest.TestCase): @@ -71,9 +75,10 @@ def test_tracking(self): np.testing.assert_allclose( quantity_true.get(), self.env.robot_state.q[2]) - gamma = - np.log(0.01) / cutoff ** 2 + gamma = - np.log(CUTOFF_ESP) / cutoff ** 2 value = np.exp(- gamma * np.sum((reward.quantity.op( quantity_true.get(), quantity_ref.get())) ** 2)) + assert value > 0.01 np.testing.assert_allclose(reward(terminated, {}), value) del reward @@ -104,3 +109,37 @@ def test_mixture(self): assert reward_sum(terminated, {}) == ( 0.5 * reward_odometry(terminated, {}) + 0.2 * reward_survive(terminated, {})) + + def test_stability(self): + CUTOFF_INNER, CUTOFF_OUTER = 0.1, 0.5 + reward_stability = MaximizeStability( + self.env, cutoff_inner=0.1, cutoff_outer=0.5) + quantity = reward_stability.quantity + + self.env.reset(seed=0) + action = self.env.action_space.sample() + _, _, terminated, _, _ = self.env.step(action) + + dist = quantity.support_polygon.get_distance_to_point(quantity.zmp) + value = tanh_normalization(dist.item(), -CUTOFF_INNER, CUTOFF_OUTER) + np.testing.assert_allclose(tanh_normalization( + -CUTOFF_INNER, -CUTOFF_INNER, CUTOFF_OUTER), 1.0 - CUTOFF_ESP) + np.testing.assert_allclose(tanh_normalization( + CUTOFF_OUTER, -CUTOFF_INNER, CUTOFF_OUTER), CUTOFF_ESP) + np.testing.assert_allclose(reward_stability(terminated, {}), value) + + def test_friction(self): + CUTOFF = 0.5 + env = gym.make("gym_jiminy.envs:atlas-pid", debug=True) + reward_friction = MinimizeFrictionReward(env, cutoff=CUTOFF) + quantity = reward_friction.quantity + + env.reset(seed=0) + _, _, terminated, _, _ = env.step(env.action) + force_tangential_rel = quantity.get() + force_tangential_rel_norm = np.sum(np.square(force_tangential_rel)) + + gamma = - np.log(CUTOFF_ESP) / CUTOFF ** 2 + value = np.exp(- gamma * force_tangential_rel_norm) + assert value > 0.01 + np.testing.assert_allclose(reward_friction(terminated, {}), value) diff --git a/python/jiminy_py/unit_py/utilities.py b/python/jiminy_py/unit_py/utilities.py index 28dd8055b..6c7a182c8 100644 --- a/python/jiminy_py/unit_py/utilities.py +++ b/python/jiminy_py/unit_py/utilities.py @@ -167,7 +167,7 @@ def integrate_dynamics(time: np.ndarray, dynamics(t: float, x: np.ndarray) -> np:ndarray - :return: 2D array for which the i-th line is the solution x at `time[i]`. + :returns: 2D array for which the i-th line is the solution x at `time[i]`. """ solver = ode(dynamics) solver.set_initial_value(x0, t=time[0])