From eca32f5624de0c7f3e6c2777aa7cd993b807ab97 Mon Sep 17 00:00:00 2001 From: Alexis DUBURCQ Date: Tue, 18 Jun 2024 18:53:53 +0200 Subject: [PATCH] [gym/common] Add foot vertical forces quantity. (#810) * [core] Rename 'runge_kutta_dopri5' in 'runge_kutta_dopri' for clarity. * [gym/common] Add foot vertical forces quantity. * [gym/common] Add foot force distribution reward. * [gym/common] Do not pre-allocate memory for tracking op as it makes no difference. --- .../double_pendulum/double_pendulum.cc | 2 +- core/include/jiminy/core/engine/engine.h | 6 +- core/src/engine/engine.cc | 2 +- .../gym_jiminy/common/bases/quantities.py | 3 +- .../common/compositions/__init__.py | 2 + .../common/compositions/locomotion.py | 68 +++- .../gym_jiminy/common/quantities/__init__.py | 2 + .../gym_jiminy/common/quantities/generic.py | 7 +- .../common/quantities/locomotion.py | 353 +++++++++++++++--- python/gym_jiminy/unit_py/test_quantities.py | 33 +- python/gym_jiminy/unit_py/test_rewards.py | 14 +- .../box_uneven_ground_impulse_contact.py | 2 +- python/jiminy_py/examples/constraint_wheel.py | 2 +- python/jiminy_py/examples/double_pendulum.py | 4 +- python/jiminy_py/examples/force_coupling.py | 2 +- python/jiminy_py/examples/sphere_rolling.py | 2 +- python/jiminy_py/src/jiminy_py/dynamics.py | 5 +- .../unit_py/data/double_pendulum_options.toml | 2 +- .../unit_py/test_double_spring_mass.py | 14 +- python/jiminy_py/unit_py/test_multi_robot.py | 2 +- .../jiminy_py/unit_py/test_simple_pendulum.py | 8 +- 21 files changed, 429 insertions(+), 106 deletions(-) diff --git a/core/examples/double_pendulum/double_pendulum.cc b/core/examples/double_pendulum/double_pendulum.cc index 5b6cf434c..491771600 100644 --- a/core/examples/double_pendulum/double_pendulum.cc +++ b/core/examples/double_pendulum/double_pendulum.cc @@ -101,7 +101,7 @@ int main(int /* argc */, char * /* argv */[]) GenericConfig & worldOptions = boost::get(simuOptions.at("world")); boost::get(worldOptions.at("gravity"))[2] = -9.81; GenericConfig & stepperOptions = boost::get(simuOptions.at("stepper")); - boost::get(stepperOptions.at("odeSolver")) = std::string("runge_kutta_dopri5"); + boost::get(stepperOptions.at("odeSolver")) = std::string("runge_kutta_dopri"); boost::get(stepperOptions.at("tolRel")) = 1.0e-5; boost::get(stepperOptions.at("tolAbs")) = 1.0e-4; boost::get(stepperOptions.at("dtMax")) = 3.0e-3; diff --git a/core/include/jiminy/core/engine/engine.h b/core/include/jiminy/core/engine/engine.h index 3aa1fa410..a5730c1d4 100644 --- a/core/include/jiminy/core/engine/engine.h +++ b/core/include/jiminy/core/engine/engine.h @@ -35,7 +35,7 @@ namespace jiminy const std::map CONSTRAINT_SOLVERS_MAP{ {"PGS", ConstraintSolverType::PGS}}; - const std::set STEPPERS{"euler_explicit", "runge_kutta_4", "runge_kutta_dopri5"}; + const std::set STEPPERS{"euler_explicit", "runge_kutta_4", "runge_kutta_dopri"}; class Robot; class AbstractConstraintSolver; @@ -306,8 +306,8 @@ namespace jiminy GenericConfig config; config["verbose"] = false; config["randomSeedSeq"] = VectorX::Zero(1).eval(); - /// \details Must be either "runge_kutta_dopri5", "runge_kutta_4" or "euler_explicit". - config["odeSolver"] = std::string{"runge_kutta_dopri5"}; + /// \details Must be either "runge_kutta_dopri", "runge_kutta_4" or "euler_explicit". + config["odeSolver"] = std::string{"runge_kutta_dopri"}; config["tolAbs"] = 1.0e-5; config["tolRel"] = 1.0e-4; config["dtMax"] = SIMULATION_MAX_TIMESTEP; diff --git a/core/src/engine/engine.cc b/core/src/engine/engine.cc index c7d5adb59..3ffb10d08 100644 --- a/core/src/engine/engine.cc +++ b/core/src/engine/engine.cc @@ -1149,7 +1149,7 @@ namespace jiminy robots_.end(), std::back_inserter(robots), [](const auto & robot) { return robot.get(); }); - if (engineOptions_->stepper.odeSolver == "runge_kutta_dopri5") + if (engineOptions_->stepper.odeSolver == "runge_kutta_dopri") { stepper_ = std::unique_ptr(new RungeKuttaDOPRIStepper( robotOde, robots, engineOptions_->stepper.tolAbs, engineOptions_->stepper.tolRel)); diff --git a/python/gym_jiminy/common/gym_jiminy/common/bases/quantities.py b/python/gym_jiminy/common/gym_jiminy/common/bases/quantities.py index 37876ae65..352b188a7 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/bases/quantities.py +++ b/python/gym_jiminy/common/gym_jiminy/common/bases/quantities.py @@ -1145,7 +1145,8 @@ def initialize(self) -> None: self.env.robot_state.a, self.env.robot_state.u, self.env.robot_state.command, - self._f_external_batch) + self._f_external_batch, + self._constraint_lambda_batch) def refresh(self) -> State: """Compute the current state depending on the mode of evaluation, and 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 a0ac75db7..c82060f4c 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/compositions/__init__.py +++ b/python/gym_jiminy/common/gym_jiminy/common/compositions/__init__.py @@ -12,6 +12,7 @@ TrackingCapturePointReward, TrackingFootPositionsReward, TrackingFootOrientationsReward, + TrackingFootForceDistributionReward, MinimizeAngularMomentumReward, MinimizeFrictionReward) @@ -27,6 +28,7 @@ "TrackingCapturePointReward", "TrackingFootPositionsReward", "TrackingFootOrientationsReward", + "TrackingFootForceDistributionReward", "MinimizeAngularMomentumReward", "MinimizeFrictionReward", "SurviveReward" 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 6932e895f..eaf10fde2 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/compositions/locomotion.py +++ b/python/gym_jiminy/common/gym_jiminy/common/compositions/locomotion.py @@ -1,7 +1,7 @@ """Rewards mainly relevant for locomotion tasks on floating-base robots. """ from functools import partial -from typing import Union, Sequence, Literal +from typing import Union, Sequence, Literal, Callable, cast import numpy as np import pinocchio as pin @@ -11,7 +11,7 @@ from ..quantities import ( MaskedQuantity, UnaryOpQuantity, AverageBaseOdometryVelocity, CapturePoint, MultiFootRelativeXYZQuat, MultiContactRelativeForceTangential, - AverageBaseMomentum) + MultiFootRelativeForceVertical, AverageBaseMomentum) from ..quantities.locomotion import sanitize_foot_frame_names from ..utils import quat_difference @@ -169,20 +169,6 @@ def __init__(self, # Sanitize frame names corresponding to the feet of the robot frame_names = tuple(sanitize_foot_frame_names(env, frame_names)) - # Buffer storing the difference before current and reference poses - # FIXME: Is it worth it to create a temporary ? - self._diff = np.zeros((3, len(frame_names) - 1)) - - # Define buffered quaternion difference operator for efficiency - def quat_difference_buffered(out: np.ndarray, - q1: np.ndarray, - q2: np.ndarray) -> np.ndarray: - """Wrapper around `quat_difference` passing buffer in and out - instead of allocating fresh memory for efficiency. - """ - quat_difference(q1, q2, out) - return out - # Call base implementation super().__init__( env, @@ -194,7 +180,55 @@ def quat_difference_buffered(out: np.ndarray, axis=0, keys=(3, 4, 5, 6))), cutoff, - op=partial(quat_difference_buffered, self._diff)) + op=cast(Callable[ + [np.ndarray, np.ndarray], np.ndarray], quat_difference)) + + +class TrackingFootForceDistributionReward(BaseTrackingReward): + """Reward the agent for tracking the relative vertical force in world frame + applied on each foot. + + .. note:: + The force is normalized by the weight of the robot rather than the + total force applied on all feet. This is important as it not only takes + into account the force distribution between the feet, but also the + overall ground contact interact force. This way, building up momentum + before jumping will be distinguished for standing still. Moreover, it + ensures that the reward is always properly defined, even if the robot + has no contact with the ground at all, which typically arises during + the flying phase of running. + + .. seealso:: + See `BaseTrackingReward` documentation for technical details. + """ + def __init__(self, + env: InterfaceJiminyEnv, + cutoff: float, + *, + frame_names: Union[Sequence[str], Literal['auto']] = 'auto' + ) -> None: + """ + :param env: Base or wrapped jiminy environment. + :param cutoff: Cutoff threshold for the RBF kernel transform. + :param frame_names: Name of the frames corresponding to the feet of the + robot. 'auto' to automatically detect them from the + set of contact and force sensors of the robot. + Optional: 'auto' by default. + """ + # Backup some user argument(s) + self.cutoff = cutoff + + # Sanitize frame names corresponding to the feet of the robot + frame_names = tuple(sanitize_foot_frame_names(env, frame_names)) + + # Call base implementation + super().__init__( + env, + "reward_tracking_foot_force_distribution", + lambda mode: (MultiFootRelativeForceVertical, dict( + frame_names=frame_names, + mode=mode)), + cutoff) class MinimizeAngularMomentumReward(BaseQuantityReward): 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 416a5836c..dbd4ae526 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/quantities/__init__.py +++ b/python/gym_jiminy/common/gym_jiminy/common/quantities/__init__.py @@ -24,6 +24,7 @@ AverageBaseMomentum, MultiFootMeanXYZQuat, MultiFootRelativeXYZQuat, + MultiFootRelativeForceVertical, MultiContactRelativeForceTangential, CenterOfMass, CapturePoint, @@ -48,6 +49,7 @@ 'MultiFootMeanXYZQuat', 'MultiFootRelativeXYZQuat', 'MultiFootMeanOdometryPose', + 'MultiFootRelativeForceVertical', 'MultiContactRelativeForceTangential', 'AverageFrameXYZQuat', 'AverageFrameRollPitch', 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 6970034b9..3e6a8699a 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/quantities/generic.py +++ b/python/gym_jiminy/common/gym_jiminy/common/quantities/generic.py @@ -18,7 +18,8 @@ import pinocchio as pin from ..bases import ( - InterfaceJiminyEnv, InterfaceQuantity, AbstractQuantity, QuantityEvalMode) + InterfaceJiminyEnv, InterfaceQuantity, AbstractQuantity, StateQuantity, + QuantityEvalMode) from ..utils import ( matrix_to_rpy, matrix_to_quat, quat_apply, remove_yaw_from_quat, quat_interpolate_middle) @@ -1404,7 +1405,9 @@ def __init__(self, super().__init__( env, parent, - requirements={}, + requirements=dict( + state=(StateQuantity, dict( + update_kinematics=False))), mode=mode, auto_refresh=False) 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 0ba1d3b20..45db50258 100644 --- a/python/gym_jiminy/common/gym_jiminy/common/quantities/locomotion.py +++ b/python/gym_jiminy/common/gym_jiminy/common/quantities/locomotion.py @@ -20,8 +20,8 @@ matrix_to_yaw, quat_to_yaw, quat_to_matrix, quat_multiply, quat_apply) from ..quantities import ( - MaskedQuantity, AverageFrameSpatialVelocity, MultiFramesXYZQuat, - MultiFramesMeanXYZQuat, AverageFrameRollPitch) + MaskedQuantity, MultiFramesXYZQuat, MultiFramesMeanXYZQuat, + AverageFrameSpatialVelocity, AverageFrameRollPitch) def sanitize_foot_frame_names( @@ -343,7 +343,7 @@ def initialize(self) -> None: self._inertia_local = self.pinocchio_model.inertias[1].inertia def refresh(self) -> np.ndarray: - # Compute the local angular angular momentum of inertia + # Compute the local angular momentum of inertia np.matmul(self._inertia_local, self.v_angular, self._h_angular) # Apply quaternion rotation of the local angular momentum of inertia @@ -676,12 +676,13 @@ def initialize(self) -> None: # Call base implementation super().initialize() - # Make sure that the state is consistent with required kinematic level + # Make sure that the state data meet requirements state = self.state if ((self.kinematic_level == pin.ACCELERATION and state.a is None) or (self.kinematic_level >= pin.VELOCITY and state.v is None)): raise RuntimeError( - "State data inconsistent with required kinematic level") + "Available state data do not meet requirements for kinematic " + f"level '{self.kinematic_level}'.") # Refresh CoM quantity proxy based on kinematic level if self.kinematic_level == pin.POSITION: @@ -771,10 +772,11 @@ def initialize(self) -> None: # Call base implementation super().initialize() - # Make sure that the state is consistent with required kinematic level - if (self.state.v is None or self.state.a is None): + # Make sure that the state data meet requirements + if self.state.v is None or self.state.a is None: raise RuntimeError( - "State data inconsistent with required kinematic level") + "State data do not meet requirements. Velocity and " + "acceleration are missing.") # Compute the weight of the robot gravity = abs(self.pinocchio_model.gravity.linear[2]) @@ -875,10 +877,10 @@ def initialize(self) -> None: # Call base implementation super().initialize() - # Make sure that the state is consistent with required kinematic level + # Make sure that the state data meet requirements if self.state.v is None: raise RuntimeError( - "State data inconsistent with required kinematic level") + "State data do not meet requirements. Velocity is missing.") # Compute the natural frequency of linear pendulum approximate model. # Note that the height of the robot is defined as the position of the @@ -911,9 +913,13 @@ def refresh(self) -> np.ndarray: @dataclass(unsafe_hash=True) -class MultiContactRelativeForceTangential(InterfaceQuantity[np.ndarray]): +class MultiContactRelativeForceTangential(AbstractQuantity[np.ndarray]): """Standardized tangential forces apply on all contact points and collision - bodies. + bodies in their respective local contact frame. + + The local contact frame is defined as the frame having the normal of the + ground as vertical axis, and the vector orthogonal to the x-axis in world + frame as y-axis. 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 @@ -930,7 +936,9 @@ class MultiContactRelativeForceTangential(InterfaceQuantity[np.ndarray]): def __init__(self, env: InterfaceJiminyEnv, - parent: Optional[InterfaceQuantity]) -> None: + 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 @@ -938,59 +946,308 @@ def __init__(self, """ # Call base implementation super().__init__( - env, parent, requirements={}, auto_refresh=False) + env, + parent, + requirements=dict( + state=(StateQuantity, dict( + update_kinematics=False))), + mode=mode, + auto_refresh=False) - # Weight of the robot - self._robot_weight: float = -1 + # Define jit-able method compute the normalized tangential forces + @nb.jit(nopython=True, cache=True, fastmath=True) + def normalize_tangential_forces(lambda_c: np.ndarray, + index_start: int, + index_end: int, + robot_weight: float, + out: np.ndarray) -> None: + """Compute the tangential forces of all the constraints associated + with contact frames and collision bodies, normalized by the total + weight of the robot. + + :param lambda_c: Stacked lambda multipliers all the constraints. + :param index_start: First index of the constraints associated with + contact frames and collisions bodies. + :param index_end: One-past-last index of the constraints associated + with contact frames and collisions bodies. + :param robot_weight: Total weight of the robot which will be used + to rescale the tangential forces. + :param out: Pre-allocated array in which to store the result. + """ + # Extract constraint lambdas of contacts and collisions from state + lambda_ = lambda_c[index_start:index_end].reshape((-1, 4)).T - # Stacked tangential forces on all contact points and collision bodies - self._force_tangential_batch = np.array([]) + # Extract references to all the tangential forces + # f_lin, f_ang = lambda_[:3], np.array([0.0, 0.0, lambda_[3]]) + forces_tangential = lambda_[:2] + + # Scale the tangential forces by the weight of the robot + np.divide(forces_tangential, robot_weight, out) + + self._normalize_tangential_forces = normalize_tangential_forces - # Define proxy for views of individual tangential forces in the stack - self._force_tangential_views: Tuple[np.ndarray, ...] = () + # Weight of the robot + self._robot_weight: float = float("nan") + + # Slice of constraint lambda multipliers for contacts and collisions + self._contact_slice: Tuple[int, int] = (0, 0) - # Define proxy for the current tangential forces - self._force_tangential_refs: Tuple[np.ndarray, ...] = () + # Stacked tangential forces on all contact points and collision bodies + self._force_tangential_rel_batch = np.array([]) def initialize(self) -> None: # Call base implementation super().initialize() + # Make sure that the state data meet requirements + if self.state.lambda_c is None: + raise RuntimeError("State data do not meet requirements. " + "Constraints lambda multipliers are missing.") + # 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] + gravity = abs(self.pinocchio_model.gravity.linear[2]) + robot_mass = self.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(): + # Extract slice of constraints associated with contacts and collisions + index_first, index_last = None, None + for i, fieldname in enumerate(self.robot.log_constraint_fieldnames): + is_contact = any( + fieldname.startswith(f"Constraint{registry_type}") + for registry_type in ("ContactFrames", "CollisionBodies")) + if index_first is None: + if is_contact: + index_first = i + elif index_last is None: # type: ignore[unreachable] + if not is_contact: + index_last = i + elif is_contact: + raise ValueError( + "Constraints associated with contacts and collisions are " + "not continuously ordered.") + if index_last is None: + index_last = i + 1 + assert index_first is not None + self._contact_slice = (index_first, index_last) + + # Make sure that all contacts and collisions constraints are supported + for constraint in self.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 constraints_body in self.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) + # Make sure that the extracted slice is consistent with the constraints + num_contraints = len(self.robot.constraints.contact_frames) + sum( + map(len, self.robot.constraints.collision_bodies)) + assert 4 * num_contraints == index_last - index_first - # 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) + # Pre-allocated memory for stacked normalized tangential forces + self._force_tangential_rel_batch = np.zeros( + (2, num_contraints), order='F') def refresh(self) -> np.ndarray: - # Copy all tangential forces in contiguous buffer - multi_array_copyto(self._force_tangential_views, - self._force_tangential_refs) + self._normalize_tangential_forces( + self.state.lambda_c, + *self._contact_slice, + self._robot_weight, + self._force_tangential_rel_batch) + + return self._force_tangential_rel_batch + + +@dataclass(unsafe_hash=True) +class MultiFootRelativeForceVertical(AbstractQuantity[np.ndarray]): + """Standardized total vertical forces apply on each foot in world frame. + + The lambda multipliers of the contact constraints are used to compute the + total forces applied on each foot. Although relying on the total wrench + acting on their respective parent joint seems enticing, it aggregates all + external forces, not just the ground contact reaction forces. Most often, + there is no difference, but not in the case of multiple robots interacting + with each others, or if user-specified external forces are manually applied + on the foot, eg to create disturbances. Relying on sensors to get the + desired information is not an option either, because they do not give + access to the ground truth. + """ - # Scale the tangential forces by the weight of the robot - self._force_tangential_batch /= self._robot_weight + frame_names: Tuple[str, ...] + """Name of the frames corresponding to some feet of the robot. - return self._force_tangential_batch + These frames must be part of the end-effectors, ie being associated with a + leaf joint in the kinematic tree of the robot. + """ + + def __init__(self, + env: InterfaceJiminyEnv, + parent: Optional[InterfaceQuantity], + frame_names: Union[Sequence[str], Literal['auto']] = 'auto', + *, + 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 frame_names: Name of the frames corresponding to some feet of + the robot. 'auto' to automatically detect them from + the set of contact and force sensors of the robot. + Optional: 'auto' by default. + :param mode: Desired mode of evaluation for this quantity. + """ + # Backup some user argument(s) + self.frame_names = tuple(sanitize_foot_frame_names(env, frame_names)) + + # Call base implementation + super().__init__( + env, + parent, + requirements=dict( + state=(StateQuantity, dict( + update_kinematics=False))), + mode=mode, + auto_refresh=False) + + # Define jit-able method compute the normalized tangential forces + @nb.jit(nopython=True, cache=True, fastmath=True) + def normalize_vertical_forces( + lambda_c: np.ndarray, + foot_slices: Tuple[Tuple[int, int], ...], + vertical_transform_batches: Tuple[np.ndarray, ...], + robot_weight: float, + out: np.ndarray) -> None: + """Compute the sum of the vertical forces in world frame of all the + constraints associated with contact frames and collision bodies, + normalized by the total weight of the robot. + + :param lambda_c: Stacked lambda multipliers all the constraints. + :param foot_slices: Slices of lambda multiplier of the constraints + associated with contact frames and collisions + bodies acting each foot, as a sequence of pairs + (index_start, index_end) corresponding to the + first and one-past-last indices respectively. + :param vertical_transform_batches: + Last row of the rotation matrices from world to local contact + frame associated with all contact and collision constraints + acting on each foot, as a list of 2D arrays. The first + dimension gathers the 3 spatial coordinates while the second + corresponds to the N individual constraints on each foot. + :param robot_weight: Total weight of the robot which will be used + to rescale the tangential forces. + :param out: Pre-allocated array in which to store the result. + """ + for i, ((index_start, index_end), vertical_transforms) in ( + enumerate(zip(foot_slices, vertical_transform_batches))): + # Extract constraint multipliers from state + lambda_ = lambda_c[index_start:index_end].reshape((-1, 4)).T + + # Extract references to all the linear forces + # f_ang = np.array([0.0, 0.0, lambda_[3]]) + f_lin = lambda_[:3] + + # Compute vertical forces in world frame and aggregate them + f_z_world = np.sum(vertical_transforms * f_lin) + + # Scale the vertical forces by the weight of the robot + out[i] = f_z_world / robot_weight + + self._normalize_vertical_forces = normalize_vertical_forces + + # Weight of the robot + self._robot_weight: float = float("nan") + + # Slice of constraint lambda multipliers associated with each foot + self._foot_slices: Tuple[Tuple[int, int], ...] = () + + # Stacked vertical forces in (world frame) on each foot + self._vertical_force_batch = np.array([]) + + # Define proxies for vertical axis transform of each frame constraint + self._vertical_transform_list: Tuple[np.ndarray, ...] = () + + # Stacked vertical axis transforms + self._vertical_transform_batches: Tuple[np.ndarray, ...] = () + + # Define proxy for views of the batch storing vertical axis transforms + self._vertical_transform_views: Tuple[Tuple[np.ndarray, ...], ...] = () + + def initialize(self) -> None: + # Call base implementation + super().initialize() + + # Make sure that the state data meet requirements + if self.state.lambda_c is None: + raise RuntimeError("State data do not meet requirements. " + "Constraints lambda multipliers are missing.") + + # Compute the weight of the robot + gravity = abs(self.pinocchio_model.gravity.linear[2]) + robot_mass = self.pinocchio_data.mass[0] + self._robot_weight = robot_mass * gravity + + # Get the joint index corresponding to each foot frame + foot_joint_indices: List[int] = [] + for frame_name in self.frame_names: + frame_index = self.pinocchio_model.getFrameId(frame_name) + joint_index = self.pinocchio_model.frames[frame_index].parent + foot_joint_indices.append(joint_index) + + # Get constraint names and vertical axis transforms for each foot + num_contraints = 0 + foot_lookup_names: List[List[str]] = [[] for _ in self.frame_names] + vertical_transforms: List[List[np.ndarray]] = [ + [] for _ in self.frame_names] + constraint_lookup_pairs = ( + ("ContactFrames", self.robot.constraints.contact_frames), + ("CollisionBodies", { + name: constraint for constraints in ( + self.robot.constraints.collision_bodies) + for name, constraint in constraints.items()})) + for registry_type, registry in constraint_lookup_pairs: + for name, constraint in registry.items(): + assert isinstance(constraint, jiminy.FrameConstraint) + frame = self.pinocchio_model.frames[constraint.frame_index] + try: + foot_index = foot_joint_indices.index(frame.parent) + foot_lookup_names[foot_index] += ( + f"Constraint{registry_type}{name}{i}" + for i in range(constraint.size)) + vertical_transforms[foot_index].append( + constraint.local_rotation[2]) + num_contraints += 1 + except IndexError: + pass + assert 4 * num_contraints == sum(map(len, foot_lookup_names)) + self._vertical_transform_list = tuple( + e for values in vertical_transforms for e in values) + + # Extract constraint lambda multiplier slices associated with each foot + self._foot_slices = tuple( + (self.robot.log_constraint_fieldnames.index(lookup_names[0]), + self.robot.log_constraint_fieldnames.index(lookup_names[-1]) + 1) + for lookup_names in foot_lookup_names) + + # Pre-allocate memory for stacked vertical forces in world frame + self._vertical_force_batch = np.zeros((len(self.frame_names),)) + + # Pre-allocate memory for stacked vertical axis transforms + self._vertical_transform_batches = tuple( + np.zeros((3, num_foot_contacts), order='F') + for num_foot_contacts in map(len, vertical_transforms)) + + # Define proxy for views of the batch storing vertical axis transforms + self._vertical_transform_views = tuple( + e for values in self._vertical_transform_batches for e in values.T) + + def refresh(self) -> np.ndarray: + # Copy all vertical axis transforms in contiguous buffer + multi_array_copyto(self._vertical_transform_views, + self._vertical_transform_list) + + # Compute the normalized sum of the vertical forces in world frame + self._normalize_vertical_forces(self.state.lambda_c, + self._foot_slices, + self._vertical_transform_batches, + self._robot_weight, + self._vertical_force_batch) + + return self._vertical_force_batch diff --git a/python/gym_jiminy/unit_py/test_quantities.py b/python/gym_jiminy/unit_py/test_quantities.py index 81a1f9ce5..95ef04bca 100644 --- a/python/gym_jiminy/unit_py/test_quantities.py +++ b/python/gym_jiminy/unit_py/test_quantities.py @@ -31,7 +31,8 @@ CenterOfMass, CapturePoint, ZeroMomentPoint, - MultiContactRelativeForceTangential) + MultiContactRelativeForceTangential, + MultiFootRelativeForceVertical) class Quantities(unittest.TestCase): @@ -541,10 +542,36 @@ def test_tangential_forces(self): env.step(env.action) gravity = abs(env.robot.pinocchio_model.gravity.linear[2]) - robot_mass = env.robot.pinocchio_data.mass[0] + robot_weight = env.robot.pinocchio_data.mass[0] * gravity force_tangential_rel = np.stack(tuple( constraint.lambda_c[:2] for constraint in env.robot.constraints.contact_frames.values()), - axis=-1) / (robot_mass * gravity) + axis=-1) / robot_weight np.testing.assert_allclose( force_tangential_rel, env.quantities["force_tangential_rel"]) + + def test_vertical_forces(self): + """ TODO: Write documentation + """ + env = gym.make("gym_jiminy.envs:atlas-pid") + + env.quantities["force_vertical_rel"] = ( + MultiFootRelativeForceVertical, {}) + + env.reset(seed=0) + for _ in range(10): + env.step(env.action) + + gravity = abs(env.robot.pinocchio_model.gravity.linear[2]) + robot_weight = env.robot.pinocchio_data.mass[0] * gravity + force_vertical_rel = np.empty((2,)) + for i, frame_name in enumerate(("l_foot", "r_foot")): + frame_index = env.robot.pinocchio_model.getFrameId(frame_name) + frame = env.robot.pinocchio_model.frames[frame_index] + transform = env.robot.pinocchio_data.oMf[frame_index] + f_external = env.robot_state.f_external[frame.parent] + f_z_world = np.dot(transform.rotation[2], f_external.linear) + force_vertical_rel[i] = f_z_world / robot_weight + np.testing.assert_allclose( + force_vertical_rel, env.quantities["force_vertical_rel"]) + np.testing.assert_allclose(np.sum(force_vertical_rel), 1.0, atol=1e-3) diff --git a/python/gym_jiminy/unit_py/test_rewards.py b/python/gym_jiminy/unit_py/test_rewards.py index 04b4158ee..03e07bc8b 100644 --- a/python/gym_jiminy/unit_py/test_rewards.py +++ b/python/gym_jiminy/unit_py/test_rewards.py @@ -16,6 +16,7 @@ TrackingCapturePointReward, TrackingFootPositionsReward, TrackingFootOrientationsReward, + TrackingFootForceDistributionReward, MinimizeFrictionReward, SurviveReward, AdditiveMixtureReward) @@ -32,7 +33,7 @@ def setUp(self): self.env.eval() self.env.reset(seed=1) - action = self.env.action_space.sample() + action = 0.5 * self.env.action_space.sample() for _ in range(10): self.env.step(action) self.env.stop() @@ -54,17 +55,18 @@ def test_deletion(self): def test_tracking(self): for reward_class, cutoff in ( (TrackingBaseOdometryVelocityReward, 20.0), - (TrackingActuatedJointPositionsReward, 20.0), + (TrackingActuatedJointPositionsReward, 10.0), (TrackingBaseHeightReward, 0.1), - (TrackingCapturePointReward, 0.5), + (TrackingCapturePointReward, 0.2), (TrackingFootPositionsReward, 1.0), - (TrackingFootOrientationsReward, 2.0)) * 20: + (TrackingFootOrientationsReward, 2.0), + (TrackingFootForceDistributionReward, 2.0)): reward = reward_class(self.env, cutoff=cutoff) quantity_true = reward.quantity.requirements['value_left'] quantity_ref = reward.quantity.requirements['value_right'] self.env.reset(seed=0) - action = self.env.action_space.sample() + action = 0.5 * self.env.action_space.sample() for _ in range(5): self.env.step(action) _, _, terminated, _, _ = self.env.step(self.env.action) @@ -83,8 +85,6 @@ def test_tracking(self): assert value > 0.01 np.testing.assert_allclose(reward(terminated, {}), value) - del reward - def test_mixture(self): reward_odometry = TrackingBaseOdometryVelocityReward( self.env, cutoff=0.3) diff --git a/python/jiminy_py/examples/box_uneven_ground_impulse_contact.py b/python/jiminy_py/examples/box_uneven_ground_impulse_contact.py index 84444fe11..2f5db4b1a 100644 --- a/python/jiminy_py/examples/box_uneven_ground_impulse_contact.py +++ b/python/jiminy_py/examples/box_uneven_ground_impulse_contact.py @@ -37,7 +37,7 @@ # Configure integrator engine_options["stepper"]['logInternalStepperSteps'] = True - engine_options['stepper']['odeSolver'] = 'runge_kutta_dopri5' + engine_options['stepper']['odeSolver'] = 'runge_kutta_dopri' engine_options['stepper']['dtMax'] = 1.0e-3 # Set the ground contact options diff --git a/python/jiminy_py/examples/constraint_wheel.py b/python/jiminy_py/examples/constraint_wheel.py index e1c2bec1c..d06f31193 100644 --- a/python/jiminy_py/examples/constraint_wheel.py +++ b/python/jiminy_py/examples/constraint_wheel.py @@ -28,7 +28,7 @@ engine_options["stepper"]["sensorsUpdatePeriod"] = 0.0 # Configure integrator - engine_options['stepper']['odeSolver'] = 'runge_kutta_dopri5' + engine_options['stepper']['odeSolver'] = 'runge_kutta_dopri' engine_options['stepper']['dtMax'] = 1.0e-3 simulator.set_options(engine_options) diff --git a/python/jiminy_py/examples/double_pendulum.py b/python/jiminy_py/examples/double_pendulum.py index 0a1cd291f..f5550181c 100644 --- a/python/jiminy_py/examples/double_pendulum.py +++ b/python/jiminy_py/examples/double_pendulum.py @@ -57,10 +57,10 @@ def internal_dynamics(self, t, q, v, u_custom): engine_options["telemetry"]["enableCommand"] = True engine_options["telemetry"]["enableEnergy"] = True engine_options["world"]["gravity"][2] = -9.81 - engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri5" # ["runge_kutta_dopri5", "euler_explicit"] + engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri" # ["runge_kutta_dopri", "euler_explicit"] engine_options["stepper"]["tolRel"] = 1.0e-5 engine_options["stepper"]["tolAbs"] = 1.0e-4 - engine_options["stepper"]["dtMax"] = 2.0e-3 # 2.0e-4 for "euler_explicit", 3.0e-3 for "runge_kutta_dopri5" + engine_options["stepper"]["dtMax"] = 2.0e-3 # 2.0e-4 for "euler_explicit", 3.0e-3 for "runge_kutta_dopri" engine_options["stepper"]["sensorsUpdatePeriod"] = 1.0e-3 engine_options["stepper"]["controllerUpdatePeriod"] = 1.0e-3 engine_options["stepper"]["logInternalStepperSteps"] = False diff --git a/python/jiminy_py/examples/force_coupling.py b/python/jiminy_py/examples/force_coupling.py index e4e30a0ef..ccbbcd6ef 100644 --- a/python/jiminy_py/examples/force_coupling.py +++ b/python/jiminy_py/examples/force_coupling.py @@ -38,7 +38,7 @@ # Remove gravity engine_options = engine.get_options() - engine_options["stepper"]["odeSolver"] = 'runge_kutta_dopri5' + engine_options["stepper"]["odeSolver"] = 'runge_kutta_dopri' engine_options["stepper"]["tolAbs"] = 1e-10 engine_options["stepper"]["tolRel"] = 1e-10 engine_options["stepper"]["dtMax"] = 0.01 diff --git a/python/jiminy_py/examples/sphere_rolling.py b/python/jiminy_py/examples/sphere_rolling.py index 2f51155dc..f3fd14f20 100644 --- a/python/jiminy_py/examples/sphere_rolling.py +++ b/python/jiminy_py/examples/sphere_rolling.py @@ -26,7 +26,7 @@ engine_options["stepper"]["sensorsUpdatePeriod"] = 0.0 # Configure integrator - engine_options['stepper']['odeSolver'] = 'runge_kutta_dopri5' + engine_options['stepper']['odeSolver'] = 'runge_kutta_dopri' engine_options['stepper']['dtMax'] = 1.0e-3 simulator.set_options(engine_options) diff --git a/python/jiminy_py/src/jiminy_py/dynamics.py b/python/jiminy_py/src/jiminy_py/dynamics.py index e3cfd4b92..5b18f6969 100644 --- a/python/jiminy_py/src/jiminy_py/dynamics.py +++ b/python/jiminy_py/src/jiminy_py/dynamics.py @@ -139,10 +139,7 @@ class State: """ lambda_c: Optional[np.ndarray] = None - """Lambda multipliers associated with all the constraints as a 2D array. - - The first dimension corresponds to the N individual constraints applied on - the robot, while the second gathers the lambda multipliers. + """Lambda multipliers associated with all the constraints as a 1D array. """ diff --git a/python/jiminy_py/unit_py/data/double_pendulum_options.toml b/python/jiminy_py/unit_py/data/double_pendulum_options.toml index 638a756ec..6c44e71a7 100644 --- a/python/jiminy_py/unit_py/data/double_pendulum_options.toml +++ b/python/jiminy_py/unit_py/data/double_pendulum_options.toml @@ -1,4 +1,4 @@ [engine.stepper] -odeSolver = "runge_kutta_dopri5" +odeSolver = "runge_kutta_dopri" sensorsUpdatePeriod = 0.001 controllerUpdatePeriod = 0.001 diff --git a/python/jiminy_py/unit_py/test_double_spring_mass.py b/python/jiminy_py/unit_py/test_double_spring_mass.py index 7cf70dec3..1b0c0561f 100644 --- a/python/jiminy_py/unit_py/test_double_spring_mass.py +++ b/python/jiminy_py/unit_py/test_double_spring_mass.py @@ -89,7 +89,7 @@ def test_continuous_simulation(self): # Configure the engine engine_options = engine.get_options() - engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri5" + engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri" engine_options["stepper"]["sensorsUpdatePeriod"] = 0.0 engine_options["stepper"]["controllerUpdatePeriod"] = 0.0 engine_options["stepper"]["tolAbs"] = TOLERANCE * 1e-1 @@ -113,7 +113,7 @@ def test_discrete_simulation(self): # Configure the engine engine_options = engine.get_options() - engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri5" + engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri" engine_options["stepper"]["sensorsUpdatePeriod"] = 1.0e-3 engine_options["stepper"]["controllerUpdatePeriod"] = 1.0e-3 engine_options["stepper"]["tolAbs"] = TOLERANCE * 1e-1 @@ -136,7 +136,7 @@ def test_profile_force(self): # Configure the engine engine_options = engine.get_options() - engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri5" + engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri" engine_options["stepper"]["tolAbs"] = TOLERANCE * 1e-1 engine_options["stepper"]["tolRel"] = TOLERANCE * 1e-1 engine.set_options(engine_options) @@ -212,7 +212,7 @@ def external_force(t, q, v, f): # Configure the engine engine_options = engine.get_options() engine_options["world"]["gravity"] = np.zeros(6) - engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri5" + engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri" engine_options["stepper"]["sensorsUpdatePeriod"] = 1e-3 engine_options["stepper"]["controllerUpdatePeriod"] = 1e-3 engine_options["stepper"]["tolAbs"] = TOLERANCE * 1e-1 @@ -232,7 +232,7 @@ def test_fixed_body_constraint(self): # Configure the engine engine_options = engine.get_options() - engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri5" + engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri" engine_options["stepper"]["tolAbs"] = TOLERANCE * 1e-1 engine_options["stepper"]["tolRel"] = TOLERANCE * 1e-1 engine_options["constraints"]["regularization"] = 0.0 @@ -274,7 +274,7 @@ def test_freeflyer_multiple_constraints(self): # Configure the engine engine_options = engine.get_options() engine_options["world"]["gravity"] = np.zeros(6) # Turn off gravity - engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri5" + engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri" engine_options["stepper"]["tolAbs"] = TOLERANCE * 1e-1 engine_options["stepper"]["tolRel"] = TOLERANCE * 1e-1 engine_options["constraints"]["regularization"] = 0.0 @@ -334,7 +334,7 @@ def test_constraint_external_force(self): # Configure the engine engine_options = engine.get_options() - engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri5" + engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri" engine_options["stepper"]["tolAbs"] = TOLERANCE * 1e-1 engine_options["stepper"]["tolRel"] = TOLERANCE * 1e-1 engine_options["constraints"]["regularization"] = 0.0 diff --git a/python/jiminy_py/unit_py/test_multi_robot.py b/python/jiminy_py/unit_py/test_multi_robot.py index 4bdd8f267..4ade7d958 100644 --- a/python/jiminy_py/unit_py/test_multi_robot.py +++ b/python/jiminy_py/unit_py/test_multi_robot.py @@ -50,7 +50,7 @@ def internal_dynamics(self, t, q, v, u_custom): # Configure the engine engine_options = engine.get_options() - engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri5" + engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri" engine_options["stepper"]["tolAbs"] = TOLERANCE * 1e-1 engine_options["stepper"]["tolRel"] = TOLERANCE * 1e-2 engine.set_options(engine_options) diff --git a/python/jiminy_py/unit_py/test_simple_pendulum.py b/python/jiminy_py/unit_py/test_simple_pendulum.py index 7995a3201..9a3439c87 100644 --- a/python/jiminy_py/unit_py/test_simple_pendulum.py +++ b/python/jiminy_py/unit_py/test_simple_pendulum.py @@ -120,7 +120,7 @@ def spring_force(t, q, v, sensor_measurements, u_custom): # Configure the engine engine_options = engine.get_options() engine_options["world"]["gravity"] = np.zeros(6) - engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri5" + engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri" engine_options["stepper"]["tolAbs"] = TOLERANCE * 1e-1 engine_options["stepper"]["tolRel"] = TOLERANCE * 1e-1 engine.set_options(engine_options) @@ -170,7 +170,7 @@ def compute_command(t, q, v, sensor_measurements, command): # Configure the engine engine_options = engine.get_options() engine_options["world"]["gravity"] = np.zeros(6) - engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri5" + engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri" engine_options["stepper"]["tolAbs"] = 1e-12 engine_options["stepper"]["tolRel"] = 1e-12 engine.set_options(engine_options) @@ -703,7 +703,7 @@ def compute_command(t, q, v, sensor_measurements, command): # Configure the engine engine_options = engine.get_options() engine_options["world"]["gravity"] = np.zeros(6) - engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri5" + engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri" engine_options["stepper"]["tolAbs"] = TOLERANCE * 1e-1 engine_options["stepper"]["tolRel"] = TOLERANCE * 1e-1 engine.set_options(engine_options) @@ -779,7 +779,7 @@ def spring_force(t, q, v, sensor_measurements, u_custom): # Configure the engine engine_options = engine.get_options() engine_options["world"]["gravity"] = np.zeros(6) - engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri5" + engine_options["stepper"]["odeSolver"] = "runge_kutta_dopri" engine_options["stepper"]["tolAbs"] = TOLERANCE * 1e-1 engine_options["stepper"]["tolRel"] = TOLERANCE * 1e-1 engine_options["constraints"]["regularization"] = 0.0