diff --git a/src/hippopt/base/opti_callback.py b/src/hippopt/base/opti_callback.py index 3cfc2576..d4c1aa17 100644 --- a/src/hippopt/base/opti_callback.py +++ b/src/hippopt/base/opti_callback.py @@ -27,6 +27,7 @@ class CallbackCriterion(abc.ABC): def __init__(self) -> None: """""" self.opti = None + self.opti_debug = None @abc.abstractmethod def satisfied(self) -> bool: @@ -69,6 +70,10 @@ def set_opti(self, opti: cs.Opti) -> None: self.opti = weakref.proxy(opti) self.reset() + def update_opti_debug(self, opti_debug: cs.OptiAdvanced) -> None: + """""" + self.opti_debug = opti_debug + class BestCost(CallbackCriterion): """""" @@ -107,7 +112,7 @@ def update(self) -> None: def _get_current_cost(self) -> float: """""" - return self.opti.debug.value(self.opti.f) + return self.opti_debug.value(self.opti.f) class AcceptableCost(CallbackCriterion): @@ -151,7 +156,7 @@ def update(self) -> None: def _get_current_cost(self) -> float: """""" - return self.opti.debug.value(self.opti.f) + return self.opti_debug.value(self.opti.f) class AcceptablePrimalInfeasibility(CallbackCriterion): @@ -199,7 +204,7 @@ def update(self) -> None: def _get_current_primal_infeasibility(self) -> float: """""" - return self.opti.debug.stats()["iterations"]["inf_pr"][-1] + return self.opti_debug.stats()["iterations"]["inf_pr"][-1] class BestPrimalInfeasibility(CallbackCriterion): @@ -240,7 +245,7 @@ def update(self) -> None: def _get_current_primal_infeasibility(self) -> float: """""" - return self.opti.debug.stats()["iterations"]["inf_pr"][-1] + return self.opti_debug.stats()["iterations"]["inf_pr"][-1] class CombinedCallbackCriterion(CallbackCriterion, abc.ABC): @@ -274,6 +279,13 @@ def set_opti(self, opti: cs.Opti) -> None: self.lhs.set_opti(opti) self.rhs.set_opti(opti) + @final + def update_opti_debug(self, opti_debug: cs.OptiAdvanced) -> None: + """""" + + self.lhs.update_opti_debug(opti_debug) + self.rhs.update_opti_debug(opti_debug) + class OrCombinedCallbackCriterion(CombinedCallbackCriterion): """""" @@ -330,6 +342,8 @@ def __init__( def call(self, i: int) -> None: """""" + opti_debug = self.opti.debug + self.criterion.update_opti_debug(opti_debug) if self.criterion.satisfied(): self.criterion.update() @@ -338,24 +352,22 @@ def call(self, i: int) -> None: _logger.info(f"[i={i}] New best intermediate variables") self.best_iteration = i - self.best_cost = self.opti.debug.value(self.opti.f) + self.best_cost = opti_debug.value(self.opti.f) for variable in self.variables: if self.ignore_map[variable]: continue try: - self.best_objects[variable] = self.opti.debug.value(variable) + self.best_objects[variable] = opti_debug.value(variable) except Exception as err: # noqa self.ignore_map[variable] = True for parameter in self.parameters: if self.ignore_map[parameter]: continue - self.best_objects[parameter] = self.opti.debug.value(parameter) + self.best_objects[parameter] = opti_debug.value(parameter) self.ignore_map[parameter] = True # Parameters are saved only once - self.best_cost_values = { - cost: self.opti.debug.value(cost) for cost in self.cost - } + self.best_cost_values = {cost: opti_debug.value(cost) for cost in self.cost} self.best_constraint_multipliers = { - constraint: self.opti.debug.value(self.opti.dual(constraint)) + constraint: opti_debug.value(self.opti.dual(constraint)) for constraint in self.constraints } diff --git a/src/hippopt/deps/__init__.py b/src/hippopt/deps/__init__.py new file mode 100644 index 00000000..d3f05101 --- /dev/null +++ b/src/hippopt/deps/__init__.py @@ -0,0 +1 @@ +from .surf2stl import surf2stl diff --git a/src/hippopt/deps/surf2stl/LICENSE.txt b/src/hippopt/deps/surf2stl/LICENSE.txt new file mode 100644 index 00000000..f5a3f030 --- /dev/null +++ b/src/hippopt/deps/surf2stl/LICENSE.txt @@ -0,0 +1,7 @@ +Copyright 2020 asahidari + +Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/src/hippopt/deps/surf2stl/__init__.py b/src/hippopt/deps/surf2stl/__init__.py new file mode 100644 index 00000000..0b8a6306 --- /dev/null +++ b/src/hippopt/deps/surf2stl/__init__.py @@ -0,0 +1 @@ +from .surf2stl import write diff --git a/src/hippopt/deps/surf2stl/surf2stl.py b/src/hippopt/deps/surf2stl/surf2stl.py new file mode 100644 index 00000000..ec7fc542 --- /dev/null +++ b/src/hippopt/deps/surf2stl/surf2stl.py @@ -0,0 +1,254 @@ +### surf2stl.py --- Write a surface to a STL format file --- + +### Copyright (c) 2020 asahidari + +### This software is released under the MIT License. +### http://opensource.org/licenses/mit-license.php + +### Functions in this script (write, tri_write) export +### a stereolithography (STL) file for a surface with geometry +### defined by three matrix arguments, X, Y and Z. + +### This idea mainly come from surf2stl.m in MATLAB(or Octave): +### https://jp.mathworks.com/matlabcentral/fileexchange/4512-surf2stl + +import datetime +import math +import struct + +import numpy as np +from scipy.spatial import Delaunay + + +def write(filename, x, y, z, mode="binary"): + """ + Write a stl file for a surface with geometry + defined from three matrix arguments, x, y, and z. + Meshes are triangulated sequencially along xyz order. + + Parameters + ---------- + filename : string + output file name + + x, y, z : ndarray + Arguments x, y can be 1-dimensional arrays or 2-dimensional grids + (usually generated by np.meshgrid(x,y)), and z must be 2-dimensional, + which must have len(x)=n, len(y)=m where z.shape[m,n]. + + mode : string + STL file format, 'ascii' or 'binary'(default). + + Examples + ---------- + import numpy as np + import surf2stl + + x = np.linspace(-6, 6, 30) + y = np.linspace(-6, 6, 30) + X, Y = np.meshgrid(x, y) + Z = np.sin(np.sqrt(X ** 2 + Y ** 2)) + + surf2stl.write('3d-sinusoidal.stl', X, Y, Z) + + """ + + if type(filename) is not str: + raise Exception("Invalid filename") + + if mode != "ascii": + mode = "binary" + + if z.ndim != 2: + raise Exception("Variable z must be a 2-dimensional array") + + ### x, y can not be used as dx, dy in Python + ### if arguments type of x(or y) is 'int', + ### type error will raise in next 'if' block + # if type(x) == int and type(y) == int: + # x = np.arange(0, z.shape[1], x) + # x = np.arange(0, z.shape[0], y) + + if ( + len(x.shape) == 1 + and x.shape[0] == z.shape[1] + and len(y.shape) == 1 + and y.shape[0] == z.shape[0] + ): + x, y = np.meshgrid(x, y) + + if ( + len(x.shape) != len(z.shape) + or len(y.shape) != len(z.shape) + or x.shape[1] != z.shape[1] + or y.shape[0] != z.shape[0] + ): + raise Exception("Unable to resolve x and y variables") + + nfacets = 0 + title_str = "Created by surf2stl.py %s" % datetime.datetime.now().strftime( + "%d-%b-%Y %H:%M:%S" + ) + + f = open(filename, "wb" if mode != "ascii" else "w") + + if mode == "ascii": + f.write("solid %s\n" % title_str) + else: + title_str_ljust = title_str.ljust(80) + # f.write(title_str_ljust.encode('utf-8')) # same as 'ascii' for alphabet characters + f.write(title_str_ljust.encode("ascii")) + f.write(struct.pack("i", 0)) + + for i in range(z.shape[0] - 1): + for j in range(z.shape[1] - 1): + p1 = np.array([x[i, j], y[i, j], z[i, j]]) + p2 = np.array([x[i, j + 1], y[i, j + 1], z[i, j + 1]]) + p3 = np.array([x[i + 1, j + 1], y[i + 1, j + 1], z[i + 1, j + 1]]) + val = local_write_facet(f, p1, p2, p3, mode) + nfacets += val + + p1 = np.array([x[i + 1, j + 1], y[i + 1, j + 1], z[i + 1, j + 1]]) + p2 = np.array([x[i + 1, j], y[i + 1, j], z[i + 1, j]]) + p3 = np.array([x[i, j], y[i, j], z[i, j]]) + val = local_write_facet(f, p1, p2, p3, mode) + nfacets += val + + if mode == "ascii": + f.write("endsolid %s\n" % title_str) + else: + f.seek(80, 0) + f.write(struct.pack("i", nfacets)) + + f.close() + + print("Wrote %d facets" % nfacets) + return + + +def tri_write(filename, x, y, z, tri, mode="binary"): + """ + Write a stl file for a surface with geometry + defined from three matrix arguments, x, y, and z + with Delaunay Triangle parameter(tri). + + Meshes are triangulated with the 'tri' parameter + usually made with parameters which determine + the sequence of vertices (like [u, v]). + + Parameters + ---------- + filename : string + output file name + + x, y, z : ndarray + Each of these arguments must be 1-dimensional. + + tri : scipy.spatial.Delaunay + Delaunay Triangulation object. + When xyz coordinates are determined from other parameters(like (u, v)), + this triangle faces are basically calculated with the parameters. + + mode : string + STL file format, 'ascii' or 'binary'(default). + + Examples + ---------- + import numpy as np + from scipy.spatial import Delaunay + import surf2stl + + u = np.linspace(0, 2.0 * np.pi, endpoint=True, num=50) + v = np.linspace(-0.5, 0.5, endpoint=True, num=10) + u, v = np.meshgrid(u, v) + u, v = u.flatten(), v.flatten() + + x = (1 + 0.5 * v * np.cos(u / 2.0)) * np.cos(u) + y = (1 + 0.5 * v * np.cos(u / 2.0)) * np.sin(u) + z = 0.5 * v * np.sin(u / 2.0) + + delaunay_tri = Delaunay(np.array([u, v]).T) + surf2stl.tri_write('mobius.stl', x, y, z, delaunay_tri) + + """ + + if type(filename) is not str: + raise Exception("Invalid filename") + + if mode != "ascii": + mode = "binary" + + if len(x.shape) != 1 or len(y.shape) != 1 or len(z.shape) != 1: + raise Exception("Each variable x,y,z must be a 1-dimensional array") + + if x.shape[0] != z.shape[0] or y.shape[0] != z.shape[0]: + raise Exception("Number of x,y,z elements must be equal") + + nfacets = 0 + title_str = "Created by surf2stl.py %s" % datetime.datetime.now().strftime( + "%d-%b-%Y %H:%M:%S" + ) + + f = open(filename, "wb" if mode != "ascii" else "w") + + if mode == "ascii": + f.write("solid %s\n" % title_str) + else: + title_str_ljust = title_str.ljust(80) + # f.write(title_str_ljust.encode('utf-8')) # same as 'ascii' for alphabet characters + f.write(title_str_ljust.encode("ascii")) + f.write(struct.pack("i", 0)) + + indices = tri.simplices + verts = tri.points[indices] + for i in range(0, indices.shape[0], 1): + p = indices[i] + p1 = np.array([x[p[0]], y[p[0]], z[p[0]]]) + p2 = np.array([x[p[1]], y[p[1]], z[p[1]]]) + p3 = np.array([x[p[2]], y[p[2]], z[p[2]]]) + val = local_write_facet(f, p1, p2, p3, mode) + nfacets += val + + if mode == "ascii": + f.write("endsolid %s\n" % title_str) + else: + f.seek(80, 0) + f.write(struct.pack("i", nfacets)) + + f.close() + + print("Wrote %d facets" % nfacets) + return + + +# Local subfunctions + + +def local_write_facet(f, p1, p2, p3, mode): + if np.isnan(p1).any() or np.isnan(p2).any() or np.isnan(p3).any(): + return 0 + + n = local_find_normal(p1, p2, p3) + if mode == "ascii": + f.write("facet normal %.7f %.7f %.7f\n" % (n[0], n[1], n[2])) + f.write("outer loop\n") + f.write("vertex %.7f %.7f %.7f\n" % (p1[0], p1[1], p1[2])) + f.write("vertex %.7f %.7f %.7f\n" % (p2[0], p2[1], p2[2])) + f.write("vertex %.7f %.7f %.7f\n" % (p3[0], p3[1], p3[2])) + f.write("endloop\n") + f.write("endfacet\n") + else: + f.write(struct.pack("%sf" % len(n), *n)) + f.write(struct.pack("%sf" % len(p1), *p1)) + f.write(struct.pack("%sf" % len(p2), *p2)) + f.write(struct.pack("%sf" % len(p3), *p3)) + f.write(struct.pack("h", 0)) + return 1 + + +def local_find_normal(p1, p2, p3): + v1 = p2 - p1 + v2 = p3 - p1 + v3 = np.cross(v1, v2) + n = v3 / math.sqrt(np.sum(v3 * v3)) + return n diff --git a/src/hippopt/robot_planning/__init__.py b/src/hippopt/robot_planning/__init__.py new file mode 100644 index 00000000..a6cbdf54 --- /dev/null +++ b/src/hippopt/robot_planning/__init__.py @@ -0,0 +1,77 @@ +from . import expressions, utilities, variables +from .expressions.centroidal import ( + centroidal_dynamics_with_point_forces, + com_dynamics_from_momentum, +) +from .expressions.complementarity import ( + dcc_complementarity_margin, + dcc_planar_complementarity, + relaxed_complementarity_margin, +) +from .expressions.contacts import ( + contact_points_centroid, + contact_points_yaw_alignment_error, + friction_cone_square_margin, + normal_force_component, + swing_height_heuristic, +) +from .expressions.kinematics import ( + center_of_mass_position_from_kinematics, + centroidal_momentum_from_kinematics, + frames_relative_position, + point_position_from_kinematics, + rotation_error_from_kinematics, +) +from .expressions.quaternion import ( + quaternion_xyzw_error, + quaternion_xyzw_normalization, + quaternion_xyzw_velocity_to_right_trivialized_angular_velocity, +) +from .utilities.foot_contact_state_plotter import ( + ContactPointStatePlotter, + ContactPointStatePlotterSettings, + FootContactStatePlotter, + FootContactStatePlotterSettings, +) +from .utilities.humanoid_state_visualizer import ( + HumanoidStateVisualizer, + HumanoidStateVisualizerSettings, +) +from .utilities.interpolators import ( + feet_contact_points_interpolator, + floating_base_system_state_interpolator, + foot_contact_state_interpolator, + free_floating_object_state_interpolator, + humanoid_state_interpolator, + kinematic_tree_state_interpolator, + linear_interpolator, + quaternion_slerp, + transform_interpolator, +) +from .utilities.planar_terrain import PlanarTerrain +from .utilities.smooth_terrain import SmoothTerrain +from .utilities.terrain_descriptor import TerrainDescriptor +from .utilities.terrain_sum import TerrainSum +from .utilities.terrain_visualizer import TerrainVisualizer, TerrainVisualizerSettings +from .variables.contacts import ( + ContactPointDescriptor, + ContactPointState, + ContactPointStateDerivative, + FeetContactPhasesDescriptor, + FeetContactPointDescriptors, + FeetContactPoints, + FootContactPhaseDescriptor, + FootContactState, +) +from .variables.floating_base import ( + FloatingBaseSystem, + FloatingBaseSystemState, + FloatingBaseSystemStateDerivative, + FreeFloatingObject, + FreeFloatingObjectState, + FreeFloatingObjectStateDerivative, + KinematicTree, + KinematicTreeState, + KinematicTreeStateDerivative, +) +from .variables.humanoid import HumanoidState diff --git a/src/hippopt/robot_planning/dynamics/__init__.py b/src/hippopt/robot_planning/dynamics/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/src/hippopt/robot_planning/expressions/__init__.py b/src/hippopt/robot_planning/expressions/__init__.py index e69de29b..9593c935 100644 --- a/src/hippopt/robot_planning/expressions/__init__.py +++ b/src/hippopt/robot_planning/expressions/__init__.py @@ -0,0 +1 @@ +from . import complementarity, contacts, kinematics diff --git a/src/hippopt/robot_planning/expressions/centroidal.py b/src/hippopt/robot_planning/expressions/centroidal.py new file mode 100644 index 00000000..998ec4ed --- /dev/null +++ b/src/hippopt/robot_planning/expressions/centroidal.py @@ -0,0 +1,96 @@ +import casadi as cs + + +def centroidal_dynamics_with_point_forces( + number_of_points: int, + assume_unitary_mass: bool = False, + mass_name: str = "m", + gravity_name: str = "g", + com_name: str = "com", + point_position_names: list[str] = None, + point_force_names: list[str] = None, + options: dict = None, + **_, +) -> cs.Function: + options = {} if options is None else options + + point_position_names = ( + point_position_names + if point_position_names is not None + else [f"p{i}" for i in range(number_of_points)] + ) + + if len(point_position_names) != number_of_points: + raise ValueError( + f"Expected {number_of_points} point position names," + f" got {len(point_position_names)}" + ) + + if point_force_names is None: + point_force_names = [] + for i in range(number_of_points): + point_force_names.append(f"f{i}") + + assert len(point_force_names) == number_of_points + + input_vars = [] + + if assume_unitary_mass: + m = 1.0 + else: + m = cs.MX.sym(mass_name, 1) + input_vars.append(m) + + g = cs.MX.sym(gravity_name, 6) + input_vars.append(g) + + x = cs.MX.sym(com_name, 3) + input_vars.append(x) + + p = [] + f = [] + for point_position_name, point_force_name in zip( + point_position_names, point_force_names + ): + p.append(cs.MX.sym(point_position_name, 3)) + input_vars.append(p[-1]) + f.append(cs.MX.sym(point_force_name, 3)) + input_vars.append(f[-1]) + + input_names = [var.name() for var in input_vars] + + h_g = m * g + cs.sum2( + cs.horzcat(*[cs.vertcat(fi, cs.cross(pi - x, fi)) for fi, pi in zip(f, p)]) + ) + + return cs.Function( + "centroidal_dynamics_with_point_forces", + input_vars, + [h_g], + input_names, + ["h_g_dot"], + options, + ) + + +def com_dynamics_from_momentum( + mass_name: str = "m", + momentum_name: str = "h_g", + options: dict = None, + **_, +) -> cs.Function: + options = {} if options is None else options + + m = cs.MX.sym(mass_name, 1) + h_g = cs.MX.sym(momentum_name, 6) + + x_dot = h_g[0:3] / m + + return cs.Function( + "com_dynamics_from_momentum", + [m, h_g], + [x_dot], + [mass_name, momentum_name], + ["x_dot"], + options, + ) diff --git a/src/hippopt/robot_planning/expressions/complementarity.py b/src/hippopt/robot_planning/expressions/complementarity.py new file mode 100644 index 00000000..591f50b0 --- /dev/null +++ b/src/hippopt/robot_planning/expressions/complementarity.py @@ -0,0 +1,155 @@ +import casadi as cs + +from hippopt.robot_planning.utilities.terrain_descriptor import TerrainDescriptor + + +def dcc_planar_complementarity( + terrain: TerrainDescriptor, + point_position_name: str = None, + height_multiplier_name: str = "kt", + point_position_control_name: str = "u_p", + options: dict = None, + **_ +) -> cs.Function: + options = {} if options is None else options + point_position_name = ( + terrain.get_point_position_name() + if point_position_name is None + else point_position_name + ) + point_position = cs.MX.sym(point_position_name, 3) + height_multiplier = cs.MX.sym(height_multiplier_name, 1) + point_control = cs.MX.sym(point_position_control_name, 3) + + terrain_orientation_fun = terrain.orientation_function() + height_function = terrain.height_function() + + height = height_function(point_position) + terrain_orientation = terrain_orientation_fun(point_position) + + planar_multiplier = cs.tanh(height_multiplier * height) + multipliers = cs.diag(cs.horzcat(planar_multiplier, planar_multiplier, 1)) + planar_complementarity = terrain_orientation @ multipliers @ point_control + + return cs.Function( + "planar_complementarity_dcc", + [point_position, height_multiplier, point_control], + [planar_complementarity], + [point_position_name, height_multiplier_name, point_position_control_name], + ["planar_complementarity"], + options, + ) + + +def dcc_complementarity_margin( + terrain: TerrainDescriptor, + point_position_name: str = None, + point_force_name: str = "point_force", + point_velocity_name: str = "point_velocity", + point_force_derivative_name: str = "point_force_derivative", + dcc_gain_name: str = "k_bs", + dcc_epsilon_name: str = "eps", + options: dict = None, + **_ +) -> cs.Function: + options = {} if options is None else options + point_position_name = ( + terrain.get_point_position_name() + if point_position_name is None + else point_position_name + ) + point_position = cs.MX.sym(point_position_name, 3) + point_force = cs.MX.sym(point_force_name, 3) + point_velocity = cs.MX.sym(point_velocity_name, 3) + point_force_derivative = cs.MX.sym(point_force_derivative_name, 3) + dcc_gain = cs.MX.sym(dcc_gain_name, 1) + eps = cs.MX.sym(dcc_epsilon_name, 1) + + normal_direction_fun = terrain.normal_direction_function() + height_function = terrain.height_function() + + height = height_function(point_position) + normal_direction = normal_direction_fun(point_position) + + # See Sec III.A of https://ieeexplore.ieee.org/abstract/document/9847574 + height_derivative = cs.jtimes(height, point_position, point_velocity) + normal_derivative = cs.jtimes(normal_direction, point_position, point_velocity) + normal_force = normal_direction.T @ point_force + normal_force_derivative = normal_direction.T @ point_force_derivative + complementarity = height * normal_force + + csi = ( + height_derivative * normal_force + + height * point_force.T @ normal_derivative + + height * normal_force_derivative + ) + + margin = eps - dcc_gain * complementarity - csi + + return cs.Function( + "dcc_complementarity_margin", + [ + point_position, + point_force, + point_velocity, + point_force_derivative, + dcc_gain, + eps, + ], + [margin], + [ + point_position_name, + point_force_name, + point_velocity_name, + point_force_derivative_name, + dcc_gain_name, + dcc_epsilon_name, + ], + ["dcc_complementarity_margin"], + options, + ) + + +def relaxed_complementarity_margin( + terrain: TerrainDescriptor, + point_position_name: str = None, + point_force_name: str = "point_force", + relaxed_complementarity_epsilon_name: str = "eps", + options: dict = None, + **_ +) -> cs.Function: + options = {} if options is None else options + point_position_name = ( + terrain.get_point_position_name() + if point_position_name is None + else point_position_name + ) + point_position = cs.MX.sym(point_position_name, 3) + point_force = cs.MX.sym(point_force_name, 3) + eps = cs.MX.sym(relaxed_complementarity_epsilon_name, 1) + + normal_direction_fun = terrain.normal_direction_function() + height_function = terrain.height_function() + + height = height_function(point_position) + normal_direction = normal_direction_fun(point_position) + normal_force = normal_direction.T @ point_force + + margin = eps - height * normal_force + + return cs.Function( + "relaxed_complementarity_margin", + [ + point_position, + point_force, + eps, + ], + [margin], + [ + point_position_name, + point_force_name, + relaxed_complementarity_epsilon_name, + ], + ["relaxed_complementarity_margin"], + options, + ) diff --git a/src/hippopt/robot_planning/expressions/contacts.py b/src/hippopt/robot_planning/expressions/contacts.py new file mode 100644 index 00000000..1d692aa8 --- /dev/null +++ b/src/hippopt/robot_planning/expressions/contacts.py @@ -0,0 +1,175 @@ +import casadi as cs + +from hippopt.robot_planning.utilities.terrain_descriptor import TerrainDescriptor + + +def normal_force_component( + terrain: TerrainDescriptor, + point_position_name: str = None, + point_force_name: str = "point_force", + options: dict = None, + **_, +) -> cs.Function: + options = {} if options is None else options + point_position_name = ( + terrain.get_point_position_name() + if point_position_name is None + else point_position_name + ) + point_position = cs.MX.sym(point_position_name, 3) + point_force = cs.MX.sym(point_force_name, 3) + + normal_direction_fun = terrain.normal_direction_function() + + normal_component = normal_direction_fun(point_position).T @ point_force + + return cs.Function( + "normal_force_component", + [point_position, point_force], + [normal_component], + [point_position_name, point_force_name], + ["normal_force"], + options, + ) + + +def friction_cone_square_margin( + terrain: TerrainDescriptor, + point_position_name: str = None, + point_force_name: str = "point_force", + static_friction_name: str = "mu_s", + options: dict = None, + **_, +) -> cs.Function: + options = {} if options is None else options + point_position_name = ( + terrain.get_point_position_name() + if point_position_name is None + else point_position_name + ) + point_position = cs.MX.sym(point_position_name, 3) + point_force = cs.MX.sym(point_force_name, 3) + static_friction = cs.MX.sym(static_friction_name, 1) + + orientation_fun = terrain.orientation_function() + terrain_orientation = orientation_fun(point_position) + force_in_contact = terrain_orientation.T @ point_force + + # In principle, it should be sqrt(fx^2 + fy^2) <= u * fz, + # but since both sides are positive, we square them both. + # Their difference needs to remain positive, i.e. + # (u * fz)^2 - (fx^2 + fy^2) >= 0 + # that is equal to + # [-1, -1, u^2] * f.^2 + margin = cs.horzcat(-1, -1, cs.constpow(static_friction, 2)) @ cs.constpow( + force_in_contact, 2 + ) + + return cs.Function( + "friction_cone_square_margin", + [point_position, point_force, static_friction], + [margin], + [point_position_name, point_force_name, static_friction_name], + ["friction_cone_square_margin"], + options, + ) + + +def contact_points_centroid( + number_of_points: int, + point_position_names: list[str] = None, + options: dict = None, + **_, +) -> cs.Function: + options = {} if options is None else options + + if point_position_names is None: + point_position_names = [] + for i in range(number_of_points): + point_position_names.append(f"p{i}") + + assert len(point_position_names) == number_of_points + + input_vars = [] + p = [] + for point_position_name in point_position_names: + p.append(cs.MX.sym(point_position_name, 3)) + input_vars.append(p[-1]) + + input_names = [var.name() for var in input_vars] + + centroid = cs.DM.zeros(3, 1) + + for point in p: + centroid = centroid + point + + if number_of_points > 0: + centroid = centroid / number_of_points + + return cs.Function( + "contact_points_centroid", + input_vars, + [centroid], + input_names, + ["centroid"], + options, + ) + + +def contact_points_yaw_alignment_error( + first_point_name: str = "p_0", + second_point_name: str = "p_1", + desired_yaw_name: str = "desired_yaw", + options: dict = None, + **_, +) -> cs.Function: + options = {} if options is None else options + + p0 = cs.MX.sym(first_point_name, 3) + p1 = cs.MX.sym(second_point_name, 3) + yaw = cs.MX.sym(desired_yaw_name, 1) + + yaw_alignment = cs.horzcat(-cs.sin(yaw), cs.cos(yaw)) @ (p1 - p0)[:2] + + return cs.Function( + "contact_points_yaw_alignment_error", + [p0, p1, yaw], + [yaw_alignment], + [first_point_name, second_point_name, desired_yaw_name], + ["yaw_alignment_error"], + options, + ) + + +def swing_height_heuristic( + terrain: TerrainDescriptor, + point_position_name: str = "p", + point_velocity_name: str = "p_dot", + desired_height_name: str = "h_desired", + options: dict = None, + **_, +) -> cs.Function: + options = {} if options is None else options + + point = cs.MX.sym(point_position_name, 3) + point_velocity = cs.MX.sym(point_velocity_name, 3) + desired_height = cs.MX.sym(desired_height_name, 1) + + height_fun = terrain.height_function() + terrain_height = height_fun(point) + terrain_orientation_fun = terrain.orientation_function() + terrain_orientation = terrain_orientation_fun(point) + + height_difference = terrain_height - desired_height + planar_velocity = (terrain_orientation.T @ point_velocity)[:2] + + heuristic = 0.5 * (cs.constpow(height_difference, 2) + cs.sumsqr(planar_velocity)) + + return cs.Function( + "swing_height_heuristic", + [point, point_velocity, desired_height], + [heuristic], + [point_position_name, point_velocity_name, desired_height_name], + ["heuristic"], + options, + ) diff --git a/src/hippopt/robot_planning/expressions/kinematics.py b/src/hippopt/robot_planning/expressions/kinematics.py new file mode 100644 index 00000000..bae777a4 --- /dev/null +++ b/src/hippopt/robot_planning/expressions/kinematics.py @@ -0,0 +1,251 @@ +import casadi as cs +import liecasadi +from adam.casadi import KinDynComputations + +from hippopt.robot_planning.expressions.quaternion import ( + quaternion_xyzw_velocity_to_right_trivialized_angular_velocity, +) + + +def centroidal_momentum_from_kinematics( + kindyn_object: KinDynComputations, + base_position_name: str = "base_position", + base_quaternion_xyzw_name: str = "base_quaternion", + joint_positions_name: str = "joint_positions", + base_position_derivative_name: str = "base_position_derivative", + base_quaternion_xyzw_derivative_name: str = "base_quaternion_derivative", + joint_velocities_name: str = "joint_velocities", + options: dict = None, + **_ +) -> cs.Function: + options = {} if options is None else options + + base_position = cs.MX.sym(base_position_name, 3) + base_quaternion = cs.MX.sym(base_quaternion_xyzw_name, 4) + joint_positions = cs.MX.sym(joint_positions_name, kindyn_object.NDoF) + base_position_derivative = cs.MX.sym(base_position_derivative_name, 3) + base_quaternion_derivative = cs.MX.sym(base_quaternion_xyzw_derivative_name, 4) + joint_velocities = cs.MX.sym(joint_velocities_name, kindyn_object.NDoF) + + cmm_function = kindyn_object.centroidal_momentum_matrix_fun() + + base_pose = liecasadi.SE3.from_position_quaternion( + base_position, base_quaternion + ).as_matrix() # The quaternion is supposed normalized + + base_angular_velocity_fun = ( + quaternion_xyzw_velocity_to_right_trivialized_angular_velocity( + quaternion_xyzw_name=base_quaternion_xyzw_name, + quaternion_xyzw_velocity_name=base_quaternion_xyzw_derivative_name, + options=options, + ) + ) + base_angular_velocity = base_angular_velocity_fun( + **{ + base_quaternion_xyzw_name: base_quaternion, + base_quaternion_xyzw_derivative_name: base_quaternion_derivative, + } + )["right_trivialized_angular_velocity"] + + robot_velocity = cs.vertcat( + base_position_derivative, base_angular_velocity, joint_velocities + ) + + momentum = cmm_function(base_pose, joint_positions) @ robot_velocity + + return cs.Function( + "centroidal_momentum_from_kinematics", + [ + base_position, + base_quaternion, + joint_positions, + base_position_derivative, + base_quaternion_derivative, + joint_velocities, + ], + [momentum], + [ + base_position_name, + base_quaternion_xyzw_name, + joint_positions_name, + base_position_derivative_name, + base_quaternion_xyzw_derivative_name, + joint_velocities_name, + ], + ["h_g"], + options, + ) + + +def center_of_mass_position_from_kinematics( + kindyn_object: KinDynComputations, + base_position_name: str = "base_position", + base_quaternion_xyzw_name: str = "base_quaternion", + joint_positions_name: str = "joint_positions", + options: dict = None, + **_ +) -> cs.Function: + options = {} if options is None else options + + base_position = cs.MX.sym(base_position_name, 3) + base_quaternion = cs.MX.sym(base_quaternion_xyzw_name, 4) + joint_positions = cs.MX.sym(joint_positions_name, kindyn_object.NDoF) + + com_function = kindyn_object.CoM_position_fun() + + base_pose = liecasadi.SE3.from_position_quaternion( + base_position, base_quaternion + ).as_matrix() # The quaternion is supposed normalized + + com_position = com_function(base_pose, joint_positions) + + return cs.Function( + "center_of_mass_position_from_kinematics", + [ + base_position, + base_quaternion, + joint_positions, + ], + [com_position], + [ + base_position_name, + base_quaternion_xyzw_name, + joint_positions_name, + ], + ["com_position"], + options, + ) + + +def point_position_from_kinematics( + kindyn_object: KinDynComputations, + frame_name: str, + point_position_in_frame_name: str = "point_position", + base_position_name: str = "base_position", + base_quaternion_xyzw_name: str = "base_quaternion", + joint_positions_name: str = "joint_positions", + options: dict = None, + **_ +) -> cs.Function: + options = {} if options is None else options + + base_position = cs.MX.sym(base_position_name, 3) + base_quaternion = cs.MX.sym(base_quaternion_xyzw_name, 4) + joint_positions = cs.MX.sym(joint_positions_name, kindyn_object.NDoF) + point_position_in_frame = cs.MX.sym(point_position_in_frame_name, 3) + + fk_function = kindyn_object.forward_kinematics_fun(frame=frame_name) + + base_pose = liecasadi.SE3.from_position_quaternion( + base_position, base_quaternion + ).as_matrix() # The quaternion is supposed normalized + + frame_pose = fk_function(base_pose, joint_positions) + + point_position = frame_pose[:3, :3] @ point_position_in_frame + frame_pose[:3, 3] + + return cs.Function( + "point_position_from_kinematics", + [ + base_position, + base_quaternion, + joint_positions, + point_position_in_frame, + ], + [point_position], + [ + base_position_name, + base_quaternion_xyzw_name, + joint_positions_name, + point_position_in_frame_name, + ], + ["point_position"], + options, + ) + + +def frames_relative_position( + kindyn_object: KinDynComputations, + reference_frame: str, + target_frame: str, + joint_positions_name: str = "joint_positions", + options: dict = None, + **_ +) -> cs.Function: + options = {} if options is None else options + joint_positions = cs.MX.sym(joint_positions_name, kindyn_object.NDoF) + + base_pose = cs.DM_eye(4) + + reference_fk_function = kindyn_object.forward_kinematics_fun(frame=reference_frame) + target_fk_function = kindyn_object.forward_kinematics_fun(frame=target_frame) + + reference_pose = reference_fk_function(base_pose, joint_positions) + reference_pose_inverse_rotation = reference_pose[:3, :3].T + reference_pose_inverse_translation = ( + -reference_pose_inverse_rotation @ reference_pose[:3, 3] + ) + + target_pose = target_fk_function(base_pose, joint_positions) + + relative_position = ( + reference_pose_inverse_rotation @ target_pose[:3, 3] + + reference_pose_inverse_translation + ) + + return cs.Function( + "frames_relative_position", + [joint_positions], + [relative_position], + [joint_positions_name], + ["relative_position"], + options, + ) + + +def rotation_error_from_kinematics( + kindyn_object: KinDynComputations, + target_frame: str, + base_position_name: str = "base_position", + base_quaternion_xyzw_name: str = "base_quaternion", + joint_positions_name: str = "joint_positions", + desired_quaternion_xyzw_name: str = "desired_quaternion", + options: dict = None, + **_ +) -> cs.Function: + options = {} if options is None else options + base_position = cs.MX.sym(base_position_name, 3) + base_quaternion = cs.MX.sym(base_quaternion_xyzw_name, 4) + joint_positions = cs.MX.sym(joint_positions_name, kindyn_object.NDoF) + desired_quaternion = cs.MX.sym(desired_quaternion_xyzw_name, 4) + + base_pose = liecasadi.SE3.from_position_quaternion( + base_position, base_quaternion + ).as_matrix() # The quaternion is supposed normalized + + target_fk_function = kindyn_object.forward_kinematics_fun(frame=target_frame) + target_pose = target_fk_function(base_pose, joint_positions) + target_orientation = target_pose[:3, :3] + + rotation_error = ( + target_orientation @ liecasadi.SO3.from_quat(desired_quaternion).as_matrix().T + ) + + return cs.Function( + "quaternion_error_from_kinematics", + [ + base_position, + base_quaternion, + joint_positions, + desired_quaternion, + ], + [rotation_error], + [ + base_position_name, + base_quaternion_xyzw_name, + joint_positions_name, + desired_quaternion_xyzw_name, + ], + ["rotation_error"], + options, + ) diff --git a/src/hippopt/robot_planning/expressions/quaternion.py b/src/hippopt/robot_planning/expressions/quaternion.py new file mode 100644 index 00000000..82ee1c43 --- /dev/null +++ b/src/hippopt/robot_planning/expressions/quaternion.py @@ -0,0 +1,85 @@ +import casadi as cs +import liecasadi + + +def quaternion_xyzw_normalization( + quaternion_xyzw_name: str = "quaternion", + options: dict = None, + **_, +) -> cs.Function: + options = {} if options is None else options + quaternion = cs.MX.sym(quaternion_xyzw_name, 4) + + normalized_quaternion = liecasadi.Quaternion(xyzw=quaternion).normalize() + + return cs.Function( + "quaternion_xyzw_normalization", + [quaternion], + [normalized_quaternion.xyzw], + [quaternion_xyzw_name], + ["quaternion_normalized"], + options, + ) + + +def quaternion_xyzw_velocity_to_right_trivialized_angular_velocity( + quaternion_xyzw_name: str = "quaternion", + quaternion_xyzw_velocity_name: str = "quaternion_velocity", + options: dict = None, + **_, +) -> cs.Function: + options = {} if options is None else options + quaternion = cs.MX.sym(quaternion_xyzw_name, 4) + quaternion_velocity = cs.MX.sym(quaternion_xyzw_velocity_name, 4) + + q_w = quaternion[3] + q_i = quaternion[:3] + + q_dot_w = quaternion_velocity[3] + q_dot_i = quaternion_velocity[:3] + + # See Sec. 1.5.3 of https://arxiv.org/pdf/0811.2889.pdf + angular_velocity = 2 * (-q_dot_w * q_i + q_w * q_dot_i - cs.cross(q_dot_i, q_i)) + + return cs.Function( + "quaternion_xyzw_velocity_to_right_trivialized_angular_velocity", + [quaternion, quaternion_velocity], + [angular_velocity], + [quaternion_xyzw_name, quaternion_xyzw_velocity_name], + ["right_trivialized_angular_velocity"], + options, + ) + + +def quaternion_xyzw_error( + quaternion_xyzw_name: str = "quaternion", + desired_quaternion_xyzw_name: str = "desired_quaternion", + options: dict = None, + **_, +): + options = {} if options is None else options + quaternion = cs.MX.sym(quaternion_xyzw_name, 4) + desired_quaternion = cs.MX.sym(desired_quaternion_xyzw_name, 4) + + target_rotation = liecasadi.SO3.from_quat(quaternion) + desired_rotation = liecasadi.SO3.from_quat(desired_quaternion) + + rotation_error = desired_rotation.inverse() * target_rotation + identity = liecasadi.SO3.Identity() + + error = rotation_error.as_quat() - identity.as_quat() + + return cs.Function( + "quaternion_error", + [ + quaternion, + desired_quaternion, + ], + [error], + [ + quaternion_xyzw_name, + desired_quaternion_xyzw_name, + ], + ["quaternion_error"], + options, + ) diff --git a/src/hippopt/robot_planning/utilities/.gitignore b/src/hippopt/robot_planning/utilities/.gitignore new file mode 100644 index 00000000..ee4c1191 --- /dev/null +++ b/src/hippopt/robot_planning/utilities/.gitignore @@ -0,0 +1,2 @@ +*.urdf +*.stl diff --git a/src/hippopt/robot_planning/utilities/__init__.py b/src/hippopt/robot_planning/utilities/__init__.py index e69de29b..49c6432d 100644 --- a/src/hippopt/robot_planning/utilities/__init__.py +++ b/src/hippopt/robot_planning/utilities/__init__.py @@ -0,0 +1,10 @@ +from . import ( + foot_contact_state_plotter, + humanoid_state_visualizer, + interpolators, + planar_terrain, + smooth_terrain, + terrain_descriptor, + terrain_sum, + terrain_visualizer, +) diff --git a/src/hippopt/robot_planning/utilities/foot_contact_state_plotter.py b/src/hippopt/robot_planning/utilities/foot_contact_state_plotter.py new file mode 100644 index 00000000..fe978282 --- /dev/null +++ b/src/hippopt/robot_planning/utilities/foot_contact_state_plotter.py @@ -0,0 +1,369 @@ +import copy +import dataclasses +import logging +import math +import multiprocessing +from typing import Callable, TypeVar + +import matplotlib.axes +import matplotlib.pyplot as plt +import numpy as np + +from hippopt.robot_planning.utilities.planar_terrain import PlanarTerrain +from hippopt.robot_planning.utilities.terrain_descriptor import TerrainDescriptor +from hippopt.robot_planning.variables.contacts import ( + ContactPointState, + FootContactState, +) + + +@dataclasses.dataclass +class ContactPointStatePlotterSettings: + complementarity_axes: list[matplotlib.axes.Axes] | None = dataclasses.field( + default=None + ) + force_axes: matplotlib.axes.Axes | None = dataclasses.field(default=None) + terrain: TerrainDescriptor = dataclasses.field(default=None) + + input_complementarity_axes: dataclasses.InitVar[ + list[matplotlib.axes.Axes] + ] = dataclasses.field(default=None) + input_force_axes: dataclasses.InitVar[matplotlib.axes.Axes] = dataclasses.field( + default=None + ) + input_terrain: dataclasses.InitVar[TerrainDescriptor] = dataclasses.field( + default=None + ) + + def __post_init__( + self, + input_complementarity_axes: list[matplotlib.axes.Axes], + input_force_axes: matplotlib.axes.Axes, + input_terrain: TerrainDescriptor, + ): + self.complementarity_axes = None + if isinstance(input_complementarity_axes, list): + if len(input_complementarity_axes) != 2: + raise ValueError("input_axes must be a list of length 2.") + + self.complementarity_axes = input_complementarity_axes + + self.force_axes = ( + input_force_axes + if isinstance(input_force_axes, matplotlib.axes.Axes) + else None + ) + + if not isinstance(self.terrain, TerrainDescriptor): + self.terrain = ( + input_terrain + if isinstance(input_terrain, TerrainDescriptor) + else PlanarTerrain() + ) + + +class ContactPointStatePlotter: + def __init__( + self, + settings: ContactPointStatePlotterSettings = ContactPointStatePlotterSettings(), + ): + self.settings = settings + self._complementarity_axes = self.settings.complementarity_axes + self._complementarity_fig = None + self._force_axes = self.settings.force_axes + self._force_fig = None + + def plot_complementarity( + self, + states: list[ContactPointState], + time_s: float | list[float] | np.ndarray = None, + title: str = "Contact Point Complementarity", + ) -> None: + _time_s = copy.deepcopy(time_s) + if _time_s is None or isinstance(_time_s, float) or _time_s.size == 1: + single_step = _time_s if _time_s is not None else 0.0 + _time_s = np.linspace(0, len(states) * single_step, len(states)) + + if len(_time_s) != len(states): + raise ValueError( + "timestep_s and foot_contact_states have different lengths." + ) + + if self._complementarity_axes is None: + self._complementarity_fig, self._complementarity_axes = plt.subplots( + nrows=1, ncols=2 + ) + plt.tight_layout() + + height_function = self.settings.terrain.height_function() + normal_direction_fun = self.settings.terrain.normal_direction_function() + + positions = np.array([height_function(s.p) for s in states]).flatten() + forces = np.array([normal_direction_fun(s.p).T @ s.f for s in states]).flatten() + complementarity_error = np.multiply(positions, forces) + self._complementarity_axes[1].plot(_time_s, complementarity_error) + self._complementarity_axes[1].set_ylabel("Complementarity Error [Nm]") + self._complementarity_axes[1].set_xlabel("Time [s]") + self._complementarity_axes[0].plot(_time_s, positions) + self._complementarity_axes[0].set_ylabel("Height [m]", color="C0") + self._complementarity_axes[0].tick_params(axis="y", color="C0", labelcolor="C0") + self._complementarity_axes[0].set_xlabel("Time [s]") + axes_force = self._complementarity_axes[0].twinx() + axes_force.plot(_time_s, forces, "C1") + axes_force.set_ylabel("Normal Force [N]", color="C1") + axes_force.tick_params(axis="y", color="C1", labelcolor="C1") + axes_force.spines["right"].set_color("C1") + axes_force.spines["left"].set_color("C0") + + if self._complementarity_fig is not None: + self._complementarity_fig.suptitle(title) + plt.draw() + plt.pause(0.001) + plt.show() + + def plot_forces( + self, + states: list[ContactPointState], + time_s: float | list[float] | np.ndarray = None, + title: str = "Contact Point Forces", + ) -> None: + _time_s = copy.deepcopy(time_s) + if _time_s is None or isinstance(_time_s, float) or _time_s.size == 1: + single_step = _time_s if _time_s is not None else 0.0 + _time_s = np.linspace(0, len(states) * single_step, len(states)) + + if len(_time_s) != len(states): + raise ValueError( + "timestep_s and foot_contact_states have different lengths." + ) + + if self._force_axes is None: + self._force_fig, self._force_axes = plt.subplots(nrows=1, ncols=1) + plt.tight_layout() + + forces = np.array([s.f for s in states]) + self._force_axes.plot(_time_s, forces) + self._force_axes.set_ylabel("Force [N]") + self._force_axes.set_xlabel("Time [s]") + + if self._force_fig is not None: + self._force_fig.suptitle(title) + plt.draw() + plt.pause(0.001) + plt.show() + + +@dataclasses.dataclass +class FootContactStatePlotterSettings: + number_of_columns: int = dataclasses.field(default=-1) + terrain: TerrainDescriptor = dataclasses.field(default=None) + + +TFootContactStatePlotter = TypeVar( + "TFootContactStatePlotter", bound="FootContactStatePlotter" +) + + +class FootContactStatePlotter: + def __init__( + self, + settings: FootContactStatePlotterSettings = FootContactStatePlotterSettings(), + ): + self._settings = settings + self._ext_proc_complementarity = None + self._ext_proc_forces = None + self._logger = logging.getLogger("[hippopt::FootContactStatePlotter]") + + def _plot( + self, + states: list[FootContactState], + time_s: float | list[float] | np.ndarray, + title: str, + blocking: bool, + plot_function: Callable[ + [list[FootContactState], np.ndarray, str, int, TerrainDescriptor], None + ], + ) -> multiprocessing.Process | None: + _time_s = copy.deepcopy(time_s) + _states = copy.deepcopy(states) + _terrain = copy.deepcopy(self._settings.terrain) + if _time_s is None or isinstance(_time_s, float) or _time_s.size == 1: + single_step = _time_s if _time_s is not None else 0.0 + _time_s = np.linspace(0, len(states) * single_step, len(states)) + + if len(_time_s) != len(_states): + raise ValueError( + "timestep_s and foot_contact_states have different lengths." + ) + + if len(_states) == 0: + return None + + process = multiprocessing.Process( + target=plot_function, + args=( + _states, + _time_s, + title, + self._settings.number_of_columns, + _terrain, + ), + ) + process.start() + + if blocking: + process.join() + process = None + + return process + + @staticmethod + def _create_complementarity_plot( + states: list[FootContactState], + time_s: np.ndarray, + title: str, + number_of_columns: int, + terrain: TerrainDescriptor, + ) -> None: + number_of_points = len(states[0]) + number_of_plots = number_of_points + 1 + _number_of_columns = ( + math.floor(math.sqrt(number_of_plots)) + if number_of_columns < 1 + else number_of_columns + ) + number_of_rows = math.ceil(number_of_plots / _number_of_columns) + + _fig, axes_list = plt.subplots( + nrows=number_of_rows, + ncols=_number_of_columns, + squeeze=False, + ) + plt.tight_layout() + last_plot_column = number_of_points - _number_of_columns * (number_of_rows - 1) + last_plot = axes_list[number_of_rows - 1][last_plot_column] + _point_plotters = [ + ContactPointStatePlotter( + ContactPointStatePlotterSettings( + input_complementarity_axes=[el, last_plot], + input_terrain=terrain, + ) + ) + for row in axes_list + for el in row + ] + for i in range(last_plot_column + 1, _number_of_columns): + axes_list[number_of_rows - 1][i].remove() + + for p in range(number_of_points): + contact_states = [state[p] for state in states] + _point_plotters[p].plot_complementarity( + states=contact_states, time_s=time_s + ) + + _fig.suptitle(title) + plt.draw() + plt.pause(0.001) + plt.show() + + @staticmethod + def _create_forces_plot( + states: list[FootContactState], + time_s: np.ndarray, + title: str, + number_of_columns: int, + _: TerrainDescriptor, + ) -> None: + number_of_points = len(states[0]) + number_of_plots = number_of_points + _number_of_columns = ( + math.floor(math.sqrt(number_of_plots)) + if number_of_columns < 1 + else number_of_columns + ) + number_of_rows = math.ceil(number_of_plots / _number_of_columns) + + _fig, axes_list = plt.subplots( + nrows=number_of_rows, + ncols=_number_of_columns, + squeeze=False, + ) + plt.tight_layout() + last_plot_column = ( + number_of_points - 1 - _number_of_columns * (number_of_rows - 1) + ) + _point_plotters = [ + ContactPointStatePlotter( + ContactPointStatePlotterSettings( + input_force_axes=el, + ) + ) + for row in axes_list + for el in row + ] + for i in range(last_plot_column + 1, _number_of_columns): + axes_list[number_of_rows - 1][i].remove() + + for p in range(number_of_points): + contact_states = [state[p] for state in states] + _point_plotters[p].plot_forces(states=contact_states, time_s=time_s) + + _fig.suptitle(title) + plt.draw() + plt.pause(0.001) + plt.show() + + def plot_complementarity( + self, + states: list[FootContactState], + time_s: float | list[float] | np.ndarray = None, + title: str = "Foot Contact Complementarity", + blocking: bool = False, + ) -> None: + if self._ext_proc_complementarity is not None: + self._logger.warning( + "A complementarity plot is already running. " + "Make sure to close the previous plot first." + ) + self._ext_proc_complementarity.join() + self._ext_proc_complementarity = None + + self._ext_proc_complementarity = self._plot( + states, + time_s, + title, + blocking, + self._create_complementarity_plot, + ) + + def plot_forces( + self, + states: list[FootContactState], + time_s: float | list[float] | np.ndarray = None, + title: str = "Foot Contact Forces", + blocking: bool = False, + ) -> None: + if self._ext_proc_forces is not None: + self._logger.warning( + "A force plot is already running. " + "Make sure to close the previous plot first." + ) + self._ext_proc_forces.join() + self._ext_proc_forces = None + + self._ext_proc_forces = self._plot( + states, + time_s, + title, + blocking, + self._create_forces_plot, + ) + + def close(self) -> None: + if self._ext_proc_complementarity is not None: + self._ext_proc_complementarity.terminate() + self._ext_proc_complementarity = None + if self._ext_proc_forces is not None: + self._ext_proc_forces.terminate() + self._ext_proc_forces = None + plt.close("all") diff --git a/src/hippopt/robot_planning/utilities/humanoid_state_visualizer.py b/src/hippopt/robot_planning/utilities/humanoid_state_visualizer.py new file mode 100644 index 00000000..f5086878 --- /dev/null +++ b/src/hippopt/robot_planning/utilities/humanoid_state_visualizer.py @@ -0,0 +1,366 @@ +import copy +import dataclasses +import logging +import pathlib +import time + +import ffmpeg +import liecasadi +import numpy as np +from idyntree.visualize import MeshcatVisualizer + +from hippopt.robot_planning.utilities.terrain_visualizer import ( + TerrainVisualizer, + TerrainVisualizerSettings, +) +from hippopt.robot_planning.variables.humanoid import ( + FeetContactPointDescriptors, + HumanoidState, +) + + +@dataclasses.dataclass +class HumanoidStateVisualizerSettings(TerrainVisualizerSettings): + robot_model: str = dataclasses.field(default=None) + considered_joints: list[str] = dataclasses.field(default=None) + contact_points: FeetContactPointDescriptors = dataclasses.field(default=None) + robot_color: list[float] = dataclasses.field(default=None) + com_color: list[float] = dataclasses.field(default=None) + com_radius: float = dataclasses.field(default=None) + contact_points_color: list[float] = dataclasses.field(default=None) + contact_points_radius: float = dataclasses.field(default=None) + contact_forces_color: list[float] = dataclasses.field(default=None) + contact_force_radius: float = dataclasses.field(default=None) + contact_force_scaling: float = dataclasses.field(default=None) + pre_allocated_clones: int = dataclasses.field(default=None) + + def __post_init__(self): + TerrainVisualizerSettings.__post_init__(self) + if self.robot_color is None: + self.robot_color = [1, 1, 1, 0.25] + + if self.com_color is None: + self.com_color = [1, 0, 0, 1] + + if self.com_radius is None: + self.com_radius = 0.02 + + if self.contact_points_color is None: + self.contact_points_color = [0, 0, 0, 1] + + if self.contact_forces_color is None: + self.contact_forces_color = [1, 0, 0, 1] + + if self.contact_points_radius is None: + self.contact_points_radius = 0.01 + + if self.contact_force_radius is None: + self.contact_force_radius = 0.005 + + if self.contact_force_scaling is None: + self.contact_force_scaling = 0.01 + + if self.pre_allocated_clones is None: + self.pre_allocated_clones = 1 + + def is_valid(self) -> bool: + ok = True + logger = logging.getLogger("[hippopt::HumanoidStateVisualizerSettings]") + if not TerrainVisualizerSettings.is_valid(self): + ok = False + if self.robot_model is None: + logger.error("robot_model is not specified.") + ok = False + if self.considered_joints is None: + logger.error("considered_joints is not specified.") + ok = False + if self.contact_points is None: + logger.error("contact_points is not specified.") + ok = False + if len(self.robot_color) != 4: + logger.error("robot_color is not specified correctly.") + ok = False + if len(self.com_color) != 4: + logger.error("com_color is not specified correctly.") + ok = False + if len(self.contact_points_color) != 4: + logger.error("contact_points_color is not specified correctly.") + ok = False + if len(self.contact_forces_color) != 4: + logger.error("contact_forces_color is not specified correctly.") + ok = False + if self.com_radius <= 0: + logger.error("com_radius is not specified correctly.") + ok = False + if self.contact_points_radius <= 0: + logger.error("contact_points_radius is not specified correctly.") + ok = False + if self.contact_force_radius <= 0: + logger.error("contact_force_radius is not specified correctly.") + ok = False + if self.contact_force_scaling <= 0: + logger.error("contact_force_scaling is not specified correctly.") + ok = False + if self.pre_allocated_clones <= 0: + logger.error("pre_allocated_clones is not specified correctly.") + ok = False + return ok + + +class HumanoidStateVisualizer: + def __init__(self, settings: HumanoidStateVisualizerSettings) -> None: + if not settings.is_valid(): + raise ValueError("Settings are not valid.") + self._logger = logging.getLogger("[hippopt::HumanoidStateVisualizer]") + self._settings = settings + self._number_of_clones = self._settings.pre_allocated_clones + self._viz = MeshcatVisualizer() + self._terrain_visualizer = TerrainVisualizer(self._settings, self._viz) + for i in range(self._number_of_clones): + self._allocate_clone(i) + if i != 0: + self._set_clone_visibility(i, False) + + def _allocate_clone(self, index: int) -> None: + self._viz.load_model_from_file( + model_path=self._settings.robot_model, + model_name=f"[{index}]robot", + considered_joints=self._settings.considered_joints, + color=self._settings.robot_color, + ) + self._viz.load_sphere( + radius=self._settings.com_radius, + shape_name=f"[{index}]CoM", + color=self._settings.com_color, + ) + for i, point in enumerate( + (self._settings.contact_points.left + self._settings.contact_points.right) + ): + self._viz.load_sphere( + radius=self._settings.contact_points_radius, + shape_name=f"[{index}]p_{i}", + color=self._settings.contact_points_color, + ) + self._viz.load_arrow( + radius=self._settings.contact_force_radius, + shape_name=f"[{index}]f_{i}", + color=self._settings.contact_forces_color, + ) + + def _set_clone_visibility(self, index: int, visible: bool) -> None: + self._viz.viewer[f"[{index}]robot"].set_property(key="visible", value=visible) + self._viz.viewer[f"[{index}]CoM"].set_property(key="visible", value=visible) + for i, point in enumerate( + (self._settings.contact_points.left + self._settings.contact_points.right) + ): + self._viz.viewer[f"[{index}]p_{i}"].set_property( + key="visible", value=visible + ) + self._viz.viewer[f"[{index}]f_{i}"].set_property( + key="visible", value=visible + ) + + def _update_clone(self, index: int, state: HumanoidState) -> None: + self._viz.set_multibody_system_state( + np.array(state.kinematics.base.position).flatten(), + liecasadi.SO3.from_quat(state.kinematics.base.quaternion_xyzw) + .as_matrix() + .full(), + np.array(state.kinematics.joints.positions).flatten(), + f"[{index}]robot", + ) + self._viz.set_primitive_geometry_transform( + np.array(state.com).flatten(), + np.eye(3), + f"[{index}]CoM", + ) + + for i, point in enumerate( + (state.contact_points.left + state.contact_points.right) + ): + self._viz.set_primitive_geometry_transform( + np.array(point.p).flatten(), + np.eye(3), + f"[{index}]p_{i}", + ) + + point_force = ( + np.array(point.f).flatten() * self._settings.contact_force_scaling + ).flatten() + + self._viz.set_arrow_transform( + origin=np.array(point.p).flatten(), + vector=point_force, + shape_name=f"[{index}]f_{i}", + ) + + def _visualize_single_state( + self, + states: list[HumanoidState], + save: bool, + file_name_stem: str, + ) -> None: + if len(states) > self._number_of_clones: + self._logger.warning( + f"Number of states ({len(states)}) is greater than the number of " + f"allocated clones ({self._number_of_clones}). " + "Creating new clones." + ) + for i in range(self._number_of_clones, len(states)): + self._allocate_clone(i) + self._set_clone_visibility(i, False) + self._number_of_clones += 1 + + for i in range(len(states)): + self._update_clone(i, states[i]) + self._set_clone_visibility(i, True) + + for i in range(len(states), self._number_of_clones): + self._set_clone_visibility(i, False) + + if save: + image = self._viz.viewer.get_image() + image.save(file_name_stem + ".png") + + def _visualize_multiple_states( + self, + states: list[HumanoidState], + timestep_s: float | list[float] | np.ndarray, + time_multiplier: float, + number_of_clones: int, + save: bool, + file_name_stem: str, + ) -> None: + _timestep_s = copy.deepcopy(timestep_s) + if ( + _timestep_s is None + or isinstance(_timestep_s, float) + or _timestep_s.size == 1 + ): + single_step = _timestep_s if _timestep_s is not None else 0.0 + _timestep_s = [single_step] * len(states) + + if len(_timestep_s) != len(states): + raise ValueError("timestep_s and states have different lengths.") + + folder_name = f"{self._settings.working_folder}/{file_name_stem}_frames" + + if save: + pathlib.Path(folder_name).mkdir(parents=True, exist_ok=True) + self._logger.info( + f"Saving visualization frames in {folder_name}. " + "Make sure to have the visualizer open, " + "otherwise the process will hang." + ) + + frame_prefix = "frame_" + + for i in range(number_of_clones, len(states) + 1): + initial_index = i - number_of_clones + visualized_states = states[initial_index:i] + if number_of_clones > 1: + self._logger.info( + f"Visualizing states [{i-number_of_clones + 1},{i}]" + f" of {len(states)}." + ) + else: + self._logger.info(f"Visualizing state {i}/{len(states)}") + start = time.time() + self._visualize_single_state( + visualized_states, + save=save, + file_name_stem=f"{folder_name}/{frame_prefix}{i-number_of_clones:03}", + ) + end = time.time() + elapsed_s = end - start + sleep_time = _timestep_s[i - 1] * time_multiplier - elapsed_s + time.sleep(max(0.0, sleep_time)) + + if save: + if timestep_s is None: + self._logger.warning("timestep_s is None. Saving video with 1.0 fps.") + fps = 1.0 + elif isinstance(timestep_s, list): + if len(timestep_s) > 1: + self._logger.warning( + "The input timestep is a list. " + "Using the average to compute the fps." + ) + fps = 1.0 / (sum(timestep_s) / len(timestep_s)) + else: + fps = 1.0 / timestep_s[0] + elif isinstance(timestep_s, np.ndarray): + if timestep_s.size > 1: + self._logger.warning( + "The input timestep is a list. " + "Using the average to compute the fps." + ) + fps = 1.0 / (sum(timestep_s) / len(timestep_s)) + else: + fps = 1.0 / timestep_s + elif isinstance(timestep_s, float): + fps = 1.0 / timestep_s + else: + self._logger.warning("Using the fps=1.0") + fps = 1.0 + + self.generate_video_from_frames( + file_name_stem=file_name_stem, + frames_folder=folder_name, + frames_prefix=frame_prefix, + fps=fps / time_multiplier, + ) + + def generate_video_from_frames( + self, file_name_stem: str, frames_folder: str, frames_prefix: str, fps: float + ) -> None: + frames = ffmpeg.input( + filename=f"{frames_folder}/{frames_prefix}%3d.png", + framerate=fps, + ) + video = ffmpeg.output( + frames, + f"{self._settings.working_folder}/{file_name_stem}.mp4", + video_bitrate="20M", + ) + try: + ffmpeg.run(video) + except ffmpeg.Error as e: + self._logger.error( + "ffmpeg failed to generate the video. " + "The following output might contain additional information: " + + str(e) + + " stderr: " + + str(e.stderr) + + " stdout: " + + str(e.stdout) + ) + + def visualize( + self, + states: HumanoidState | list[HumanoidState], + timestep_s: float | list[float] | np.ndarray = None, + time_multiplier: float = 1.0, + number_of_clones: int = 1, + save: bool = False, + file_name_stem: str = "humanoid_state_visualization", + ) -> None: + if number_of_clones < 1: + raise ValueError("number_of_clones must be greater than 0.") + + if not isinstance(states, list): + states = [states] + + if number_of_clones < len(states): + self._visualize_multiple_states( + states=states, + timestep_s=timestep_s, + time_multiplier=time_multiplier, + number_of_clones=number_of_clones, + save=save, + file_name_stem=file_name_stem, + ) + else: + self._visualize_single_state( + states=states, save=save, file_name_stem=file_name_stem + ) diff --git a/src/hippopt/robot_planning/utilities/interpolators.py b/src/hippopt/robot_planning/utilities/interpolators.py new file mode 100644 index 00000000..d639e04b --- /dev/null +++ b/src/hippopt/robot_planning/utilities/interpolators.py @@ -0,0 +1,443 @@ +import copy +import math + +import liecasadi +import numpy as np + +from hippopt import StorageType +from hippopt.robot_planning.variables.contacts import ( + ContactPointDescriptor, + FeetContactPhasesDescriptor, + FeetContactPointDescriptors, + FeetContactPoints, + FootContactPhaseDescriptor, + FootContactState, +) +from hippopt.robot_planning.variables.floating_base import ( + FloatingBaseSystemState, + FreeFloatingObjectState, + KinematicTreeState, +) +from hippopt.robot_planning.variables.humanoid import HumanoidState + + +def linear_interpolator( + initial: StorageType, final: StorageType, number_of_points: int +) -> list[StorageType]: + assert not isinstance(initial, list) and not isinstance(final, list) + + initial = np.array(initial) + final = np.array(final) + + if len(initial.shape) < 2: + initial = np.expand_dims(initial, axis=1) + + if len(final.shape) < 2: + final = np.expand_dims(final, axis=1) + + if ( + hasattr(initial, "shape") + and hasattr(final, "shape") + and initial.shape != final.shape + ): + raise ValueError( + f"Initial value has shape {initial.shape}, " + f"but final value has shape {final.shape}." + ) + + t = np.linspace(start=0.0, stop=1.0, num=number_of_points) + output = [] + for t_i in t: + output.append((1 - t_i) * initial + t_i * final) + return output + + +def quaternion_slerp( + initial: StorageType, final: StorageType, number_of_points: int +) -> list[StorageType]: + assert not isinstance(initial, list) and not isinstance(final, list) + + initial = np.array(initial) + final = np.array(final) + + if len(initial.shape) < 2: + initial = np.expand_dims(initial, axis=1) + + if len(final.shape) < 2: + final = np.expand_dims(final, axis=1) + + dot = initial.T @ final + angle = math.acos(dot) + + t = np.linspace(start=0.0, stop=1.0, num=number_of_points) + output = [] + for t_i in t: + output.append( + liecasadi.Quaternion.slerp_step(initial, final, t_i).coeffs() + if abs(angle) > 1e-6 + else initial + ) + return output + + +def transform_interpolator( + initial: liecasadi.SE3, final: liecasadi.SE3, number_of_points: int +) -> list[liecasadi.SE3]: + linear_interpolation = linear_interpolator( + initial=initial.translation(), + final=final.translation(), + number_of_points=number_of_points, + ) + quaternion_interpolation = quaternion_slerp( + initial=initial.rotation().as_quat().coeffs(), + final=final.rotation().as_quat().coeffs(), + number_of_points=number_of_points, + ) + output = [] + for i in range(number_of_points): + output.append( + liecasadi.SE3.from_position_quaternion( + linear_interpolation[i], quaternion_interpolation[i] + ) + ) + return output + + +def foot_contact_state_interpolator( + phases: list[FootContactPhaseDescriptor], + descriptor: list[ContactPointDescriptor], + number_of_points: int, + dt: float, + t0: float = 0.0, +) -> list[FootContactState]: + assert len(phases) > 0 + assert number_of_points > 0 + assert dt > 0.0 + + end_time = t0 + dt * number_of_points + + phases_copy = copy.deepcopy(phases) + + if phases_copy[0].activation_time is None: + deactivation_time = ( + phases_copy[0].deactivation_time + if phases_copy[0].deactivation_time is not None + else t0 + ) + phases_copy[0].activation_time = min(deactivation_time, t0) - dt + + if phases_copy[0].activation_time > t0: + raise ValueError( + f"The first phase activation time " + f"({phases_copy[0].activation_time}) is after " + f"the start time ({t0})." + ) + + for i, phase in enumerate(phases_copy): + if phase.activation_time is None: + raise ValueError( + f"Phase {i} has no activation time, but is not the first phase." + ) + + last = len(phases_copy) - 1 + if phases_copy[last].deactivation_time is None: + phases_copy[last].deactivation_time = ( + max(end_time, phases_copy[last].activation_time) + dt + ) + + if phases_copy[last].deactivation_time < end_time: + raise ValueError( + f"The Last phase deactivation time " + f"({phases_copy[len(phases_copy) - 1].deactivation_time}) is before " + f"the end time ({end_time}, computed from the inputs)." + ) + + for i, phase in enumerate(phases_copy): + if phase.deactivation_time is None: + raise ValueError( + f"Phase {i} has no deactivation time, but is not the last phase." + ) + if phase.activation_time > phase.deactivation_time: + raise ValueError( + f"Phase {i} has an activation time ({phase.activation_time}) " + f"greater than its deactivation time ({phase.deactivation_time})." + ) + + if i < last: + if phase.deactivation_time > phases_copy[i + 1].activation_time: + raise ValueError( + f"Phase {i} has a deactivation time ({phase.deactivation_time}) " + f"greater than the activation time of the next phase " + f"({phases_copy[i + 1].activation_time})." + ) + + output = [] + + def append_stance_phase( + stance_phase: FootContactPhaseDescriptor, + points: int, + ) -> None: + for _ in range(points): + foot_state = FootContactState.from_parent_frame_transform( + descriptor=descriptor, transform=stance_phase.transform + ) + for point in foot_state: + point.f = stance_phase.force + output.append(foot_state) + + def append_swing_phase( + start_phase: FootContactPhaseDescriptor, + end_phase: FootContactPhaseDescriptor, + points: int, + ): + full_swing_points = int( + np.ceil((end_phase.activation_time - start_phase.deactivation_time) / dt) + ) + + if start_phase.mid_swing_transform is None: + start_phase.mid_swing_transform = ( + liecasadi.SE3.from_translation_and_rotation( + ( + start_phase.transform.translation() + + end_phase.transform.translation() + ) + / 2, + end_phase.transform.rotation(), + ) + ) + + mid_swing_points = min(round(full_swing_points / 2), points) + mid_swing_transforms = transform_interpolator( + start_phase.transform, start_phase.mid_swing_transform, mid_swing_points + ) + for transform in mid_swing_transforms: + foot_state = FootContactState.from_parent_frame_transform( + descriptor=descriptor, transform=transform + ) + for point in foot_state: + point.f = np.zeros((3, 1)) + output.append(foot_state) + second_half_points = points - mid_swing_points + if second_half_points == 0: + return + second_half_transforms = transform_interpolator( + start_phase.mid_swing_transform, end_phase.transform, second_half_points + ) + for transform in second_half_transforms: + foot_state = FootContactState.from_parent_frame_transform( + descriptor=descriptor, transform=transform + ) + for point in foot_state: + point.f = np.zeros((3, 1)) + output.append(foot_state) + + if len(phases_copy) == 1 or phases_copy[0].deactivation_time >= end_time: + append_stance_phase(phases_copy[0], number_of_points) + return output + + i = 0 + activation_time = phases_copy[0].activation_time + while activation_time < t0: + if phases_copy[i].deactivation_time > t0: + break + i += 1 + activation_time = phases_copy[i].activation_time + + if activation_time > t0: + previous_active_phase = phases_copy[i - 1] + new_t0 = previous_active_phase.deactivation_time - dt + advance_points = int(np.ceil((t0 - new_t0) / dt)) + increased_output = foot_contact_state_interpolator( + phases=phases_copy, + descriptor=descriptor, + number_of_points=number_of_points + advance_points, + dt=dt, + t0=new_t0, + ) + return increased_output[advance_points:] + + remaining_points = number_of_points + while i < len(phases_copy) - 1: + phase = phases_copy[i] + next_phase = phases_copy[i + 1] + + stance_points = int( + np.ceil((phase.deactivation_time - max(phase.activation_time, t0)) / dt) + ) + stance_points = min(stance_points, remaining_points) + + append_stance_phase(phase, stance_points) + remaining_points -= stance_points + + if remaining_points == 0: + return output + + swing_points = int( + np.ceil((next_phase.activation_time - phase.deactivation_time) / dt) + ) + + swing_points = min(swing_points, remaining_points) + + if swing_points == 0: + continue + + append_swing_phase(phase, next_phase, swing_points) + remaining_points -= swing_points + + if remaining_points == 0: + return output + i += 1 + + last_phase = phases_copy[len(phases_copy) - 1] + append_stance_phase(last_phase, remaining_points) + return output + + +def feet_contact_points_interpolator( + phases: FeetContactPhasesDescriptor, + descriptor: FeetContactPointDescriptors, + number_of_points: int, + dt: float, + t0: float = 0.0, +) -> list[FeetContactPoints]: + left_output = foot_contact_state_interpolator( + phases=phases.left, + descriptor=descriptor.left, + number_of_points=number_of_points, + dt=dt, + t0=t0, + ) + right_output = foot_contact_state_interpolator( + phases=phases.right, + descriptor=descriptor.right, + number_of_points=number_of_points, + dt=dt, + t0=t0, + ) + + assert len(left_output) == len(right_output) == number_of_points + + return [ + FeetContactPoints(left=left, right=right) + for left, right in zip(left_output, right_output) + ] + + +def free_floating_object_state_interpolator( + initial_state: FreeFloatingObjectState, + final_state: FreeFloatingObjectState, + number_of_points: int, +) -> list[FreeFloatingObjectState]: + position_interpolation = linear_interpolator( + initial=initial_state.position, + final=final_state.position, + number_of_points=number_of_points, + ) + quaternion_interpolation = quaternion_slerp( + initial=initial_state.quaternion_xyzw, + final=final_state.quaternion_xyzw, + number_of_points=number_of_points, + ) + assert ( + len(position_interpolation) == len(quaternion_interpolation) == number_of_points + ) + return [ + FreeFloatingObjectState(position=position, quaternion_xyzw=quaternion) + for position, quaternion in zip( + position_interpolation, quaternion_interpolation + ) + ] + + +def kinematic_tree_state_interpolator( + initial_state: KinematicTreeState, + final_state: KinematicTreeState, + number_of_points: int, +) -> list[KinematicTreeState]: + if len(initial_state.positions) != len(final_state.positions): + raise ValueError( + f"Initial state has {len(initial_state.positions)} joints, " + f"but final state has {len(final_state.positions)} joints." + ) + + positions_interpolation = linear_interpolator( + initial=initial_state.positions, + final=final_state.positions, + number_of_points=number_of_points, + ) + return [ + KinematicTreeState(positions=positions) for positions in positions_interpolation + ] + + +def floating_base_system_state_interpolator( + initial_state: FloatingBaseSystemState, + final_state: FloatingBaseSystemState, + number_of_points: int, +) -> list[FloatingBaseSystemState]: + base_interpolation = free_floating_object_state_interpolator( + initial_state=initial_state.base, + final_state=final_state.base, + number_of_points=number_of_points, + ) + joints_interpolation = kinematic_tree_state_interpolator( + initial_state=initial_state.joints, + final_state=final_state.joints, + number_of_points=number_of_points, + ) + + assert len(base_interpolation) == len(joints_interpolation) == number_of_points + + return [ + FloatingBaseSystemState(base=base, joints=joints) + for base, joints in zip(base_interpolation, joints_interpolation) + ] + + +def humanoid_state_interpolator( + initial_state: HumanoidState, + final_state: HumanoidState, + contact_phases: FeetContactPhasesDescriptor, + contact_descriptor: FeetContactPointDescriptors, + number_of_points: int, + dt: float, + t0: float = 0.0, +): + contact_points_interpolation = feet_contact_points_interpolator( + phases=contact_phases, + descriptor=contact_descriptor, + number_of_points=number_of_points, + dt=dt, + t0=t0, + ) + kinematics_interpolation = floating_base_system_state_interpolator( + initial_state=initial_state.kinematics, + final_state=final_state.kinematics, + number_of_points=number_of_points, + ) + com_interpolation = linear_interpolator( + initial=initial_state.com, + final=final_state.com, + number_of_points=number_of_points, + ) + + assert ( + len(contact_points_interpolation) + == len(kinematics_interpolation) + == len(com_interpolation) + == number_of_points + ) + + output = [] + for points, kin, com in zip( + contact_points_interpolation, kinematics_interpolation, com_interpolation + ): + output_state = HumanoidState( + contact_point_descriptors=contact_descriptor, + number_of_joints=len(initial_state.kinematics.joints.positions), + ) + output_state.contact_points = points + output_state.kinematics = kin + output_state.com = com + output.append(output_state) + return output diff --git a/src/hippopt/robot_planning/utilities/planar_terrain.py b/src/hippopt/robot_planning/utilities/planar_terrain.py new file mode 100644 index 00000000..6f5d73e5 --- /dev/null +++ b/src/hippopt/robot_planning/utilities/planar_terrain.py @@ -0,0 +1,41 @@ +import casadi as cs + +from hippopt.robot_planning.utilities.terrain_descriptor import TerrainDescriptor + + +class PlanarTerrain(TerrainDescriptor): + def create_height_function(self) -> cs.Function: + point_position = cs.MX.sym(self.get_point_position_name(), 3) + + return cs.Function( + "planar_terrain_height", + [point_position], + [point_position[2]], + [self.get_point_position_name()], + ["point_height"], + self._options, + ) + + def create_normal_direction_function(self) -> cs.Function: + point_position = cs.MX.sym(self.get_point_position_name(), 3) + + return cs.Function( + "planar_terrain_normal", + [point_position], + [cs.MX.eye(3)[:, 2]], + [self.get_point_position_name()], + ["normal_direction"], + self._options, + ) + + def create_orientation_function(self) -> cs.Function: + point_position = cs.MX.sym(self.get_point_position_name(), 3) + + return cs.Function( + "planar_terrain_orientation", + [point_position], + [cs.MX.eye(3)], + [self.get_point_position_name()], + ["plane_rotation"], + self._options, + ) diff --git a/src/hippopt/robot_planning/utilities/smooth_terrain.py b/src/hippopt/robot_planning/utilities/smooth_terrain.py new file mode 100644 index 00000000..ceef2530 --- /dev/null +++ b/src/hippopt/robot_planning/utilities/smooth_terrain.py @@ -0,0 +1,475 @@ +import dataclasses +from typing import TypeVar + +import casadi as cs +import numpy as np + +from hippopt.robot_planning.utilities.terrain_descriptor import TerrainDescriptor +from hippopt.robot_planning.utilities.terrain_sum import TerrainSum +from hippopt.robot_planning.utilities.terrain_visualizer import ( + TerrainVisualizer, + TerrainVisualizerSettings, +) + +TSmoothTerrain = TypeVar("TSmoothTerrain", bound="SmoothTerrain") + + +@dataclasses.dataclass +class SmoothTerrain(TerrainDescriptor): + """ + Smooth terrain is a terrain with a smooth height function. + The height is defined as follows: + h(x, y) = exp(−g(x, y)^(2s)) π(x, y). + + Here, g(x, y) is the equation of a curve in the xy-plane, and π(x, y) is the + equation of a plane in the xy-plane, defining the shape of the terrain when + exp(−g(x, y)^(2s)) = 1, i.e. when g(x, y) = 0. The parameter s ≥ 1 controls the + smoothness of the terrain. + + Independently of the value of s, the value of exp(−g(.)^(2s)) is always passing + through 1 when g(x,y) = 0, and through 1/e when |g(x,y)| = 1. Then, it will + tend to zero as |g(x, y)| → ∞. The parameter s controls how fast exp(−g(.)^(2s)) + tends to zero as |g(x, y)| grows. By multiplying times the equation of a plane, i.e. + π(x, y), we can control the inclination of the top surface when g(x, y) = 0. + Instead, |g(x, y)| = 1 controls the shape of the terrain at height 1/e * π(x, y). + + For example, to define a classical step with a square base of dimension l, + we can use: + g(x, y) = |2/l * x|^p + |2/l * y|^p + π(x, y) = 1 + where p ≥ 1 is a parameter controlling the sharpness of the edges of the square + at height 1/e. + Here g(x, y) = 0 only at the origin, hence, the top surface is flat at the origin, + and the parameter s determines the sharpness of the top face. + + It is then possible to specify a position offset and a transformation matrix to + move, rotate, and scale the terrain. + In particular, we have that: + [x, y, z]^T = R^T * ([x_i, y_i, z_i]^T - [x_offset, y_offset, z_offset]^T) + where [x_i, y_i, z_i]^T is the position of the point in the inertial frame, R is + the transformation matrix, and [x_offset, y_offset, z_offset]^T is the offset. + When applying a transformation matrix, it is also possible to rotate the terrain. + Because of the choice of the projection method, the modified z axis needs to be + parallel to the original one (i.e. no rotations around x or y are allowed), + otherwise there are issues in the computation of the normal direction. + + """ + + _shape_function: cs.Function = dataclasses.field(default=None) + _top_surface_function: cs.Function = dataclasses.field(default=None) + _sharpness: float = dataclasses.field(default=None) + _offset: np.ndarray = dataclasses.field(default=None) + _transformation_matrix: np.ndarray = dataclasses.field(default=None) + + shape_function: dataclasses.InitVar[cs.Function] = dataclasses.field(default=None) + top_surface_function: dataclasses.InitVar[cs.Function] = dataclasses.field( + default=None + ) + sharpness: dataclasses.InitVar[float] = dataclasses.field(default=None) + offset: dataclasses.InitVar[np.ndarray] = dataclasses.field(default=None) + transformation_matrix: dataclasses.InitVar[np.ndarray] = dataclasses.field( + default=None + ) + + def __post_init__( + self, + point_position_name: str = None, + options: dict = None, + name: str = None, + shape_function: cs.Function = None, + top_surface_function: cs.Function = None, + sharpness: float = None, + offset: np.ndarray = None, + transformation_matrix: np.ndarray = None, + ): + TerrainDescriptor.__post_init__(self, point_position_name, options, name) + + if self._sharpness is None: + self._sharpness = 1.0 + + if self._offset is None: + self._offset = np.zeros(3) + + if self._transformation_matrix is None: + self._transformation_matrix = np.eye(3) + + point_position_xy = cs.MX.sym(self.get_point_position_name() + "_xy", 2) + if self._shape_function is None: + self._shape_function = cs.Function( + "smooth_terrain_shape", + [point_position_xy], + [0], + [point_position_xy.name()], + ["g"], + self._options, + ) + if self._top_surface_function is None: + self._top_surface_function = cs.Function( + "smooth_terrain_top_surface", + [point_position_xy], + [0], + [point_position_xy.name()], + ["pi"], + self._options, + ) + + self.set_terrain( + shape_function=shape_function, + top_surface_function=top_surface_function, + sharpness=sharpness, + offset=offset, + transformation_matrix=transformation_matrix, + ) + + def set_terrain( + self, + shape_function: cs.Function = None, + top_surface_function: cs.Function = None, + sharpness: float = None, + offset: np.ndarray = None, + transformation_matrix: np.ndarray = None, + ) -> None: + if isinstance(shape_function, cs.Function): + if shape_function.n_in() != 1: + raise ValueError( + "The shape function must have exactly one input argument." + ) + if shape_function.n_out() != 1: + raise ValueError( + "The shape function must have exactly one output argument." + ) + if shape_function.numel_in() != 2: + raise ValueError( + "The input argument of the shape function must be a 2D vector." + ) + self._shape_function = shape_function + + if isinstance(top_surface_function, cs.Function): + if top_surface_function.n_in() != 1: + raise ValueError( + "The top surface function must have exactly one input argument." + ) + if top_surface_function.n_out() != 1: + raise ValueError( + "The top surface function must have exactly one output argument." + ) + if top_surface_function.numel_in() != 2: + raise ValueError( + "The input argument of the top surface function" + " must be a 2D vector." + ) + self._top_surface_function = top_surface_function + + if sharpness is not None: + if sharpness < 1: + raise ValueError( + "The sharpness parameter must be greater than or equal to 1." + ) + self._sharpness = sharpness + + if offset is not None: + if not isinstance(offset, np.ndarray): + raise TypeError("The offset must be a numpy array.") + if offset.size != 3: + raise ValueError("The offset must be a 3D vector.") + self._offset = offset + + if transformation_matrix is not None: + if not isinstance(transformation_matrix, np.ndarray): + raise TypeError("The transformation matrix must be a numpy matrix.") + if transformation_matrix.shape != (3, 3): + raise ValueError("The transformation matrix must be a 3x3 matrix.") + if ( + np.abs(np.linalg.det(transformation_matrix)) < 1e-6 + or (np.linalg.norm(transformation_matrix, axis=0) < 1e-6).any() + ): + raise ValueError( + "The transformation matrix must be invertible and have a non-zero" + " norm for each column." + ) + if np.abs(np.dot(transformation_matrix[:, 2], [0, 0, 1])) < ( + 1 - 1e-6 + ) * np.linalg.norm(transformation_matrix[:, 2]): + raise ValueError( + "The transformation matrix should not change" + " the z axis orientation." + ) + self._transformation_matrix = transformation_matrix + + self.invalidate_functions() + + def create_height_function(self) -> cs.Function: + point_position = cs.MX.sym(self.get_point_position_name(), 3) + + position_in_terrain_frame = np.linalg.inv(self._transformation_matrix) @ cs.MX( + point_position - self._offset, + ) + + shape = self._shape_function(position_in_terrain_frame[:2]) + top_surface = self._top_surface_function(position_in_terrain_frame[:2]) + + z_terrain = cs.exp(-(shape ** (2 * self._sharpness))) * top_surface + terrain_position = ( + self._transformation_matrix + @ cs.vertcat(position_in_terrain_frame[:2], z_terrain) + + self._offset + ) + + height = point_position[2] - terrain_position[2] + + return cs.Function( + "smooth_terrain_height", + [point_position], + [height], + [self.get_point_position_name()], + ["point_height"], + self._options, + ) + + def __add__(self, other: TerrainDescriptor) -> TerrainSum: + return TerrainSum.add(self, other) + + def __radd__(self, other: TerrainDescriptor) -> TerrainSum: + return TerrainSum.add(other, self) + + @staticmethod + def _top_expression_from_normal( + height: float, point_position_xy: cs.MX, top_normal_direction: np.ndarray + ): + top_expression = cs.MX(height) + if top_normal_direction is not None: + if not isinstance(top_normal_direction, np.ndarray): + raise TypeError("The top normal direction must be a numpy array.") + if top_normal_direction.size != 3: + raise ValueError("The top normal direction must be a 3D vector.") + + norm = np.linalg.norm(top_normal_direction) + if norm < 1e-6: + raise ValueError("The top normal direction must be non-zero.") + + top_normal_direction = top_normal_direction / norm + if np.abs(top_normal_direction[2]) < 1e-6: + raise ValueError( + "The top normal direction must not be parallel to the xy-plane." + ) + top_expression = cs.MX( + -top_normal_direction[0] + / top_normal_direction[2] + * point_position_xy[0] + - top_normal_direction[1] + / top_normal_direction[2] + * point_position_xy[1] + + height + ) + return top_expression + + @classmethod + def step( + cls, + length: float, + width: float, + height: float, + edge_sharpness: float = 5, + side_sharpness: float = 10, + position: np.ndarray = None, + orientation: float = 0.0, + top_normal_direction: np.ndarray = None, + point_position_name: str = None, + options: dict = None, + name: str = None, + ) -> TSmoothTerrain: + if edge_sharpness < 1: + raise ValueError("The edge sharpness must be greater than or equal to 1.") + + name = ( + name + if isinstance(name, str) + else f"step_{length:.2f}x{width:.2f}x{height:.2f}" + ) + options = {} if options is None else options + + output = cls( + point_position_name=point_position_name, + options=options, + name=name, + ) + + rotation_z = np.array( + [ + [np.cos(orientation), -np.sin(orientation), 0], + [np.sin(orientation), np.cos(orientation), 0], + [0, 0, 1], + ] + ) + point_position_xy = cs.MX.sym(output.get_point_position_name() + "_xy", 2) + + top_expression = cls._top_expression_from_normal( + height, point_position_xy, top_normal_direction + ) + + shape_function = cs.Function( + "step_shape", + [point_position_xy], + [ + (2 / length * point_position_xy[0]) ** (2 * edge_sharpness) + + (2 / width * point_position_xy[1]) ** (2 * edge_sharpness) + ], + [point_position_xy.name()], + ["g"], + options, + ) + top_surface_function = cs.Function( + "step_top_surface", + [point_position_xy], + [top_expression], + [point_position_xy.name()], + ["pi"], + options, + ) + output.set_terrain( + shape_function=shape_function, + top_surface_function=top_surface_function, + sharpness=side_sharpness, + offset=position, + transformation_matrix=rotation_z, + ) + return output + + @classmethod + def cylinder( + cls, + radius: float, + height: float, + side_sharpness: float = 10, + top_normal_direction: np.ndarray = None, + position: np.ndarray = None, + point_position_name: str = None, + options: dict = None, + name: str = None, + ) -> TSmoothTerrain: + name = name if isinstance(name, str) else f"cylinder_{radius:.2f}x{height:.2f}" + options = {} if options is None else options + + output = cls( + point_position_name=point_position_name, + options=options, + name=name, + ) + + point_position_xy = cs.MX.sym(output.get_point_position_name() + "_xy", 2) + + top_expression = cls._top_expression_from_normal( + height, point_position_xy, top_normal_direction + ) + + shape_function = cs.Function( + "cylinder_shape", + [point_position_xy], + [ + (1 / radius * point_position_xy[0]) ** 2 + + (1 / radius * point_position_xy[1]) ** 2 + ], + [point_position_xy.name()], + ["g"], + options, + ) + top_surface_function = cs.Function( + "cylinder_top_surface", + [point_position_xy], + [top_expression], + [point_position_xy.name()], + ["pi"], + options, + ) + output.set_terrain( + shape_function=shape_function, + top_surface_function=top_surface_function, + sharpness=side_sharpness, + offset=position, + ) + return output + + @classmethod + def plane( + cls, + normal_direction: np.ndarray = None, + position: np.ndarray = None, + point_position_name: str = None, + options: dict = None, + name: str = None, + ) -> TSmoothTerrain: + name = ( + name + if isinstance(name, str) + else f"plane_{normal_direction[0]:.2f}x{normal_direction[1]:.2f}" + f"x{normal_direction[2]:.2f}" + ) + options = {} if options is None else options + + output = cls( + point_position_name=point_position_name, + options=options, + name=name, + ) + + point_position_xy = cs.MX.sym(output.get_point_position_name() + "_xy", 2) + + top_expression = cls._top_expression_from_normal( + 0, point_position_xy, normal_direction + ) + + shape_function = cs.Function( + "plane_shape", + [point_position_xy], + [0], + [point_position_xy.name()], + ["g"], + options, + ) + top_surface_function = cs.Function( + "plane_top_surface", + [point_position_xy], + [top_expression], + [point_position_xy.name()], + ["pi"], + options, + ) + output.set_terrain( + shape_function=shape_function, + top_surface_function=top_surface_function, + sharpness=1, + offset=position, + ) + return output + + +if __name__ == "__main__": + viz_settings = TerrainVisualizerSettings() + viz_settings.terrain = ( + SmoothTerrain.step( + length=1.0, + width=0.5, + height=0.5, + position=np.array([-0.5, -0.5, 0.0]), + orientation=np.pi / 4, + top_normal_direction=np.array([0.5, 1.0, 1.0]), + ) + + SmoothTerrain.cylinder( + radius=0.5, height=1.0, position=np.array([0.8, 0.5, 0.0]) + ) + + SmoothTerrain.step( + length=0.5, + width=0.5, + height=0.5, + edge_sharpness=2, + position=np.array([-0.5, 0.5, 0.0]), + ) + + SmoothTerrain.plane( + normal_direction=np.array([0.5, 0.1, 1.0]), + position=np.array([0.0, 0.0, 0.0]), + ) + ) + viz_settings.overwrite_terrain_files = True + viz_settings.draw_terrain_frames = True + viz = TerrainVisualizer(viz_settings) + input("Press Enter to exit.") diff --git a/src/hippopt/robot_planning/utilities/terrain_descriptor.py b/src/hippopt/robot_planning/utilities/terrain_descriptor.py new file mode 100644 index 00000000..a1198bd0 --- /dev/null +++ b/src/hippopt/robot_planning/utilities/terrain_descriptor.py @@ -0,0 +1,111 @@ +import abc +import dataclasses + +import casadi as cs + + +@dataclasses.dataclass +class TerrainDescriptor(abc.ABC): + _height_function: cs.Function | None = dataclasses.field(default=None) + _normal_direction_function: cs.Function | None = dataclasses.field(default=None) + _orientation_function: cs.Function | None = dataclasses.field(default=None) + _point_position_name: str = dataclasses.field(default="point_position") + _name: str = dataclasses.field(default=None) + _options: dict = dataclasses.field(default_factory=dict) + point_position_name: dataclasses.InitVar[str] = dataclasses.field( + default="point_position" + ) + options: dataclasses.InitVar[dict] = dataclasses.field(default=None) + name: dataclasses.InitVar[str] = dataclasses.field(default=None) + + def __post_init__( + self, + point_position_name: str = None, + options: dict = None, + name: str = None, + ): + self.change_options(point_position_name, options) + if name is not None: + self.set_name(name) + + def change_options( + self, point_position_name: str = None, options: dict = None, **_ + ) -> None: + if options is not None: + self._options = options + if point_position_name is not None: + self._point_position_name = point_position_name + self.invalidate_functions() + + @abc.abstractmethod + def create_height_function(self) -> cs.Function: + pass + + def create_normal_direction_function(self) -> cs.Function: + point_position = cs.MX.sym(self.get_point_position_name(), 3) + + # The normal direction is the gradient of the implicit function h(x, y, z) = 0 + height_gradient = cs.gradient( + self.height_function()(point_position), point_position + ) + normal_direction = height_gradient / cs.norm_2(height_gradient) + + return cs.Function( + "terrain_normal", + [point_position], + [normal_direction], + [self.get_point_position_name()], + ["normal_direction"], + self._options, + ) + + def create_orientation_function(self) -> cs.Function: + point_position = cs.MX.sym(self.get_point_position_name(), 3) + + normal_direction = self.normal_direction_function()(point_position) + y_direction = cs.cross(normal_direction, cs.DM([1, 0, 0])) + x_direction = cs.cross(y_direction, normal_direction) + x_direction = x_direction / cs.norm_2(x_direction) + # Make sure the y direction is orthogonal even after the transformation + y_direction = cs.cross(normal_direction, x_direction) + + return cs.Function( + "terrain_orientation", + [point_position], + [cs.horzcat(x_direction, y_direction, normal_direction)], + [self.get_point_position_name()], + ["plane_rotation"], + self._options, + ) + + def height_function(self) -> cs.Function: + if not isinstance(self._height_function, cs.Function): + self._height_function = self.create_height_function() + + return self._height_function + + def normal_direction_function(self) -> cs.Function: + if not isinstance(self._normal_direction_function, cs.Function): + self._normal_direction_function = self.create_normal_direction_function() + + return self._normal_direction_function + + def orientation_function(self) -> cs.Function: + if not isinstance(self._orientation_function, cs.Function): + self._orientation_function = self.create_orientation_function() + + return self._orientation_function + + def get_point_position_name(self) -> str: + return self._point_position_name + + def set_name(self, name: str) -> None: + self._name = name + + def get_name(self) -> str: + return self._name if isinstance(self._name, str) else self.__class__.__name__ + + def invalidate_functions(self) -> None: + self._height_function = None + self._normal_direction_function = None + self._orientation_function = None diff --git a/src/hippopt/robot_planning/utilities/terrain_sum.py b/src/hippopt/robot_planning/utilities/terrain_sum.py new file mode 100644 index 00000000..bf813663 --- /dev/null +++ b/src/hippopt/robot_planning/utilities/terrain_sum.py @@ -0,0 +1,55 @@ +import dataclasses +from typing import TypeVar + +import casadi as cs + +from hippopt.robot_planning.utilities.terrain_descriptor import TerrainDescriptor + +TTerrainSum = TypeVar("TTerrainSum", bound="TerrainSum") + + +class TerrainSum(TerrainDescriptor): + _lhs: TerrainDescriptor = dataclasses.field(default=None) + _rhs: TerrainDescriptor = dataclasses.field(default=None) + + def set_terms(self, lhs: TerrainDescriptor, rhs: TerrainDescriptor): + self._lhs = lhs + self._rhs = rhs + + def create_height_function(self) -> cs.Function: + assert ( + self._lhs is not None and self._rhs is not None + ), "Both lhs and rhs must be provided" + + point_position = cs.MX.sym(self.get_point_position_name(), 3) + + # We need to subtract the point height because otherwise is counted twice + return cs.Function( + "terrain_sum_height", + [point_position], + [ + self._lhs.height_function()(point_position) + + self._rhs.height_function()(point_position) + - point_position[2] + ], + [self.get_point_position_name()], + ["point_height"], + self._options, + ) + + @classmethod + def add(cls, lhs: TerrainDescriptor, rhs: TerrainDescriptor) -> TTerrainSum: + output = cls( + point_position_name=lhs.get_point_position_name(), + options=lhs._options, + name=f"{lhs.get_name()}+{rhs.get_name()}", + ) + output.set_terms(lhs, rhs) + + return output + + def __add__(self, other: TTerrainSum) -> TTerrainSum: + return TerrainSum.add(self, other) + + def __radd__(self, other: TTerrainSum) -> TTerrainSum: + return TerrainSum.add(other, self) diff --git a/src/hippopt/robot_planning/utilities/terrain_visualizer.py b/src/hippopt/robot_planning/utilities/terrain_visualizer.py new file mode 100644 index 00000000..bf5ce0f9 --- /dev/null +++ b/src/hippopt/robot_planning/utilities/terrain_visualizer.py @@ -0,0 +1,312 @@ +import dataclasses +import logging +import os + +import numpy as np +from idyntree.visualize import MeshcatVisualizer + +import hippopt.deps.surf2stl as surf2stl +from hippopt.robot_planning.utilities.terrain_descriptor import TerrainDescriptor + + +@dataclasses.dataclass +class TerrainVisualizerSettings: + terrain: TerrainDescriptor = dataclasses.field(default=None) + terrain_color: list[float] = dataclasses.field(default=None) + working_folder: str = dataclasses.field(default=None) + terrain_mesh_axis_points: int = dataclasses.field(default=None) + terrain_x_limits: list[float] = dataclasses.field(default=None) + terrain_y_limits: list[float] = dataclasses.field(default=None) + overwrite_terrain_files: bool = dataclasses.field(default=None) + draw_terrain_normals: bool = dataclasses.field(default=None) + terrain_normals_color: list[float] = dataclasses.field(default=None) + terrain_normals_radius: float = dataclasses.field(default=None) + terrain_normal_axis_points: int = dataclasses.field(default=None) + terrain_normal_length: float = dataclasses.field(default=None) + draw_terrain_frames: bool = dataclasses.field(default=None) + terrain_frames_opacity: float = dataclasses.field(default=None) + terrain_frames_axis_length: float = dataclasses.field(default=None) + terrain_frames_axis_radius: float = dataclasses.field(default=None) + terrain_frames_axis_points: int = dataclasses.field(default=None) + + def __post_init__(self): + if self.terrain_color is None: + self.terrain_color = [0.5, 0.5, 0.5, 0.75] + + if self.working_folder is None: + self.working_folder = "./" + + if self.terrain_mesh_axis_points is None: + self.terrain_mesh_axis_points = 200 + + if self.terrain_x_limits is None: + self.terrain_x_limits = [-1.5, 1.5] + + if self.terrain_y_limits is None: + self.terrain_y_limits = [-1.5, 1.5] + + if self.overwrite_terrain_files is None: + self.overwrite_terrain_files = False + + if self.draw_terrain_normals is None: + self.draw_terrain_normals = False + + if self.terrain_normals_color is None: + self.terrain_normals_color = [1.0, 0.0, 0.0, 1.0] + + if self.terrain_normals_radius is None: + self.terrain_normals_radius = 0.01 + + if self.terrain_normal_axis_points is None: + self.terrain_normal_axis_points = 20 + + if self.terrain_normal_length is None: + self.terrain_normal_length = 0.1 + + if self.draw_terrain_frames is None: + self.draw_terrain_frames = False + + if self.terrain_frames_opacity is None: + self.terrain_frames_opacity = 0.5 + + if self.terrain_frames_axis_length is None: + self.terrain_frames_axis_length = 0.1 + + if self.terrain_frames_axis_radius is None: + self.terrain_frames_axis_radius = 0.01 + + if self.terrain_frames_axis_points is None: + self.terrain_frames_axis_points = 20 + + def is_valid(self) -> bool: + ok = True + logger = logging.getLogger("[hippopt::TerrainVisualizerSettings]") + if not os.access(self.working_folder, os.W_OK): + logger.error("working_folder is not writable.") + ok = False + if len(self.terrain_color) != 4: + logger.error("terrain_color is not specified correctly.") + ok = False + if self.terrain_mesh_axis_points <= 0: + logger.error("terrain_mesh_axis_points is not specified correctly.") + ok = False + if self.terrain is None: + logger.error("terrain is not specified.") + ok = False + if len(self.terrain_x_limits) != 2 or ( + self.terrain_x_limits[0] >= self.terrain_x_limits[1] + ): + logger.error("terrain_x_limits are not specified correctly.") + ok = False + if len(self.terrain_y_limits) != 2 or ( + self.terrain_y_limits[0] >= self.terrain_y_limits[1] + ): + logger.error("terrain_y_limits are not specified correctly.") + ok = False + if len(self.terrain_normals_color) != 4: + logger.error("terrain_normals_color is not specified correctly.") + ok = False + if self.terrain_normals_radius <= 0: + logger.error("terrain_normals_radius is not specified correctly.") + ok = False + if self.terrain_normal_axis_points <= 0: + logger.error("terrain_normal_axis_points is not specified correctly.") + ok = False + if self.terrain_frames_axis_points <= 0: + logger.error("terrain_frames_axis_points is not specified correctly.") + ok = False + if self.terrain_frames_opacity < 0 or self.terrain_frames_opacity > 1: + logger.error("terrain_frames_opacity is not specified correctly.") + ok = False + if self.terrain_frames_axis_radius <= 0: + logger.error("terrain_frames_axis_radius is not specified correctly.") + ok = False + + return ok + + +class TerrainVisualizer: + def __init__( + self, settings: TerrainVisualizerSettings, viz: MeshcatVisualizer = None + ) -> None: + if not settings.is_valid(): + raise ValueError("Settings are not valid.") + self._logger = logging.getLogger("[hippopt::TerrainVisualizer]") + self._settings = settings + mesh_file = self._create_terrain_mesh() + self._create_terrain_urdf(mesh_file) + self._viz = MeshcatVisualizer() if viz is None else viz + self._viz.load_model_from_file( + model_path=os.path.join( + self._settings.working_folder, + self._settings.terrain.get_name() + ".urdf", + ), + model_name="terrain", + color=self._settings.terrain_color, + ) + if self._settings.draw_terrain_normals: + self._draw_terrain_normals() + if self._settings.draw_terrain_frames: + self._draw_terrain_frames() + + def _create_terrain_urdf(self, mesh_filename: str) -> None: + filename = self._settings.terrain.get_name() + ".urdf" + full_filename = os.path.join(self._settings.working_folder, filename) + if os.path.exists(full_filename): + if self._settings.overwrite_terrain_files: + self._logger.info(f"Overwriting {filename}") + else: + self._logger.info(f"{filename} already exists. Skipping creation.") + return + + with open(full_filename, "w") as f: + f.write( + f""" + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + """ + ) + + def _create_terrain_mesh(self) -> str: + filename = self._settings.terrain.get_name() + ".stl" + full_filename = os.path.join(self._settings.working_folder, filename) + if os.path.exists(full_filename): + if self._settings.overwrite_terrain_files: + self._logger.info(f"Overwriting {filename}") + else: + self._logger.info(f"{filename} already exists. Skipping creation.") + return full_filename + x_step = ( + self._settings.terrain_x_limits[1] - self._settings.terrain_x_limits[0] + ) / self._settings.terrain_mesh_axis_points + y_step = ( + self._settings.terrain_y_limits[1] - self._settings.terrain_y_limits[0] + ) / self._settings.terrain_mesh_axis_points + x = np.arange( + self._settings.terrain_x_limits[0], + self._settings.terrain_x_limits[1] + x_step, + x_step, + ) + y = np.arange( + self._settings.terrain_y_limits[0], + self._settings.terrain_y_limits[1] + y_step, + y_step, + ) + x, y = np.meshgrid(x, y) + assert x.shape == y.shape + points = np.array([x.flatten(), y.flatten(), np.zeros(x.size)]) + height_function_map = self._settings.terrain.height_function().map(x.size) + z = -np.array(height_function_map(points).full()).reshape(x.shape) + surf2stl.write(full_filename, x, y, z) + return full_filename + + def _draw_terrain_normals(self) -> None: + x_step = ( + self._settings.terrain_x_limits[1] - self._settings.terrain_x_limits[0] + ) / self._settings.terrain_normal_axis_points + y_step = ( + self._settings.terrain_y_limits[1] - self._settings.terrain_y_limits[0] + ) / self._settings.terrain_normal_axis_points + x = np.arange( + self._settings.terrain_x_limits[0], + self._settings.terrain_x_limits[1] + x_step, + x_step, + ) + y = np.arange( + self._settings.terrain_y_limits[0], + self._settings.terrain_y_limits[1] + y_step, + y_step, + ) + x, y = np.meshgrid(x, y) + assert x.shape == y.shape + points = np.array([x.flatten(), y.flatten(), np.zeros(x.size)]) + height_function_map = self._settings.terrain.height_function().map(x.size) + z = -np.array(height_function_map(points).full()).reshape(x.shape) + points = np.array([x.flatten(), y.flatten(), z.flatten()]) + normal_function_map = self._settings.terrain.normal_direction_function().map( + x.size + ) + normals = normal_function_map(points).full() + + for i in range(normals.shape[1]): + self._viz.load_arrow( + radius=self._settings.terrain_normals_radius, + color=self._settings.terrain_normals_color, + shape_name=f"normal_{i}", + ) + self._viz.set_arrow_transform( + origin=points[:, i], + vector=self._settings.terrain_normal_length * normals[:, i], + shape_name=f"normal_{i}", + ) + + def _draw_terrain_frames(self) -> None: + x_step = ( + self._settings.terrain_x_limits[1] - self._settings.terrain_x_limits[0] + ) / self._settings.terrain_frames_axis_points + y_step = ( + self._settings.terrain_y_limits[1] - self._settings.terrain_y_limits[0] + ) / self._settings.terrain_frames_axis_points + x = np.arange( + self._settings.terrain_x_limits[0], + self._settings.terrain_x_limits[1] + x_step, + x_step, + ) + y = np.arange( + self._settings.terrain_y_limits[0], + self._settings.terrain_y_limits[1] + y_step, + y_step, + ) + x, y = np.meshgrid(x, y) + assert x.shape == y.shape + points = np.array([x.flatten(), y.flatten(), np.zeros(x.size)]) + height_function_map = self._settings.terrain.height_function().map(x.size) + z = -np.array(height_function_map(points).full()).reshape(x.shape) + points = np.array([x.flatten(), y.flatten(), z.flatten()]) + + for i in range(x.size): + frame = self._settings.terrain.orientation_function()(points[:, i]).full() + assert np.abs(np.linalg.det(frame) - 1.0) < 1e-6 + for el in range(3): + color = [0.0, 0.0, 0.0, 0.0] + color[el] = 1.0 + color[3] = self._settings.terrain_frames_opacity + self._viz.load_arrow( + radius=self._settings.terrain_frames_axis_radius, + color=color, + shape_name=f"frame_{i}_{el}", + ) + self._viz.set_arrow_transform( + origin=points[:, i], + vector=self._settings.terrain_frames_axis_length * frame[:, el], + shape_name=f"frame_{i}_{el}", + ) diff --git a/src/hippopt/robot_planning/variables/__init__.py b/src/hippopt/robot_planning/variables/__init__.py index e69de29b..aaa3d8a8 100644 --- a/src/hippopt/robot_planning/variables/__init__.py +++ b/src/hippopt/robot_planning/variables/__init__.py @@ -0,0 +1 @@ +from . import contacts, floating_base, humanoid diff --git a/src/hippopt/robot_planning/variables/contacts.py b/src/hippopt/robot_planning/variables/contacts.py new file mode 100644 index 00000000..f1e63a3d --- /dev/null +++ b/src/hippopt/robot_planning/variables/contacts.py @@ -0,0 +1,165 @@ +import dataclasses +from typing import TypeVar + +import casadi as cs +import liecasadi +import numpy as np + +from hippopt import ( + CompositeType, + OptimizationObject, + OverridableVariable, + Parameter, + StorageType, + default_composite_field, + default_storage_field, +) + + +@dataclasses.dataclass +class ContactPointDescriptor(OptimizationObject): + position_in_foot_frame: StorageType = default_storage_field(Parameter) + foot_frame: str = dataclasses.field(default=None) + + input_foot_frame: dataclasses.InitVar[str] = dataclasses.field(default=None) + input_position_in_foot_frame: dataclasses.InitVar[np.ndarray] = dataclasses.field( + default=None + ) + + def __post_init__( + self, input_foot_frame: str, input_position_in_foot_frame: np.ndarray + ) -> None: + if input_foot_frame is not None: + self.foot_frame = input_foot_frame + if input_position_in_foot_frame is not None: + self.position_in_foot_frame = input_position_in_foot_frame + + @staticmethod + def rectangular_foot( + foot_frame: str, + x_length: float, + y_length: float, + top_left_point_position: np.array, + ): + return [ + ContactPointDescriptor( + input_foot_frame=foot_frame, + input_position_in_foot_frame=top_left_point_position, + ), + ContactPointDescriptor( + input_foot_frame=foot_frame, + input_position_in_foot_frame=top_left_point_position + + np.array([-x_length, 0.0, 0.0]), + ), + ContactPointDescriptor( + input_foot_frame=foot_frame, + input_position_in_foot_frame=top_left_point_position + + np.array([-x_length, -y_length, 0.0]), + ), + ContactPointDescriptor( + input_foot_frame=foot_frame, + input_position_in_foot_frame=top_left_point_position + + np.array([0.0, -y_length, 0.0]), + ), + ] + + +@dataclasses.dataclass +class ContactPointState(OptimizationObject): + p: StorageType = default_storage_field(OverridableVariable) + f: StorageType = default_storage_field(OverridableVariable) + + descriptor: CompositeType[ContactPointDescriptor] = default_composite_field( + factory=ContactPointDescriptor, time_varying=False + ) + + input_descriptor: dataclasses.InitVar[ContactPointDescriptor] = dataclasses.field( + default=None + ) + + def __post_init__(self, input_descriptor: ContactPointDescriptor) -> None: + self.p = np.zeros(3) if self.p is None else self.p + self.f = np.zeros(3) if self.f is None else self.f + if input_descriptor is not None: + self.descriptor = input_descriptor + + +@dataclasses.dataclass +class ContactPointStateDerivative(OptimizationObject): + v: StorageType = default_storage_field(OverridableVariable) + f_dot: StorageType = default_storage_field(OverridableVariable) + + def __post_init__(self) -> None: + self.v = np.zeros(3) if self.v is None else self.v + self.f_dot = np.zeros(3) if self.f_dot is None else self.f_dot + + +TFootContactState = TypeVar("TFootContactState", bound="FootContactState") + + +@dataclasses.dataclass +class FootContactState(list[ContactPointState], OptimizationObject): + def set_from_parent_frame_transform(self, transform: liecasadi.SE3): + for contact_point in self: + contact_point.p = transform.translation() + transform.rotation().act( + contact_point.descriptor.position_in_foot_frame + ) + + @staticmethod + def from_list(input_list: list[ContactPointState]) -> TFootContactState: + output = FootContactState() + for contact_point in input_list: + output.append(contact_point) + return output + + @staticmethod + def from_parent_frame_transform( + descriptor: list[ContactPointDescriptor], transform: liecasadi.SE3 + ) -> TFootContactState: + foot_contact_state = FootContactState() + for contact_point_descriptor in descriptor: + foot_contact_state.append( + ContactPointState(input_descriptor=contact_point_descriptor) + ) + + foot_contact_state.set_from_parent_frame_transform(transform) + return foot_contact_state + + +@dataclasses.dataclass +class FeetContactPointDescriptors: + left: list[ContactPointDescriptor] = dataclasses.field(default_factory=list) + right: list[ContactPointDescriptor] = dataclasses.field(default_factory=list) + + +@dataclasses.dataclass +class FeetContactPoints(OptimizationObject): + left: FootContactState = default_composite_field(factory=FootContactState) + right: FootContactState = default_composite_field(factory=FootContactState) + + +@dataclasses.dataclass +class FootContactPhaseDescriptor: + transform: liecasadi.SE3 = dataclasses.field(default=None) + mid_swing_transform: liecasadi.SE3 | None = dataclasses.field(default=None) + force: np.ndarray = dataclasses.field(default=None) + activation_time: float | None = dataclasses.field(default=None) + deactivation_time: float | None = dataclasses.field(default=None) + + def __post_init__(self) -> None: + self.transform = ( + liecasadi.SE3.from_translation_and_rotation( + cs.DM.zeros(3), liecasadi.SO3.Identity() + ) + if self.transform is None + else self.transform + ) + if self.force is None: + self.force = np.zeros(3) + self.force[2] = 100 + + +@dataclasses.dataclass +class FeetContactPhasesDescriptor: + left: list[FootContactPhaseDescriptor] = dataclasses.field(default_factory=list) + right: list[FootContactPhaseDescriptor] = dataclasses.field(default_factory=list) diff --git a/src/hippopt/robot_planning/variables/floating_base.py b/src/hippopt/robot_planning/variables/floating_base.py new file mode 100644 index 00000000..30b4c2f3 --- /dev/null +++ b/src/hippopt/robot_planning/variables/floating_base.py @@ -0,0 +1,179 @@ +import dataclasses +from typing import TypeVar + +import numpy as np + +from hippopt import ( + CompositeType, + OptimizationObject, + OverridableVariable, + StorageType, + default_composite_field, + default_storage_field, +) + + +@dataclasses.dataclass +class FreeFloatingObjectState(OptimizationObject): + position: StorageType = default_storage_field(OverridableVariable) + quaternion_xyzw: StorageType = default_storage_field(OverridableVariable) + + def __post_init__(self): + self.position = np.zeros(3) if self.position is None else self.position + if self.quaternion_xyzw is None: + self.quaternion_xyzw = np.zeros(4) + self.quaternion_xyzw[3] = 1.0 + + +@dataclasses.dataclass +class FreeFloatingObjectStateDerivative(OptimizationObject): + linear_velocity: StorageType = default_storage_field(OverridableVariable) + quaternion_velocity_xyzw: StorageType = default_storage_field(OverridableVariable) + + def __post_init__(self): + self.linear_velocity = ( + np.zeros(3) if self.linear_velocity is None else self.linear_velocity + ) + + self.quaternion_velocity_xyzw = ( + np.zeros(4) + if self.quaternion_velocity_xyzw is None + else self.quaternion_velocity_xyzw + ) + + +@dataclasses.dataclass +class FreeFloatingObject(FreeFloatingObjectState, FreeFloatingObjectStateDerivative): + def __post_init__(self): + FreeFloatingObjectState.__post_init__(self) + FreeFloatingObjectStateDerivative.__post_init__(self) + + +@dataclasses.dataclass +class KinematicTreeState(OptimizationObject): + positions: StorageType = default_storage_field(OverridableVariable) + + number_of_joints_state: dataclasses.InitVar[int] = dataclasses.field(default=0) + + def __post_init__(self, number_of_joints_state: int): + if number_of_joints_state is not None and self.positions is None: + self.positions = np.zeros(number_of_joints_state) + + +@dataclasses.dataclass +class KinematicTreeStateDerivative(OptimizationObject): + velocities: StorageType = default_storage_field(OverridableVariable) + + number_of_joints_derivative: dataclasses.InitVar[int] = dataclasses.field( + default=None + ) + + def __post_init__(self, number_of_joints_derivative: int): + if number_of_joints_derivative is not None: + self.velocities = np.zeros(number_of_joints_derivative) + + +@dataclasses.dataclass +class KinematicTree(KinematicTreeState, KinematicTreeStateDerivative): + def __post_init__( + self, + number_of_joints_derivative: int = None, + number_of_joints_state: int = None, + ): + if ( + number_of_joints_derivative is not None + or number_of_joints_state is not None + ): + number_of_joints_state = ( + number_of_joints_derivative + if number_of_joints_state is None + else number_of_joints_state + ) + number_of_joints_derivative = ( + number_of_joints_state + if number_of_joints_derivative is None + else number_of_joints_derivative + ) + KinematicTreeState.__post_init__( + self, number_of_joints_state=number_of_joints_state + ) + KinematicTreeStateDerivative.__post_init__( + self, number_of_joints_derivative=number_of_joints_derivative + ) + + +@dataclasses.dataclass +class FloatingBaseSystemState(OptimizationObject): + base: CompositeType[FreeFloatingObjectState] = default_composite_field( + factory=FreeFloatingObjectState + ) + joints: CompositeType[KinematicTreeState] = default_composite_field( + factory=KinematicTreeState + ) + + number_of_joints_state: dataclasses.InitVar[int] = dataclasses.field(default=None) + + def __post_init__(self, number_of_joints_state: int): + if number_of_joints_state is not None: + self.joints = KinematicTreeState( + number_of_joints_state=number_of_joints_state + ) + + +@dataclasses.dataclass +class FloatingBaseSystemStateDerivative(OptimizationObject): + base: CompositeType[FreeFloatingObjectStateDerivative] = default_composite_field( + factory=FreeFloatingObjectStateDerivative + ) + joints: CompositeType[KinematicTreeStateDerivative] = default_composite_field( + factory=KinematicTreeStateDerivative + ) + + number_of_joints_derivative: dataclasses.InitVar[int] = dataclasses.field( + default=None + ) + + def __post_init__(self, number_of_joints_derivative: int): + if number_of_joints_derivative is not None: + self.joints = KinematicTreeStateDerivative( + number_of_joints_derivative=number_of_joints_derivative + ) + + +TFloatingBaseSystem = TypeVar("TFloatingBaseSystem", bound="FloatingBaseSystem") + + +@dataclasses.dataclass +class FloatingBaseSystem(OptimizationObject): + base: CompositeType[FreeFloatingObject] = default_composite_field( + factory=FreeFloatingObject + ) + joints: CompositeType[KinematicTree] = default_composite_field( + factory=KinematicTree + ) + + number_of_joints: dataclasses.InitVar[int] = dataclasses.field(default=None) + + def __post_init__(self, number_of_joints: int): + if number_of_joints is not None: + self.joints = KinematicTree(number_of_joints_state=number_of_joints) + + def to_floating_base_system_state(self): + output = FloatingBaseSystemState() + output.base.position = self.base.position + output.base.quaternion_xyzw = self.base.quaternion_xyzw + output.joints.positions = self.joints.positions + return output + + @staticmethod + def from_floating_base_system_state( + state: FloatingBaseSystemState, + ) -> TFloatingBaseSystem: + output = FloatingBaseSystem(number_of_joints=len(state.joints.positions)) + output.base.position = state.base.position + output.base.quaternion_xyzw = state.base.quaternion_xyzw + output.base.linear_velocity = None + output.base.quaternion_velocity_xyzw = None + output.joints.positions = state.joints.positions + output.joints.velocities = None + return output diff --git a/src/hippopt/robot_planning/variables/humanoid.py b/src/hippopt/robot_planning/variables/humanoid.py new file mode 100644 index 00000000..884adb26 --- /dev/null +++ b/src/hippopt/robot_planning/variables/humanoid.py @@ -0,0 +1,56 @@ +import dataclasses + +import numpy as np + +from hippopt import ( + CompositeType, + OptimizationObject, + OverridableVariable, + StorageType, + default_composite_field, + default_storage_field, +) +from hippopt.robot_planning.variables.contacts import ( + ContactPointState, + FeetContactPointDescriptors, + FeetContactPoints, +) +from hippopt.robot_planning.variables.floating_base import FloatingBaseSystemState + + +@dataclasses.dataclass +class HumanoidState(OptimizationObject): + contact_points: CompositeType[FeetContactPoints] = default_composite_field( + factory=FeetContactPoints, time_varying=False + ) + + kinematics: CompositeType[FloatingBaseSystemState] = default_composite_field( + factory=FloatingBaseSystemState, time_varying=False + ) + + com: StorageType = default_storage_field(OverridableVariable) + + contact_point_descriptors: dataclasses.InitVar[ + FeetContactPointDescriptors + ] = dataclasses.field(default=None) + number_of_joints: dataclasses.InitVar[int] = dataclasses.field(default=None) + + def __post_init__( + self, + contact_point_descriptors: FeetContactPointDescriptors, + number_of_joints: int, + ) -> None: + if contact_point_descriptors is not None: + self.contact_points.left = [ + ContactPointState(input_descriptor=point) + for point in contact_point_descriptors.left + ] + self.contact_points.right = [ + ContactPointState(input_descriptor=point) + for point in contact_point_descriptors.right + ] + if number_of_joints is not None: + self.kinematics = FloatingBaseSystemState( + number_of_joints_state=number_of_joints + ) + self.com = np.zeros(3) if self.com is None else self.com diff --git a/test/test_planar_terrain.py b/test/test_planar_terrain.py new file mode 100644 index 00000000..4ac8a0ba --- /dev/null +++ b/test/test_planar_terrain.py @@ -0,0 +1,38 @@ +import casadi as cs +import numpy as np + +import hippopt.robot_planning + + +def test_planar_terrain(): + planar_terrain = hippopt.robot_planning.PlanarTerrain(point_position_name="p") + + dummy_point = cs.DM.zeros(3) + dummy_point[2] = 0.5 + + height_function = planar_terrain.height_function() + + assert next(iter(height_function(p=dummy_point).values())) == 0.5 + + assert ( + height_function is planar_terrain.height_function() + ) # Check that the function is created only once + + normal_function = planar_terrain.normal_direction_function() + normal_direction = np.zeros((3, 1)) + normal_direction[2] = 1.0 + output = normal_function(dummy_point).full() + + assert (normal_direction == output).all() # noqa + + orientation_fun = planar_terrain.orientation_function() + expected_orientation = np.eye(3) + + output = orientation_fun(dummy_point).full() + + assert (expected_orientation == output).all() # noqa + + assert planar_terrain.get_name() == "PlanarTerrain" + + other = hippopt.robot_planning.PlanarTerrain(name="other") + assert other.get_name() == "other"