From a3ee6974ad53d6773ccdfe06994f207383472186 Mon Sep 17 00:00:00 2001 From: Daniel Ordonez Date: Wed, 26 Mar 2025 10:41:41 -0500 Subject: [PATCH 1/2] Ruff format and auto linter --- quadruped_pympc/config.py | 279 +-- .../centroidal_model_collaborative.py | 359 ++-- .../centroidal_nmpc_collaborative.py | 1685 +++++++++-------- .../centroidal_model_input_rates.py | 170 +- .../centroidal_nmpc_input_rates.py | 1658 ++++++++-------- .../gradient/kinodynamic/kinodynamic_model.py | 401 ++-- .../gradient/kinodynamic/kinodynamic_nmpc.py | 1467 +++++++------- .../lyapunov/centroidal_model_lyapunov.py | 349 ++-- .../lyapunov/centroidal_nmpc_lyapunov.py | 1532 ++++++++------- .../nominal/centroidal_model_nominal.py | 308 +-- .../nominal/centroidal_nmpc_gait_adaptive.py | 560 +++--- .../nominal/centroidal_nmpc_nominal.py | 902 +++++---- .../sampling/centroidal_model_jax.py | 261 +-- .../sampling/centroidal_nmpc_jax.py | 1501 ++++++++------- .../centroidal_nmpc_jax_gait_adaptive.py | 1562 ++++++++------- .../helpers/foothold_reference_generator.py | 127 +- .../inverse_kinematics_numeric.py | 242 ++- .../inverse_kinematics_numeric_mujoco.py | 176 +- .../inverse_kinematics_qp.py | 118 +- quadruped_pympc/helpers/math_utils.py | 4 +- .../helpers/periodic_gait_generator.py | 98 +- .../helpers/periodic_gait_generator_jax.py | 80 +- quadruped_pympc/helpers/quadruped_utils.py | 106 +- .../helpers/srb_inertia_computation.py | 49 +- .../explicit_swing_trajectory_generator.py | 126 +- .../ndcurves_swing_trajectory_generator.py | 87 +- .../scipy_swing_trajectory_generator.py | 118 +- .../helpers/swing_trajectory_controller.py | 374 ++-- quadruped_pympc/helpers/terrain_estimator.py | 65 +- quadruped_pympc/helpers/velocity_modulator.py | 37 +- .../helpers/visual_foothold_adaptation.py | 121 +- .../srbd_batched_controller_interface.py | 86 +- .../interfaces/srbd_controller_interface.py | 294 +-- quadruped_pympc/interfaces/wb_interface.py | 605 +++--- quadruped_pympc/quadruped_pympc_wrapper.py | 340 ++-- simulation/batched_simulations.py | 52 +- simulation/simulation.py | 303 +-- 37 files changed, 8763 insertions(+), 7839 deletions(-) diff --git a/quadruped_pympc/config.py b/quadruped_pympc/config.py index 439b506..fefaf5f 100644 --- a/quadruped_pympc/config.py +++ b/quadruped_pympc/config.py @@ -1,70 +1,114 @@ """This file includes all of the configuration parameters for the MPC controllers and of the internal simulations that can be launch from the folder /simulation. """ + import numpy as np + from quadruped_pympc.helpers.quadruped_utils import GaitType # These are used both for a real experiment and a simulation ----------- # These are the only attributes needed per quadruped, the rest can be computed automatically ---------------------- -robot = 'aliengo' # 'go1', 'go2', 'b2', 'aliengo', 'hyqreal', 'mini_cheetah' # TODO: Load from robot_descriptions.py -robot_leg_joints = dict(FL=['FL_hip_joint', 'FL_thigh_joint', 'FL_calf_joint',], # TODO: Make configs per robot. - FR=['FR_hip_joint', 'FR_thigh_joint', 'FR_calf_joint',], - RL=['RL_hip_joint', 'RL_thigh_joint', 'RL_calf_joint',], - RR=['RR_hip_joint', 'RR_thigh_joint', 'RR_calf_joint',]) -robot_feet_geom_names = dict(FL='FL', FR='FR', RL='RL', RR='RR') +robot = "aliengo" # 'go1', 'go2', 'b2', 'aliengo', 'hyqreal', 'mini_cheetah' # TODO: Load from robot_descriptions.py +robot_leg_joints = dict( + FL=[ + "FL_hip_joint", + "FL_thigh_joint", + "FL_calf_joint", + ], # TODO: Make configs per robot. + FR=[ + "FR_hip_joint", + "FR_thigh_joint", + "FR_calf_joint", + ], + RL=[ + "RL_hip_joint", + "RL_thigh_joint", + "RL_calf_joint", + ], + RR=[ + "RR_hip_joint", + "RR_thigh_joint", + "RR_calf_joint", + ], +) +robot_feet_geom_names = dict(FL="FL", FR="FR", RL="RL", RR="RR") qpos0_js = None # Zero joint-space configuration. If None it will be extracted from the URDF. # ---------------------------------------------------------------------------------------------------------------- -if (robot == 'go1'): +if robot == "go1": mass = 12.019 - inertia = np.array([[1.58460467e-01, 1.21660000e-04, -1.55444692e-02], - [1.21660000e-04, 4.68645637e-01, -3.12000000e-05], - [-1.55444692e-02, -3.12000000e-05, 5.24474661e-01]]) + inertia = np.array( + [ + [1.58460467e-01, 1.21660000e-04, -1.55444692e-02], + [1.21660000e-04, 4.68645637e-01, -3.12000000e-05], + [-1.55444692e-02, -3.12000000e-05, 5.24474661e-01], + ] + ) urdf_filename = "go1.urdf" hip_height = 0.3 -elif (robot == 'go2'): +elif robot == "go2": mass = 15.019 - inertia = np.array([[1.58460467e-01, 1.21660000e-04, -1.55444692e-02], - [1.21660000e-04, 4.68645637e-01, -3.12000000e-05], - [-1.55444692e-02, -3.12000000e-05, 5.24474661e-01]]) + inertia = np.array( + [ + [1.58460467e-01, 1.21660000e-04, -1.55444692e-02], + [1.21660000e-04, 4.68645637e-01, -3.12000000e-05], + [-1.55444692e-02, -3.12000000e-05, 5.24474661e-01], + ] + ) urdf_filename = "go2.urdf" hip_height = 0.28 -elif (robot == 'aliengo'): +elif robot == "aliengo": mass = 24.637 - inertia = np.array([[0.2310941359705289, -0.0014987128245817424, -0.021400468992761768], - [-0.0014987128245817424, 1.4485084687476608, 0.0004641447134275615], - [-0.021400468992761768, 0.0004641447134275615, 1.503217877350808]]) + inertia = np.array( + [ + [0.2310941359705289, -0.0014987128245817424, -0.021400468992761768], + [-0.0014987128245817424, 1.4485084687476608, 0.0004641447134275615], + [-0.021400468992761768, 0.0004641447134275615, 1.503217877350808], + ] + ) urdf_filename = "aliengo.urdf" hip_height = 0.35 -elif (robot == 'b2'): +elif robot == "b2": mass = 83.49 - inertia = np.array([[0.2310941359705289, -0.0014987128245817424, -0.021400468992761768], - [-0.0014987128245817424, 1.4485084687476608, 0.0004641447134275615], - [-0.021400468992761768, 0.0004641447134275615, 1.503217877350808]]) + inertia = np.array( + [ + [0.2310941359705289, -0.0014987128245817424, -0.021400468992761768], + [-0.0014987128245817424, 1.4485084687476608, 0.0004641447134275615], + [-0.021400468992761768, 0.0004641447134275615, 1.503217877350808], + ] + ) urdf_filename = "b2.urdf" hip_height = 0.485 -elif (robot == 'hyqreal'): +elif robot == "hyqreal": mass = 108.40 - inertia = np.array([[4.55031444e+00, 2.75249434e-03, -5.11957307e-01], - [2.75249434e-03, 2.02411774e+01, -7.38560592e-04], - [-5.11957307e-01, -7.38560592e-04, 2.14269772e+01]]) + inertia = np.array( + [ + [4.55031444e00, 2.75249434e-03, -5.11957307e-01], + [2.75249434e-03, 2.02411774e01, -7.38560592e-04], + [-5.11957307e-01, -7.38560592e-04, 2.14269772e01], + ] + ) urdf_filename = "hyqreal.urdf" hip_height = 0.5 -elif (robot == 'mini_cheetah'): +elif robot == "mini_cheetah": mass = 12.5 - inertia = np.array([[1.58460467e-01, 1.21660000e-04, -1.55444692e-02], - [1.21660000e-04, 4.68645637e-01, -3.12000000e-05], - [-1.55444692e-02, -3.12000000e-05, 5.24474661e-01]]) + inertia = np.array( + [ + [1.58460467e-01, 1.21660000e-04, -1.55444692e-02], + [1.21660000e-04, 4.68645637e-01, -3.12000000e-05], + [-1.55444692e-02, -3.12000000e-05, 5.24474661e-01], + ] + ) urdf_filename = "mini_cheetah.urdf" hip_height = 0.225 - qpos0_js = np.concatenate((np.array([0, -np.pi/2, 0] * 2), np.array([0, np.pi/2, 0] * 2))) + qpos0_js = np.concatenate((np.array([0, -np.pi / 2, 0] * 2), np.array([0, np.pi / 2, 0] * 2))) mpc_params = { # 'nominal' optimized directly the GRF @@ -73,157 +117,122 @@ # 'collaborative' optimized directly the GRF and has a passive arm model inside # 'lyapunov' optimized directly the GRF and has a Lyapunov-based stability constraint # 'kynodynamic' sbrd with joints - experimental - 'type': 'nominal', - + "type": "nominal", # print the mpc info - 'verbose': False, - + "verbose": False, # horizon is the number of timesteps in the future that the mpc will optimize # dt is the discretization time used in the mpc - 'horizon': 12, - 'dt': 0.02, - + "horizon": 12, + "dt": 0.02, # GRF limits for each single leg - "grf_max": mass * 9.81, - "grf_min": 0, - 'mu': 0.5, - + "grf_max": mass * 9.81, + "grf_min": 0, + "mu": 0.5, # this is used to have a smaller dt near the start of the horizon - 'use_nonuniform_discretization': False, - 'horizon_fine_grained': 2, - 'dt_fine_grained': 0.01, - + "use_nonuniform_discretization": False, + "horizon_fine_grained": 2, + "dt_fine_grained": 0.01, # if this is true, we optimize the step frequency as well # for the sampling controller, this is done in the rollout # for the gradient-based controller, this is done with a batched version of the ocp - 'optimize_step_freq': False, - 'step_freq_available': [1.4, 2.0, 2.4], - + "optimize_step_freq": False, + "step_freq_available": [1.4, 2.0, 2.4], # ----- START properties only for the gradient-based mpc ----- - # this is used if you want to manually warm start the mpc - 'use_warm_start': False, - + "use_warm_start": False, # this enables integrators for height, linear velocities, roll and pitch - 'use_integrators': False, - 'alpha_integrator': 0.1, - 'integrator_cap': [0.5, 0.2, 0.2, 0.0, 0.0, 1.0], - + "use_integrators": False, + "alpha_integrator": 0.1, + "integrator_cap": [0.5, 0.2, 0.2, 0.0, 0.0, 1.0], # if this is off, the mpc will not optimize the footholds and will # use only the ones provided in the reference - 'use_foothold_optimization': True, - + "use_foothold_optimization": True, # this is set to false automatically is use_foothold_optimization is false # because in that case we cannot chose the footholds and foothold # constraints do not any make sense - 'use_foothold_constraints': False, - + "use_foothold_constraints": False, # works with all the mpc types except 'sampling'. In sim does not do much for now, # but in real it minizimes the delay between the mpc control and the state - 'use_RTI': False, + "use_RTI": False, # If RTI is used, we can set the advance RTI-step! (Standard is the simpler RTI) # See https://arxiv.org/pdf/2403.07101.pdf - 'as_rti_type': "Standard", # "AS-RTI-A", "AS-RTI-B", "AS-RTI-C", "AS-RTI-D", "Standard" - 'as_rti_iter': 1, # > 0, the higher the better, but slower computation! - + "as_rti_type": "Standard", # "AS-RTI-A", "AS-RTI-B", "AS-RTI-C", "AS-RTI-D", "Standard" + "as_rti_iter": 1, # > 0, the higher the better, but slower computation! # This will force to use DDP instead of SQP, based on https://arxiv.org/abs/2403.10115. # Note that RTI is not compatible with DDP, and no state costraints for now are considered - 'use_DDP': False, - + "use_DDP": False, # this is used only in the case 'use_RTI' is false in a single mpc feedback loop. # More is better, but slower computation! - 'num_qp_iterations': 1, - + "num_qp_iterations": 1, # this is used to speeding up or robustify acados' solver (hpipm). - 'solver_mode': 'balance', # balance, robust, speed, crazy_speed - - + "solver_mode": "balance", # balance, robust, speed, crazy_speed # these is used only for the case 'input_rates', using as GRF not the actual state # of the robot of the predicted one. Can be activated to compensate # for the delay in the control loop on the real robot - 'use_input_prediction': False, - + "use_input_prediction": False, # ONLY ONE CAN BE TRUE AT A TIME (only gradient) - 'use_static_stability': False, - 'use_zmp_stability': False, - 'trot_stability_margin': 0.04, - 'pace_stability_margin': 0.1, - 'crawl_stability_margin': 0.04, # in general, 0.02 is a good value - + "use_static_stability": False, + "use_zmp_stability": False, + "trot_stability_margin": 0.04, + "pace_stability_margin": 0.1, + "crawl_stability_margin": 0.04, # in general, 0.02 is a good value # this is used to compensate for the external wrenches # you should provide explicitly this value in compute_control - 'external_wrenches_compensation': True, - 'external_wrenches_compensation_num_step': 15, - + "external_wrenches_compensation": True, + "external_wrenches_compensation_num_step": 15, # this is used only in the case of collaborative mpc, to # compensate for the external wrench in the prediction (only collaborative) - 'passive_arm_compensation': True, - - + "passive_arm_compensation": True, # Gain for Lyapunov-based MPC - 'K_z1': np.array([1, 1, 10]), - 'K_z2': np.array([1, 4, 10]), - 'residual_dynamics_upper_bound': 30, - 'use_residual_dynamics_decay': False, - + "K_z1": np.array([1, 1, 10]), + "K_z2": np.array([1, 4, 10]), + "residual_dynamics_upper_bound": 30, + "use_residual_dynamics_decay": False, # ----- END properties for the gradient-based mpc ----- - - # ----- START properties only for the sampling-based mpc ----- - # this is used only in the case 'sampling'. - 'sampling_method': 'random_sampling', # 'random_sampling', 'mppi', 'cem_mppi' - 'control_parametrization': 'cubic_spline_1', + "sampling_method": "random_sampling", # 'random_sampling', 'mppi', 'cem_mppi' + "control_parametrization": "cubic_spline_1", # 'cubic_spline_1', 'cubic_spline_2', 'linear_spline_1', 'linear_spline_2', 'zero_order' - 'num_parallel_computations': 10000, # More is better, but slower computation! - 'num_sampling_iterations': 1, # More is better, but slower computation! + "num_parallel_computations": 10000, # More is better, but slower computation! + "num_sampling_iterations": 1, # More is better, but slower computation! # convariances for the sampling methods - 'sigma_cem_mppi': 3, - 'sigma_mppi': 3, - 'sigma_random_sampling': [0.2, 3, 10], - 'shift_solution': False, - + "sigma_cem_mppi": 3, + "sigma_mppi": 3, + "sigma_random_sampling": [0.2, 3, 10], + "shift_solution": False, # ----- END properties for the sampling-based mpc ----- - - } +} # ----------------------------------------------------------------------- simulation_params = { - 'swing_generator': 'scipy', # 'scipy', 'explicit', 'ndcurves' - 'swing_position_gain_fb': 500, - 'swing_velocity_gain_fb': 10, - 'swing_integral_gain_fb': 0, - 'impedence_joint_position_gain': 10.0, - 'impedence_joint_velocity_gain': 1.0, - 'step_height': 0.3 * hip_height, - + "swing_generator": "scipy", # 'scipy', 'explicit', 'ndcurves' + "swing_position_gain_fb": 500, + "swing_velocity_gain_fb": 10, + "swing_integral_gain_fb": 0, + "impedence_joint_position_gain": 10.0, + "impedence_joint_velocity_gain": 1.0, + "step_height": 0.3 * hip_height, # Visual Foothold adapatation - "visual_foothold_adaptation": 'blind', #'blind', 'height', 'vfa' - + "visual_foothold_adaptation": "blind", #'blind', 'height', 'vfa' # this is the integration time used in the simulator - 'dt': 0.002, - - 'gait': 'trot', # 'trot', 'pace', 'crawl', 'bound', 'full_stance' - 'gait_params': {'trot': {'step_freq': 1.4, 'duty_factor': 0.65, 'type': GaitType.TROT.value}, - 'crawl': {'step_freq': 0.5, 'duty_factor': 0.8, 'type': GaitType.BACKDIAGONALCRAWL.value}, - 'pace': {'step_freq': 1.4, 'duty_factor': 0.7, 'type': GaitType.PACE.value}, - 'bound': {'step_freq': 1.8, 'duty_factor': 0.65, 'type': GaitType.BOUNDING.value}, - 'full_stance': {'step_freq': 2, 'duty_factor': 0.65, 'type': GaitType.FULL_STANCE.value}, - }, - + "dt": 0.002, + "gait": "trot", # 'trot', 'pace', 'crawl', 'bound', 'full_stance' + "gait_params": { + "trot": {"step_freq": 1.4, "duty_factor": 0.65, "type": GaitType.TROT.value}, + "crawl": {"step_freq": 0.5, "duty_factor": 0.8, "type": GaitType.BACKDIAGONALCRAWL.value}, + "pace": {"step_freq": 1.4, "duty_factor": 0.7, "type": GaitType.PACE.value}, + "bound": {"step_freq": 1.8, "duty_factor": 0.65, "type": GaitType.BOUNDING.value}, + "full_stance": {"step_freq": 2, "duty_factor": 0.65, "type": GaitType.FULL_STANCE.value}, + }, # velocity mode: human will give you the possibility to use the keyboard, the other are # forward only random linear-velocity, random will give you random linear-velocity and yaw-velocity - 'mode': 'human', # 'human', 'forward', 'random' - 'ref_z': hip_height, - - + "mode": "human", # 'human', 'forward', 'random' + "ref_z": hip_height, # the MPC will be called every 1/(mpc_frequency*dt) timesteps # this helps to evaluate more realistically the performance of the controller - 'mpc_frequency': 100, - - 'use_inertia_recomputation': True, - - 'scene': 'flat', # flat, random_boxes, random_pyramids, perlin - - } + "mpc_frequency": 100, + "use_inertia_recomputation": True, + "scene": "perlin", # flat, random_boxes, random_pyramids, perlin +} # ----------------------------------------------------------------------- diff --git a/quadruped_pympc/controllers/gradient/collaborative/centroidal_model_collaborative.py b/quadruped_pympc/controllers/gradient/collaborative/centroidal_model_collaborative.py index 5243371..f0c44b4 100644 --- a/quadruped_pympc/controllers/gradient/collaborative/centroidal_model_collaborative.py +++ b/quadruped_pympc/controllers/gradient/collaborative/centroidal_model_collaborative.py @@ -3,39 +3,36 @@ # Authors: Giulio Turrisi - -import time -import unittest -import casadi as cs +import os +import casadi as cs import numpy as np from acados_template import AcadosModel -import os dir_path = os.path.dirname(os.path.realpath(__file__)) import sys + sys.path.append(dir_path) -sys.path.append(dir_path + '/../../') +sys.path.append(dir_path + "/../../") from quadruped_pympc import config # Class that defines the prediction model of the NMPC class Centroidal_Model_Collaborative: - def __init__(self,) -> None: + def __init__( + self, + ) -> None: """ This method initializes the foothold generator Centroidal_Model, which creates the prediction model of the NMPC. """ - # Set mass and inertia from config file self.mass = config.mass self.inertia = config.inertia - - - # Define state and its casadi variables com_position_x = cs.SX.sym("com_position_x") com_position_y = cs.SX.sym("com_position_y") @@ -44,20 +41,19 @@ def __init__(self,) -> None: com_velocity_x = cs.SX.sym("com_velocity_x") com_velocity_y = cs.SX.sym("com_velocity_y") com_velocity_z = cs.SX.sym("com_velocity_z") - + roll = cs.SX.sym("roll", 1, 1) pitch = cs.SX.sym("pitch", 1, 1) yaw = cs.SX.sym("yaw", 1, 1) omega_x = cs.SX.sym("omega_x", 1, 1) omega_y = cs.SX.sym("omega_y", 1, 1) omega_z = cs.SX.sym("omega_z", 1, 1) - + foot_position_fl = cs.SX.sym("foot_position_fl", 3, 1) foot_position_fr = cs.SX.sym("foot_position_fr", 3, 1) foot_position_rl = cs.SX.sym("foot_position_rl", 3, 1) foot_position_rr = cs.SX.sym("foot_position_rr", 3, 1) - com_position_z_integral = cs.SX.sym("com_position_z_integral") com_velocity_x_integral = cs.SX.sym("com_velocity_x_integral") com_velocity_y_integral = cs.SX.sym("com_velocity_y_integral") @@ -65,55 +61,50 @@ def __init__(self,) -> None: roll_integral = cs.SX.sym("roll_integral") pitch_integral = cs.SX.sym("pitch_integral") - - passive_arm_force = cs.SX.sym("passive_arm_force", 6, 1) - - - - self.states = cs.vertcat(com_position_x, - com_position_y, - com_position_z, - com_velocity_x, - com_velocity_y, - com_velocity_z, - roll, - pitch, - yaw, - omega_x, - omega_y, - omega_z, - foot_position_fl, - foot_position_fr, - foot_position_rl, - foot_position_rr, - com_position_z_integral, - com_velocity_x_integral, - com_velocity_y_integral, - com_velocity_z_integral, - roll_integral, - pitch_integral, - passive_arm_force) - - - - # Define state dot - self.states_dot = cs.vertcat(cs.SX.sym("linear_com_vel", 3, 1), - cs.SX.sym("linear_com_acc", 3, 1), - cs.SX.sym("euler_rates_base", 3, 1), - cs.SX.sym("angular_acc_base", 3, 1), - cs.SX.sym("linear_vel_foot_FL", 3, 1), - cs.SX.sym("linear_vel_foot_FR", 3, 1), - cs.SX.sym("linear_vel_foot_RL", 3, 1), - cs.SX.sym("linear_vel_foot_RR", 3, 1), - cs.SX.sym("linear_com_vel_z_integral", 1, 1), - cs.SX.sym("linear_com_acc_integral", 3, 1), - cs.SX.sym("euler_rates_roll_integral", 1, 1), - cs.SX.sym("euler_rates_pitch_integral", 1, 1), - cs.SX.sym("passive_arm_force_dot", 6, 1)) - - + self.states = cs.vertcat( + com_position_x, + com_position_y, + com_position_z, + com_velocity_x, + com_velocity_y, + com_velocity_z, + roll, + pitch, + yaw, + omega_x, + omega_y, + omega_z, + foot_position_fl, + foot_position_fr, + foot_position_rl, + foot_position_rr, + com_position_z_integral, + com_velocity_x_integral, + com_velocity_y_integral, + com_velocity_z_integral, + roll_integral, + pitch_integral, + passive_arm_force, + ) + + # Define state dot + self.states_dot = cs.vertcat( + cs.SX.sym("linear_com_vel", 3, 1), + cs.SX.sym("linear_com_acc", 3, 1), + cs.SX.sym("euler_rates_base", 3, 1), + cs.SX.sym("angular_acc_base", 3, 1), + cs.SX.sym("linear_vel_foot_FL", 3, 1), + cs.SX.sym("linear_vel_foot_FR", 3, 1), + cs.SX.sym("linear_vel_foot_RL", 3, 1), + cs.SX.sym("linear_vel_foot_RR", 3, 1), + cs.SX.sym("linear_com_vel_z_integral", 1, 1), + cs.SX.sym("linear_com_acc_integral", 3, 1), + cs.SX.sym("euler_rates_roll_integral", 1, 1), + cs.SX.sym("euler_rates_pitch_integral", 1, 1), + cs.SX.sym("passive_arm_force_dot", 6, 1), + ) # Define input and its casadi variables foot_velocity_fl = cs.SX.sym("foot_velocity_fl", 3, 1) @@ -126,28 +117,27 @@ def __init__(self,) -> None: foot_force_rl = cs.SX.sym("foot_force_rl", 3, 1) foot_force_rr = cs.SX.sym("foot_force_rr", 3, 1) - self.inputs = cs.vertcat(foot_velocity_fl, - foot_velocity_fr, - foot_velocity_rl, - foot_velocity_rr, - foot_force_fl, - foot_force_fr, - foot_force_rl, - foot_force_rr) - - + self.inputs = cs.vertcat( + foot_velocity_fl, + foot_velocity_fr, + foot_velocity_rl, + foot_velocity_rr, + foot_force_fl, + foot_force_fr, + foot_force_rl, + foot_force_rr, + ) + # Usefull for debug what things goes where in y_ref in the compute_control function, # because there are a lot of variables self.y_ref = cs.vertcat(self.states, self.inputs) - # Define acados parameters that can be changed at runtine self.stanceFL = cs.SX.sym("stanceFL", 1, 1) self.stanceFR = cs.SX.sym("stanceFR", 1, 1) self.stanceRL = cs.SX.sym("stanceRL", 1, 1) self.stanceRR = cs.SX.sym("stanceRR", 1, 1) - self.stance_param = cs.vertcat(self.stanceFL , self.stanceFR , self.stanceRL , self.stanceRR) - + self.stance_param = cs.vertcat(self.stanceFL, self.stanceFR, self.stanceRL, self.stanceRR) self.mu_friction = cs.SX.sym("mu_friction", 1, 1) self.stance_proximity = cs.SX.sym("stanceProximity", 4, 1) @@ -160,23 +150,26 @@ def __init__(self,) -> None: self.end_effector_jacobian_transpose_inv = cs.SX.sym("end_effector_jacobian_transpose_inv", 9, 1) self.end_effector_gain = cs.SX.sym("end_effector_gain", 3, 1) - # Not so useful, i can instantiate a casadi function for the fd - param = cs.vertcat(self.stance_param, self.mu_friction, self.stance_proximity, - self.base_position, self.base_yaw, self.external_wrench, - self.end_effector_position, self.end_effector_jacobian_inv, - self.end_effector_jacobian_transpose_inv, self.end_effector_gain) + param = cs.vertcat( + self.stance_param, + self.mu_friction, + self.stance_proximity, + self.base_position, + self.base_yaw, + self.external_wrench, + self.end_effector_position, + self.end_effector_jacobian_inv, + self.end_effector_jacobian_transpose_inv, + self.end_effector_gain, + ) fd = self.forward_dynamics(self.states, self.inputs, param) - self.fun_forward_dynamics = cs.Function('fun_forward_dynamics', [self.states, self.inputs, param], [fd]) - - - - + self.fun_forward_dynamics = cs.Function("fun_forward_dynamics", [self.states, self.inputs, param], [fd]) def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.ndarray) -> cs.SX: """ This method computes the symbolic forward dynamics of the robot. It is used inside - Acados to compute the prediction model. It uses the same variable of self.states, self.param, and + Acados to compute the prediction model. It uses the same variable of self.states, self.param, and fill the same variables as the one in self.states_dot. Args: @@ -188,7 +181,7 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda Returns: A CasADi SX object of shape (36,1) representing the predicted state of the robot. """ - + # Saving for clarity a bunch of variables foot_velocity_fl = inputs[0:3] foot_velocity_fr = inputs[3:6] @@ -214,138 +207,116 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda stance_proximity_RL = param[7] stance_proximity_RR = param[8] - # Rotation matrix from world to base. Usefull for later roll = states[6] pitch = states[7] yaw = states[8] Rx = cs.SX.eye(3) - Rx[0,0] = 1 - Rx[0,1] = 0 - Rx[0,2] = 0 - Rx[1,0] = 0 - Rx[1,1] = cs.cos(roll) - Rx[1,2] = cs.sin(roll) - Rx[2,0] = 0 - Rx[2,1] = -cs.sin(roll) - Rx[2,2] = cs.cos(roll) - + Rx[0, 0] = 1 + Rx[0, 1] = 0 + Rx[0, 2] = 0 + Rx[1, 0] = 0 + Rx[1, 1] = cs.cos(roll) + Rx[1, 2] = cs.sin(roll) + Rx[2, 0] = 0 + Rx[2, 1] = -cs.sin(roll) + Rx[2, 2] = cs.cos(roll) + Ry = cs.SX.eye(3) - Ry[0,0] = cs.cos(pitch) - Ry[0,1] = 0 - Ry[0,2] = -cs.sin(pitch) - Ry[1,0] = 0 - Ry[1,1] = 1 - Ry[1,2] = 0 - Ry[2,0] = cs.sin(pitch) - Ry[2,1] = 0 - Ry[2,2] = cs.cos(pitch) + Ry[0, 0] = cs.cos(pitch) + Ry[0, 1] = 0 + Ry[0, 2] = -cs.sin(pitch) + Ry[1, 0] = 0 + Ry[1, 1] = 1 + Ry[1, 2] = 0 + Ry[2, 0] = cs.sin(pitch) + Ry[2, 1] = 0 + Ry[2, 2] = cs.cos(pitch) Rz = cs.SX.eye(3) - Rz[0,0] = cs.cos(yaw) - Rz[0,1] = cs.sin(yaw) - Rz[0,2] = 0 - Rz[1,0] = -cs.sin(yaw) - Rz[1,1] = cs.cos(yaw) - Rz[1,2] = 0 - Rz[2,0] = 0 - Rz[2,1] = 0 - Rz[2,2] = 1 - - b_R_w = Rx@Ry@Rz - - - - # Passive Arm equation + Rz[0, 0] = cs.cos(yaw) + Rz[0, 1] = cs.sin(yaw) + Rz[0, 2] = 0 + Rz[1, 0] = -cs.sin(yaw) + Rz[1, 1] = cs.cos(yaw) + Rz[1, 2] = 0 + Rz[2, 0] = 0 + Rz[2, 1] = 0 + Rz[2, 2] = 1 + + b_R_w = Rx @ Ry @ Rz + + # Passive Arm equation end_effector_initial_position = param[19:22] end_effector_to_base = end_effector_initial_position - com_position - K_bar = cs.diag(param[40:43]) - K_bar_yaw = end_effector_to_base.T@end_effector_to_base@K_bar[1, 1] + K_bar_yaw = end_effector_to_base.T @ end_effector_to_base @ K_bar[1, 1] - passive_arm_force_dot = cs.SX.zeros(6,1) - passive_arm_force_dot[0] = -K_bar[0, 0]@(states[3]) #-? - passive_arm_force_dot[1] = -K_bar[1, 1]@(states[4]) #-? + passive_arm_force_dot = cs.SX.zeros(6, 1) + passive_arm_force_dot[0] = -K_bar[0, 0] @ (states[3]) # -? + passive_arm_force_dot[1] = -K_bar[1, 1] @ (states[4]) # -? passive_arm_force_dot[2] = 0.0 passive_arm_force_dot[3] = 0.0 passive_arm_force_dot[4] = 0.0 - passive_arm_force_dot[5] = -K_bar_yaw@states[11] - - - - + passive_arm_force_dot[5] = -K_bar_yaw @ states[11] # We can choose between a fixed external compensation or a dynamic one - if(config.mpc_params['passive_arm_compensation']): + if config.mpc_params["passive_arm_compensation"]: external_wrench_linear = states[30:33] external_wrench_angular = states[33:36] else: external_wrench_linear = param[13:16] external_wrench_angular = param[16:19] - - - - # FINAL linear_com_vel STATE (1) linear_com_vel = states[3:6] - # FINAL linear_com_acc STATE (2) - temp = foot_force_fl@stanceFL - temp += foot_force_fr@stanceFR - temp += foot_force_rl@stanceRL - temp += foot_force_rr@stanceRR + temp = foot_force_fl @ stanceFL + temp += foot_force_fr @ stanceFR + temp += foot_force_rl @ stanceRL + temp += foot_force_rr @ stanceRR temp += external_wrench_linear gravity = np.array([0, 0, -9.81]) - linear_com_acc = (1/self.mass)@temp + gravity + linear_com_acc = (1 / self.mass) @ temp + gravity - - - # Start to write the component of euler_rates_base and angular_acc_base STATES w = states[9:12] roll = states[6] pitch = states[7] yaw = states[8] - + conj_euler_rates = cs.SX.eye(3) conj_euler_rates[1, 1] = cs.cos(roll) - conj_euler_rates[2, 2] = cs.cos(pitch)*cs.cos(roll) + conj_euler_rates[2, 2] = cs.cos(pitch) * cs.cos(roll) conj_euler_rates[2, 1] = -cs.sin(roll) conj_euler_rates[0, 2] = -cs.sin(pitch) - conj_euler_rates[1, 2] = cs.cos(pitch)*cs.sin(roll) - + conj_euler_rates[1, 2] = cs.cos(pitch) * cs.sin(roll) - temp2 = cs.skew(foot_position_fl - com_position)@foot_force_fl@stanceFL - temp2 += cs.skew(foot_position_fr - com_position)@foot_force_fr@stanceFR - temp2 += cs.skew(foot_position_rl - com_position)@foot_force_rl@stanceRL - temp2 += cs.skew(foot_position_rr - com_position)@foot_force_rr@stanceRR + temp2 = cs.skew(foot_position_fl - com_position) @ foot_force_fl @ stanceFL + temp2 += cs.skew(foot_position_fr - com_position) @ foot_force_fr @ stanceFR + temp2 += cs.skew(foot_position_rl - com_position) @ foot_force_rl @ stanceRL + temp2 += cs.skew(foot_position_rr - com_position) @ foot_force_rr @ stanceRR - # FINAL euler_rates_base STATE (3) - euler_rates_base = cs.inv(conj_euler_rates)@w - + euler_rates_base = cs.inv(conj_euler_rates) @ w # FINAL angular_acc_base STATE (4) - #Z Y X rotations are calculated at the top of this function! + # Z Y X rotations are calculated at the top of this function! temp2 = temp2 + external_wrench_angular - angular_acc_base = cs.inv(self.inertia)@(b_R_w@temp2 - cs.skew(w)@self.inertia@w) - - + angular_acc_base = cs.inv(self.inertia) @ (b_R_w @ temp2 - cs.skew(w) @ self.inertia @ w) # FINAL linear_foot_vel STATES (5,6,7,8) - if(not config.mpc_params["use_foothold_optimization"]): - foot_velocity_fl = foot_velocity_fl@0.0 - foot_velocity_fr = foot_velocity_fr@0.0 - foot_velocity_rl = foot_velocity_rl@0.0 - foot_velocity_rr = foot_velocity_rr@0.0 - linear_foot_vel_FL = foot_velocity_fl@(1-stanceFL)@(1-stance_proximity_FL) - linear_foot_vel_FR = foot_velocity_fr@(1-stanceFR)@(1-stance_proximity_FR) - linear_foot_vel_RL = foot_velocity_rl@(1-stanceRL)@(1-stance_proximity_RL) - linear_foot_vel_RR = foot_velocity_rr@(1-stanceRR)@(1-stance_proximity_RR) - + if not config.mpc_params["use_foothold_optimization"]: + foot_velocity_fl = foot_velocity_fl @ 0.0 + foot_velocity_fr = foot_velocity_fr @ 0.0 + foot_velocity_rl = foot_velocity_rl @ 0.0 + foot_velocity_rr = foot_velocity_rr @ 0.0 + linear_foot_vel_FL = foot_velocity_fl @ (1 - stanceFL) @ (1 - stance_proximity_FL) + linear_foot_vel_FR = foot_velocity_fr @ (1 - stanceFR) @ (1 - stance_proximity_FR) + linear_foot_vel_RL = foot_velocity_rl @ (1 - stanceRL) @ (1 - stance_proximity_RL) + linear_foot_vel_RR = foot_velocity_rr @ (1 - stanceRR) @ (1 - stance_proximity_RR) # Integral states integral_states = states[24:30] @@ -356,30 +327,42 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda integral_states[4] += roll integral_states[5] += pitch - - - - # The order of the return should be equal to the order of the states_dot - return cs.vertcat(linear_com_vel, linear_com_acc, euler_rates_base, angular_acc_base, - linear_foot_vel_FL,linear_foot_vel_FR, linear_foot_vel_RL, linear_foot_vel_RR, - integral_states, passive_arm_force_dot) - - - - - def export_robot_model(self,) -> AcadosModel: + return cs.vertcat( + linear_com_vel, + linear_com_acc, + euler_rates_base, + angular_acc_base, + linear_foot_vel_FL, + linear_foot_vel_FR, + linear_foot_vel_RL, + linear_foot_vel_RR, + integral_states, + passive_arm_force_dot, + ) + + def export_robot_model( + self, + ) -> AcadosModel: """ This method set some general properties of the NMPC, such as the params, prediction mode, etc...! It will be called in centroidal_nmpc.py """ - + # dynamics - self.param = cs.vertcat(self.stance_param, self.mu_friction, self.stance_proximity, self.base_position, - self.base_yaw, self.external_wrench, - self.end_effector_position, self.end_effector_jacobian_inv, - self.end_effector_jacobian_transpose_inv, self.end_effector_gain) - + self.param = cs.vertcat( + self.stance_param, + self.mu_friction, + self.stance_proximity, + self.base_position, + self.base_yaw, + self.external_wrench, + self.end_effector_position, + self.end_effector_jacobian_inv, + self.end_effector_jacobian_transpose_inv, + self.end_effector_gain, + ) + f_expl = self.forward_dynamics(self.states, self.inputs, self.param) f_impl = self.states_dot - f_expl @@ -392,6 +375,4 @@ def export_robot_model(self,) -> AcadosModel: acados_model.p = self.param acados_model.name = "centroidal_model" - return acados_model - diff --git a/quadruped_pympc/controllers/gradient/collaborative/centroidal_nmpc_collaborative.py b/quadruped_pympc/controllers/gradient/collaborative/centroidal_nmpc_collaborative.py index 63d7882..655c21a 100644 --- a/quadruped_pympc/controllers/gradient/collaborative/centroidal_nmpc_collaborative.py +++ b/quadruped_pympc/controllers/gradient/collaborative/centroidal_nmpc_collaborative.py @@ -1,23 +1,25 @@ # Description: This file contains the class for the NMPC controller -# Authors: Giulio Turrisi - +# Authors: Giulio Turrisi - from acados_template import AcadosOcp, AcadosOcpSolver + ACADOS_INFTY = 1000 -from .centroidal_model_collaborative import Centroidal_Model_Collaborative +import copy +import os + +import casadi as cs import numpy as np import scipy.linalg -import casadi as cs -import copy -import time +from .centroidal_model_collaborative import Centroidal_Model_Collaborative -import os dir_path = os.path.dirname(os.path.realpath(__file__)) import sys + sys.path.append(dir_path) -sys.path.append(dir_path + '/../../') +sys.path.append(dir_path + "/../../") from quadruped_pympc import config @@ -25,37 +27,28 @@ # Class for the Acados NMPC, the model is in another file! class Acados_NMPC_Collaborative: def __init__(self): - - self.horizon = config.mpc_params['horizon'] # Define the number of discretization steps - self.dt = config.mpc_params['dt'] - self.T_horizon = self.horizon*self.dt - self.use_RTI = config.mpc_params['use_RTI'] - self.use_integrators = config.mpc_params['use_integrators'] - self.use_warm_start = config.mpc_params['use_warm_start'] - self.use_foothold_constraints = config.mpc_params['use_foothold_constraints'] - - - self.use_static_stability = config.mpc_params['use_static_stability'] - self.use_zmp_stability = config.mpc_params['use_zmp_stability'] + self.horizon = config.mpc_params["horizon"] # Define the number of discretization steps + self.dt = config.mpc_params["dt"] + self.T_horizon = self.horizon * self.dt + self.use_RTI = config.mpc_params["use_RTI"] + self.use_integrators = config.mpc_params["use_integrators"] + self.use_warm_start = config.mpc_params["use_warm_start"] + self.use_foothold_constraints = config.mpc_params["use_foothold_constraints"] + + self.use_static_stability = config.mpc_params["use_static_stability"] + self.use_zmp_stability = config.mpc_params["use_zmp_stability"] self.use_stability_constraints = self.use_static_stability or self.use_zmp_stability - - - - self.previous_status = -1 self.previous_contact_sequence = np.zeros((4, self.horizon)) self.optimal_next_state = np.zeros((24,)) self.previous_optimal_GRF = np.zeros((12,)) - self.integral_errors = np.zeros((6,)) - # For centering the variable around 0, 0, 0 (World frame) self.initial_base_position = np.array([0, 0, 0]) - # Create the class of the centroidal model and instantiate the acados model self.centroidal_model = Centroidal_Model_Collaborative() acados_model = self.centroidal_model.export_robot_model() @@ -64,15 +57,15 @@ def __init__(self): # Create the acados ocp solver self.ocp = self.create_ocp_solver_description(acados_model) - - self.ocp.code_export_directory = dir_path + '/c_generated_code' - #if (not os.path.isdir(dir_path + "/c_generated_code") or os.listdir(dir_path + "/c_generated_code") == []): - self.acados_ocp_solver = AcadosOcpSolver(self.ocp, json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json") - - #else : - # self.acados_ocp_solver = AcadosOcpSolver(self.ocp, json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json", build = False, generate = True) + self.ocp.code_export_directory = dir_path + "/c_generated_code" + # if (not os.path.isdir(dir_path + "/c_generated_code") or os.listdir(dir_path + "/c_generated_code") == []): + self.acados_ocp_solver = AcadosOcpSolver( + self.ocp, json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json" + ) + # else : + # self.acados_ocp_solver = AcadosOcpSolver(self.ocp, json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json", build = False, generate = True) # Initialize solver for stage in range(self.horizon + 1): @@ -80,17 +73,12 @@ def __init__(self): for stage in range(self.horizon): self.acados_ocp_solver.set(stage, "u", np.zeros((self.inputs_dim,))) - if(self.use_RTI): + if self.use_RTI: # first preparation phase - self.acados_ocp_solver.options_set('rti_phase', 1) - status = self.acados_ocp_solver.solve() - - - - - + self.acados_ocp_solver.options_set("rti_phase", 1) + status = self.acados_ocp_solver.solve() - # Set cost, constraints and options + # Set cost, constraints and options def create_ocp_solver_description(self, acados_model) -> AcadosOcp: # Create ocp object to formulate the OCP ocp = AcadosOcp() @@ -126,39 +114,33 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: ocp.cost.yref_e = np.zeros((ny_e,)) # Set friction and foothold constraints - expr_h_friction, \ - self.constr_uh_friction, \ - self.constr_lh_friction = self.create_friction_cone_constraints() - + expr_h_friction, self.constr_uh_friction, self.constr_lh_friction = self.create_friction_cone_constraints() + ocp.model.con_h_expr = expr_h_friction ocp.constraints.uh = self.constr_uh_friction ocp.constraints.lh = self.constr_lh_friction ocp.model.con_h_expr_0 = expr_h_friction ocp.constraints.uh_0 = self.constr_uh_friction ocp.constraints.lh_0 = self.constr_lh_friction - nsh = expr_h_friction.shape[0] + nsh = expr_h_friction.shape[0] nsh_state_constraint_start = copy.copy(nsh) - - if(self.use_foothold_constraints): - expr_h_foot, \ - self.constr_uh_foot, \ - self.constr_lh_foot = self.create_foothold_constraints() + if self.use_foothold_constraints: + expr_h_foot, self.constr_uh_foot, self.constr_lh_foot = self.create_foothold_constraints() - ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_foot) + ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_foot) ocp.constraints.uh = np.concatenate((ocp.constraints.uh, self.constr_uh_foot)) ocp.constraints.lh = np.concatenate((ocp.constraints.lh, self.constr_lh_foot)) nsh += expr_h_foot.shape[0] - - # Set stability constraints - if(self.use_stability_constraints): + # Set stability constraints + if self.use_stability_constraints: self.nsh_stability_start = copy.copy(nsh) - expr_h_support_polygon, \ - self.constr_uh_support_polygon, \ - self.constr_lh_support_polygon = self.create_stability_constraints() + expr_h_support_polygon, self.constr_uh_support_polygon, self.constr_lh_support_polygon = ( + self.create_stability_constraints() + ) - ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_support_polygon) + ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_support_polygon) ocp.constraints.uh = np.concatenate((ocp.constraints.uh, self.constr_uh_support_polygon)) ocp.constraints.lh = np.concatenate((ocp.constraints.lh, self.constr_lh_support_polygon)) nsh += expr_h_support_polygon.shape[0] @@ -166,19 +148,23 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: nsh_state_constraint_end = copy.copy(nsh) - # Set slack variable configuration: num_state_cstr = nsh_state_constraint_end - nsh_state_constraint_start - if(num_state_cstr > 0): - ocp.constraints.lsh = np.zeros(num_state_cstr) # Lower bounds on slacks corresponding to soft lower bounds for nonlinear constraints - ocp.constraints.ush = np.zeros(num_state_cstr) # Lower bounds on slacks corresponding to soft upper bounds for nonlinear constraints - ocp.constraints.idxsh = np.array(range(nsh_state_constraint_start, nsh_state_constraint_end)) # Jsh + if num_state_cstr > 0: + ocp.constraints.lsh = np.zeros( + num_state_cstr + ) # Lower bounds on slacks corresponding to soft lower bounds for nonlinear constraints + ocp.constraints.ush = np.zeros( + num_state_cstr + ) # Lower bounds on slacks corresponding to soft upper bounds for nonlinear constraints + ocp.constraints.idxsh = np.array(range(nsh_state_constraint_start, nsh_state_constraint_end)) # Jsh ns = num_state_cstr - ocp.cost.zl = 100 * np.ones((ns,)) # gradient wrt lower slack at intermediate shooting nodes (1 to N-1) - ocp.cost.Zl = 1 * np.ones((ns,)) # diagonal of Hessian wrt lower slack at intermediate shooting nodes (1 to N-1) - ocp.cost.zu = 100 * np.ones((ns,)) - ocp.cost.Zu = 1 * np.ones((ns,)) - + ocp.cost.zl = 100 * np.ones((ns,)) # gradient wrt lower slack at intermediate shooting nodes (1 to N-1) + ocp.cost.Zl = 1 * np.ones( + (ns,) + ) # diagonal of Hessian wrt lower slack at intermediate shooting nodes (1 to N-1) + ocp.cost.zu = 100 * np.ones((ns,)) + ocp.cost.Zu = 1 * np.ones((ns,)) # Variables to save the upper and lower bound of the constraints applied list_upper_bound = [] @@ -189,13 +175,12 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: self.upper_bound = np.array(list_upper_bound, dtype=object) self.lower_bound = np.array(list_lower_bound, dtype=object) - # Set initial state constraint X0 = np.zeros(shape=(nx,)) ocp.constraints.x0 = X0 # Set initialize parameters - init_contact_status = np.array([1., 1., 1., 1.]) + init_contact_status = np.array([1.0, 1.0, 1.0, 1.0]) init_mu = np.array([0.5]) init_stance_proximity = np.array([0, 0, 0, 0]) init_base_position = np.array([0, 0, 0]) @@ -205,136 +190,155 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: init_end_effector_jacobian_inv = np.array([1, 1, 1, 1, 1, 1, 1, 1, 1]) init_end_effector_jacobian_transpose_inv = np.array([1, 1, 1, 1, 1, 1, 1, 1, 1]) init_end_effector_gain = np.array([10, 10, 10]) - ocp.parameter_values = np.concatenate((init_contact_status, init_mu, init_stance_proximity, - init_base_position, init_base_yaw, init_external_wrench, - init_end_effector_position, init_end_effector_jacobian_inv, - init_end_effector_jacobian_transpose_inv, init_end_effector_gain)) - - - + ocp.parameter_values = np.concatenate( + ( + init_contact_status, + init_mu, + init_stance_proximity, + init_base_position, + init_base_yaw, + init_external_wrench, + init_end_effector_position, + init_end_effector_jacobian_inv, + init_end_effector_jacobian_transpose_inv, + init_end_effector_gain, + ) + ) # Set options - ocp.solver_options.qp_solver = "PARTIAL_CONDENSING_HPIPM" # FULL_CONDENSING_QPOASES PARTIAL_CONDENSING_OSQP PARTIAL_CONDENSING_HPIPM + ocp.solver_options.qp_solver = ( + "PARTIAL_CONDENSING_HPIPM" # FULL_CONDENSING_QPOASES PARTIAL_CONDENSING_OSQP PARTIAL_CONDENSING_HPIPM + ) ocp.solver_options.hessian_approx = "GAUSS_NEWTON" # 'GAUSS_NEWTON', 'EXACT' - ocp.solver_options.integrator_type = "ERK" #ERK IRK GNSF DISCRETE - if(self.use_RTI): - ocp.solver_options.nlp_solver_type = "SQP_RTI" + ocp.solver_options.integrator_type = "ERK" # ERK IRK GNSF DISCRETE + if self.use_RTI: + ocp.solver_options.nlp_solver_type = "SQP_RTI" ocp.solver_options.nlp_solver_max_iter = 1 - # Set the RTI type for the advanced RTI method + # Set the RTI type for the advanced RTI method # (see https://arxiv.org/pdf/2403.07101.pdf) - if(config.mpc_params['as_rti_type'] == "AS-RTI-A"): + if config.mpc_params["as_rti_type"] == "AS-RTI-A": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 0 - elif(config.mpc_params['as_rti_type'] == "AS-RTI-B"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-B": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 1 - elif(config.mpc_params['as_rti_type'] == "AS-RTI-C"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-C": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 2 - elif(config.mpc_params['as_rti_type'] == "AS-RTI-D"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-D": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 3 else: - ocp.solver_options.nlp_solver_type = "SQP" - ocp.solver_options.nlp_solver_max_iter = config.mpc_params['num_qp_iterations'] - #ocp.solver_options.globalization = "MERIT_BACKTRACKING" # FIXED_STEP, MERIT_BACKTRACKING - + ocp.solver_options.nlp_solver_type = "SQP" + ocp.solver_options.nlp_solver_max_iter = config.mpc_params["num_qp_iterations"] + # ocp.solver_options.globalization = "MERIT_BACKTRACKING" # FIXED_STEP, MERIT_BACKTRACKING - if(config.mpc_params['solver_mode'] == "balance"): + if config.mpc_params["solver_mode"] == "balance": ocp.solver_options.hpipm_mode = "BALANCE" - elif(config.mpc_params['solver_mode'] == "robust"): + elif config.mpc_params["solver_mode"] == "robust": ocp.solver_options.hpipm_mode = "ROBUST" - elif(config.mpc_params['solver_mode'] == "fast"): + elif config.mpc_params["solver_mode"] == "fast": ocp.solver_options.qp_solver_iter_max = 10 ocp.solver_options.hpipm_mode = "SPEED" - elif(config.mpc_params['solver_mode'] == "crazy_speed"): + elif config.mpc_params["solver_mode"] == "crazy_speed": ocp.solver_options.qp_solver_iter_max = 5 ocp.solver_options.hpipm_mode = "SPEED_ABS" - #ocp.solver_options.line_search_use_sufficient_descent = 1 - #ocp.solver_options.regularize_method = "PROJECT_REDUC_HESS" - #ocp.solver_options.nlp_solver_ext_qp_res = 1 + # ocp.solver_options.line_search_use_sufficient_descent = 1 + # ocp.solver_options.regularize_method = "PROJECT_REDUC_HESS" + # ocp.solver_options.nlp_solver_ext_qp_res = 1 ocp.solver_options.levenberg_marquardt = 1e-3 - # Set prediction horizon ocp.solver_options.tf = self.T_horizon - # Nonuniform discretization - if(config.mpc_params['use_nonuniform_discretization']): - time_steps_fine_grained = np.tile(config.mpc_params['dt_fine_grained'], config.mpc_params['horizon_fine_grained']) - time_steps = np.concatenate((time_steps_fine_grained, np.tile(self.dt, self.horizon - config.mpc_params['horizon_fine_grained']))) - shooting_nodes = np.zeros((self.horizon+1,)) + if config.mpc_params["use_nonuniform_discretization"]: + time_steps_fine_grained = np.tile( + config.mpc_params["dt_fine_grained"], config.mpc_params["horizon_fine_grained"] + ) + time_steps = np.concatenate( + (time_steps_fine_grained, np.tile(self.dt, self.horizon - config.mpc_params["horizon_fine_grained"])) + ) + shooting_nodes = np.zeros((self.horizon + 1,)) for i in range(len(time_steps)): - shooting_nodes[i+1] = shooting_nodes[i] + time_steps[i] + shooting_nodes[i + 1] = shooting_nodes[i] + time_steps[i] ocp.solver_options.shooting_nodes = shooting_nodes - return ocp - # Create a constraint for stability (COM, ZMP or CP inside support polygon) - def create_stability_constraints(self,) -> None: - + def create_stability_constraints( + self, + ) -> None: base_w = self.centroidal_model.states[0:3] base_vel_w = self.centroidal_model.states[3:6] - FL = self.centroidal_model.states[12:15] FR = self.centroidal_model.states[15:18] RL = self.centroidal_model.states[18:21] RR = self.centroidal_model.states[21:24] - #yaw = self.centroidal_model.base_yaw[0] + # yaw = self.centroidal_model.base_yaw[0] yaw = self.centroidal_model.states[8] h_R_w = cs.SX.zeros(2, 2) - h_R_w[0,0] = cs.cos(yaw) - h_R_w[0,1] = cs.sin(yaw) - h_R_w[1,0] = -cs.sin(yaw) - h_R_w[1,1] = cs.cos(yaw) - - FL[0:2] = h_R_w@(FL[0:2] - base_w[0:2]) - FR[0:2] = h_R_w@(FR[0:2] - base_w[0:2]) - RL[0:2] = h_R_w@(RL[0:2] - base_w[0:2]) - RR[0:2] = h_R_w@(RR[0:2] - base_w[0:2]) - + h_R_w[0, 0] = cs.cos(yaw) + h_R_w[0, 1] = cs.sin(yaw) + h_R_w[1, 0] = -cs.sin(yaw) + h_R_w[1, 1] = cs.cos(yaw) - if(self.use_static_stability): + FL[0:2] = h_R_w @ (FL[0:2] - base_w[0:2]) + FR[0:2] = h_R_w @ (FR[0:2] - base_w[0:2]) + RL[0:2] = h_R_w @ (RL[0:2] - base_w[0:2]) + RR[0:2] = h_R_w @ (RR[0:2] - base_w[0:2]) + + if self.use_static_stability: x = 0.0 y = 0.0 else: # Compute the ZMP robotHeight = base_w[2] - foot_force_fl = self.centroidal_model.inputs[12:15]@self.centroidal_model.param[0] - foot_force_fr = self.centroidal_model.inputs[15:18]@self.centroidal_model.param[1] - foot_force_rl = self.centroidal_model.inputs[18:21]@self.centroidal_model.param[2] - foot_force_rr = self.centroidal_model.inputs[21:24]@self.centroidal_model.param[3] + foot_force_fl = self.centroidal_model.inputs[12:15] @ self.centroidal_model.param[0] + foot_force_fr = self.centroidal_model.inputs[15:18] @ self.centroidal_model.param[1] + foot_force_rl = self.centroidal_model.inputs[18:21] @ self.centroidal_model.param[2] + foot_force_rr = self.centroidal_model.inputs[21:24] @ self.centroidal_model.param[3] temp = foot_force_fl + foot_force_fr + foot_force_rl + foot_force_rr - - gravity = np.array([0, 0, -9.81]) - linear_com_acc = (1/self.centroidal_model.mass)@temp + gravity - zmp = base_w[0:2] - linear_com_acc[0:2]*(robotHeight/(-gravity[2])) + gravity = np.array([0, 0, -9.81]) + linear_com_acc = (1 / self.centroidal_model.mass) @ temp + gravity + zmp = base_w[0:2] - linear_com_acc[0:2] * (robotHeight / (-gravity[2])) external_forces_linear = self.centroidal_model.states[30:33] end_effector_position = self.centroidal_model.param[19:22] - zmp_x = self.centroidal_model.mass*-gravity[2]*base_w[0] - base_w[2]*self.centroidal_model.mass*linear_com_acc[0]*0 - zmp_x = zmp_x - end_effector_position[0]*external_forces_linear[2] + end_effector_position[2]*external_forces_linear[0] - zmp_x = zmp_x/(self.centroidal_model.mass*-gravity[2] - external_forces_linear[2]) - - zmp_y = self.centroidal_model.mass*-gravity[2]*base_w[1] - base_w[2]*self.centroidal_model.mass*linear_com_acc[1]*0 - zmp_y = zmp_y - end_effector_position[1]*external_forces_linear[2] + end_effector_position[2]*external_forces_linear[1] - zmp_y = zmp_y/(self.centroidal_model.mass*-gravity[2] - external_forces_linear[2]) + zmp_x = ( + self.centroidal_model.mass * -gravity[2] * base_w[0] + - base_w[2] * self.centroidal_model.mass * linear_com_acc[0] * 0 + ) + zmp_x = ( + zmp_x + - end_effector_position[0] * external_forces_linear[2] + + end_effector_position[2] * external_forces_linear[0] + ) + zmp_x = zmp_x / (self.centroidal_model.mass * -gravity[2] - external_forces_linear[2]) + + zmp_y = ( + self.centroidal_model.mass * -gravity[2] * base_w[1] + - base_w[2] * self.centroidal_model.mass * linear_com_acc[1] * 0 + ) + zmp_y = ( + zmp_y + - end_effector_position[1] * external_forces_linear[2] + + end_effector_position[2] * external_forces_linear[1] + ) + zmp_y = zmp_y / (self.centroidal_model.mass * -gravity[2] - external_forces_linear[2]) zmp = cs.vertcat(zmp_x, zmp_y) - zmp = h_R_w@(zmp - base_w[0:2]) + zmp = h_R_w @ (zmp - base_w[0:2]) x = zmp[0] y = zmp[1] - y_FL = FL[1] y_FR = FR[1] y_RL = RL[1] @@ -345,64 +349,60 @@ def create_stability_constraints(self,) -> None: x_RL = RL[0] x_RR = RR[0] - #LF - RF : x < (x2 - x1) (y - y1) / (y2 - y1) + x1 - #RF - RH: y > (y2 - y1) (x - x1) / (x2 - x1) + y1 - #RH - LH : x > (x2 - x1) (y - y1) / (y2 - y1) + x1 - #LH - LF: y < (y2 - y1) (x - x1) / (x2 - x1) + y1 - - #FL and FR cannot stay at the same x! #constrint should be less than zero - constraint_FL_FR = x - (x_FR - x_FL)*(y - y_FL) / (y_FR - y_FL + 0.001) - x_FL + # LF - RF : x < (x2 - x1) (y - y1) / (y2 - y1) + x1 + # RF - RH: y > (y2 - y1) (x - x1) / (x2 - x1) + y1 + # RH - LH : x > (x2 - x1) (y - y1) / (y2 - y1) + x1 + # LH - LF: y < (y2 - y1) (x - x1) / (x2 - x1) + y1 + + # FL and FR cannot stay at the same x! #constrint should be less than zero + constraint_FL_FR = x - (x_FR - x_FL) * (y - y_FL) / (y_FR - y_FL + 0.001) - x_FL + + # FR and RR cannot stay at the same y! #constraint should be bigger than zero + constraint_FR_RR = y - (y_RR - y_FR) * (x - x_FR) / (x_RR - x_FR + 0.001) - y_FR + + # RL and RR cannot stay at the same x! #constraint should be bigger than zero + constraint_RR_RL = x - (x_RL - x_RR) * (y - y_RR) / (y_RL - y_RR + 0.001) - x_RR + + # FL and RL cannot stay at the same y! #constraint should be less than zero + constraint_RL_FL = y - (y_FL - y_RL) * (x - x_RL) / (x_FL - x_RL + 0.001) - y_RL - #FR and RR cannot stay at the same y! #constraint should be bigger than zero - constraint_FR_RR = y - (y_RR - y_FR)*(x - x_FR) / (x_RR - x_FR + 0.001) - y_FR - - #RL and RR cannot stay at the same x! #constraint should be bigger than zero - constraint_RR_RL = x - (x_RL - x_RR)*(y - y_RR) / (y_RL - y_RR + 0.001) - x_RR - - #FL and RL cannot stay at the same y! #constraint should be less than zero - constraint_RL_FL = y - (y_FL - y_RL)*(x - x_RL) / (x_FL - x_RL + 0.001) - y_RL - # the diagonal stuff can be at any point... - constraint_FL_RR = y - (y_RR - y_FL)*(x - x_FL) / (x_RR - x_FL + 0.001) - y_FL #bigger than zero - constraint_FR_RL = y - (y_RL - y_FR)*(x - x_FR) / (x_RL - x_FR + 0.001) - y_FR #bigger than zero + constraint_FL_RR = y - (y_RR - y_FL) * (x - x_FL) / (x_RR - x_FL + 0.001) - y_FL # bigger than zero + constraint_FR_RL = y - (y_RL - y_FR) * (x - x_FR) / (x_RL - x_FR + 0.001) - y_FR # bigger than zero # upper and lower bound - ub = np.ones(6)*ACADOS_INFTY - lb = -np.ones(6)*ACADOS_INFTY + ub = np.ones(6) * ACADOS_INFTY + lb = -np.ones(6) * ACADOS_INFTY - #constraint_FL_FR + # constraint_FL_FR ub[0] = 0 lb[0] = -ACADOS_INFTY - - #constraint_FR_RR + + # constraint_FR_RR ub[1] = ACADOS_INFTY lb[1] = 0 - #constraint_RR_RL + # constraint_RR_RL ub[2] = ACADOS_INFTY lb[2] = 0 - #constraint_RL_FL + # constraint_RL_FL ub[3] = 0 lb[3] = -ACADOS_INFTY - - #constraint_FL_RR + + # constraint_FL_RR ub[4] = 0 lb[4] = -ACADOS_INFTY - - #constraint_FR_RL + + # constraint_FR_RL ub[5] = 0 lb[5] = -ACADOS_INFTY + Jb = cs.vertcat( + constraint_FL_FR, constraint_FR_RR, constraint_RR_RL, constraint_RL_FL, constraint_FL_RR, constraint_FR_RL + ) - - Jb = cs.vertcat(constraint_FL_FR, constraint_FR_RR, - constraint_RR_RL, constraint_RL_FL, - constraint_FL_RR, constraint_FR_RL) - - - - #create some casadi function for the derivative of the constraint if needed + # create some casadi function for the derivative of the constraint if needed """temp = cs.vertcat(self.centroidal_model.states, self.centroidal_model.inputs, self.centroidal_model.stanceFL, self.centroidal_model.stanceFR, self.centroidal_model.stanceRL, self.centroidal_model.stanceRR) @@ -423,13 +423,13 @@ def create_stability_constraints(self,) -> None: constraint_FR_RL_jac = cs.jacobian(constraint_FR_RL, temp) self.constraint_FR_RL_jac_fun = cs.Function('constraint_FR_RL_jac_fun', [temp], [constraint_FR_RL_jac])""" - - return Jb, ub, lb - + return Jb, ub, lb # Create a standard foothold box constraint - def create_foothold_constraints(self,): + def create_foothold_constraints( + self, + ): """ This function calculates the symbolic foothold constraints for the centroidal NMPC problem. @@ -439,52 +439,45 @@ def create_foothold_constraints(self,): lbu: A numpy array of shape (8,) representing the lower bounds of the foothold constraints. """ - ubu = np.ones(12)*ACADOS_INFTY - lbu = -np.ones(12)*ACADOS_INFTY + ubu = np.ones(12) * ACADOS_INFTY + lbu = -np.ones(12) * ACADOS_INFTY - - # The foothold constraint in acados are in the horizontal frame, + # The foothold constraint in acados are in the horizontal frame, # but they arrive to us in the world frame. We need to rotate them # using the robot yaw yaw = self.centroidal_model.base_yaw[0] h_R_w = cs.SX.zeros(2, 2) - h_R_w[0,0] = cs.cos(yaw) - h_R_w[0,1] = cs.sin(yaw) - h_R_w[1,0] = -cs.sin(yaw) - h_R_w[1,1] = cs.cos(yaw) - + h_R_w[0, 0] = cs.cos(yaw) + h_R_w[0, 1] = cs.sin(yaw) + h_R_w[1, 0] = -cs.sin(yaw) + h_R_w[1, 1] = cs.cos(yaw) # and translate them using the robot base position base = self.centroidal_model.base_position[0:3] - - - foot_position_fl = cs.SX.zeros(3,1) - foot_position_fl[0:2] = h_R_w@cs.vertcat(self.centroidal_model.states[12:14] - base[0:2]) + foot_position_fl = cs.SX.zeros(3, 1) + foot_position_fl[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.states[12:14] - base[0:2]) foot_position_fl[2] = self.centroidal_model.states[14] - - foot_position_fr = cs.SX.zeros(3,1) - foot_position_fr[0:2] = h_R_w@cs.vertcat(self.centroidal_model.states[15:17] - base[0:2]) + + foot_position_fr = cs.SX.zeros(3, 1) + foot_position_fr[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.states[15:17] - base[0:2]) foot_position_fr[2] = self.centroidal_model.states[17] - - foot_position_rl = cs.SX.zeros(3,1) - foot_position_rl[0:2] = h_R_w@cs.vertcat(self.centroidal_model.states[18:20] - base[0:2]) + + foot_position_rl = cs.SX.zeros(3, 1) + foot_position_rl[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.states[18:20] - base[0:2]) foot_position_rl[2] = self.centroidal_model.states[20] - - foot_position_rr = cs.SX.zeros(3,1) - foot_position_rr[0:2] = h_R_w@cs.vertcat(self.centroidal_model.states[21:23] - base[0:2]) + + foot_position_rr = cs.SX.zeros(3, 1) + foot_position_rr[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.states[21:23] - base[0:2]) foot_position_rr[2] = self.centroidal_model.states[23] - - - Jbu = cs.vertcat(foot_position_fl, foot_position_fr, - foot_position_rl, foot_position_rr) + Jbu = cs.vertcat(foot_position_fl, foot_position_fr, foot_position_rl, foot_position_rr) return Jbu, ubu, lbu - - # Create the friction cone constraint - def create_friction_cone_constraints(self,) ->None: + def create_friction_cone_constraints( + self, + ) -> None: """ This function calculates the symbolic friction cone constraints for the centroidal NMPC problem. @@ -497,47 +490,45 @@ def create_friction_cone_constraints(self,) ->None: t = np.array([1, 0, 0]) b = np.array([0, 1, 0]) mu = self.centroidal_model.mu_friction - f_max = config.mpc_params['grf_max'] - f_min = config.mpc_params['grf_min'] + f_max = config.mpc_params["grf_max"] + f_min = config.mpc_params["grf_min"] # Derivation can be found in the paper # "High-slope terrain locomotion for torque-controlled quadruped robots", # M Focchi, A Del Prete, I Havoutis, R Featherstone, DG Caldwell, C Semini Jbu = cs.SX.zeros(20, 12) - Jbu[0, :3] = -n*mu + t - Jbu[1, :3] = -n*mu + b - Jbu[2, :3] = n*mu + b - Jbu[3, :3] = n*mu + t + Jbu[0, :3] = -n * mu + t + Jbu[1, :3] = -n * mu + b + Jbu[2, :3] = n * mu + b + Jbu[3, :3] = n * mu + t Jbu[4, :3] = n - Jbu[5, 3:6] = -n*mu + t - Jbu[6, 3:6] = -n*mu + b - Jbu[7, 3:6] = n*mu + b - Jbu[8, 3:6] = n*mu + t + Jbu[5, 3:6] = -n * mu + t + Jbu[6, 3:6] = -n * mu + b + Jbu[7, 3:6] = n * mu + b + Jbu[8, 3:6] = n * mu + t Jbu[9, 3:6] = n - - Jbu[10, 6:9] = -n*mu + t - Jbu[11, 6:9] = -n*mu + b - Jbu[12, 6:9] = n*mu + b - Jbu[13, 6:9] = n*mu + t + + Jbu[10, 6:9] = -n * mu + t + Jbu[11, 6:9] = -n * mu + b + Jbu[12, 6:9] = n * mu + b + Jbu[13, 6:9] = n * mu + t Jbu[14, 6:9] = n - - Jbu[15, 9:12] = -n*mu + t - Jbu[16, 9:12] = -n*mu + b - Jbu[17, 9:12] = n*mu + b - Jbu[18, 9:12] = n*mu + t + + Jbu[15, 9:12] = -n * mu + t + Jbu[16, 9:12] = -n * mu + b + Jbu[17, 9:12] = n * mu + b + Jbu[18, 9:12] = n * mu + t Jbu[19, 9:12] = n # C matrix - Jbu = Jbu@cs.vertcat(self.centroidal_model.inputs[12:24]) - - + Jbu = Jbu @ cs.vertcat(self.centroidal_model.inputs[12:24]) # lower bound (-ACADOS_INFTY is "almost" -inf) lbu = np.zeros(20) lbu[0] = -ACADOS_INFTY lbu[1] = -ACADOS_INFTY - lbu[2] = 0 + lbu[2] = 0 lbu[3] = 0 lbu[4] = f_min lbu[5:10] = lbu[0:5] @@ -546,7 +537,7 @@ def create_friction_cone_constraints(self,) ->None: # upper bound (ACADOS_INFTY is "almost" inf) ubu = np.zeros(20) - ubu[0] = 0 + ubu[0] = 0 ubu[1] = 0 ubu[2] = ACADOS_INFTY ubu[3] = ACADOS_INFTY @@ -557,57 +548,73 @@ def create_friction_cone_constraints(self,) ->None: return Jbu, ubu, lbu - - def set_weight(self, nx, nu): # Define the weight matrices for the cost function - Q_position = np.array([0, 0, 1500]) #x, y, z - Q_velocity = np.array([200, 200, 200]) #x_vel, y_vel, z_vel - Q_base_angle = np.array([1000, 1200, 0]) #roll, pitch, yaw - Q_base_angle_rates = np.array([20, 20, 600]) #roll_rate, pitch_rate, yaw_rate - Q_foot_pos = np.array([300, 300, 300]) #f_x, f_y, f_z (should be 4 times this, once per foot) - Q_com_position_z_integral = np.array([50]) #integral of z_com - Q_com_velocity_x_integral = np.array([10]) #integral of x_com - Q_com_velocity_y_integral = np.array([10]) #integral of y_com - Q_com_velocity_z_integral = np.array([10]) #integral of z_com_vel - Q_roll_integral_integral = np.array([10]) #integral of roll - Q_pitch_integral_integral = np.array([10]) #integral of pitch - Q_passive_arm_force = np.array([0, 0, 0, 0, 0, 0]) #end_effector_position_z - - R_foot_vel = np.array([0.0001, 0.0001, 0.00001]) #v_x, v_y, v_z (should be 4 times this, once per foot) - if(config.robot == "hyqreal"): - R_foot_force = np.array([0.00001, 0.00001, 0.00001])#f_x, f_y, f_z (should be 4 times this, once per foot) + Q_position = np.array([0, 0, 1500]) # x, y, z + Q_velocity = np.array([200, 200, 200]) # x_vel, y_vel, z_vel + Q_base_angle = np.array([1000, 1200, 0]) # roll, pitch, yaw + Q_base_angle_rates = np.array([20, 20, 600]) # roll_rate, pitch_rate, yaw_rate + Q_foot_pos = np.array([300, 300, 300]) # f_x, f_y, f_z (should be 4 times this, once per foot) + Q_com_position_z_integral = np.array([50]) # integral of z_com + Q_com_velocity_x_integral = np.array([10]) # integral of x_com + Q_com_velocity_y_integral = np.array([10]) # integral of y_com + Q_com_velocity_z_integral = np.array([10]) # integral of z_com_vel + Q_roll_integral_integral = np.array([10]) # integral of roll + Q_pitch_integral_integral = np.array([10]) # integral of pitch + Q_passive_arm_force = np.array([0, 0, 0, 0, 0, 0]) # end_effector_position_z + + R_foot_vel = np.array([0.0001, 0.0001, 0.00001]) # v_x, v_y, v_z (should be 4 times this, once per foot) + if config.robot == "hyqreal": + R_foot_force = np.array( + [0.00001, 0.00001, 0.00001] + ) # f_x, f_y, f_z (should be 4 times this, once per foot) else: R_foot_force = np.array([0.001, 0.001, 0.001]) - Q_mat = np.diag(np.concatenate((Q_position, Q_velocity, - Q_base_angle, Q_base_angle_rates, - Q_foot_pos, Q_foot_pos, Q_foot_pos, Q_foot_pos, - Q_com_position_z_integral, Q_com_velocity_x_integral, - Q_com_velocity_y_integral, Q_com_velocity_z_integral, - Q_roll_integral_integral, Q_pitch_integral_integral, - Q_passive_arm_force))) - - R_mat = np.diag(np.concatenate((R_foot_vel, R_foot_vel, R_foot_vel, R_foot_vel, - R_foot_force, R_foot_force, R_foot_force, R_foot_force))) + Q_mat = np.diag( + np.concatenate( + ( + Q_position, + Q_velocity, + Q_base_angle, + Q_base_angle_rates, + Q_foot_pos, + Q_foot_pos, + Q_foot_pos, + Q_foot_pos, + Q_com_position_z_integral, + Q_com_velocity_x_integral, + Q_com_velocity_y_integral, + Q_com_velocity_z_integral, + Q_roll_integral_integral, + Q_pitch_integral_integral, + Q_passive_arm_force, + ) + ) + ) + + R_mat = np.diag( + np.concatenate( + (R_foot_vel, R_foot_vel, R_foot_vel, R_foot_vel, R_foot_force, R_foot_force, R_foot_force, R_foot_force) + ) + ) return Q_mat, R_mat - - def reset(self): self.acados_ocp_solver.reset() - self.acados_ocp_solver = AcadosOcpSolver(self.ocp, - json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json", - build = False, generate = False) - - + self.acados_ocp_solver = AcadosOcpSolver( + self.ocp, + json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json", + build=False, + generate=False, + ) def set_stage_constraint(self, constraint, state, reference, contact_sequence, h_R_w, stance_proximity): """ Set the stage constraint for the centroidal NMPC problem. We only consider the stance constraint, and the swing - constraint up to 2 maximum references. + constraint up to 2 maximum references. Args: constraint (numpy.ndarray or None): Constraint passed from outside (e.g. vision). If None, nominal foothold is used. @@ -621,9 +628,7 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h None """ try: - - - # Take the array of the contact sequence and split + # Take the array of the contact sequence and split # it in 4 arrays for clarity FL_contact_sequence = contact_sequence[0] FR_contact_sequence = contact_sequence[1] @@ -641,74 +646,67 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h RL_reference_foot = reference["ref_foot_RL"] RR_reference_foot = reference["ref_foot_RR"] - - # Take the base position and the yaw rotation matrix. This is needed to + # Take the base position and the yaw rotation matrix. This is needed to # express the foothold constraint in the horizontal frame base = state["position"] - - h_R_w = h_R_w.reshape((2,2)) - + h_R_w = h_R_w.reshape((2, 2)) # Divide the constraint in upper and lower bound. The constraint are - # represented by 4 vertex, but we only use box constraint hence + # represented by 4 vertex, but we only use box constraint hence # we need only 2 vertex for each constraint (first an last) # For the leg in stance now, we simply enlarge the actual position as a box constraint # maybe this is not needed, but we cannot disable constraint!! - + # FL Stance Constraint stance_up_constraint_FL = np.array([FL_actual_foot[0], FL_actual_foot[1], FL_actual_foot[2] + 0.002]) - stance_up_constraint_FL[0:2] = h_R_w@(stance_up_constraint_FL[0:2] - base[0:2]) + stance_up_constraint_FL[0:2] = h_R_w @ (stance_up_constraint_FL[0:2] - base[0:2]) stance_up_constraint_FL[0:2] = stance_up_constraint_FL[0:2] + 0.1 stance_up_constraint_FL[2] = stance_up_constraint_FL[2] + 0.01 - + stance_low_constraint_FL = np.array([FL_actual_foot[0], FL_actual_foot[1], FL_actual_foot[2] - 0.002]) - stance_low_constraint_FL[0:2] = h_R_w@(stance_low_constraint_FL[0:2] - base[0:2]) + stance_low_constraint_FL[0:2] = h_R_w @ (stance_low_constraint_FL[0:2] - base[0:2]) stance_low_constraint_FL[0:2] = stance_low_constraint_FL[0:2] - 0.1 stance_low_constraint_FL[2] = stance_low_constraint_FL[2] - 0.01 - # FR Stance Constraint stance_up_constraint_FR = np.array([FR_actual_foot[0], FR_actual_foot[1], FR_actual_foot[2] + 0.002]) - stance_up_constraint_FR[0:2] = h_R_w@(stance_up_constraint_FR[0:2] - base[0:2]) + stance_up_constraint_FR[0:2] = h_R_w @ (stance_up_constraint_FR[0:2] - base[0:2]) stance_up_constraint_FR[0:2] = stance_up_constraint_FR[0:2] + 0.1 stance_up_constraint_FR[2] = stance_up_constraint_FR[2] + 0.01 stance_low_constraint_FR = np.array([FR_actual_foot[0], FR_actual_foot[1], FR_actual_foot[2] - 0.002]) - stance_low_constraint_FR[0:2] = h_R_w@(stance_low_constraint_FR[0:2] - base[0:2]) + stance_low_constraint_FR[0:2] = h_R_w @ (stance_low_constraint_FR[0:2] - base[0:2]) stance_low_constraint_FR[0:2] = stance_low_constraint_FR[0:2] - 0.1 stance_low_constraint_FR[2] = stance_low_constraint_FR[2] - 0.01 - # RL Stance Constraint stance_up_constraint_RL = np.array([RL_actual_foot[0], RL_actual_foot[1], RL_actual_foot[2] + 0.002]) - stance_up_constraint_RL[0:2] = h_R_w@(stance_up_constraint_RL[0:2]- base[0:2]) + stance_up_constraint_RL[0:2] = h_R_w @ (stance_up_constraint_RL[0:2] - base[0:2]) stance_up_constraint_RL[0:2] = stance_up_constraint_RL[0:2] + 0.1 stance_up_constraint_RL[2] = stance_up_constraint_RL[2] + 0.01 - + stance_low_constraint_RL = np.array([RL_actual_foot[0], RL_actual_foot[1], RL_actual_foot[2] - 0.002]) - stance_low_constraint_RL[0:2] = h_R_w@(stance_low_constraint_RL[0:2]- base[0:2]) + stance_low_constraint_RL[0:2] = h_R_w @ (stance_low_constraint_RL[0:2] - base[0:2]) stance_low_constraint_RL[0:2] = stance_low_constraint_RL[0:2] - 0.1 stance_low_constraint_RL[2] = stance_low_constraint_RL[2] - 0.01 - # RR Stance Constraint stance_up_constraint_RR = np.array([RR_actual_foot[0], RR_actual_foot[1], RR_actual_foot[2] + 0.002]) - stance_up_constraint_RR[0:2] = h_R_w@(stance_up_constraint_RR[0:2]- base[0:2]) + stance_up_constraint_RR[0:2] = h_R_w @ (stance_up_constraint_RR[0:2] - base[0:2]) stance_up_constraint_RR[0:2] = stance_up_constraint_RR[0:2] + 0.1 stance_up_constraint_RR[2] = stance_up_constraint_RR[2] + 0.01 - + stance_low_constraint_RR = np.array([RR_actual_foot[0], RR_actual_foot[1], RR_actual_foot[2] - 0.002]) - stance_low_constraint_RR[0:2] = h_R_w@(stance_low_constraint_RR[0:2]- base[0:2]) + stance_low_constraint_RR[0:2] = h_R_w @ (stance_low_constraint_RR[0:2] - base[0:2]) stance_low_constraint_RR[0:2] = stance_low_constraint_RR[0:2] - 0.1 stance_low_constraint_RR[2] = stance_low_constraint_RR[2] - 0.01 - # Constraint for the first footholds at the next touchdown. If constraint == True # we have some constraint passed from outside (e.g. vision), otherwise we use the - # nominal foothold enlarged as we do previously - if(constraint is not None): + # nominal foothold enlarged as we do previously + if constraint is not None: # From the VFA first_up_constraint_FL = np.array([constraint[0][0], constraint[1][0], constraint[2][0] + 0.002]) first_up_constraint_FR = np.array([constraint[0][1], constraint[1][1], constraint[2][1] + 0.002]) @@ -719,75 +717,85 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h first_low_constraint_FR = np.array([constraint[9][1], constraint[10][1], constraint[11][1] - 0.002]) first_low_constraint_RL = np.array([constraint[9][2], constraint[10][2], constraint[11][2] - 0.002]) first_low_constraint_RR = np.array([constraint[9][3], constraint[10][3], constraint[11][3] - 0.002]) - + # Rotate the constraint in the horizontal frame - first_up_constraint_FL[0:2] = h_R_w@(first_up_constraint_FL[0:2] - base[0:2]) + first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - base[0:2]) first_up_constraint_FL = first_up_constraint_FL + 0.005 - #first_up_constraint_FL[2] = first_up_constraint_FL[2] + 0.5 - - first_up_constraint_FR[0:2] = h_R_w@(first_up_constraint_FR[0:2] - base[0:2]) + # first_up_constraint_FL[2] = first_up_constraint_FL[2] + 0.5 + + first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - base[0:2]) first_up_constraint_FR = first_up_constraint_FR + 0.005 - #first_up_constraint_FR[2] = first_up_constraint_FR[2] + 0.5 + # first_up_constraint_FR[2] = first_up_constraint_FR[2] + 0.5 - first_up_constraint_RL[0:2] = h_R_w@(first_up_constraint_RL[0:2] - base[0:2]) + first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - base[0:2]) first_up_constraint_RL = first_up_constraint_RL + 0.005 - #first_up_constraint_RL[2] = first_up_constraint_RL[2] + 0.5 + # first_up_constraint_RL[2] = first_up_constraint_RL[2] + 0.5 - first_up_constraint_RR[0:2] = h_R_w@(first_up_constraint_RR[0:2] - base[0:2]) + first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - base[0:2]) first_up_constraint_RR = first_up_constraint_RR + 0.005 - #first_up_constraint_RR[2] = first_up_constraint_RR[2] + 0.5 + # first_up_constraint_RR[2] = first_up_constraint_RR[2] + 0.5 - first_low_constraint_FL[0:2] = h_R_w@(first_low_constraint_FL[0:2] - base[0:2]) + first_low_constraint_FL[0:2] = h_R_w @ (first_low_constraint_FL[0:2] - base[0:2]) first_low_constraint_FL = first_low_constraint_FL - 0.005 - #first_low_constraint_FL[2] = first_low_constraint_FL[2] - 0.5 + # first_low_constraint_FL[2] = first_low_constraint_FL[2] - 0.5 - first_low_constraint_FR[0:2] = h_R_w@(first_low_constraint_FR[0:2] - base[0:2]) + first_low_constraint_FR[0:2] = h_R_w @ (first_low_constraint_FR[0:2] - base[0:2]) first_low_constraint_FR = first_low_constraint_FR - 0.005 - #first_low_constraint_FR[2] = first_low_constraint_FR[2] - 0.5 - - first_low_constraint_RL[0:2] = h_R_w@(first_low_constraint_RL[0:2] - base[0:2]) + # first_low_constraint_FR[2] = first_low_constraint_FR[2] - 0.5 + + first_low_constraint_RL[0:2] = h_R_w @ (first_low_constraint_RL[0:2] - base[0:2]) first_low_constraint_RL = first_low_constraint_RL - 0.005 - #first_low_constraint_RL[2] = first_low_constraint_RL[2] - 0.5 + # first_low_constraint_RL[2] = first_low_constraint_RL[2] - 0.5 - first_low_constraint_RR[0:2] = h_R_w@(first_low_constraint_RR[0:2] - base[0:2]) + first_low_constraint_RR[0:2] = h_R_w @ (first_low_constraint_RR[0:2] - base[0:2]) first_low_constraint_RR = first_low_constraint_RR - 0.005 - #first_low_constraint_RR[2] = first_low_constraint_RR[2] - 0.5 + # first_low_constraint_RR[2] = first_low_constraint_RR[2] - 0.5 else: # Constrain taken from the nominal foothold (NO VISION) # FL first touchdown constraint - first_up_constraint_FL = np.array([FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] + 0.002]) - first_up_constraint_FL[0:2] = h_R_w@(first_up_constraint_FL[0:2] - base[0:2]) + 0.15 - - first_low_constraint_FL = np.array([FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] - 0.002]) - first_low_constraint_FL[0:2] = h_R_w@(first_low_constraint_FL[0:2] - base[0:2]) - 0.15 - + first_up_constraint_FL = np.array( + [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] + 0.002] + ) + first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - base[0:2]) + 0.15 - # FR first touchdown constraint - first_up_constraint_FR = np.array([FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] + 0.002]) - first_up_constraint_FR[0:2] = h_R_w@(first_up_constraint_FR[0:2] - base[0:2]) + 0.15 + first_low_constraint_FL = np.array( + [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] - 0.002] + ) + first_low_constraint_FL[0:2] = h_R_w @ (first_low_constraint_FL[0:2] - base[0:2]) - 0.15 - first_low_constraint_FR = np.array([FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] - 0.002]) - first_low_constraint_FR[0:2] = h_R_w@(first_low_constraint_FR[0:2] - base[0:2]) - 0.15 + # FR first touchdown constraint + first_up_constraint_FR = np.array( + [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] + 0.002] + ) + first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - base[0:2]) + 0.15 + first_low_constraint_FR = np.array( + [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] - 0.002] + ) + first_low_constraint_FR[0:2] = h_R_w @ (first_low_constraint_FR[0:2] - base[0:2]) - 0.15 # RL first touchdown constraint - first_up_constraint_RL = np.array([RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] + 0.002]) - first_up_constraint_RL[0:2] = h_R_w@(first_up_constraint_RL[0:2] - base[0:2]) + 0.15 - - first_low_constraint_RL = np.array([RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] - 0.002]) - first_low_constraint_RL[0:2] = h_R_w@(first_low_constraint_RL[0:2] - base[0:2]) - 0.15 + first_up_constraint_RL = np.array( + [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] + 0.002] + ) + first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - base[0:2]) + 0.15 + first_low_constraint_RL = np.array( + [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] - 0.002] + ) + first_low_constraint_RL[0:2] = h_R_w @ (first_low_constraint_RL[0:2] - base[0:2]) - 0.15 # RR first touchdown constraint - first_up_constraint_RR = np.array([RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] + 0.002]) - first_up_constraint_RR[0:2] = h_R_w@(first_up_constraint_RR[0:2] - base[0:2]) + 0.15 + first_up_constraint_RR = np.array( + [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] + 0.002] + ) + first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - base[0:2]) + 0.15 - first_low_constraint_RR = np.array([RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] - 0.002]) - first_low_constraint_RR[0:2] = h_R_w@(first_low_constraint_RR[0:2] - base[0:2]) - 0.15 - - - + first_low_constraint_RR = np.array( + [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] - 0.002] + ) + first_low_constraint_RR[0:2] = h_R_w @ (first_low_constraint_RR[0:2] - base[0:2]) - 0.15 # we stack all the constraint we have for now up_constraint_FL = np.vstack((stance_up_constraint_FL, first_up_constraint_FL)) @@ -800,64 +808,69 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h low_constraint_RL = np.vstack((stance_low_constraint_RL, first_low_constraint_RL)) low_constraint_RR = np.vstack((stance_low_constraint_RR, first_low_constraint_RR)) - - # If there are more than 1 nominal foothold per leg, we create additional constraints # We do not expect more than two reference footholds... # FL second touchdown constraint - if(FL_reference_foot.shape[0] == 2): - second_up_constraint_FL = np.array([FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] + 0.002]) - second_up_constraint_FL[0:2] = h_R_w@(second_up_constraint_FL[0:2] - base[0:2]) + 0.15 - - second_low_constraint_FL = np.array([FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] - 0.002]) - second_low_constraint_FL[0:2] = h_R_w@(second_low_constraint_FL[0:2] - base[0:2]) - 0.15 + if FL_reference_foot.shape[0] == 2: + second_up_constraint_FL = np.array( + [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] + 0.002] + ) + second_up_constraint_FL[0:2] = h_R_w @ (second_up_constraint_FL[0:2] - base[0:2]) + 0.15 + second_low_constraint_FL = np.array( + [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] - 0.002] + ) + second_low_constraint_FL[0:2] = h_R_w @ (second_low_constraint_FL[0:2] - base[0:2]) - 0.15 up_constraint_FL = np.vstack((up_constraint_FL, second_up_constraint_FL)) low_constraint_FL = np.vstack((low_constraint_FL, second_low_constraint_FL)) # FR second touchdown constraint - if(FR_reference_foot.shape[0] == 2): - second_up_constraint_FR = np.array([FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] + 0.002]) - second_up_constraint_FR[0:2] = h_R_w@(second_up_constraint_FR[0:2] - base[0:2]) + 0.15 - - - second_low_constraint_FR = np.array([FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] - 0.002]) - second_low_constraint_FR[0:2] = h_R_w@(second_low_constraint_FR[0:2]- base[0:2]) - 0.15 + if FR_reference_foot.shape[0] == 2: + second_up_constraint_FR = np.array( + [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] + 0.002] + ) + second_up_constraint_FR[0:2] = h_R_w @ (second_up_constraint_FR[0:2] - base[0:2]) + 0.15 + second_low_constraint_FR = np.array( + [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] - 0.002] + ) + second_low_constraint_FR[0:2] = h_R_w @ (second_low_constraint_FR[0:2] - base[0:2]) - 0.15 up_constraint_FR = np.vstack((up_constraint_FR, second_up_constraint_FR)) low_constraint_FR = np.vstack((low_constraint_FR, second_low_constraint_FR)) - - # RL second touchdown constraint - if(RL_reference_foot.shape[0] == 2): - second_up_constraint_RL = np.array([RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] + 0.002]) - second_up_constraint_RL[0:2] = h_R_w@(second_up_constraint_RL[0:2]- base[0:2]) + 0.15 - - second_low_constraint_RL = np.array([RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] - 0.002]) - second_low_constraint_RL[0:2] = h_R_w@(second_low_constraint_RL[0:2] - base[0:2]) - 0.15 + # RL second touchdown constraint + if RL_reference_foot.shape[0] == 2: + second_up_constraint_RL = np.array( + [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] + 0.002] + ) + second_up_constraint_RL[0:2] = h_R_w @ (second_up_constraint_RL[0:2] - base[0:2]) + 0.15 + second_low_constraint_RL = np.array( + [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] - 0.002] + ) + second_low_constraint_RL[0:2] = h_R_w @ (second_low_constraint_RL[0:2] - base[0:2]) - 0.15 up_constraint_RL = np.vstack((up_constraint_RL, second_up_constraint_RL)) low_constraint_RL = np.vstack((low_constraint_RL, second_low_constraint_RL)) - - # RR second touchdown constraint - if(RR_reference_foot.shape[0] == 2): - second_up_constraint_RR = np.array([RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] + 0.002]) - second_up_constraint_RR[0:2] = h_R_w@(second_up_constraint_RR[0:2] - base[0:2]) + 0.15 - - second_low_constraint_RR = np.array([RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] - 0.002]) - second_low_constraint_RR[0:2] = h_R_w@(second_low_constraint_RR[0:2] - base[0:2]) - 0.15 + # RR second touchdown constraint + if RR_reference_foot.shape[0] == 2: + second_up_constraint_RR = np.array( + [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] + 0.002] + ) + second_up_constraint_RR[0:2] = h_R_w @ (second_up_constraint_RR[0:2] - base[0:2]) + 0.15 + second_low_constraint_RR = np.array( + [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] - 0.002] + ) + second_low_constraint_RR[0:2] = h_R_w @ (second_low_constraint_RR[0:2] - base[0:2]) - 0.15 up_constraint_RR = np.vstack((up_constraint_RR, second_up_constraint_RR)) low_constraint_RR = np.vstack((low_constraint_RR, second_low_constraint_RR)) - - # We pass all the constraints (foothold and friction conte) to acados. # Then one regarding friction are precomputed ub_friction = self.constr_uh_friction @@ -865,48 +878,42 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h # If the foothold is in swing, the idx of the constraint start from 1 idx_constraint = np.array([0, 0, 0, 0]) - if(FL_contact_sequence[0] == 0): + if FL_contact_sequence[0] == 0: idx_constraint[0] = 1 - if(FR_contact_sequence[0] == 0): + if FR_contact_sequence[0] == 0: idx_constraint[1] = 1 - if(RL_contact_sequence[0] == 0): + if RL_contact_sequence[0] == 0: idx_constraint[2] = 1 - if(RR_contact_sequence[0] == 0): + if RR_contact_sequence[0] == 0: idx_constraint[3] = 1 - - for j in range(0, self.horizon): - # take the constraint for the current timestep - ub_foot_FL = up_constraint_FL[idx_constraint[0]] - lb_foot_FL = low_constraint_FL[idx_constraint[0]] + for j in range(0, self.horizon): + # take the constraint for the current timestep + ub_foot_FL = up_constraint_FL[idx_constraint[0]] + lb_foot_FL = low_constraint_FL[idx_constraint[0]] - ub_foot_FR = up_constraint_FR[idx_constraint[1]] + ub_foot_FR = up_constraint_FR[idx_constraint[1]] lb_foot_FR = low_constraint_FR[idx_constraint[1]] - ub_foot_RL = up_constraint_RL[idx_constraint[2]] + ub_foot_RL = up_constraint_RL[idx_constraint[2]] lb_foot_RL = low_constraint_RL[idx_constraint[2]] - ub_foot_RR = up_constraint_RR[idx_constraint[3]] + ub_foot_RR = up_constraint_RR[idx_constraint[3]] lb_foot_RR = low_constraint_RR[idx_constraint[3]] - - - # Concatenate the friction and foothold constraint - ub_foot = copy.deepcopy(np.concatenate((ub_foot_FL, ub_foot_FR, - ub_foot_RL, ub_foot_RR))) - lb_foot = copy.deepcopy(np.concatenate((lb_foot_FL, lb_foot_FR, - lb_foot_RL, lb_foot_RR))) - if(self.use_foothold_constraints): + # Concatenate the friction and foothold constraint + ub_foot = copy.deepcopy(np.concatenate((ub_foot_FL, ub_foot_FR, ub_foot_RL, ub_foot_RR))) + lb_foot = copy.deepcopy(np.concatenate((lb_foot_FL, lb_foot_FR, lb_foot_RL, lb_foot_RR))) + if self.use_foothold_constraints: ub_total = np.concatenate((ub_friction, ub_foot)) lb_total = np.concatenate((lb_friction, lb_foot)) else: ub_total = ub_friction lb_total = lb_friction - - #Constraints for the support polygon depending on the leg in stance + # Constraints for the support polygon depending on the leg in stance # all disabled at the beginning!! - if(self.use_stability_constraints): + if self.use_stability_constraints: ub_support_FL_FR = ACADOS_INFTY lb_support_FL_FR = -ACADOS_INFTY @@ -918,20 +925,21 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h ub_support_RL_FL = ACADOS_INFTY lb_support_RL_FL = -ACADOS_INFTY - + ub_support_FL_RR = ACADOS_INFTY lb_support_FL_RR = -ACADOS_INFTY ub_support_FR_RL = ACADOS_INFTY - lb_support_FR_RL = -ACADOS_INFTY - - + lb_support_FR_RL = -ACADOS_INFTY + # We have 4 cases for the stability constraint: trot, pace, crawl, full stance - if(FL_contact_sequence[j] == 1 and - FR_contact_sequence[j] == 1 and - RL_contact_sequence[j] == 1 and - RR_contact_sequence[j] == 1): - #FULL STANCE TODO + if ( + FL_contact_sequence[j] == 1 + and FR_contact_sequence[j] == 1 + and RL_contact_sequence[j] == 1 + and RR_contact_sequence[j] == 1 + ): + # FULL STANCE TODO ub_support_FL_FR = ACADOS_INFTY lb_support_FL_FR = -ACADOS_INFTY @@ -943,148 +951,161 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h ub_support_RL_FL = ACADOS_INFTY lb_support_RL_FL = -ACADOS_INFTY - + ub_support_FL_RR = ACADOS_INFTY lb_support_FL_RR = -ACADOS_INFTY ub_support_FR_RL = ACADOS_INFTY lb_support_FR_RL = -ACADOS_INFTY - elif(np.array_equal(FL_contact_sequence, RR_contact_sequence) - and np.array_equal(FR_contact_sequence, RL_contact_sequence)): - #TROT - stability_margin = config.mpc_params['trot_stability_margin'] - if(FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0): + elif np.array_equal(FL_contact_sequence, RR_contact_sequence) and np.array_equal( + FR_contact_sequence, RL_contact_sequence + ): + # TROT + stability_margin = config.mpc_params["trot_stability_margin"] + if FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0: ub_support_FL_RR = 0 + stability_margin lb_support_FL_RR = 0 - stability_margin - - if(FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0): + + if FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0: ub_support_FR_RL = 0 + stability_margin lb_support_FR_RL = 0 - stability_margin - elif(np.array_equal(FL_contact_sequence, RL_contact_sequence) - and np.array_equal(FR_contact_sequence, RR_contact_sequence)): - #PACE - stability_margin = config.mpc_params['pace_stability_margin'] - if(FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0): + elif np.array_equal(FL_contact_sequence, RL_contact_sequence) and np.array_equal( + FR_contact_sequence, RR_contact_sequence + ): + # PACE + stability_margin = config.mpc_params["pace_stability_margin"] + if FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0: ub_support_RL_FL = 0 + stability_margin lb_support_RL_FL = 0 - stability_margin - - if(FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0): + if FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0: ub_support_FR_RR = 0 + stability_margin lb_support_FR_RR = 0 - stability_margin - + else: - #CRAWL BACKDIAGONALCRAWL ONLY - stability_margin = config.mpc_params['crawl_stability_margin'] + # CRAWL BACKDIAGONALCRAWL ONLY + stability_margin = config.mpc_params["crawl_stability_margin"] - if(FL_contact_sequence[j] == 1): - if(FR_contact_sequence[j] == 1): - ub_support_FL_FR = -0.0 - stability_margin + if FL_contact_sequence[j] == 1: + if FR_contact_sequence[j] == 1: + ub_support_FL_FR = -0.0 - stability_margin lb_support_FL_FR = -ACADOS_INFTY else: ub_support_FL_RR = ACADOS_INFTY - lb_support_FL_RR = 0.0 + stability_margin - + lb_support_FL_RR = 0.0 + stability_margin - if(FR_contact_sequence[j] == 1): - if(RR_contact_sequence[j] == 1): + if FR_contact_sequence[j] == 1: + if RR_contact_sequence[j] == 1: ub_support_FR_RR = ACADOS_INFTY lb_support_FR_RR = 0.0 + stability_margin else: ub_support_FR_RL = ACADOS_INFTY lb_support_FR_RL = 0.0 + stability_margin - - if(RR_contact_sequence[j] == 1): - if(RL_contact_sequence[j] == 1): + if RR_contact_sequence[j] == 1: + if RL_contact_sequence[j] == 1: ub_support_RR_RL = ACADOS_INFTY - lb_support_RR_RL = 0.0 + stability_margin + lb_support_RR_RL = 0.0 + stability_margin else: - ub_support_FL_RR = -0.0 - stability_margin + ub_support_FL_RR = -0.0 - stability_margin lb_support_FL_RR = -ACADOS_INFTY - - if(RL_contact_sequence[j] == 1): - if(FL_contact_sequence[j] == 1): - ub_support_RL_FL = -0.0 - stability_margin + if RL_contact_sequence[j] == 1: + if FL_contact_sequence[j] == 1: + ub_support_RL_FL = -0.0 - stability_margin lb_support_RL_FL = -ACADOS_INFTY else: ub_support_FR_RL = -0.0 - stability_margin lb_support_FR_RL = -ACADOS_INFTY + ub_support = np.array( + [ + ub_support_FL_FR, + ub_support_FR_RR, + ub_support_RR_RL, + ub_support_RL_FL, + ub_support_FL_RR, + ub_support_FR_RL, + ] + ) + lb_support = np.array( + [ + lb_support_FL_FR, + lb_support_FR_RR, + lb_support_RR_RL, + lb_support_RL_FL, + lb_support_FL_RR, + lb_support_FR_RL, + ] + ) - - - ub_support = np.array([ub_support_FL_FR, ub_support_FR_RR, ub_support_RR_RL, ub_support_RL_FL, - ub_support_FL_RR, ub_support_FR_RL]) - lb_support = np.array([lb_support_FL_FR, lb_support_FR_RR, lb_support_RR_RL, lb_support_RL_FL, - lb_support_FL_RR, lb_support_FR_RL]) - ub_total = np.concatenate((ub_total, ub_support)) lb_total = np.concatenate((lb_total, lb_support)) - # No friction constraint at the end! we don't have u_N - if(j == self.horizon): - if(self.use_foothold_constraints): - if(self.use_stability_constraints): + if j == self.horizon: + if self.use_foothold_constraints: + if self.use_stability_constraints: ub_total = np.concatenate((ub_foot, ub_support)) lb_total = np.concatenate((lb_foot, lb_support)) else: ub_total = ub_foot lb_total = lb_foot else: - if(self.use_stability_constraints): + if self.use_stability_constraints: ub_total = ub_support lb_total = lb_support else: continue - + # Only friction costraints at the beginning - if(j == 0): + if j == 0: self.acados_ocp_solver.constraints_set(j, "uh", ub_friction) self.acados_ocp_solver.constraints_set(j, "lh", lb_friction) - if(j > 0): + if j > 0: self.acados_ocp_solver.constraints_set(j, "uh", ub_total) self.acados_ocp_solver.constraints_set(j, "lh", lb_total) - - #save the constraint for logging + # save the constraint for logging self.upper_bound[j] = ub_total.tolist() self.lower_bound[j] = lb_total.tolist() - - # ugly procedure to update the idx of the constraint - if(j>=1): - if(FL_contact_sequence[j] == 0 and FL_contact_sequence[j-1] == 1): - if(idx_constraint[0] < up_constraint_FL.shape[0] - 1): + if j >= 1: + if FL_contact_sequence[j] == 0 and FL_contact_sequence[j - 1] == 1: + if idx_constraint[0] < up_constraint_FL.shape[0] - 1: idx_constraint[0] += 1 - - if(FR_contact_sequence[j] == 0 and FR_contact_sequence[j-1] == 1): - if(idx_constraint[1] < up_constraint_FR.shape[0] - 1): + + if FR_contact_sequence[j] == 0 and FR_contact_sequence[j - 1] == 1: + if idx_constraint[1] < up_constraint_FR.shape[0] - 1: idx_constraint[1] += 1 - - if(RL_contact_sequence[j] == 0 and RL_contact_sequence[j-1] == 1): - if(idx_constraint[2] < up_constraint_RL.shape[0] - 1): + + if RL_contact_sequence[j] == 0 and RL_contact_sequence[j - 1] == 1: + if idx_constraint[2] < up_constraint_RL.shape[0] - 1: idx_constraint[2] += 1 - - if(RR_contact_sequence[j] == 0 and RR_contact_sequence[j-1] == 1): - if(idx_constraint[3] < up_constraint_RR.shape[0] - 1): + + if RR_contact_sequence[j] == 0 and RR_contact_sequence[j - 1] == 1: + if idx_constraint[3] < up_constraint_RR.shape[0] - 1: idx_constraint[3] += 1 - + except: print("###WARNING: error in setting the constraints") return - - - def set_warm_start(self, state_acados, reference, FL_contact_sequence, FR_contact_sequence, RL_contact_sequence, RR_contact_sequence): + def set_warm_start( + self, + state_acados, + reference, + FL_contact_sequence, + FR_contact_sequence, + RL_contact_sequence, + RR_contact_sequence, + ): """ - Sets the warm start for the optimization problem. This could be useful in the case + Sets the warm start for the optimization problem. This could be useful in the case some old guess is outside some new constraints.. Maybe we could have some infeasibility problem. Still to investigate.. @@ -1097,69 +1118,62 @@ def set_warm_start(self, state_acados, reference, FL_contact_sequence, FR_contac RR_contact_sequence (np.ndarray): Contact sequence for the rear right leg. """ idx_ref_foot_to_assign = np.array([0, 0, 0, 0]) - + for j in range(self.horizon): # We only want to modify the warm start of the footholds... - # hence i take the current warm start from acados and + # hence i take the current warm start from acados and # modify onlya subset of it warm_start = copy.deepcopy(self.acados_ocp_solver.get(j, "x")) - # Update the idx_ref_foot_to_assign. Every time there is a change in the contact phase # between 1 and 0, it means that the leg go into swing and a new reference is needed!!! - if(j > 1 and j idx_ref_foot_to_assign[0]+1): + if j > 1 and j < self.horizon - 1: + if FL_contact_sequence[j] == 0 and FL_contact_sequence[j - 1] == 1: + if reference["ref_foot_FL"].shape[0] > idx_ref_foot_to_assign[0] + 1: idx_ref_foot_to_assign[0] += 1 - if(FR_contact_sequence[j] == 0 and FR_contact_sequence[j-1] == 1): - if(reference['ref_foot_FR'].shape[0] > idx_ref_foot_to_assign[1]+1): + if FR_contact_sequence[j] == 0 and FR_contact_sequence[j - 1] == 1: + if reference["ref_foot_FR"].shape[0] > idx_ref_foot_to_assign[1] + 1: idx_ref_foot_to_assign[1] += 1 - if(RL_contact_sequence[j] == 0 and RL_contact_sequence[j-1] == 1): - if(reference['ref_foot_RL'].shape[0] > idx_ref_foot_to_assign[2]+1): + if RL_contact_sequence[j] == 0 and RL_contact_sequence[j - 1] == 1: + if reference["ref_foot_RL"].shape[0] > idx_ref_foot_to_assign[2] + 1: idx_ref_foot_to_assign[2] += 1 - if(RR_contact_sequence[j] == 0 and RR_contact_sequence[j-1] == 1): - if(reference['ref_foot_RR'].shape[0] > idx_ref_foot_to_assign[3]+1): - idx_ref_foot_to_assign[3] += 1 - + if RR_contact_sequence[j] == 0 and RR_contact_sequence[j - 1] == 1: + if reference["ref_foot_RR"].shape[0] > idx_ref_foot_to_assign[3] + 1: + idx_ref_foot_to_assign[3] += 1 + warm_start[8] = state_acados[8] - - if(idx_ref_foot_to_assign[0] == 0): + + if idx_ref_foot_to_assign[0] == 0: warm_start[12:15] = state_acados[12:15].reshape((3,)) else: warm_start[12:15] = reference["ref_foot_FL"][0] - if(idx_ref_foot_to_assign[1] == 0): + if idx_ref_foot_to_assign[1] == 0: warm_start[15:18] = state_acados[15:18].reshape((3,)) else: warm_start[15:18] = reference["ref_foot_FR"][0] - if(idx_ref_foot_to_assign[2] == 0): + if idx_ref_foot_to_assign[2] == 0: warm_start[18:21] = state_acados[18:21].reshape((3,)) else: warm_start[18:21] = reference["ref_foot_RL"][0] - if(idx_ref_foot_to_assign[3] == 0): + if idx_ref_foot_to_assign[3] == 0: warm_start[21:24] = state_acados[21:24].reshape((3,)) else: warm_start[21:24] = reference["ref_foot_RR"][0] - self.acados_ocp_solver.set(j, "x", warm_start) - - # Method to perform the centering of the states and the reference around (0, 0, 0) - def perform_scaling(self, state, reference, constraint = None): - - + def perform_scaling(self, state, reference, constraint=None): self.initial_base_position = copy.deepcopy(state["position"]) reference = copy.deepcopy(reference) state = copy.deepcopy(state) - + reference["ref_position"] = reference["ref_position"] - state["position"] reference["ref_foot_FL"] = reference["ref_foot_FL"] - state["position"] reference["ref_foot_FR"] = reference["ref_foot_FR"] - state["position"] reference["ref_foot_RL"] = reference["ref_foot_RL"] - state["position"] reference["ref_foot_RR"] = reference["ref_foot_RR"] - state["position"] - state["foot_FL"] = state["foot_FL"] - state["position"] state["foot_FR"] = state["foot_FR"] - state["position"] state["foot_RL"] = state["foot_RL"] - state["position"] @@ -1168,15 +1182,19 @@ def perform_scaling(self, state, reference, constraint = None): return state, reference, constraint - # Main loop for computing the control - def compute_control(self, state, reference, contact_sequence, constraint = None, external_wrenches = np.zeros((6,)), - end_effector_position = np.zeros((3,)), - end_effector_jacobian = np.ones((9,)), - end_effector_gain = np.zeros((3,))): - - - # Take the array of the contact sequence and split it in 4 arrays, + def compute_control( + self, + state, + reference, + contact_sequence, + constraint=None, + external_wrenches=np.zeros((6,)), + end_effector_position=np.zeros((3,)), + end_effector_jacobian=np.ones((9,)), + end_effector_gain=np.zeros((3,)), + ): + # Take the array of the contact sequence and split it in 4 arrays, # one for each leg FL_contact_sequence = contact_sequence[0] FR_contact_sequence = contact_sequence[1] @@ -1188,254 +1206,265 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, RL_previous_contact_sequence = self.previous_contact_sequence[2] RR_previous_contact_sequence = self.previous_contact_sequence[3] - # Perform the scaling of the states and the reference - state, \ - reference, \ - constraint = self.perform_scaling(state, reference, constraint) + state, reference, constraint = self.perform_scaling(state, reference, constraint) - - desired_force = np.array([0, 0]) #x, y, to be put in the reference + desired_force = np.array([0, 0]) # x, y, to be put in the reference # Fill reference (self.states_dim+self.inputs_dim) idx_ref_foot_to_assign = np.array([0, 0, 0, 0]) for j in range(self.horizon): - yref = np.zeros(shape=(self.states_dim + self.inputs_dim,)) - + # If and Else to check if the reference is a single value or a sequence - if(reference['ref_position'].shape[0] == 3): - yref[0:3] = reference['ref_position'] + if reference["ref_position"].shape[0] == 3: + yref[0:3] = reference["ref_position"] else: - yref[0:3] = reference['ref_position'][j] - - if(reference['ref_linear_velocity'].shape[0] == 3): - yref[3:6] = reference['ref_linear_velocity'] + yref[0:3] = reference["ref_position"][j] + + if reference["ref_linear_velocity"].shape[0] == 3: + yref[3:6] = reference["ref_linear_velocity"] else: - yref[3:6] = reference['ref_linear_velocity'][j] + yref[3:6] = reference["ref_linear_velocity"][j] - if(reference['ref_orientation'].shape[0] == 3): - yref[6:9] = reference['ref_orientation'] + if reference["ref_orientation"].shape[0] == 3: + yref[6:9] = reference["ref_orientation"] else: - yref[6:9] = reference['ref_orientation'][j] + yref[6:9] = reference["ref_orientation"][j] - if(reference['ref_angular_velocity'].shape[0] == 3): - yref[9:12] = reference['ref_angular_velocity'] + if reference["ref_angular_velocity"].shape[0] == 3: + yref[9:12] = reference["ref_angular_velocity"] else: - yref[9:12] = reference['ref_angular_velocity'][j] + yref[9:12] = reference["ref_angular_velocity"][j] - yref[12:15] = reference['ref_foot_FL'][idx_ref_foot_to_assign[0]] - yref[15:18] = reference['ref_foot_FR'][idx_ref_foot_to_assign[1]] - yref[18:21] = reference['ref_foot_RL'][idx_ref_foot_to_assign[2]] - yref[21:24] = reference['ref_foot_RR'][idx_ref_foot_to_assign[3]] + yref[12:15] = reference["ref_foot_FL"][idx_ref_foot_to_assign[0]] + yref[15:18] = reference["ref_foot_FR"][idx_ref_foot_to_assign[1]] + yref[18:21] = reference["ref_foot_RL"][idx_ref_foot_to_assign[2]] + yref[21:24] = reference["ref_foot_RR"][idx_ref_foot_to_assign[3]] yref[30:32] = np.array([desired_force[0], desired_force[1]]) - - - - # Update the idx_ref_foot_to_assign. Every time there is a change in the contact phase # between 1 and 0, it means that the leg go into swing and a new reference is needed!!! - if(j > 1 and j idx_ref_foot_to_assign[0]+1): + if j > 1 and j < self.horizon - 1: + if FL_contact_sequence[j + 1] == 0 and FL_contact_sequence[j] == 1: + if reference["ref_foot_FL"].shape[0] > idx_ref_foot_to_assign[0] + 1: idx_ref_foot_to_assign[0] += 1 - if(FR_contact_sequence[j+1] == 0 and FR_contact_sequence[j] == 1): - if(reference['ref_foot_FR'].shape[0] > idx_ref_foot_to_assign[1]+1): + if FR_contact_sequence[j + 1] == 0 and FR_contact_sequence[j] == 1: + if reference["ref_foot_FR"].shape[0] > idx_ref_foot_to_assign[1] + 1: idx_ref_foot_to_assign[1] += 1 - if(RL_contact_sequence[j+1] == 0 and RL_contact_sequence[j] == 1): - if(reference['ref_foot_RL'].shape[0] > idx_ref_foot_to_assign[2]+1): + if RL_contact_sequence[j + 1] == 0 and RL_contact_sequence[j] == 1: + if reference["ref_foot_RL"].shape[0] > idx_ref_foot_to_assign[2] + 1: idx_ref_foot_to_assign[2] += 1 - if(RR_contact_sequence[j+1] == 0 and RR_contact_sequence[j] == 1): - if(reference['ref_foot_RR'].shape[0] > idx_ref_foot_to_assign[3]+1): - idx_ref_foot_to_assign[3] += 1 - + if RR_contact_sequence[j + 1] == 0 and RR_contact_sequence[j] == 1: + if reference["ref_foot_RR"].shape[0] > idx_ref_foot_to_assign[3] + 1: + idx_ref_foot_to_assign[3] += 1 # Calculate the reference force z for the leg in stance # It's simply mass*acc/number_of_legs_in_stance!! # Force x and y are always 0 - number_of_legs_in_stance = np.array([FL_contact_sequence[j], FR_contact_sequence[j], - RL_contact_sequence [j], RR_contact_sequence[j]]).sum() - if(number_of_legs_in_stance == 0): + number_of_legs_in_stance = np.array( + [FL_contact_sequence[j], FR_contact_sequence[j], RL_contact_sequence[j], RR_contact_sequence[j]] + ).sum() + if number_of_legs_in_stance == 0: reference_force_stance_legs = 0 else: reference_force_stance_legs = (self.centroidal_model.mass * 9.81) / number_of_legs_in_stance - + reference_force_fl_z = reference_force_stance_legs * FL_contact_sequence[j] reference_force_fr_z = reference_force_stance_legs * FR_contact_sequence[j] reference_force_rl_z = reference_force_stance_legs * RL_contact_sequence[j] reference_force_rr_z = reference_force_stance_legs * RR_contact_sequence[j] - + yref[50] = reference_force_fl_z yref[53] = reference_force_fr_z yref[56] = reference_force_rl_z yref[59] = reference_force_rr_z - - # Setting the reference to acados self.acados_ocp_solver.set(j, "yref", yref) - # Fill last step horizon reference (self.states_dim - no control action!!) - yref_N = np.zeros(shape=(self.states_dim ,)) - if(reference['ref_position'].shape[0] == 3): - yref_N[0:3] = reference['ref_position'] + yref_N = np.zeros(shape=(self.states_dim,)) + if reference["ref_position"].shape[0] == 3: + yref_N[0:3] = reference["ref_position"] else: - yref_N[0:3] = reference['ref_position'][-1] - - if(reference['ref_linear_velocity'].shape[0] == 3): - yref_N[3:6] = reference['ref_linear_velocity'] + yref_N[0:3] = reference["ref_position"][-1] + + if reference["ref_linear_velocity"].shape[0] == 3: + yref_N[3:6] = reference["ref_linear_velocity"] else: - yref_N[3:6] = reference['ref_linear_velocity'][-1] - - if(reference['ref_orientation'].shape[0] == 3): - yref_N[6:9] = reference['ref_orientation'] + yref_N[3:6] = reference["ref_linear_velocity"][-1] + + if reference["ref_orientation"].shape[0] == 3: + yref_N[6:9] = reference["ref_orientation"] else: - yref_N[6:9] = reference['ref_orientation'][-1] - - if(reference['ref_angular_velocity'].shape[0] == 3): - yref_N[9:12] = reference['ref_angular_velocity'] + yref_N[6:9] = reference["ref_orientation"][-1] + + if reference["ref_angular_velocity"].shape[0] == 3: + yref_N[9:12] = reference["ref_angular_velocity"] else: - yref_N[9:12] = reference['ref_angular_velocity'][-1] - - yref_N[12:15] = reference['ref_foot_FL'][idx_ref_foot_to_assign[0]] - yref_N[15:18] = reference['ref_foot_FR'][idx_ref_foot_to_assign[1]] - yref_N[18:21] = reference['ref_foot_RL'][idx_ref_foot_to_assign[2]] - yref_N[21:24] = reference['ref_foot_RR'][idx_ref_foot_to_assign[3]] + yref_N[9:12] = reference["ref_angular_velocity"][-1] + + yref_N[12:15] = reference["ref_foot_FL"][idx_ref_foot_to_assign[0]] + yref_N[15:18] = reference["ref_foot_FR"][idx_ref_foot_to_assign[1]] + yref_N[18:21] = reference["ref_foot_RL"][idx_ref_foot_to_assign[2]] + yref_N[21:24] = reference["ref_foot_RR"][idx_ref_foot_to_assign[3]] yref_N[30:32] = np.array([desired_force[0], desired_force[1]]) # Setting the reference to acados self.acados_ocp_solver.set(self.horizon, "yref", yref_N) - - - - # Fill stance param, friction and stance proximity + # Fill stance param, friction and stance proximity # (stance proximity will disable foothold optimization near a stance!!) - mu = config.mpc_params['mu'] + mu = config.mpc_params["mu"] yaw = state["orientation"][2] - + # Stance Proximity ugly routine. Basically we disable foothold optimization # in the proximity of a stance phase (the real foot cannot travel too fast in # a small time!!) - stance_proximity_FL = np.zeros((self.horizon, )) - stance_proximity_FR = np.zeros((self.horizon, )) - stance_proximity_RL = np.zeros((self.horizon, )) - stance_proximity_RR = np.zeros((self.horizon, )) + stance_proximity_FL = np.zeros((self.horizon,)) + stance_proximity_FR = np.zeros((self.horizon,)) + stance_proximity_RL = np.zeros((self.horizon,)) + stance_proximity_RR = np.zeros((self.horizon,)) for j in range(self.horizon): - if(FL_contact_sequence[j] == 0): - if(j+1) < self.horizon: - if(FL_contact_sequence[j+1] == 1): - stance_proximity_FL[j] = 1*0 - if(j+2) < self.horizon: - if(FL_contact_sequence[j+1] == 0 and FL_contact_sequence[j+2] == 1): - stance_proximity_FL[j] = 1*0 - - if(FR_contact_sequence[j] == 0): - if(j+1) < self.horizon: - if(FR_contact_sequence[j+1] == 1): - stance_proximity_FR[j] = 1*0 - if(j+2) < self.horizon: - if(FR_contact_sequence[j+1] == 0 and FR_contact_sequence[j+2] == 1): - stance_proximity_FR[j] = 1*0 - - if(RL_contact_sequence[j] == 0): - if(j+1) < self.horizon: - if(RL_contact_sequence[j+1] == 1): - stance_proximity_RL[j] = 1*0 - if(j+2) < self.horizon: - if(RL_contact_sequence[j+1] == 0 and RL_contact_sequence[j+2] == 1): - stance_proximity_RL[j] = 1*0 - - if(RR_contact_sequence[j] == 0): - if(j+1) < self.horizon: - if(RR_contact_sequence[j+1] == 1): - stance_proximity_RR[j] = 1*0 - if(j+2) < self.horizon: - if(RR_contact_sequence[j+1] == 0 and RR_contact_sequence[j+2] == 1): - stance_proximity_RR[j] = 1*0 - - + if FL_contact_sequence[j] == 0: + if (j + 1) < self.horizon: + if FL_contact_sequence[j + 1] == 1: + stance_proximity_FL[j] = 1 * 0 + if (j + 2) < self.horizon: + if FL_contact_sequence[j + 1] == 0 and FL_contact_sequence[j + 2] == 1: + stance_proximity_FL[j] = 1 * 0 + + if FR_contact_sequence[j] == 0: + if (j + 1) < self.horizon: + if FR_contact_sequence[j + 1] == 1: + stance_proximity_FR[j] = 1 * 0 + if (j + 2) < self.horizon: + if FR_contact_sequence[j + 1] == 0 and FR_contact_sequence[j + 2] == 1: + stance_proximity_FR[j] = 1 * 0 + + if RL_contact_sequence[j] == 0: + if (j + 1) < self.horizon: + if RL_contact_sequence[j + 1] == 1: + stance_proximity_RL[j] = 1 * 0 + if (j + 2) < self.horizon: + if RL_contact_sequence[j + 1] == 0 and RL_contact_sequence[j + 2] == 1: + stance_proximity_RL[j] = 1 * 0 + + if RR_contact_sequence[j] == 0: + if (j + 1) < self.horizon: + if RR_contact_sequence[j + 1] == 1: + stance_proximity_RR[j] = 1 * 0 + if (j + 2) < self.horizon: + if RR_contact_sequence[j + 1] == 0 and RR_contact_sequence[j + 2] == 1: + stance_proximity_RR[j] = 1 * 0 # Set the parameters to acados for j in range(self.horizon): # If we have estimated an external wrench, we can compensate it for all steps # or less (maybe the disturbance is not costant along the horizon!) - if(config.mpc_params['external_wrenches_compensation'] and - j < config.mpc_params['external_wrenches_compensation_num_step']): + if ( + config.mpc_params["external_wrenches_compensation"] + and j < config.mpc_params["external_wrenches_compensation_num_step"] + ): external_wrenches_estimated_param = copy.deepcopy(external_wrenches) - external_wrenches_estimated_param = external_wrenches_estimated_param.reshape((6, )) + external_wrenches_estimated_param = external_wrenches_estimated_param.reshape((6,)) else: external_wrenches_estimated_param = np.zeros((6,)) - + # To be filled with an estimation of the ee position in the WORLD end_effector_position = end_effector_position.reshape((3,)) end_effector_jacobian = end_effector_jacobian.reshape((3, 3)) - + end_effector_jacobian_inv = np.linalg.pinv(end_effector_jacobian) end_effector_jacobian_inv = end_effector_jacobian.reshape((9,)) - + end_effector_jacobian_transpose_inv = np.linalg.pinv(end_effector_jacobian.T) end_effector_jacobian_transpose_inv = end_effector_jacobian_transpose_inv.reshape((9,)) - - end_effector_gain = end_effector_gain.reshape((3,)) + end_effector_gain = end_effector_gain.reshape((3,)) - param = np.array([FL_contact_sequence[j], FR_contact_sequence[j], - RL_contact_sequence[j], RR_contact_sequence[j], mu, - stance_proximity_FL[j], - stance_proximity_FR[j], - stance_proximity_RL[j], - stance_proximity_RR[j], - state["position"][0], state["position"][1], - state["position"][2], state["orientation"][2], - external_wrenches_estimated_param[0], external_wrenches_estimated_param[1], - external_wrenches_estimated_param[2], external_wrenches_estimated_param[3], - external_wrenches_estimated_param[4], external_wrenches_estimated_param[5], - end_effector_position[0], end_effector_position[1], end_effector_position[2], - end_effector_jacobian_inv[0], end_effector_jacobian_inv[1], end_effector_jacobian_inv[2], - end_effector_jacobian_inv[3], end_effector_jacobian_inv[4], end_effector_jacobian_inv[5], - end_effector_jacobian_inv[6], end_effector_jacobian_inv[7], end_effector_jacobian_inv[8], - end_effector_jacobian_transpose_inv[0], end_effector_jacobian_transpose_inv[1], end_effector_jacobian_transpose_inv[2], - end_effector_jacobian_transpose_inv[3], end_effector_jacobian_transpose_inv[4], end_effector_jacobian_transpose_inv[5], - end_effector_jacobian_transpose_inv[6], end_effector_jacobian_transpose_inv[7], end_effector_jacobian_transpose_inv[8], - end_effector_gain[0], end_effector_gain[1], end_effector_gain[2]]) + param = np.array( + [ + FL_contact_sequence[j], + FR_contact_sequence[j], + RL_contact_sequence[j], + RR_contact_sequence[j], + mu, + stance_proximity_FL[j], + stance_proximity_FR[j], + stance_proximity_RL[j], + stance_proximity_RR[j], + state["position"][0], + state["position"][1], + state["position"][2], + state["orientation"][2], + external_wrenches_estimated_param[0], + external_wrenches_estimated_param[1], + external_wrenches_estimated_param[2], + external_wrenches_estimated_param[3], + external_wrenches_estimated_param[4], + external_wrenches_estimated_param[5], + end_effector_position[0], + end_effector_position[1], + end_effector_position[2], + end_effector_jacobian_inv[0], + end_effector_jacobian_inv[1], + end_effector_jacobian_inv[2], + end_effector_jacobian_inv[3], + end_effector_jacobian_inv[4], + end_effector_jacobian_inv[5], + end_effector_jacobian_inv[6], + end_effector_jacobian_inv[7], + end_effector_jacobian_inv[8], + end_effector_jacobian_transpose_inv[0], + end_effector_jacobian_transpose_inv[1], + end_effector_jacobian_transpose_inv[2], + end_effector_jacobian_transpose_inv[3], + end_effector_jacobian_transpose_inv[4], + end_effector_jacobian_transpose_inv[5], + end_effector_jacobian_transpose_inv[6], + end_effector_jacobian_transpose_inv[7], + end_effector_jacobian_transpose_inv[8], + end_effector_gain[0], + end_effector_gain[1], + end_effector_gain[2], + ] + ) self.acados_ocp_solver.set(j, "p", copy.deepcopy(param)) - - - - - # Set initial state constraint. We teleported the robot foothold # to the previous optimal foothold. This is done to avoid the optimization # of a foothold that is not considered at all at touchdown! In any case, - # the height cames always from the VFA - if(FL_contact_sequence[0] == 0): + # the height cames always from the VFA + if FL_contact_sequence[0] == 0: state["foot_FL"] = reference["ref_foot_FL"][0] - if(FR_contact_sequence[0] == 0): + if FR_contact_sequence[0] == 0: state["foot_FR"] = reference["ref_foot_FR"][0] - if(RL_contact_sequence[0] == 0): + if RL_contact_sequence[0] == 0: state["foot_RL"] = reference["ref_foot_RL"][0] - - if(RR_contact_sequence[0] == 0): - state["foot_RR"] = reference["ref_foot_RR"][0] + if RR_contact_sequence[0] == 0: + state["foot_RR"] = reference["ref_foot_RR"][0] - if(self.use_integrators): + if self.use_integrators: # Compute error for integral action alpha_integrator = config.mpc_params["alpha_integrator"] - self.integral_errors[0] += (state["position"][2] - reference["ref_position"][2])*alpha_integrator - self.integral_errors[1] += (state["linear_velocity"][0] - reference["ref_linear_velocity"][0])*alpha_integrator - self.integral_errors[2] += (state["linear_velocity"][1] - reference["ref_linear_velocity"][1])*alpha_integrator - self.integral_errors[3] += (state["linear_velocity"][2] - reference["ref_linear_velocity"][2])*alpha_integrator - self.integral_errors[4] += (state["orientation"][0] - reference["ref_orientation"][0])*(alpha_integrator) - self.integral_errors[5] += (state["orientation"][1] - reference["ref_orientation"][1])*alpha_integrator - + self.integral_errors[0] += (state["position"][2] - reference["ref_position"][2]) * alpha_integrator + self.integral_errors[1] += ( + state["linear_velocity"][0] - reference["ref_linear_velocity"][0] + ) * alpha_integrator + self.integral_errors[2] += ( + state["linear_velocity"][1] - reference["ref_linear_velocity"][1] + ) * alpha_integrator + self.integral_errors[3] += ( + state["linear_velocity"][2] - reference["ref_linear_velocity"][2] + ) * alpha_integrator + self.integral_errors[4] += (state["orientation"][0] - reference["ref_orientation"][0]) * (alpha_integrator) + self.integral_errors[5] += (state["orientation"][1] - reference["ref_orientation"][1]) * alpha_integrator cap_integrator_z = config.mpc_params["integrator_cap"][0] cap_integrator_x_dot = config.mpc_params["integrator_cap"][1] @@ -1444,67 +1473,99 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, cap_integrator_roll = config.mpc_params["integrator_cap"][4] cap_integrator_pitch = config.mpc_params["integrator_cap"][5] - self.integral_errors[0] = np.where(np.abs(self.integral_errors[0]) > cap_integrator_z, cap_integrator_z*np.sign(self.integral_errors[0]), self.integral_errors[0]) - self.integral_errors[1] = np.where(np.abs(self.integral_errors[1]) > cap_integrator_x_dot, cap_integrator_x_dot*np.sign(self.integral_errors[1]), self.integral_errors[1]) - self.integral_errors[2] = np.where(np.abs(self.integral_errors[2]) > cap_integrator_y_dot, cap_integrator_y_dot*np.sign(self.integral_errors[2]), self.integral_errors[2]) - self.integral_errors[3] = np.where(np.abs(self.integral_errors[3]) > cap_integrator_z_dot, cap_integrator_z_dot*np.sign(self.integral_errors[3]), self.integral_errors[3]) - self.integral_errors[4] = np.where(np.abs(self.integral_errors[4]) > cap_integrator_roll, cap_integrator_roll*np.sign(self.integral_errors[4]), self.integral_errors[4]) - self.integral_errors[5] = np.where(np.abs(self.integral_errors[5]) > cap_integrator_pitch, cap_integrator_pitch*np.sign(self.integral_errors[5]), self.integral_errors[5]) + self.integral_errors[0] = np.where( + np.abs(self.integral_errors[0]) > cap_integrator_z, + cap_integrator_z * np.sign(self.integral_errors[0]), + self.integral_errors[0], + ) + self.integral_errors[1] = np.where( + np.abs(self.integral_errors[1]) > cap_integrator_x_dot, + cap_integrator_x_dot * np.sign(self.integral_errors[1]), + self.integral_errors[1], + ) + self.integral_errors[2] = np.where( + np.abs(self.integral_errors[2]) > cap_integrator_y_dot, + cap_integrator_y_dot * np.sign(self.integral_errors[2]), + self.integral_errors[2], + ) + self.integral_errors[3] = np.where( + np.abs(self.integral_errors[3]) > cap_integrator_z_dot, + cap_integrator_z_dot * np.sign(self.integral_errors[3]), + self.integral_errors[3], + ) + self.integral_errors[4] = np.where( + np.abs(self.integral_errors[4]) > cap_integrator_roll, + cap_integrator_roll * np.sign(self.integral_errors[4]), + self.integral_errors[4], + ) + self.integral_errors[5] = np.where( + np.abs(self.integral_errors[5]) > cap_integrator_pitch, + cap_integrator_pitch * np.sign(self.integral_errors[5]), + self.integral_errors[5], + ) print("self.integral_errors: ", self.integral_errors) - # Set initial state constraint acados, converting first the dictionary to np array - state_acados = np.concatenate((state["position"], state["linear_velocity"], - state["orientation"], state["angular_velocity"], - state["foot_FL"], state["foot_FR"], - state["foot_RL"], state["foot_RR"], - self.integral_errors, - state["passive_arm_force"])).reshape((self.states_dim, 1)) + state_acados = np.concatenate( + ( + state["position"], + state["linear_velocity"], + state["orientation"], + state["angular_velocity"], + state["foot_FL"], + state["foot_FR"], + state["foot_RL"], + state["foot_RR"], + self.integral_errors, + state["passive_arm_force"], + ) + ).reshape((self.states_dim, 1)) self.acados_ocp_solver.set(0, "lbx", state_acados) self.acados_ocp_solver.set(0, "ubx", state_acados) - # Set Warm start in case... - if(self.use_warm_start): - self.set_warm_start(state_acados, reference, FL_contact_sequence, FR_contact_sequence, RL_contact_sequence, RR_contact_sequence) - + if self.use_warm_start: + self.set_warm_start( + state_acados, + reference, + FL_contact_sequence, + FR_contact_sequence, + RL_contact_sequence, + RR_contact_sequence, + ) # Set stage constraint - if(self.use_foothold_constraints or self.use_stability_constraints): - - h_R_w = np.array([np.cos(yaw), np.sin(yaw), - -np.sin(yaw), np.cos(yaw)]) - stance_proximity = np.vstack((stance_proximity_FL, stance_proximity_FR, - stance_proximity_RL, stance_proximity_RR)) + if self.use_foothold_constraints or self.use_stability_constraints: + h_R_w = np.array([np.cos(yaw), np.sin(yaw), -np.sin(yaw), np.cos(yaw)]) + stance_proximity = np.vstack( + (stance_proximity_FL, stance_proximity_FR, stance_proximity_RL, stance_proximity_RR) + ) self.set_stage_constraint(constraint, state, reference, contact_sequence, h_R_w, stance_proximity) - # Solve ocp via RTI or normal ocp if self.use_RTI: # feedback phase - self.acados_ocp_solver.options_set('rti_phase', 2) + self.acados_ocp_solver.options_set("rti_phase", 2) status = self.acados_ocp_solver.solve() - print("feedback phase time: ", self.acados_ocp_solver.get_stats('time_tot')) + print("feedback phase time: ", self.acados_ocp_solver.get_stats("time_tot")) else: status = self.acados_ocp_solver.solve() - print("ocp time: ", self.acados_ocp_solver.get_stats('time_tot')) - + print("ocp time: ", self.acados_ocp_solver.get_stats("time_tot")) - # Take the solution + # Take the solution control = self.acados_ocp_solver.get(0, "u") optimal_GRF = control[12:] optimal_foothold = np.zeros((4, 3)) - optimal_footholds_assigned = np.zeros((4, ), dtype='bool') + optimal_footholds_assigned = np.zeros((4,), dtype="bool") # Saturation of the GRF - optimal_GRF = np.clip(optimal_GRF, 0.0, config.mpc_params['grf_max']) - + optimal_GRF = np.clip(optimal_GRF, 0.0, config.mpc_params["grf_max"]) # We need to provide the next touchdown foothold position. # We first take the foothold in stance now (they are not optimized!) - # and flag them as True (aka "assigned") + # and flag them as True (aka "assigned") if FL_contact_sequence[0] == 1: optimal_foothold[0] = state["foot_FL"] optimal_footholds_assigned[0] = True @@ -1518,116 +1579,173 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, optimal_foothold[3] = state["foot_RR"] optimal_footholds_assigned[3] = True - - - - # Then we take the foothold at the next touchdown from the one - # that are not flagged as True from before, and saturate them!! + # Then we take the foothold at the next touchdown from the one + # that are not flagged as True from before, and saturate them!! # P.S. The saturation is in the horizontal frame - h_R_w = h_R_w.reshape((2, 2)) + h_R_w = h_R_w.reshape((2, 2)) for j in range(1, self.horizon): - if(FL_contact_sequence[j] != FL_contact_sequence[j-1] and not optimal_footholds_assigned[0]): + if FL_contact_sequence[j] != FL_contact_sequence[j - 1] and not optimal_footholds_assigned[0]: optimal_foothold[0] = self.acados_ocp_solver.get(j, "x")[12:15] optimal_footholds_assigned[0] = True - - if(constraint is None): - first_up_constraint_FL = np.array([reference["ref_foot_FL"][0][0], reference["ref_foot_FL"][0][1], reference["ref_foot_FL"][0][2] + 0.002]) - first_low_constraint_FL = np.array([reference["ref_foot_FL"][0][0], reference["ref_foot_FL"][0][1], reference["ref_foot_FL"][0][2] - 0.002]) - first_up_constraint_FL[0:2] = h_R_w@(first_up_constraint_FL[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_FL[0:2] = h_R_w@(first_low_constraint_FL[0:2] - state["position"][0:2]) - 0.15 + if constraint is None: + first_up_constraint_FL = np.array( + [ + reference["ref_foot_FL"][0][0], + reference["ref_foot_FL"][0][1], + reference["ref_foot_FL"][0][2] + 0.002, + ] + ) + first_low_constraint_FL = np.array( + [ + reference["ref_foot_FL"][0][0], + reference["ref_foot_FL"][0][1], + reference["ref_foot_FL"][0][2] - 0.002, + ] + ) + + first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - state["position"][0:2]) + 0.15 + first_low_constraint_FL[0:2] = ( + h_R_w @ (first_low_constraint_FL[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_FL = np.array([constraint[0][0], constraint[1][0], constraint[2][0] + 0.002]) first_low_constraint_FL = np.array([constraint[9][0], constraint[10][0], constraint[11][0] - 0.002]) - first_up_constraint_FL[0:2] = h_R_w@(first_up_constraint_FL[0:2] - state["position"][0:2]) - first_low_constraint_FL[0:2] = h_R_w@(first_low_constraint_FL[0:2] - state["position"][0:2]) + first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - state["position"][0:2]) + first_low_constraint_FL[0:2] = h_R_w @ (first_low_constraint_FL[0:2] - state["position"][0:2]) - optimal_foothold[0][0:2] = h_R_w@(optimal_foothold[0][0:2] - state["position"][0:2]) - optimal_foothold[0][0:2] = np.clip(optimal_foothold[0][0:2], first_low_constraint_FL[0:2], first_up_constraint_FL[0:2]) - optimal_foothold[0][0:2] = h_R_w.T@optimal_foothold[0][0:2] + state["position"][0:2] + optimal_foothold[0][0:2] = h_R_w @ (optimal_foothold[0][0:2] - state["position"][0:2]) + optimal_foothold[0][0:2] = np.clip( + optimal_foothold[0][0:2], first_low_constraint_FL[0:2], first_up_constraint_FL[0:2] + ) + optimal_foothold[0][0:2] = h_R_w.T @ optimal_foothold[0][0:2] + state["position"][0:2] - - if(FR_contact_sequence[j] != FR_contact_sequence[j-1] and not optimal_footholds_assigned[1]): + if FR_contact_sequence[j] != FR_contact_sequence[j - 1] and not optimal_footholds_assigned[1]: optimal_foothold[1] = self.acados_ocp_solver.get(j, "x")[15:18] optimal_footholds_assigned[1] = True - if(constraint is None): - first_up_constraint_FR = np.array([reference["ref_foot_FR"][0][0], reference["ref_foot_FR"][0][1], reference["ref_foot_FR"][0][2] + 0.002]) - first_low_constraint_FR = np.array([reference["ref_foot_FR"][0][0], reference["ref_foot_FR"][0][1], reference["ref_foot_FR"][0][2] - 0.002]) - - first_up_constraint_FR[0:2] = h_R_w@(first_up_constraint_FR[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_FR[0:2] = h_R_w@(first_low_constraint_FR[0:2] - state["position"][0:2]) - 0.15 + if constraint is None: + first_up_constraint_FR = np.array( + [ + reference["ref_foot_FR"][0][0], + reference["ref_foot_FR"][0][1], + reference["ref_foot_FR"][0][2] + 0.002, + ] + ) + first_low_constraint_FR = np.array( + [ + reference["ref_foot_FR"][0][0], + reference["ref_foot_FR"][0][1], + reference["ref_foot_FR"][0][2] - 0.002, + ] + ) + + first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - state["position"][0:2]) + 0.15 + first_low_constraint_FR[0:2] = ( + h_R_w @ (first_low_constraint_FR[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_FR = np.array([constraint[0][1], constraint[1][1], constraint[2][1] + 0.002]) first_low_constraint_FR = np.array([constraint[9][1], constraint[10][1], constraint[11][1] - 0.002]) - first_up_constraint_FR[0:2] = h_R_w@(first_up_constraint_FR[0:2] - state["position"][0:2]) - first_low_constraint_FR[0:2] = h_R_w@(first_low_constraint_FR[0:2] - state["position"][0:2]) + first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - state["position"][0:2]) + first_low_constraint_FR[0:2] = h_R_w @ (first_low_constraint_FR[0:2] - state["position"][0:2]) - optimal_foothold[1][0:2] = h_R_w@(optimal_foothold[1][0:2] - state["position"][0:2]) - optimal_foothold[1][0:2] = np.clip(optimal_foothold[1][0:2], first_low_constraint_FR[0:2], first_up_constraint_FR[0:2]) - optimal_foothold[1][0:2] = h_R_w.T@optimal_foothold[1][0:2] + state["position"][0:2] + optimal_foothold[1][0:2] = h_R_w @ (optimal_foothold[1][0:2] - state["position"][0:2]) + optimal_foothold[1][0:2] = np.clip( + optimal_foothold[1][0:2], first_low_constraint_FR[0:2], first_up_constraint_FR[0:2] + ) + optimal_foothold[1][0:2] = h_R_w.T @ optimal_foothold[1][0:2] + state["position"][0:2] - - if(RL_contact_sequence[j] != RL_contact_sequence[j-1] and not optimal_footholds_assigned[2]): + if RL_contact_sequence[j] != RL_contact_sequence[j - 1] and not optimal_footholds_assigned[2]: optimal_foothold[2] = self.acados_ocp_solver.get(j, "x")[18:21] optimal_footholds_assigned[2] = True - if(constraint is None): - first_up_constraint_RL = np.array([reference["ref_foot_RL"][0][0], reference["ref_foot_RL"][0][1], reference["ref_foot_RL"][0][2] + 0.002]) - first_low_constraint_RL = np.array([reference["ref_foot_RL"][0][0], reference["ref_foot_RL"][0][1], reference["ref_foot_RL"][0][2] - 0.002]) - - first_up_constraint_RL[0:2] = h_R_w@(first_up_constraint_RL[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_RL[0:2] = h_R_w@(first_low_constraint_RL[0:2] - state["position"][0:2]) - 0.15 + if constraint is None: + first_up_constraint_RL = np.array( + [ + reference["ref_foot_RL"][0][0], + reference["ref_foot_RL"][0][1], + reference["ref_foot_RL"][0][2] + 0.002, + ] + ) + first_low_constraint_RL = np.array( + [ + reference["ref_foot_RL"][0][0], + reference["ref_foot_RL"][0][1], + reference["ref_foot_RL"][0][2] - 0.002, + ] + ) + + first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - state["position"][0:2]) + 0.15 + first_low_constraint_RL[0:2] = ( + h_R_w @ (first_low_constraint_RL[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_RL = np.array([constraint[0][2], constraint[1][2], constraint[2][2] + 0.002]) first_low_constraint_RL = np.array([constraint[9][2], constraint[10][2], constraint[11][2] - 0.002]) - first_up_constraint_RL[0:2] = h_R_w@(first_up_constraint_RL[0:2] - state["position"][0:2]) - first_low_constraint_RL[0:2] = h_R_w@(first_low_constraint_RL[0:2] - state["position"][0:2]) - - optimal_foothold[2][0:2] = h_R_w@(optimal_foothold[2][0:2] - state["position"][0:2]) - optimal_foothold[2][0:2] = np.clip(optimal_foothold[2][0:2], first_low_constraint_RL[0:2], first_up_constraint_RL[0:2]) - optimal_foothold[2][0:2] = h_R_w.T@optimal_foothold[2][0:2] + state["position"][0:2] - - if(RR_contact_sequence[j] != RR_contact_sequence[j-1] and not optimal_footholds_assigned[3]): + first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - state["position"][0:2]) + first_low_constraint_RL[0:2] = h_R_w @ (first_low_constraint_RL[0:2] - state["position"][0:2]) + + optimal_foothold[2][0:2] = h_R_w @ (optimal_foothold[2][0:2] - state["position"][0:2]) + optimal_foothold[2][0:2] = np.clip( + optimal_foothold[2][0:2], first_low_constraint_RL[0:2], first_up_constraint_RL[0:2] + ) + optimal_foothold[2][0:2] = h_R_w.T @ optimal_foothold[2][0:2] + state["position"][0:2] + + if RR_contact_sequence[j] != RR_contact_sequence[j - 1] and not optimal_footholds_assigned[3]: optimal_foothold[3] = self.acados_ocp_solver.get(j, "x")[21:24] optimal_footholds_assigned[3] = True - if(constraint is None): - first_up_constraint_RR = np.array([reference["ref_foot_RR"][0][0], reference["ref_foot_RR"][0][1], reference["ref_foot_RR"][0][2] + 0.002]) - first_low_constraint_RR = np.array([reference["ref_foot_RR"][0][0], reference["ref_foot_RR"][0][1], reference["ref_foot_RR"][0][2] - 0.002]) - - first_up_constraint_RR[0:2] = h_R_w@(first_up_constraint_RR[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_RR[0:2] = h_R_w@(first_low_constraint_RR[0:2] - state["position"][0:2]) - 0.15 + if constraint is None: + first_up_constraint_RR = np.array( + [ + reference["ref_foot_RR"][0][0], + reference["ref_foot_RR"][0][1], + reference["ref_foot_RR"][0][2] + 0.002, + ] + ) + first_low_constraint_RR = np.array( + [ + reference["ref_foot_RR"][0][0], + reference["ref_foot_RR"][0][1], + reference["ref_foot_RR"][0][2] - 0.002, + ] + ) + + first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - state["position"][0:2]) + 0.15 + first_low_constraint_RR[0:2] = ( + h_R_w @ (first_low_constraint_RR[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_RR = np.array([constraint[0][3], constraint[1][3], constraint[2][3] + 0.002]) first_low_constraint_RR = np.array([constraint[9][3], constraint[10][3], constraint[11][3] - 0.002]) - first_up_constraint_RR[0:2] = h_R_w@(first_up_constraint_RR[0:2] - state["position"][0:2]) - first_low_constraint_RR[0:2] = h_R_w@(first_low_constraint_RR[0:2] - state["position"][0:2]) + first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - state["position"][0:2]) + first_low_constraint_RR[0:2] = h_R_w @ (first_low_constraint_RR[0:2] - state["position"][0:2]) - optimal_foothold[3][0:2] = h_R_w@(optimal_foothold[3][0:2] - state["position"][0:2]) - optimal_foothold[3][0:2] = np.clip(optimal_foothold[3][0:2], first_low_constraint_RR[0:2], first_up_constraint_RR[0:2]) - optimal_foothold[3][0:2] = h_R_w.T@optimal_foothold[3][0:2] + state["position"][0:2] - + optimal_foothold[3][0:2] = h_R_w @ (optimal_foothold[3][0:2] - state["position"][0:2]) + optimal_foothold[3][0:2] = np.clip( + optimal_foothold[3][0:2], first_low_constraint_RR[0:2], first_up_constraint_RR[0:2] + ) + optimal_foothold[3][0:2] = h_R_w.T @ optimal_foothold[3][0:2] + state["position"][0:2] - - # If in the prediction horizon, the foot is never in stance, we replicate the reference - #to not confuse the swing controller - if(optimal_footholds_assigned[0] == False): + # If in the prediction horizon, the foot is never in stance, we replicate the reference + # to not confuse the swing controller + if optimal_footholds_assigned[0] == False: optimal_foothold[0] = reference["ref_foot_FL"][0] - if(optimal_footholds_assigned[1] == False): + if optimal_footholds_assigned[1] == False: optimal_foothold[1] = reference["ref_foot_FR"][0] - if(optimal_footholds_assigned[2] == False): + if optimal_footholds_assigned[2] == False: optimal_foothold[2] = reference["ref_foot_RL"][0] - if(optimal_footholds_assigned[3] == False): + if optimal_footholds_assigned[3] == False: optimal_foothold[3] = reference["ref_foot_RR"][0] - - if(config.mpc_params['dt'] <= 0.02 or ( - config.mpc_params['use_nonuniform_discretization'] and config.mpc_params['dt_fine_grained'] <= 0.02)): + if config.mpc_params["dt"] <= 0.02 or ( + config.mpc_params["use_nonuniform_discretization"] and config.mpc_params["dt_fine_grained"] <= 0.02 + ): optimal_next_state_index = 2 else: optimal_next_state_index = 1 @@ -1636,12 +1754,9 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, self.optimal_next_state = optimal_next_state self.acados_ocp_solver.print_statistics() - - - # Check if QPs converged, if not just use the reference footholds # and a GRF over Z distribuited between the leg in stance - if(status == 1 or status == 4): + if status == 1 or status == 4: print("status", status) if FL_contact_sequence[0] == 0: optimal_foothold[0] = reference["ref_foot_FL"][0] @@ -1651,38 +1766,33 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, optimal_foothold[2] = reference["ref_foot_RL"][0] if RR_contact_sequence[0] == 0: optimal_foothold[3] = reference["ref_foot_RR"][0] - - number_of_legs_in_stance = np.array([FL_contact_sequence[0], FR_contact_sequence[0], - RL_contact_sequence [0], RR_contact_sequence[0]]).sum() - if(number_of_legs_in_stance == 0): + + number_of_legs_in_stance = np.array( + [FL_contact_sequence[0], FR_contact_sequence[0], RL_contact_sequence[0], RR_contact_sequence[0]] + ).sum() + if number_of_legs_in_stance == 0: reference_force_stance_legs = 0 else: reference_force_stance_legs = (self.centroidal_model.mass * 9.81) / number_of_legs_in_stance - + reference_force_fl_z = reference_force_stance_legs * FL_contact_sequence[0] reference_force_fr_z = reference_force_stance_legs * FR_contact_sequence[0] reference_force_rl_z = reference_force_stance_legs * RL_contact_sequence[0] reference_force_rr_z = reference_force_stance_legs * RR_contact_sequence[0] - optimal_GRF = np.zeros((12, )) + optimal_GRF = np.zeros((12,)) optimal_GRF[2] = reference_force_fl_z optimal_GRF[5] = reference_force_fr_z optimal_GRF[8] = reference_force_rl_z optimal_GRF[11] = reference_force_rr_z - optimal_GRF = self.previous_optimal_GRF self.reset() - - # Save the previous optimal GRF, the previous status and the previous contact sequence self.previous_optimal_GRF = optimal_GRF self.previous_status = status self.previous_contact_sequence = contact_sequence - - - # Decenter the optimal foothold and the next state (they were centered around zero at the beginning) optimal_foothold[0] = optimal_foothold[0] + self.initial_base_position optimal_foothold[1] = optimal_foothold[1] + self.initial_base_position @@ -1695,6 +1805,5 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, optimal_next_state[18:21] = optimal_foothold[2] optimal_next_state[21:24] = optimal_foothold[3] - # Return the optimal GRF, the optimal foothold, the next state and the status of the optimization return optimal_GRF, optimal_foothold, optimal_next_state, status diff --git a/quadruped_pympc/controllers/gradient/input_rates/centroidal_model_input_rates.py b/quadruped_pympc/controllers/gradient/input_rates/centroidal_model_input_rates.py index 9e3861d..7d9d2a8 100644 --- a/quadruped_pympc/controllers/gradient/input_rates/centroidal_model_input_rates.py +++ b/quadruped_pympc/controllers/gradient/input_rates/centroidal_model_input_rates.py @@ -3,15 +3,18 @@ # Authors: Giulio Turrisi - -import numpy as np import casadi as cs +import numpy as np from acados_template import AcadosModel + from quadruped_pympc import config # Class that defines the prediction model of the NMPC class Centroidal_Model_InputRates: - def __init__(self, ) -> None: + def __init__( + self, + ) -> None: """ This method initializes the foothold generator Centroidal_Model, which creates the prediction model of the NMPC. @@ -53,50 +56,54 @@ def __init__(self, ) -> None: foot_force_rl = cs.SX.sym("foot_force_rl", 3, 1) foot_force_rr = cs.SX.sym("foot_force_rr", 3, 1) - self.states = cs.vertcat(com_position_x, - com_position_y, - com_position_z, - com_velocity_x, - com_velocity_y, - com_velocity_z, - roll, - pitch, - yaw, - omega_x, - omega_y, - omega_z, - foot_position_fl, - foot_position_fr, - foot_position_rl, - foot_position_rr, - com_position_z_integral, - com_velocity_x_integral, - com_velocity_y_integral, - com_velocity_z_integral, - roll_integral, - pitch_integral, - foot_force_fl, - foot_force_fr, - foot_force_rl, - foot_force_rr) - - # Define state dot - self.states_dot = cs.vertcat(cs.SX.sym("linear_com_vel", 3, 1), - cs.SX.sym("linear_com_acc", 3, 1), - cs.SX.sym("euler_rates_base", 3, 1), - cs.SX.sym("angular_acc_base", 3, 1), - cs.SX.sym("linear_vel_foot_FL", 3, 1), - cs.SX.sym("linear_vel_foot_FR", 3, 1), - cs.SX.sym("linear_vel_foot_RL", 3, 1), - cs.SX.sym("linear_vel_foot_RR", 3, 1), - cs.SX.sym("linear_com_vel_z_integral", 1, 1), - cs.SX.sym("linear_com_acc_integral", 3, 1), - cs.SX.sym("euler_rates_roll_integral", 1, 1), - cs.SX.sym("euler_rates_pitch_integral", 1, 1), - cs.SX.sym("foot_force_rate_FL", 3, 1), - cs.SX.sym("foot_force_rate_FR", 3, 1), - cs.SX.sym("foot_force_rate_RL", 3, 1), - cs.SX.sym("foot_force_rate_RR", 3, 1), ) + self.states = cs.vertcat( + com_position_x, + com_position_y, + com_position_z, + com_velocity_x, + com_velocity_y, + com_velocity_z, + roll, + pitch, + yaw, + omega_x, + omega_y, + omega_z, + foot_position_fl, + foot_position_fr, + foot_position_rl, + foot_position_rr, + com_position_z_integral, + com_velocity_x_integral, + com_velocity_y_integral, + com_velocity_z_integral, + roll_integral, + pitch_integral, + foot_force_fl, + foot_force_fr, + foot_force_rl, + foot_force_rr, + ) + + # Define state dot + self.states_dot = cs.vertcat( + cs.SX.sym("linear_com_vel", 3, 1), + cs.SX.sym("linear_com_acc", 3, 1), + cs.SX.sym("euler_rates_base", 3, 1), + cs.SX.sym("angular_acc_base", 3, 1), + cs.SX.sym("linear_vel_foot_FL", 3, 1), + cs.SX.sym("linear_vel_foot_FR", 3, 1), + cs.SX.sym("linear_vel_foot_RL", 3, 1), + cs.SX.sym("linear_vel_foot_RR", 3, 1), + cs.SX.sym("linear_com_vel_z_integral", 1, 1), + cs.SX.sym("linear_com_acc_integral", 3, 1), + cs.SX.sym("euler_rates_roll_integral", 1, 1), + cs.SX.sym("euler_rates_pitch_integral", 1, 1), + cs.SX.sym("foot_force_rate_FL", 3, 1), + cs.SX.sym("foot_force_rate_FR", 3, 1), + cs.SX.sym("foot_force_rate_RL", 3, 1), + cs.SX.sym("foot_force_rate_RR", 3, 1), + ) # Define input and its casadi variables foot_velocity_fl = cs.SX.sym("foot_velocity_fl", 3, 1) @@ -109,14 +116,16 @@ def __init__(self, ) -> None: foot_force_rate_rl = cs.SX.sym("foot_force_rate_rl", 3, 1) foot_force_rate_rr = cs.SX.sym("foot_force_rate_rr", 3, 1) - self.inputs = cs.vertcat(foot_velocity_fl, - foot_velocity_fr, - foot_velocity_rl, - foot_velocity_rr, - foot_force_rate_fl, - foot_force_rate_fr, - foot_force_rate_rl, - foot_force_rate_rr) + self.inputs = cs.vertcat( + foot_velocity_fl, + foot_velocity_fr, + foot_velocity_rl, + foot_velocity_rr, + foot_force_rate_fl, + foot_force_rate_fr, + foot_force_rate_rl, + foot_force_rate_rr, + ) # Usefull for debug what things goes where in y_ref in the compute_control function, # because there are a lot of variables @@ -140,10 +149,18 @@ def __init__(self, ) -> None: self.mass = cs.SX.sym("mass", 1, 1) # Not so useful, i can instantiate a casadi function for the fd - param = cs.vertcat(self.stance_param, self.mu_friction, self.stance_proximity, - self.base_initial, self.base_yaw, self.external_wrench, self.inertia, self.mass) + param = cs.vertcat( + self.stance_param, + self.mu_friction, + self.stance_proximity, + self.base_initial, + self.base_yaw, + self.external_wrench, + self.inertia, + self.mass, + ) fd = self.forward_dynamics(self.states, self.inputs, param) - self.fun_forward_dynamics = cs.Function('fun_forward_dynamics', [self.states, self.inputs, param], [fd]) + self.fun_forward_dynamics = cs.Function("fun_forward_dynamics", [self.states, self.inputs, param], [fd]) def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.ndarray) -> cs.SX: """ @@ -276,7 +293,7 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda # external_wrench_angular # FINAL linear_foot_vel STATES (5,6,7,8) - if (not config.mpc_params["use_foothold_optimization"]): + if not config.mpc_params["use_foothold_optimization"]: foot_velocity_fl = foot_velocity_fl @ 0.0 foot_velocity_fr = foot_velocity_fr @ 0.0 foot_velocity_rl = foot_velocity_rl @ 0.0 @@ -302,20 +319,41 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda foot_force_rate_RR = foot_force_rate_rr # The order of the return should be equal to the order of the states_dot - return cs.vertcat(linear_com_vel, linear_com_acc, euler_rates_base, angular_acc_base, - linear_foot_vel_FL, linear_foot_vel_FR, linear_foot_vel_RL, linear_foot_vel_RR, - integral_states, - foot_force_rate_FL, foot_force_rate_FR, foot_force_rate_RL, foot_force_rate_RR) - - def export_robot_model(self, ) -> AcadosModel: + return cs.vertcat( + linear_com_vel, + linear_com_acc, + euler_rates_base, + angular_acc_base, + linear_foot_vel_FL, + linear_foot_vel_FR, + linear_foot_vel_RL, + linear_foot_vel_RR, + integral_states, + foot_force_rate_FL, + foot_force_rate_FR, + foot_force_rate_RL, + foot_force_rate_RR, + ) + + def export_robot_model( + self, + ) -> AcadosModel: """ This method set some general properties of the NMPC, such as the params, prediction mode, etc...! It will be called in centroidal_nmpc.py """ # dynamics - self.param = cs.vertcat(self.stance_param, self.mu_friction, self.stance_proximity, - self.base_initial, self.base_yaw, self.external_wrench, self.inertia, self.mass) + self.param = cs.vertcat( + self.stance_param, + self.mu_friction, + self.stance_proximity, + self.base_initial, + self.base_yaw, + self.external_wrench, + self.inertia, + self.mass, + ) f_expl = self.forward_dynamics(self.states, self.inputs, self.param) f_impl = self.states_dot - f_expl diff --git a/quadruped_pympc/controllers/gradient/input_rates/centroidal_nmpc_input_rates.py b/quadruped_pympc/controllers/gradient/input_rates/centroidal_nmpc_input_rates.py index 1cbe783..548f293 100644 --- a/quadruped_pympc/controllers/gradient/input_rates/centroidal_nmpc_input_rates.py +++ b/quadruped_pympc/controllers/gradient/input_rates/centroidal_nmpc_input_rates.py @@ -1,19 +1,23 @@ # Description: This file contains the class for the NMPC controller -# Authors: Giulio Turrisi - +# Authors: Giulio Turrisi - +import copy import os + +import casadi as cs import numpy as np import scipy.linalg -import casadi as cs -import copy -from typing import Tuple + # TODO: Check acados installation and trow error + instructions if needed -> pip install quadruped_pympc[acados] from acados_template import AcadosOcp, AcadosOcpSolver + ACADOS_INFTY = ACADOS_INFTY = 1000 from quadruped_pympc import config + from .centroidal_model_input_rates import Centroidal_Model_InputRates + # Class for the Acados NMPC, the model is in another file! class Acados_NMPC_InputRates: def __init__(self): @@ -28,32 +32,28 @@ def __init__(self): use_integrators (bool): Flag indicating whether to use integrators. """ - self.horizon = config.mpc_params['horizon'] # Define the number of discretization steps - self.dt = config.mpc_params['dt'] - self.T_horizon = self.horizon*self.dt - self.use_RTI = config.mpc_params['use_RTI'] - self.use_integrators = config.mpc_params['use_integrators'] - self.use_warm_start = config.mpc_params['use_warm_start'] - self.use_foothold_constraints = config.mpc_params['use_foothold_constraints'] + self.horizon = config.mpc_params["horizon"] # Define the number of discretization steps + self.dt = config.mpc_params["dt"] + self.T_horizon = self.horizon * self.dt + self.use_RTI = config.mpc_params["use_RTI"] + self.use_integrators = config.mpc_params["use_integrators"] + self.use_warm_start = config.mpc_params["use_warm_start"] + self.use_foothold_constraints = config.mpc_params["use_foothold_constraints"] - - self.use_static_stability = config.mpc_params['use_static_stability'] - self.use_zmp_stability = config.mpc_params['use_zmp_stability'] + self.use_static_stability = config.mpc_params["use_static_stability"] + self.use_zmp_stability = config.mpc_params["use_zmp_stability"] self.use_stability_constraints = self.use_static_stability or self.use_zmp_stability - - self.use_input_prediction = config.mpc_params['use_input_prediction'] + self.use_input_prediction = config.mpc_params["use_input_prediction"] - self.use_DDP = config.mpc_params['use_DDP'] + self.use_DDP = config.mpc_params["use_DDP"] - self.verbose = config.mpc_params['verbose'] + self.verbose = config.mpc_params["verbose"] - self.previous_status = -1 self.previous_contact_sequence = np.zeros((4, self.horizon)) self.optimal_next_state = np.zeros((24,)) self.previous_optimal_GRF = np.zeros((12,)) - self.integral_errors = np.zeros((6,)) @@ -62,11 +62,9 @@ def __init__(self): self.force_RL = np.zeros((3,)) self.force_RR = np.zeros((3,)) - # For centering the variable around 0, 0, 0 (World frame) self.initial_base_position = np.array([0, 0, 0]) - # Create the class of the centroidal model and instantiate the acados model self.centroidal_model = Centroidal_Model_InputRates() acados_model = self.centroidal_model.export_robot_model() @@ -77,9 +75,11 @@ def __init__(self): self.ocp = self.create_ocp_solver_description(acados_model) dir_path = os.path.dirname(os.path.realpath(__file__)) - self.ocp.code_export_directory = dir_path + '/c_generated_code' - - self.acados_ocp_solver = AcadosOcpSolver(self.ocp, json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json") + self.ocp.code_export_directory = dir_path + "/c_generated_code" + + self.acados_ocp_solver = AcadosOcpSolver( + self.ocp, json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json" + ) # Initialize solver for stage in range(self.horizon + 1): @@ -87,13 +87,12 @@ def __init__(self): for stage in range(self.horizon): self.acados_ocp_solver.set(stage, "u", np.zeros((self.inputs_dim,))) - if(self.use_RTI): + if self.use_RTI: # first preparation phase - self.acados_ocp_solver.options_set('rti_phase', 1) - status = self.acados_ocp_solver.solve() - + self.acados_ocp_solver.options_set("rti_phase", 1) + status = self.acados_ocp_solver.solve() - # Set cost, constraints and options + # Set cost, constraints and options def create_ocp_solver_description(self, acados_model) -> AcadosOcp: """ Creates and configures an optimal control problem (OCP) solver description. Acados STUFF! @@ -138,41 +137,35 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: ocp.cost.yref_e = np.zeros((ny_e,)) # Set friction and foothold constraints - expr_h_friction, \ - self.constr_uh_friction, \ - self.constr_lh_friction = self.create_friction_cone_constraints() + expr_h_friction, self.constr_uh_friction, self.constr_lh_friction = self.create_friction_cone_constraints() ocp.model.con_h_expr = expr_h_friction ocp.constraints.uh = self.constr_uh_friction ocp.constraints.lh = self.constr_lh_friction - nsh = expr_h_friction.shape[0] + nsh = expr_h_friction.shape[0] nsh_state_constraint_start = copy.copy(nsh) - if(self.use_foothold_constraints): - expr_h_foot, \ - self.constr_uh_foot, \ - self.constr_lh_foot = self.create_foothold_constraints() + if self.use_foothold_constraints: + expr_h_foot, self.constr_uh_foot, self.constr_lh_foot = self.create_foothold_constraints() - ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_foot) + ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_foot) ocp.constraints.uh = np.concatenate((ocp.constraints.uh, self.constr_uh_foot)) ocp.constraints.lh = np.concatenate((ocp.constraints.lh, self.constr_lh_foot)) nsh += expr_h_foot.shape[0] - - # Set stability constraints - if(self.use_stability_constraints): + # Set stability constraints + if self.use_stability_constraints: self.nsh_stability_start = copy.copy(nsh) - expr_h_support_polygon, \ - self.constr_uh_support_polygon, \ - self.constr_lh_support_polygon = self.create_stability_constraints() + expr_h_support_polygon, self.constr_uh_support_polygon, self.constr_lh_support_polygon = ( + self.create_stability_constraints() + ) - ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_support_polygon) + ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_support_polygon) ocp.constraints.uh = np.concatenate((ocp.constraints.uh, self.constr_uh_support_polygon)) ocp.constraints.lh = np.concatenate((ocp.constraints.lh, self.constr_lh_support_polygon)) nsh += expr_h_support_polygon.shape[0] self.nsh_stability_end = copy.copy(nsh) - nsh_state_constraint_end = copy.copy(nsh) # Set slack variable configuration: @@ -191,18 +184,23 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: ocp.cost.zu = 1000 * np.ones((ns,)) ocp.cost.Zu = 1 * np.ones((ns,)) """ - # Set slack variable configuration: num_state_cstr = nsh_state_constraint_end - nsh_state_constraint_start - if(num_state_cstr > 0): - ocp.constraints.lsh = np.zeros(num_state_cstr) # Lower bounds on slacks corresponding to soft lower bounds for nonlinear constraints - ocp.constraints.ush = np.zeros(num_state_cstr) # Lower bounds on slacks corresponding to soft upper bounds for nonlinear constraints - ocp.constraints.idxsh = np.array(range(nsh_state_constraint_start, nsh_state_constraint_end)) # Jsh + if num_state_cstr > 0: + ocp.constraints.lsh = np.zeros( + num_state_cstr + ) # Lower bounds on slacks corresponding to soft lower bounds for nonlinear constraints + ocp.constraints.ush = np.zeros( + num_state_cstr + ) # Lower bounds on slacks corresponding to soft upper bounds for nonlinear constraints + ocp.constraints.idxsh = np.array(range(nsh_state_constraint_start, nsh_state_constraint_end)) # Jsh ns = num_state_cstr - ocp.cost.zl = 1000 * np.ones((ns,)) # gradient wrt lower slack at intermediate shooting nodes (1 to N-1) - ocp.cost.Zl = 1 * np.ones((ns,)) # diagonal of Hessian wrt lower slack at intermediate shooting nodes (1 to N-1) - ocp.cost.zu = 1000 * np.ones((ns,)) - ocp.cost.Zu = 1 * np.ones((ns,)) + ocp.cost.zl = 1000 * np.ones((ns,)) # gradient wrt lower slack at intermediate shooting nodes (1 to N-1) + ocp.cost.Zl = 1 * np.ones( + (ns,) + ) # diagonal of Hessian wrt lower slack at intermediate shooting nodes (1 to N-1) + ocp.cost.zu = 1000 * np.ones((ns,)) + ocp.cost.Zu = 1 * np.ones((ns,)) # Variables to save the upper and lower bound of the constraints applied list_upper_bound = [] @@ -213,13 +211,12 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: self.upper_bound = np.array(list_upper_bound, dtype=object) self.lower_bound = np.array(list_lower_bound, dtype=object) - # Set initial state constraint X0 = np.zeros(shape=(nx,)) ocp.constraints.x0 = X0 # Set initialize parameters - init_contact_status = np.array([1., 1., 1., 1.]) + init_contact_status = np.array([1.0, 1.0, 1.0, 1.0]) init_mu = np.array([0.5]) init_stance_proximity = np.array([0, 0, 0, 0]) init_base = np.array([0, 0, 0]) @@ -227,136 +224,143 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: init_external_wrench = np.array([0, 0, 0, 0, 0, 0]) init_inertia = config.inertia.reshape((9,)) init_mass = np.array([config.mass]) - - ocp.parameter_values = np.concatenate((init_contact_status, init_mu, init_stance_proximity, - init_base, init_base_yaw, init_external_wrench, - init_inertia, init_mass)) + ocp.parameter_values = np.concatenate( + ( + init_contact_status, + init_mu, + init_stance_proximity, + init_base, + init_base_yaw, + init_external_wrench, + init_inertia, + init_mass, + ) + ) # Set options - ocp.solver_options.qp_solver = "PARTIAL_CONDENSING_HPIPM" # FULL_CONDENSING_QPOASES PARTIAL_CONDENSING_OSQP PARTIAL_CONDENSING_HPIPM + ocp.solver_options.qp_solver = ( + "PARTIAL_CONDENSING_HPIPM" # FULL_CONDENSING_QPOASES PARTIAL_CONDENSING_OSQP PARTIAL_CONDENSING_HPIPM + ) ocp.solver_options.hessian_approx = "GAUSS_NEWTON" # 'GAUSS_NEWTON', 'EXACT' - ocp.solver_options.integrator_type = "ERK" #ERK IRK GNSF DISCRETE - if(self.use_DDP): - ocp.solver_options.nlp_solver_type = 'DDP' - ocp.solver_options.nlp_solver_max_iter = config.mpc_params['num_qp_iterations'] - #ocp.solver_options.globalization = 'MERIT_BACKTRACKING' + ocp.solver_options.integrator_type = "ERK" # ERK IRK GNSF DISCRETE + if self.use_DDP: + ocp.solver_options.nlp_solver_type = "DDP" + ocp.solver_options.nlp_solver_max_iter = config.mpc_params["num_qp_iterations"] + # ocp.solver_options.globalization = 'MERIT_BACKTRACKING' ocp.solver_options.with_adaptive_levenberg_marquardt = True - ocp.cost.cost_type = 'NONLINEAR_LS' - ocp.cost.cost_type_e = 'NONLINEAR_LS' + ocp.cost.cost_type = "NONLINEAR_LS" + ocp.cost.cost_type_e = "NONLINEAR_LS" ocp.model.cost_y_expr = cs.vertcat(ocp.model.x, ocp.model.u) ocp.model.cost_y_expr_e = ocp.model.x - + ocp.translate_to_feasibility_problem(keep_x0=True, keep_cost=True) - elif(self.use_RTI): - ocp.solver_options.nlp_solver_type = "SQP_RTI" + elif self.use_RTI: + ocp.solver_options.nlp_solver_type = "SQP_RTI" ocp.solver_options.nlp_solver_max_iter = 1 - # Set the RTI type for the advanced RTI method + # Set the RTI type for the advanced RTI method # (see https://arxiv.org/pdf/2403.07101.pdf) - if(config.mpc_params['as_rti_type'] == "AS-RTI-A"): + if config.mpc_params["as_rti_type"] == "AS-RTI-A": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 0 - elif(config.mpc_params['as_rti_type'] == "AS-RTI-B"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-B": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 1 - elif(config.mpc_params['as_rti_type'] == "AS-RTI-C"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-C": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 2 - elif(config.mpc_params['as_rti_type'] == "AS-RTI-D"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-D": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 3 - + else: - ocp.solver_options.nlp_solver_type = "SQP" - ocp.solver_options.nlp_solver_max_iter = config.mpc_params['num_qp_iterations'] - #ocp.solver_options.globalization = "MERIT_BACKTRACKING" # FIXED_STEP, MERIT_BACKTRACKING - - if(config.mpc_params['solver_mode'] == "balance"): + ocp.solver_options.nlp_solver_type = "SQP" + ocp.solver_options.nlp_solver_max_iter = config.mpc_params["num_qp_iterations"] + # ocp.solver_options.globalization = "MERIT_BACKTRACKING" # FIXED_STEP, MERIT_BACKTRACKING + + if config.mpc_params["solver_mode"] == "balance": ocp.solver_options.hpipm_mode = "BALANCE" - elif(config.mpc_params['solver_mode'] == "robust"): + elif config.mpc_params["solver_mode"] == "robust": ocp.solver_options.hpipm_mode = "ROBUST" - elif(config.mpc_params['solver_mode'] == "fast"): + elif config.mpc_params["solver_mode"] == "fast": ocp.solver_options.qp_solver_iter_max = 10 ocp.solver_options.hpipm_mode = "SPEED" - elif(config.mpc_params['solver_mode'] == "crazy_speed"): + elif config.mpc_params["solver_mode"] == "crazy_speed": ocp.solver_options.qp_solver_iter_max = 5 ocp.solver_options.hpipm_mode = "SPEED_ABS" - #ocp.solver_options.globalization = "MERIT_BACKTRACKING" # FIXED_STEP, MERIT_BACKTRACKING - #ocp.solver_options.qp_solver_iter_max = 40 - #ocp.solver_options.line_search_use_sufficient_descent = 1 - #ocp.solver_options.hpipm_mode = "ROBUST" - #ocp.solver_options.regularize_method = "PROJECT_REDUC_HESS" - #ocp.solver_options.nlp_solver_ext_qp_res = 1 + # ocp.solver_options.globalization = "MERIT_BACKTRACKING" # FIXED_STEP, MERIT_BACKTRACKING + # ocp.solver_options.qp_solver_iter_max = 40 + # ocp.solver_options.line_search_use_sufficient_descent = 1 + # ocp.solver_options.hpipm_mode = "ROBUST" + # ocp.solver_options.regularize_method = "PROJECT_REDUC_HESS" + # ocp.solver_options.nlp_solver_ext_qp_res = 1 ocp.solver_options.levenberg_marquardt = 1e-3 # Set prediction horizon ocp.solver_options.tf = self.T_horizon - # Nonuniform discretization - if(config.mpc_params['use_nonuniform_discretization']): - time_steps_fine_grained = np.tile(config.mpc_params['dt_fine_grained'], config.mpc_params['horizon_fine_grained']) - time_steps = np.concatenate((time_steps_fine_grained, np.tile(self.dt, self.horizon - config.mpc_params['horizon_fine_grained']))) - shooting_nodes = np.zeros((self.horizon+1,)) + if config.mpc_params["use_nonuniform_discretization"]: + time_steps_fine_grained = np.tile( + config.mpc_params["dt_fine_grained"], config.mpc_params["horizon_fine_grained"] + ) + time_steps = np.concatenate( + (time_steps_fine_grained, np.tile(self.dt, self.horizon - config.mpc_params["horizon_fine_grained"])) + ) + shooting_nodes = np.zeros((self.horizon + 1,)) for i in range(len(time_steps)): - shooting_nodes[i+1] = shooting_nodes[i] + time_steps[i] + shooting_nodes[i + 1] = shooting_nodes[i] + time_steps[i] ocp.solver_options.shooting_nodes = shooting_nodes - - return ocp - - # Create a constraint for stability (COM, ZMP or CP inside support polygon) - def create_stability_constraints(self,) -> None: - + def create_stability_constraints( + self, + ) -> None: base_w = self.centroidal_model.states[0:3] base_vel_w = self.centroidal_model.states[3:6] - FL = self.centroidal_model.states[12:15] FR = self.centroidal_model.states[15:18] RL = self.centroidal_model.states[18:21] RR = self.centroidal_model.states[21:24] - #yaw = self.centroidal_model.base_yaw[0] + # yaw = self.centroidal_model.base_yaw[0] yaw = self.centroidal_model.states[8] h_R_w = cs.SX.zeros(2, 2) - h_R_w[0,0] = cs.cos(yaw) - h_R_w[0,1] = cs.sin(yaw) - h_R_w[1,0] = -cs.sin(yaw) - h_R_w[1,1] = cs.cos(yaw) - - FL[0:2] = h_R_w@(FL[0:2] - base_w[0:2]) - FR[0:2] = h_R_w@(FR[0:2] - base_w[0:2]) - RL[0:2] = h_R_w@(RL[0:2] - base_w[0:2]) - RR[0:2] = h_R_w@(RR[0:2] - base_w[0:2]) - - - if(self.use_static_stability): + h_R_w[0, 0] = cs.cos(yaw) + h_R_w[0, 1] = cs.sin(yaw) + h_R_w[1, 0] = -cs.sin(yaw) + h_R_w[1, 1] = cs.cos(yaw) + + FL[0:2] = h_R_w @ (FL[0:2] - base_w[0:2]) + FR[0:2] = h_R_w @ (FR[0:2] - base_w[0:2]) + RL[0:2] = h_R_w @ (RL[0:2] - base_w[0:2]) + RR[0:2] = h_R_w @ (RR[0:2] - base_w[0:2]) + + if self.use_static_stability: x = 0.0 y = 0.0 else: # Compute the ZMP robotHeight = base_w[2] - foot_force_fl = self.centroidal_model.states[30:33]#@self.centroidal_model.param[0] - foot_force_fr = self.centroidal_model.states[33:36]#@self.centroidal_model.param[1] - foot_force_rl = self.centroidal_model.states[36:39]#@self.centroidal_model.param[2] - foot_force_rr = self.centroidal_model.states[39:42]#@self.centroidal_model.param[3] + foot_force_fl = self.centroidal_model.states[30:33] # @self.centroidal_model.param[0] + foot_force_fr = self.centroidal_model.states[33:36] # @self.centroidal_model.param[1] + foot_force_rl = self.centroidal_model.states[36:39] # @self.centroidal_model.param[2] + foot_force_rr = self.centroidal_model.states[39:42] # @self.centroidal_model.param[3] temp = foot_force_fl + foot_force_fr + foot_force_rl + foot_force_rr gravity = np.array([0, 0, -9.81]) - linear_com_acc = (1/self.centroidal_model.mass)@temp + gravity - zmp = base_w[0:2] - linear_com_acc[0:2]*(robotHeight/(-gravity[2])) - #zmp = base_w[0:2] - base_vel_w[0:2]*(robotHeight/gravity[2]) - zmp = h_R_w@(zmp - base_w[0:2]) + linear_com_acc = (1 / self.centroidal_model.mass) @ temp + gravity + zmp = base_w[0:2] - linear_com_acc[0:2] * (robotHeight / (-gravity[2])) + # zmp = base_w[0:2] - base_vel_w[0:2]*(robotHeight/gravity[2]) + zmp = h_R_w @ (zmp - base_w[0:2]) x = zmp[0] y = zmp[1] - y_FL = FL[1] y_FR = FR[1] y_RL = RL[1] @@ -367,67 +371,65 @@ def create_stability_constraints(self,) -> None: x_RL = RL[0] x_RR = RR[0] - #LF - RF : x < (x2 - x1) (y - y1) / (y2 - y1) + x1 - #RF - RH: y > (y2 - y1) (x - x1) / (x2 - x1) + y1 - #RH - LH : x > (x2 - x1) (y - y1) / (y2 - y1) + x1 - #LH - LF: y < (y2 - y1) (x - x1) / (x2 - x1) + y1 - - #FL and FR cannot stay at the same x! #constrint should be less than zero - constraint_FL_FR = x - (x_FR - x_FL)*(y - y_FL) / (y_FR - y_FL + 0.001) - x_FL - - #FR and RR cannot stay at the same y! #constraint should be bigger than zero - constraint_FR_RR = y - (y_RR - y_FR)*(x - x_FR) / (x_RR - x_FR + 0.001) - y_FR - - #RL and RR cannot stay at the same x! #constraint should be bigger than zero - constraint_RR_RL = x - (x_RL - x_RR)*(y - y_RR) / (y_RL - y_RR + 0.001) - x_RR - - #FL and RL cannot stay at the same y! #constraint should be less than zero - constraint_RL_FL = y - (y_FL - y_RL)*(x - x_RL) / (x_FL - x_RL + 0.001) - y_RL - + # LF - RF : x < (x2 - x1) (y - y1) / (y2 - y1) + x1 + # RF - RH: y > (y2 - y1) (x - x1) / (x2 - x1) + y1 + # RH - LH : x > (x2 - x1) (y - y1) / (y2 - y1) + x1 + # LH - LF: y < (y2 - y1) (x - x1) / (x2 - x1) + y1 + + # FL and FR cannot stay at the same x! #constrint should be less than zero + constraint_FL_FR = x - (x_FR - x_FL) * (y - y_FL) / (y_FR - y_FL + 0.001) - x_FL + + # FR and RR cannot stay at the same y! #constraint should be bigger than zero + constraint_FR_RR = y - (y_RR - y_FR) * (x - x_FR) / (x_RR - x_FR + 0.001) - y_FR + + # RL and RR cannot stay at the same x! #constraint should be bigger than zero + constraint_RR_RL = x - (x_RL - x_RR) * (y - y_RR) / (y_RL - y_RR + 0.001) - x_RR + + # FL and RL cannot stay at the same y! #constraint should be less than zero + constraint_RL_FL = y - (y_FL - y_RL) * (x - x_RL) / (x_FL - x_RL + 0.001) - y_RL + # the diagonal stuff can be at any point... - constraint_FL_RR = y - (y_RR - y_FL)*(x - x_FL) / (x_RR - x_FL + 0.001) - y_FL #bigger than zero - constraint_FR_RL = y - (y_RL - y_FR)*(x - x_FR) / (x_RL - x_FR + 0.001) - y_FR #bigger than zero + constraint_FL_RR = y - (y_RR - y_FL) * (x - x_FL) / (x_RR - x_FL + 0.001) - y_FL # bigger than zero + constraint_FR_RL = y - (y_RL - y_FR) * (x - x_FR) / (x_RL - x_FR + 0.001) - y_FR # bigger than zero # upper and lower bound - ub = np.ones(6)*ACADOS_INFTY - lb = -np.ones(6)*ACADOS_INFTY + ub = np.ones(6) * ACADOS_INFTY + lb = -np.ones(6) * ACADOS_INFTY - #constraint_FL_FR + # constraint_FL_FR ub[0] = 0 lb[0] = -ACADOS_INFTY - - #constraint_FR_RR + + # constraint_FR_RR ub[1] = ACADOS_INFTY lb[1] = 0 - #constraint_RR_RL + # constraint_RR_RL ub[2] = ACADOS_INFTY lb[2] = 0 - #constraint_RL_FL + # constraint_RL_FL ub[3] = 0 lb[3] = -ACADOS_INFTY - - #constraint_FL_RR + + # constraint_FL_RR ub[4] = 0 lb[4] = -ACADOS_INFTY - - #constraint_FR_RL + + # constraint_FR_RL ub[5] = 0 lb[5] = -ACADOS_INFTY - - Jb = cs.vertcat(constraint_FL_FR, constraint_FR_RR, - constraint_RR_RL, constraint_RL_FL, - constraint_FL_RR, constraint_FR_RL) - - return Jb, ub, lb - - + Jb = cs.vertcat( + constraint_FL_FR, constraint_FR_RR, constraint_RR_RL, constraint_RL_FL, constraint_FL_RR, constraint_FR_RL + ) + return Jb, ub, lb # Create a standard foothold box constraint - def create_foothold_constraints(self,): + def create_foothold_constraints( + self, + ): """ This function calculates the symbolic foothold constraints for the centroidal NMPC problem. @@ -437,53 +439,45 @@ def create_foothold_constraints(self,): lbu: A numpy array of shape (8,) representing the lower bounds of the foothold constraints. """ - ubu = np.ones(12)*ACADOS_INFTY - lbu = -np.ones(12)*ACADOS_INFTY + ubu = np.ones(12) * ACADOS_INFTY + lbu = -np.ones(12) * ACADOS_INFTY - - # The foothold constraint in acados are in the horizontal frame, + # The foothold constraint in acados are in the horizontal frame, # but they arrive to us in the world frame. We need to rotate them # using the robot yaw yaw = self.centroidal_model.base_yaw[0] h_R_w = cs.SX.zeros(2, 2) - h_R_w[0,0] = cs.cos(yaw) - h_R_w[0,1] = cs.sin(yaw) - h_R_w[1,0] = -cs.sin(yaw) - h_R_w[1,1] = cs.cos(yaw) - + h_R_w[0, 0] = cs.cos(yaw) + h_R_w[0, 1] = cs.sin(yaw) + h_R_w[1, 0] = -cs.sin(yaw) + h_R_w[1, 1] = cs.cos(yaw) # and translate them using the robot base position base = self.centroidal_model.base_initial[0:3] - - - foot_position_fl = cs.SX.zeros(3,1) - foot_position_fl[0:2] = h_R_w@cs.vertcat(self.centroidal_model.states[12:14] - base[0:2]) + foot_position_fl = cs.SX.zeros(3, 1) + foot_position_fl[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.states[12:14] - base[0:2]) foot_position_fl[2] = self.centroidal_model.states[14] - - foot_position_fr = cs.SX.zeros(3,1) - foot_position_fr[0:2] = h_R_w@cs.vertcat(self.centroidal_model.states[15:17] - base[0:2]) + + foot_position_fr = cs.SX.zeros(3, 1) + foot_position_fr[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.states[15:17] - base[0:2]) foot_position_fr[2] = self.centroidal_model.states[17] - - foot_position_rl = cs.SX.zeros(3,1) - foot_position_rl[0:2] = h_R_w@cs.vertcat(self.centroidal_model.states[18:20] - base[0:2]) + + foot_position_rl = cs.SX.zeros(3, 1) + foot_position_rl[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.states[18:20] - base[0:2]) foot_position_rl[2] = self.centroidal_model.states[20] - - foot_position_rr = cs.SX.zeros(3,1) - foot_position_rr[0:2] = h_R_w@cs.vertcat(self.centroidal_model.states[21:23] - base[0:2]) + + foot_position_rr = cs.SX.zeros(3, 1) + foot_position_rr[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.states[21:23] - base[0:2]) foot_position_rr[2] = self.centroidal_model.states[23] - - - Jbu = cs.vertcat(foot_position_fl, foot_position_fr, - foot_position_rl, foot_position_rr) + Jbu = cs.vertcat(foot_position_fl, foot_position_fr, foot_position_rl, foot_position_rr) return Jbu, ubu, lbu - - - # Create the friction cone constraint - def create_friction_cone_constraints(self,): + def create_friction_cone_constraints( + self, + ): """ This function calculates the symbolic friction cone constraints for the centroidal NMPC problem. @@ -496,46 +490,45 @@ def create_friction_cone_constraints(self,): t = np.array([1, 0, 0]) b = np.array([0, 1, 0]) mu = self.centroidal_model.mu_friction - f_max = config.mpc_params['grf_max'] - f_min = config.mpc_params['grf_min'] + f_max = config.mpc_params["grf_max"] + f_min = config.mpc_params["grf_min"] # Derivation can be found in the paper # "High-slope terrain locomotion for torque-controlled quadruped robots", # M Focchi, A Del Prete, I Havoutis, R Featherstone, DG Caldwell, C Semini Jbu = cs.SX.zeros(20, 12) - Jbu[0, :3] = -n*mu + t - Jbu[1, :3] = -n*mu + b - Jbu[2, :3] = n*mu + b - Jbu[3, :3] = n*mu + t + Jbu[0, :3] = -n * mu + t + Jbu[1, :3] = -n * mu + b + Jbu[2, :3] = n * mu + b + Jbu[3, :3] = n * mu + t Jbu[4, :3] = n - Jbu[5, 3:6] = -n*mu + t - Jbu[6, 3:6] = -n*mu + b - Jbu[7, 3:6] = n*mu + b - Jbu[8, 3:6] = n*mu + t + Jbu[5, 3:6] = -n * mu + t + Jbu[6, 3:6] = -n * mu + b + Jbu[7, 3:6] = n * mu + b + Jbu[8, 3:6] = n * mu + t Jbu[9, 3:6] = n - - Jbu[10, 6:9] = -n*mu + t - Jbu[11, 6:9] = -n*mu + b - Jbu[12, 6:9] = n*mu + b - Jbu[13, 6:9] = n*mu + t + + Jbu[10, 6:9] = -n * mu + t + Jbu[11, 6:9] = -n * mu + b + Jbu[12, 6:9] = n * mu + b + Jbu[13, 6:9] = n * mu + t Jbu[14, 6:9] = n - - Jbu[15, 9:12] = -n*mu + t - Jbu[16, 9:12] = -n*mu + b - Jbu[17, 9:12] = n*mu + b - Jbu[18, 9:12] = n*mu + t + + Jbu[15, 9:12] = -n * mu + t + Jbu[16, 9:12] = -n * mu + b + Jbu[17, 9:12] = n * mu + b + Jbu[18, 9:12] = n * mu + t Jbu[19, 9:12] = n # C matrix - Jbu = Jbu@cs.vertcat(self.centroidal_model.states[30:42]) + Jbu = Jbu @ cs.vertcat(self.centroidal_model.states[30:42]) - # lower bound (-ACADOS_INFTY is "almost" -inf) lbu = np.zeros(20) lbu[0] = -ACADOS_INFTY lbu[1] = -ACADOS_INFTY - lbu[2] = 0 + lbu[2] = 0 lbu[3] = 0 lbu[4] = f_min lbu[5:10] = lbu[0:5] @@ -544,7 +537,7 @@ def create_friction_cone_constraints(self,): # upper bound (ACADOS_INFTY is "almost" inf) ubu = np.zeros(20) - ubu[0] = 0 + ubu[0] = 0 ubu[1] = 0 ubu[2] = ACADOS_INFTY ubu[3] = ACADOS_INFTY @@ -555,10 +548,6 @@ def create_friction_cone_constraints(self,): return Jbu, ubu, lbu - - - - def set_weight(self, nx: int, nu: int): """ Set the weight matrices for the optimization problem. @@ -572,61 +561,91 @@ def set_weight(self, nx: int, nu: int): """ - Q_position = np.array([0, 0, 1500]) #x, y, z - Q_velocity = np.array([200, 200, 200]) #x_vel, y_vel, z_vel - Q_base_angle = np.array([500, 500, 0]) #roll, pitch, yaw - Q_base_angle_rates = np.array([20, 20, 600]) #roll_rate, pitch_rate, yaw_rate - Q_foot_pos = np.array([300, 300, 300]) #f_x, f_y, f_z (should be 4 times this, once per foot) - Q_com_position_z_integral = np.array([50]) #integral of z_com - Q_com_velocity_x_integral = np.array([10]) #integral of x_com - Q_com_velocity_y_integral = np.array([10]) #integral of y_com - Q_com_velocity_z_integral = np.array([10]) #integral of z_com_vel - Q_roll_integral_integral = np.array([10]) #integral of roll - Q_pitch_integral_integral = np.array([10]) #integral of pitch - - - if(config.robot == "hyqreal"): - Q_foot_force = np.array([0.00001, 0.00001, 0.00001]) #f_x, f_y, f_z (should be 4 times this, once per foot) + Q_position = np.array([0, 0, 1500]) # x, y, z + Q_velocity = np.array([200, 200, 200]) # x_vel, y_vel, z_vel + Q_base_angle = np.array([500, 500, 0]) # roll, pitch, yaw + Q_base_angle_rates = np.array([20, 20, 600]) # roll_rate, pitch_rate, yaw_rate + Q_foot_pos = np.array([300, 300, 300]) # f_x, f_y, f_z (should be 4 times this, once per foot) + Q_com_position_z_integral = np.array([50]) # integral of z_com + Q_com_velocity_x_integral = np.array([10]) # integral of x_com + Q_com_velocity_y_integral = np.array([10]) # integral of y_com + Q_com_velocity_z_integral = np.array([10]) # integral of z_com_vel + Q_roll_integral_integral = np.array([10]) # integral of roll + Q_pitch_integral_integral = np.array([10]) # integral of pitch + + if config.robot == "hyqreal": + Q_foot_force = np.array( + [0.00001, 0.00001, 0.00001] + ) # f_x, f_y, f_z (should be 4 times this, once per foot) else: - Q_foot_force = np.array([0.0001, 0.0001, 0.0001]) #f_x, f_y, f_z (should be 4 times this, once per foot) - - R_foot_vel = np.array([0.0001, 0.0001, 0.00001]) #v_x, v_y, v_z (should be 4 times this, once per foot) - R_foot_rate_force = np.array([0.001, 0.001, 0.001]) #delta_f_x, delta_f_y, deltaf_z (should be 4 times this, once per foot) - - Q_mat = np.diag(np.concatenate((Q_position, Q_velocity, - Q_base_angle, Q_base_angle_rates, - Q_foot_pos, Q_foot_pos, Q_foot_pos, Q_foot_pos, - Q_com_position_z_integral, Q_com_velocity_x_integral, - Q_com_velocity_y_integral, Q_com_velocity_z_integral, - Q_roll_integral_integral, Q_pitch_integral_integral, - Q_foot_force, Q_foot_force, Q_foot_force, Q_foot_force))) - - R_mat = np.diag(np.concatenate((R_foot_vel, R_foot_vel, R_foot_vel, R_foot_vel, - R_foot_rate_force, R_foot_rate_force, R_foot_rate_force, R_foot_rate_force))) + Q_foot_force = np.array([0.0001, 0.0001, 0.0001]) # f_x, f_y, f_z (should be 4 times this, once per foot) + + R_foot_vel = np.array([0.0001, 0.0001, 0.00001]) # v_x, v_y, v_z (should be 4 times this, once per foot) + R_foot_rate_force = np.array( + [0.001, 0.001, 0.001] + ) # delta_f_x, delta_f_y, deltaf_z (should be 4 times this, once per foot) + + Q_mat = np.diag( + np.concatenate( + ( + Q_position, + Q_velocity, + Q_base_angle, + Q_base_angle_rates, + Q_foot_pos, + Q_foot_pos, + Q_foot_pos, + Q_foot_pos, + Q_com_position_z_integral, + Q_com_velocity_x_integral, + Q_com_velocity_y_integral, + Q_com_velocity_z_integral, + Q_roll_integral_integral, + Q_pitch_integral_integral, + Q_foot_force, + Q_foot_force, + Q_foot_force, + Q_foot_force, + ) + ) + ) + + R_mat = np.diag( + np.concatenate( + ( + R_foot_vel, + R_foot_vel, + R_foot_vel, + R_foot_vel, + R_foot_rate_force, + R_foot_rate_force, + R_foot_rate_force, + R_foot_rate_force, + ) + ) + ) return Q_mat, R_mat - - def reset(self): """ Resets the solver by calling the reset method of the acados_ocp_solver. - + Returns: None """ self.acados_ocp_solver.reset() - self.acados_ocp_solver = AcadosOcpSolver(self.ocp, - json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json", - build = False, generate = False) - - - + self.acados_ocp_solver = AcadosOcpSolver( + self.ocp, + json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json", + build=False, + generate=False, + ) def set_stage_constraint(self, constraint, state, reference, contact_sequence, h_R_w, stance_proximity): """ Set the stage constraint for the centroidal NMPC problem. We only consider the stance constraint, and the swing - constraint up to 2 maximum references. + constraint up to 2 maximum references. Args: constraint (numpy.ndarray or None): Constraint passed from outside (e.g. vision). If None, nominal foothold is used. @@ -640,9 +659,7 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h None """ try: - - - # Take the array of the contact sequence and split + # Take the array of the contact sequence and split # it in 4 arrays for clarity FL_contact_sequence = contact_sequence[0] FR_contact_sequence = contact_sequence[1] @@ -660,74 +677,67 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h RL_reference_foot = reference["ref_foot_RL"] RR_reference_foot = reference["ref_foot_RR"] - - # Take the base position and the yaw rotation matrix. This is needed to + # Take the base position and the yaw rotation matrix. This is needed to # express the foothold constraint in the horizontal frame base = state["position"] - - h_R_w = h_R_w.reshape((2,2)) - + h_R_w = h_R_w.reshape((2, 2)) # Divide the constraint in upper and lower bound. The constraint are - # represented by 4 vertex, but we only use box constraint hence + # represented by 4 vertex, but we only use box constraint hence # we need only 2 vertex for each constraint (first an last) # For the leg in stance now, we simply enlarge the actual position as a box constraint # maybe this is not needed, but we cannot disable constraint!! - + # FL Stance Constraint stance_up_constraint_FL = np.array([FL_actual_foot[0], FL_actual_foot[1], FL_actual_foot[2] + 0.002]) - stance_up_constraint_FL[0:2] = h_R_w@(stance_up_constraint_FL[0:2] - base[0:2]) + stance_up_constraint_FL[0:2] = h_R_w @ (stance_up_constraint_FL[0:2] - base[0:2]) stance_up_constraint_FL[0:2] = stance_up_constraint_FL[0:2] + 0.1 stance_up_constraint_FL[2] = stance_up_constraint_FL[2] + 0.01 - + stance_low_constraint_FL = np.array([FL_actual_foot[0], FL_actual_foot[1], FL_actual_foot[2] - 0.002]) - stance_low_constraint_FL[0:2] = h_R_w@(stance_low_constraint_FL[0:2] - base[0:2]) + stance_low_constraint_FL[0:2] = h_R_w @ (stance_low_constraint_FL[0:2] - base[0:2]) stance_low_constraint_FL[0:2] = stance_low_constraint_FL[0:2] - 0.1 stance_low_constraint_FL[2] = stance_low_constraint_FL[2] - 0.01 - # FR Stance Constraint stance_up_constraint_FR = np.array([FR_actual_foot[0], FR_actual_foot[1], FR_actual_foot[2] + 0.002]) - stance_up_constraint_FR[0:2] = h_R_w@(stance_up_constraint_FR[0:2] - base[0:2]) + stance_up_constraint_FR[0:2] = h_R_w @ (stance_up_constraint_FR[0:2] - base[0:2]) stance_up_constraint_FR[0:2] = stance_up_constraint_FR[0:2] + 0.1 stance_up_constraint_FR[2] = stance_up_constraint_FR[2] + 0.01 stance_low_constraint_FR = np.array([FR_actual_foot[0], FR_actual_foot[1], FR_actual_foot[2] - 0.002]) - stance_low_constraint_FR[0:2] = h_R_w@(stance_low_constraint_FR[0:2] - base[0:2]) + stance_low_constraint_FR[0:2] = h_R_w @ (stance_low_constraint_FR[0:2] - base[0:2]) stance_low_constraint_FR[0:2] = stance_low_constraint_FR[0:2] - 0.1 stance_low_constraint_FR[2] = stance_low_constraint_FR[2] - 0.01 - # RL Stance Constraint stance_up_constraint_RL = np.array([RL_actual_foot[0], RL_actual_foot[1], RL_actual_foot[2] + 0.002]) - stance_up_constraint_RL[0:2] = h_R_w@(stance_up_constraint_RL[0:2]- base[0:2]) + stance_up_constraint_RL[0:2] = h_R_w @ (stance_up_constraint_RL[0:2] - base[0:2]) stance_up_constraint_RL[0:2] = stance_up_constraint_RL[0:2] + 0.1 stance_up_constraint_RL[2] = stance_up_constraint_RL[2] + 0.01 - + stance_low_constraint_RL = np.array([RL_actual_foot[0], RL_actual_foot[1], RL_actual_foot[2] - 0.002]) - stance_low_constraint_RL[0:2] = h_R_w@(stance_low_constraint_RL[0:2]- base[0:2]) + stance_low_constraint_RL[0:2] = h_R_w @ (stance_low_constraint_RL[0:2] - base[0:2]) stance_low_constraint_RL[0:2] = stance_low_constraint_RL[0:2] - 0.1 stance_low_constraint_RL[2] = stance_low_constraint_RL[2] - 0.01 - # RR Stance Constraint stance_up_constraint_RR = np.array([RR_actual_foot[0], RR_actual_foot[1], RR_actual_foot[2] + 0.002]) - stance_up_constraint_RR[0:2] = h_R_w@(stance_up_constraint_RR[0:2]- base[0:2]) + stance_up_constraint_RR[0:2] = h_R_w @ (stance_up_constraint_RR[0:2] - base[0:2]) stance_up_constraint_RR[0:2] = stance_up_constraint_RR[0:2] + 0.1 stance_up_constraint_RR[2] = stance_up_constraint_RR[2] + 0.01 - + stance_low_constraint_RR = np.array([RR_actual_foot[0], RR_actual_foot[1], RR_actual_foot[2] - 0.002]) - stance_low_constraint_RR[0:2] = h_R_w@(stance_low_constraint_RR[0:2]- base[0:2]) + stance_low_constraint_RR[0:2] = h_R_w @ (stance_low_constraint_RR[0:2] - base[0:2]) stance_low_constraint_RR[0:2] = stance_low_constraint_RR[0:2] - 0.1 stance_low_constraint_RR[2] = stance_low_constraint_RR[2] - 0.01 - # Constraint for the first footholds at the next touchdown. If constraint == True # we have some constraint passed from outside (e.g. vision), otherwise we use the - # nominal foothold enlarged as we do previously - if(constraint is not None): + # nominal foothold enlarged as we do previously + if constraint is not None: # From the VFA first_up_constraint_FL = np.array([constraint[0][0], constraint[1][0], constraint[2][0] + 0.002]) first_up_constraint_FR = np.array([constraint[0][1], constraint[1][1], constraint[2][1] + 0.002]) @@ -738,75 +748,85 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h first_low_constraint_FR = np.array([constraint[9][1], constraint[10][1], constraint[11][1] - 0.002]) first_low_constraint_RL = np.array([constraint[9][2], constraint[10][2], constraint[11][2] - 0.002]) first_low_constraint_RR = np.array([constraint[9][3], constraint[10][3], constraint[11][3] - 0.002]) - + # Rotate the constraint in the horizontal frame - first_up_constraint_FL[0:2] = h_R_w@(first_up_constraint_FL[0:2] - base[0:2]) + first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - base[0:2]) first_up_constraint_FL = first_up_constraint_FL + 0.005 - #first_up_constraint_FL[2] = first_up_constraint_FL[2] + 0.5 - - first_up_constraint_FR[0:2] = h_R_w@(first_up_constraint_FR[0:2] - base[0:2]) + # first_up_constraint_FL[2] = first_up_constraint_FL[2] + 0.5 + + first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - base[0:2]) first_up_constraint_FR = first_up_constraint_FR + 0.005 - #first_up_constraint_FR[2] = first_up_constraint_FR[2] + 0.5 + # first_up_constraint_FR[2] = first_up_constraint_FR[2] + 0.5 - first_up_constraint_RL[0:2] = h_R_w@(first_up_constraint_RL[0:2] - base[0:2]) + first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - base[0:2]) first_up_constraint_RL = first_up_constraint_RL + 0.005 - #first_up_constraint_RL[2] = first_up_constraint_RL[2] + 0.5 + # first_up_constraint_RL[2] = first_up_constraint_RL[2] + 0.5 - first_up_constraint_RR[0:2] = h_R_w@(first_up_constraint_RR[0:2] - base[0:2]) + first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - base[0:2]) first_up_constraint_RR = first_up_constraint_RR + 0.005 - #first_up_constraint_RR[2] = first_up_constraint_RR[2] + 0.5 + # first_up_constraint_RR[2] = first_up_constraint_RR[2] + 0.5 - first_low_constraint_FL[0:2] = h_R_w@(first_low_constraint_FL[0:2] - base[0:2]) + first_low_constraint_FL[0:2] = h_R_w @ (first_low_constraint_FL[0:2] - base[0:2]) first_low_constraint_FL = first_low_constraint_FL - 0.005 - #first_low_constraint_FL[2] = first_low_constraint_FL[2] - 0.5 + # first_low_constraint_FL[2] = first_low_constraint_FL[2] - 0.5 - first_low_constraint_FR[0:2] = h_R_w@(first_low_constraint_FR[0:2] - base[0:2]) + first_low_constraint_FR[0:2] = h_R_w @ (first_low_constraint_FR[0:2] - base[0:2]) first_low_constraint_FR = first_low_constraint_FR - 0.005 - #first_low_constraint_FR[2] = first_low_constraint_FR[2] - 0.5 - - first_low_constraint_RL[0:2] = h_R_w@(first_low_constraint_RL[0:2] - base[0:2]) + # first_low_constraint_FR[2] = first_low_constraint_FR[2] - 0.5 + + first_low_constraint_RL[0:2] = h_R_w @ (first_low_constraint_RL[0:2] - base[0:2]) first_low_constraint_RL = first_low_constraint_RL - 0.005 - #first_low_constraint_RL[2] = first_low_constraint_RL[2] - 0.5 + # first_low_constraint_RL[2] = first_low_constraint_RL[2] - 0.5 - first_low_constraint_RR[0:2] = h_R_w@(first_low_constraint_RR[0:2] - base[0:2]) + first_low_constraint_RR[0:2] = h_R_w @ (first_low_constraint_RR[0:2] - base[0:2]) first_low_constraint_RR = first_low_constraint_RR - 0.005 - #first_low_constraint_RR[2] = first_low_constraint_RR[2] - 0.5 + # first_low_constraint_RR[2] = first_low_constraint_RR[2] - 0.5 else: # Constrain taken from the nominal foothold (NO VISION) # FL first touchdown constraint - first_up_constraint_FL = np.array([FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] + 0.002]) - first_up_constraint_FL[0:2] = h_R_w@(first_up_constraint_FL[0:2] - base[0:2]) + 0.15 - - first_low_constraint_FL = np.array([FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] - 0.002]) - first_low_constraint_FL[0:2] = h_R_w@(first_low_constraint_FL[0:2] - base[0:2]) - 0.15 - + first_up_constraint_FL = np.array( + [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] + 0.002] + ) + first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - base[0:2]) + 0.15 - # FR first touchdown constraint - first_up_constraint_FR = np.array([FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] + 0.002]) - first_up_constraint_FR[0:2] = h_R_w@(first_up_constraint_FR[0:2] - base[0:2]) + 0.15 + first_low_constraint_FL = np.array( + [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] - 0.002] + ) + first_low_constraint_FL[0:2] = h_R_w @ (first_low_constraint_FL[0:2] - base[0:2]) - 0.15 - first_low_constraint_FR = np.array([FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] - 0.002]) - first_low_constraint_FR[0:2] = h_R_w@(first_low_constraint_FR[0:2] - base[0:2]) - 0.15 + # FR first touchdown constraint + first_up_constraint_FR = np.array( + [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] + 0.002] + ) + first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - base[0:2]) + 0.15 + first_low_constraint_FR = np.array( + [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] - 0.002] + ) + first_low_constraint_FR[0:2] = h_R_w @ (first_low_constraint_FR[0:2] - base[0:2]) - 0.15 # RL first touchdown constraint - first_up_constraint_RL = np.array([RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] + 0.002]) - first_up_constraint_RL[0:2] = h_R_w@(first_up_constraint_RL[0:2] - base[0:2]) + 0.15 - - first_low_constraint_RL = np.array([RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] - 0.002]) - first_low_constraint_RL[0:2] = h_R_w@(first_low_constraint_RL[0:2] - base[0:2]) - 0.15 + first_up_constraint_RL = np.array( + [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] + 0.002] + ) + first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - base[0:2]) + 0.15 + first_low_constraint_RL = np.array( + [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] - 0.002] + ) + first_low_constraint_RL[0:2] = h_R_w @ (first_low_constraint_RL[0:2] - base[0:2]) - 0.15 # RR first touchdown constraint - first_up_constraint_RR = np.array([RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] + 0.002]) - first_up_constraint_RR[0:2] = h_R_w@(first_up_constraint_RR[0:2] - base[0:2]) + 0.15 - - first_low_constraint_RR = np.array([RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] - 0.002]) - first_low_constraint_RR[0:2] = h_R_w@(first_low_constraint_RR[0:2] - base[0:2]) - 0.15 + first_up_constraint_RR = np.array( + [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] + 0.002] + ) + first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - base[0:2]) + 0.15 - - + first_low_constraint_RR = np.array( + [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] - 0.002] + ) + first_low_constraint_RR[0:2] = h_R_w @ (first_low_constraint_RR[0:2] - base[0:2]) - 0.15 # we stack all the constraint we have for now up_constraint_FL = np.vstack((stance_up_constraint_FL, first_up_constraint_FL)) @@ -819,64 +839,69 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h low_constraint_RL = np.vstack((stance_low_constraint_RL, first_low_constraint_RL)) low_constraint_RR = np.vstack((stance_low_constraint_RR, first_low_constraint_RR)) - - # If there are more than 1 nominal foothold per leg, we create additional constraints # We do not expect more than two reference footholds... # FL second touchdown constraint - if(FL_reference_foot.shape[0] == 2): - second_up_constraint_FL = np.array([FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] + 0.002]) - second_up_constraint_FL[0:2] = h_R_w@(second_up_constraint_FL[0:2] - base[0:2]) + 0.15 - - second_low_constraint_FL = np.array([FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] - 0.002]) - second_low_constraint_FL[0:2] = h_R_w@(second_low_constraint_FL[0:2] - base[0:2]) - 0.15 + if FL_reference_foot.shape[0] == 2: + second_up_constraint_FL = np.array( + [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] + 0.002] + ) + second_up_constraint_FL[0:2] = h_R_w @ (second_up_constraint_FL[0:2] - base[0:2]) + 0.15 + second_low_constraint_FL = np.array( + [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] - 0.002] + ) + second_low_constraint_FL[0:2] = h_R_w @ (second_low_constraint_FL[0:2] - base[0:2]) - 0.15 up_constraint_FL = np.vstack((up_constraint_FL, second_up_constraint_FL)) low_constraint_FL = np.vstack((low_constraint_FL, second_low_constraint_FL)) # FR second touchdown constraint - if(FR_reference_foot.shape[0] == 2): - second_up_constraint_FR = np.array([FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] + 0.002]) - second_up_constraint_FR[0:2] = h_R_w@(second_up_constraint_FR[0:2] - base[0:2]) + 0.15 - - - second_low_constraint_FR = np.array([FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] - 0.002]) - second_low_constraint_FR[0:2] = h_R_w@(second_low_constraint_FR[0:2]- base[0:2]) - 0.15 + if FR_reference_foot.shape[0] == 2: + second_up_constraint_FR = np.array( + [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] + 0.002] + ) + second_up_constraint_FR[0:2] = h_R_w @ (second_up_constraint_FR[0:2] - base[0:2]) + 0.15 + second_low_constraint_FR = np.array( + [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] - 0.002] + ) + second_low_constraint_FR[0:2] = h_R_w @ (second_low_constraint_FR[0:2] - base[0:2]) - 0.15 up_constraint_FR = np.vstack((up_constraint_FR, second_up_constraint_FR)) low_constraint_FR = np.vstack((low_constraint_FR, second_low_constraint_FR)) - - # RL second touchdown constraint - if(RL_reference_foot.shape[0] == 2): - second_up_constraint_RL = np.array([RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] + 0.002]) - second_up_constraint_RL[0:2] = h_R_w@(second_up_constraint_RL[0:2]- base[0:2]) + 0.15 - - second_low_constraint_RL = np.array([RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] - 0.002]) - second_low_constraint_RL[0:2] = h_R_w@(second_low_constraint_RL[0:2] - base[0:2]) - 0.15 + # RL second touchdown constraint + if RL_reference_foot.shape[0] == 2: + second_up_constraint_RL = np.array( + [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] + 0.002] + ) + second_up_constraint_RL[0:2] = h_R_w @ (second_up_constraint_RL[0:2] - base[0:2]) + 0.15 + second_low_constraint_RL = np.array( + [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] - 0.002] + ) + second_low_constraint_RL[0:2] = h_R_w @ (second_low_constraint_RL[0:2] - base[0:2]) - 0.15 up_constraint_RL = np.vstack((up_constraint_RL, second_up_constraint_RL)) low_constraint_RL = np.vstack((low_constraint_RL, second_low_constraint_RL)) - - # RR second touchdown constraint - if(RR_reference_foot.shape[0] == 2): - second_up_constraint_RR = np.array([RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] + 0.002]) - second_up_constraint_RR[0:2] = h_R_w@(second_up_constraint_RR[0:2] - base[0:2]) + 0.15 - - second_low_constraint_RR = np.array([RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] - 0.002]) - second_low_constraint_RR[0:2] = h_R_w@(second_low_constraint_RR[0:2] - base[0:2]) - 0.15 + # RR second touchdown constraint + if RR_reference_foot.shape[0] == 2: + second_up_constraint_RR = np.array( + [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] + 0.002] + ) + second_up_constraint_RR[0:2] = h_R_w @ (second_up_constraint_RR[0:2] - base[0:2]) + 0.15 + second_low_constraint_RR = np.array( + [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] - 0.002] + ) + second_low_constraint_RR[0:2] = h_R_w @ (second_low_constraint_RR[0:2] - base[0:2]) - 0.15 up_constraint_RR = np.vstack((up_constraint_RR, second_up_constraint_RR)) low_constraint_RR = np.vstack((low_constraint_RR, second_low_constraint_RR)) - - # We pass all the constraints (foothold and friction conte) to acados. # The one regarding friction are precomputed ub_friction = self.constr_uh_friction @@ -884,48 +909,42 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h # If the foothold is in swing, the idx of the constraint start from 1 idx_constraint = np.array([0, 0, 0, 0]) - if(FL_contact_sequence[0] == 0): + if FL_contact_sequence[0] == 0: idx_constraint[0] = 1 - if(FR_contact_sequence[0] == 0): + if FR_contact_sequence[0] == 0: idx_constraint[1] = 1 - if(RL_contact_sequence[0] == 0): + if RL_contact_sequence[0] == 0: idx_constraint[2] = 1 - if(RR_contact_sequence[0] == 0): + if RR_contact_sequence[0] == 0: idx_constraint[3] = 1 - - for j in range(0, self.horizon): - # take the constraint for the current timestep - ub_foot_FL = up_constraint_FL[idx_constraint[0]] - lb_foot_FL = low_constraint_FL[idx_constraint[0]] + for j in range(0, self.horizon): + # take the constraint for the current timestep + ub_foot_FL = up_constraint_FL[idx_constraint[0]] + lb_foot_FL = low_constraint_FL[idx_constraint[0]] - ub_foot_FR = up_constraint_FR[idx_constraint[1]] + ub_foot_FR = up_constraint_FR[idx_constraint[1]] lb_foot_FR = low_constraint_FR[idx_constraint[1]] - ub_foot_RL = up_constraint_RL[idx_constraint[2]] + ub_foot_RL = up_constraint_RL[idx_constraint[2]] lb_foot_RL = low_constraint_RL[idx_constraint[2]] - ub_foot_RR = up_constraint_RR[idx_constraint[3]] + ub_foot_RR = up_constraint_RR[idx_constraint[3]] lb_foot_RR = low_constraint_RR[idx_constraint[3]] - - - # Concatenate the friction and foothold constraint - ub_foot = copy.deepcopy(np.concatenate((ub_foot_FL, ub_foot_FR, - ub_foot_RL, ub_foot_RR))) - lb_foot = copy.deepcopy(np.concatenate((lb_foot_FL, lb_foot_FR, - lb_foot_RL, lb_foot_RR))) - if(self.use_foothold_constraints): + # Concatenate the friction and foothold constraint + ub_foot = copy.deepcopy(np.concatenate((ub_foot_FL, ub_foot_FR, ub_foot_RL, ub_foot_RR))) + lb_foot = copy.deepcopy(np.concatenate((lb_foot_FL, lb_foot_FR, lb_foot_RL, lb_foot_RR))) + if self.use_foothold_constraints: ub_total = np.concatenate((ub_friction, ub_foot)) lb_total = np.concatenate((lb_friction, lb_foot)) else: ub_total = ub_friction lb_total = lb_friction - - #Constraints for the support polygon depending on the leg in stance + # Constraints for the support polygon depending on the leg in stance # all disabled at the beginning!! - if(self.use_stability_constraints): + if self.use_stability_constraints: ub_support_FL_FR = ACADOS_INFTY lb_support_FL_FR = -ACADOS_INFTY @@ -937,20 +956,21 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h ub_support_RL_FL = ACADOS_INFTY lb_support_RL_FL = -ACADOS_INFTY - + ub_support_FL_RR = ACADOS_INFTY lb_support_FL_RR = -ACADOS_INFTY ub_support_FR_RL = ACADOS_INFTY - lb_support_FR_RL = -ACADOS_INFTY - - + lb_support_FR_RL = -ACADOS_INFTY + # We have 4 cases for the stability constraint: trot, pace, crawl, full stance - if(FL_contact_sequence[j] == 1 and - FR_contact_sequence[j] == 1 and - RL_contact_sequence[j] == 1 and - RR_contact_sequence[j] == 1): - #FULL STANCE TODO + if ( + FL_contact_sequence[j] == 1 + and FR_contact_sequence[j] == 1 + and RL_contact_sequence[j] == 1 + and RR_contact_sequence[j] == 1 + ): + # FULL STANCE TODO ub_support_FL_FR = 0 lb_support_FL_FR = -ACADOS_INFTY @@ -962,148 +982,158 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h ub_support_RL_FL = 0 lb_support_RL_FL = -ACADOS_INFTY - + ub_support_FL_RR = ACADOS_INFTY lb_support_FL_RR = -ACADOS_INFTY ub_support_FR_RL = ACADOS_INFTY lb_support_FR_RL = -ACADOS_INFTY - elif(np.array_equal(FL_contact_sequence, RR_contact_sequence) - and np.array_equal(FR_contact_sequence, RL_contact_sequence)): - #TROT - stability_margin = config.mpc_params['trot_stability_margin'] - if(FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0): + elif np.array_equal(FL_contact_sequence, RR_contact_sequence) and np.array_equal( + FR_contact_sequence, RL_contact_sequence + ): + # TROT + stability_margin = config.mpc_params["trot_stability_margin"] + if FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0: ub_support_FL_RR = 0 + stability_margin lb_support_FL_RR = 0 - stability_margin - - if(FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0): + + if FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0: ub_support_FR_RL = 0 + stability_margin lb_support_FR_RL = 0 - stability_margin - elif(np.array_equal(FL_contact_sequence, RL_contact_sequence) - and np.array_equal(FR_contact_sequence, RR_contact_sequence)): - #PACE - stability_margin = config.mpc_params['pace_stability_margin'] - if(FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0): + elif np.array_equal(FL_contact_sequence, RL_contact_sequence) and np.array_equal( + FR_contact_sequence, RR_contact_sequence + ): + # PACE + stability_margin = config.mpc_params["pace_stability_margin"] + if FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0: ub_support_RL_FL = 0 + stability_margin lb_support_RL_FL = 0 - stability_margin - - if(FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0): + if FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0: ub_support_FR_RR = 0 + stability_margin lb_support_FR_RR = 0 - stability_margin - + else: - #CRAWL BACKDIAGONALCRAWL ONLY - stability_margin = config.mpc_params['crawl_stability_margin'] + # CRAWL BACKDIAGONALCRAWL ONLY + stability_margin = config.mpc_params["crawl_stability_margin"] - if(FL_contact_sequence[j] == 1): - if(FR_contact_sequence[j] == 1): - ub_support_FL_FR = -0.0 - stability_margin + if FL_contact_sequence[j] == 1: + if FR_contact_sequence[j] == 1: + ub_support_FL_FR = -0.0 - stability_margin lb_support_FL_FR = -ACADOS_INFTY else: ub_support_FL_RR = ACADOS_INFTY - lb_support_FL_RR = 0.0 + stability_margin + lb_support_FL_RR = 0.0 + stability_margin - - if(FR_contact_sequence[j] == 1): - if(RR_contact_sequence[j] == 1): + if FR_contact_sequence[j] == 1: + if RR_contact_sequence[j] == 1: ub_support_FR_RR = ACADOS_INFTY lb_support_FR_RR = 0.0 + stability_margin else: ub_support_FR_RL = ACADOS_INFTY lb_support_FR_RL = 0.0 + stability_margin - - if(RR_contact_sequence[j] == 1): - if(RL_contact_sequence[j] == 1): + if RR_contact_sequence[j] == 1: + if RL_contact_sequence[j] == 1: ub_support_RR_RL = ACADOS_INFTY - lb_support_RR_RL = 0.0 + stability_margin + lb_support_RR_RL = 0.0 + stability_margin else: - ub_support_FL_RR = -0.0 - stability_margin + ub_support_FL_RR = -0.0 - stability_margin lb_support_FL_RR = -ACADOS_INFTY - - if(RL_contact_sequence[j] == 1): - if(FL_contact_sequence[j] == 1): - ub_support_RL_FL = -0.0 - stability_margin + if RL_contact_sequence[j] == 1: + if FL_contact_sequence[j] == 1: + ub_support_RL_FL = -0.0 - stability_margin lb_support_RL_FL = -ACADOS_INFTY else: ub_support_FR_RL = -0.0 - stability_margin lb_support_FR_RL = -ACADOS_INFTY + ub_support = np.array( + [ + ub_support_FL_FR, + ub_support_FR_RR, + ub_support_RR_RL, + ub_support_RL_FL, + ub_support_FL_RR, + ub_support_FR_RL, + ] + ) + lb_support = np.array( + [ + lb_support_FL_FR, + lb_support_FR_RR, + lb_support_RR_RL, + lb_support_RL_FL, + lb_support_FL_RR, + lb_support_FR_RL, + ] + ) - - - - ub_support = np.array([ub_support_FL_FR, ub_support_FR_RR, ub_support_RR_RL, ub_support_RL_FL, - ub_support_FL_RR, ub_support_FR_RL]) - lb_support = np.array([lb_support_FL_FR, lb_support_FR_RR, lb_support_RR_RL, lb_support_RL_FL, - lb_support_FL_RR, lb_support_FR_RL]) - ub_total = np.concatenate((ub_total, ub_support)) lb_total = np.concatenate((lb_total, lb_support)) - # No friction constraint at the end! we don't have u_N - if(j == self.horizon): - if(self.use_foothold_constraints): - if(self.use_stability_constraints): + if j == self.horizon: + if self.use_foothold_constraints: + if self.use_stability_constraints: ub_total = np.concatenate((ub_foot, ub_support)) lb_total = np.concatenate((lb_foot, lb_support)) else: ub_total = ub_foot lb_total = lb_foot else: - if(self.use_stability_constraints): + if self.use_stability_constraints: ub_total = ub_support lb_total = lb_support else: continue - - if(j > 0): + + if j > 0: self.acados_ocp_solver.constraints_set(j, "uh", ub_total) self.acados_ocp_solver.constraints_set(j, "lh", lb_total) - - #save the constraint for logging + # save the constraint for logging self.upper_bound[j] = ub_total.tolist() self.lower_bound[j] = lb_total.tolist() - - # ugly procedure to update the idx of the constraint - if(j>=1): - if(FL_contact_sequence[j] == 0 and FL_contact_sequence[j-1] == 1): - if(idx_constraint[0] < up_constraint_FL.shape[0] - 1): + if j >= 1: + if FL_contact_sequence[j] == 0 and FL_contact_sequence[j - 1] == 1: + if idx_constraint[0] < up_constraint_FL.shape[0] - 1: idx_constraint[0] += 1 - - if(FR_contact_sequence[j] == 0 and FR_contact_sequence[j-1] == 1): - if(idx_constraint[1] < up_constraint_FR.shape[0] - 1): + + if FR_contact_sequence[j] == 0 and FR_contact_sequence[j - 1] == 1: + if idx_constraint[1] < up_constraint_FR.shape[0] - 1: idx_constraint[1] += 1 - - if(RL_contact_sequence[j] == 0 and RL_contact_sequence[j-1] == 1): - if(idx_constraint[2] < up_constraint_RL.shape[0] - 1): + + if RL_contact_sequence[j] == 0 and RL_contact_sequence[j - 1] == 1: + if idx_constraint[2] < up_constraint_RL.shape[0] - 1: idx_constraint[2] += 1 - - if(RR_contact_sequence[j] == 0 and RR_contact_sequence[j-1] == 1): - if(idx_constraint[3] < up_constraint_RR.shape[0] - 1): + + if RR_contact_sequence[j] == 0 and RR_contact_sequence[j - 1] == 1: + if idx_constraint[3] < up_constraint_RR.shape[0] - 1: idx_constraint[3] += 1 - + except: - if(self.verbose): + if self.verbose: print("###WARNING: error in setting the constraints") return - - - - - def set_warm_start(self, state_acados, reference, FL_contact_sequence, FR_contact_sequence, RL_contact_sequence, RR_contact_sequence): + def set_warm_start( + self, + state_acados, + reference, + FL_contact_sequence, + FR_contact_sequence, + RL_contact_sequence, + RR_contact_sequence, + ): """ - Sets the warm start for the optimization problem. This could be useful in the case + Sets the warm start for the optimization problem. This could be useful in the case some old guess is outside some new constraints.. Maybe we could have some infeasibility problem. Still to investigate.. @@ -1116,83 +1146,82 @@ def set_warm_start(self, state_acados, reference, FL_contact_sequence, FR_contac RR_contact_sequence (np.ndarray): Contact sequence for the rear right leg. """ idx_ref_foot_to_assign = np.array([0, 0, 0, 0]) - + for j in range(self.horizon): # We only want to modify the warm start of the footholds... - # hence i take the current warm start from acados and + # hence i take the current warm start from acados and # modify onlya subset of it warm_start = copy.deepcopy(self.acados_ocp_solver.get(j, "x")) - # Update the idx_ref_foot_to_assign. Every time there is a change in the contact phase # between 1 and 0, it means that the leg go into swing and a new reference is needed!!! - if(j > 1 and j idx_ref_foot_to_assign[0]+1): + if j > 1 and j < self.horizon - 1: + if FL_contact_sequence[j] == 0 and FL_contact_sequence[j - 1] == 1: + if reference["ref_foot_FL"].shape[0] > idx_ref_foot_to_assign[0] + 1: idx_ref_foot_to_assign[0] += 1 - if(FR_contact_sequence[j] == 0 and FR_contact_sequence[j-1] == 1): - if(reference['ref_foot_FR'].shape[0] > idx_ref_foot_to_assign[1]+1): + if FR_contact_sequence[j] == 0 and FR_contact_sequence[j - 1] == 1: + if reference["ref_foot_FR"].shape[0] > idx_ref_foot_to_assign[1] + 1: idx_ref_foot_to_assign[1] += 1 - if(RL_contact_sequence[j] == 0 and RL_contact_sequence[j-1] == 1): - if(reference['ref_foot_RL'].shape[0] > idx_ref_foot_to_assign[2]+1): + if RL_contact_sequence[j] == 0 and RL_contact_sequence[j - 1] == 1: + if reference["ref_foot_RL"].shape[0] > idx_ref_foot_to_assign[2] + 1: idx_ref_foot_to_assign[2] += 1 - if(RR_contact_sequence[j] == 0 and RR_contact_sequence[j-1] == 1): - if(reference['ref_foot_RR'].shape[0] > idx_ref_foot_to_assign[3]+1): - idx_ref_foot_to_assign[3] += 1 - + if RR_contact_sequence[j] == 0 and RR_contact_sequence[j - 1] == 1: + if reference["ref_foot_RR"].shape[0] > idx_ref_foot_to_assign[3] + 1: + idx_ref_foot_to_assign[3] += 1 + warm_start[8] = state_acados[8] - - if(idx_ref_foot_to_assign[0] == 0): + + if idx_ref_foot_to_assign[0] == 0: warm_start[12:15] = state_acados[12:15].reshape((3,)) else: warm_start[12:15] = reference["ref_foot_FL"][0] - if(idx_ref_foot_to_assign[1] == 0): + if idx_ref_foot_to_assign[1] == 0: warm_start[15:18] = state_acados[15:18].reshape((3,)) else: warm_start[15:18] = reference["ref_foot_FR"][0] - if(idx_ref_foot_to_assign[2] == 0): + if idx_ref_foot_to_assign[2] == 0: warm_start[18:21] = state_acados[18:21].reshape((3,)) else: warm_start[18:21] = reference["ref_foot_RL"][0] - if(idx_ref_foot_to_assign[3] == 0): + if idx_ref_foot_to_assign[3] == 0: warm_start[21:24] = state_acados[21:24].reshape((3,)) else: warm_start[21:24] = reference["ref_foot_RR"][0] - self.acados_ocp_solver.set(j, "x", warm_start) - # Method to perform the centering of the states and the reference around (0, 0, 0) - def perform_scaling(self, state, reference, constraint = None): - - + def perform_scaling(self, state, reference, constraint=None): self.initial_base_position = copy.deepcopy(state["position"]) reference = copy.deepcopy(reference) state = copy.deepcopy(state) - + reference["ref_position"] = reference["ref_position"] - state["position"] reference["ref_foot_FL"] = reference["ref_foot_FL"] - state["position"] reference["ref_foot_FR"] = reference["ref_foot_FR"] - state["position"] reference["ref_foot_RL"] = reference["ref_foot_RL"] - state["position"] reference["ref_foot_RR"] = reference["ref_foot_RR"] - state["position"] - state["foot_FL"] = state["foot_FL"] - state["position"] state["foot_FR"] = state["foot_FR"] - state["position"] state["foot_RL"] = state["foot_RL"] - state["position"] state["foot_RR"] = state["foot_RR"] - state["position"] state["position"] = np.array([0, 0, 0]) - return state, reference, constraint - + return state, reference, constraint # Main loop for computing the control - def compute_control(self, state, reference, contact_sequence, constraint = None, external_wrenches = np.zeros((6,)), - inertia = config.inertia.reshape((9,)), mass = config.mass): - - - # Take the array of the contact sequence and split it in 4 arrays, + def compute_control( + self, + state, + reference, + contact_sequence, + constraint=None, + external_wrenches=np.zeros((6,)), + inertia=config.inertia.reshape((9,)), + mass=config.mass, + ): + # Take the array of the contact sequence and split it in 4 arrays, # one for each leg FL_contact_sequence = contact_sequence[0] FR_contact_sequence = contact_sequence[1] @@ -1204,55 +1233,49 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, RL_previous_contact_sequence = self.previous_contact_sequence[2] RR_previous_contact_sequence = self.previous_contact_sequence[3] - # Perform the scaling of the states and the reference - state, \ - reference, \ - constraint = self.perform_scaling(state, reference, constraint) - + state, reference, constraint = self.perform_scaling(state, reference, constraint) # Fill reference (self.states_dim+self.inputs_dim) idx_ref_foot_to_assign = np.array([0, 0, 0, 0]) for j in range(self.horizon): - yref = np.zeros(shape=(self.states_dim + self.inputs_dim,)) - yref[0:3] = reference['ref_position'] - yref[3:6] = reference['ref_linear_velocity'] - yref[6:9] = reference['ref_orientation'] - yref[9:12] = reference['ref_angular_velocity'] - yref[12:15] = reference['ref_foot_FL'][idx_ref_foot_to_assign[0]] - yref[15:18] = reference['ref_foot_FR'][idx_ref_foot_to_assign[1]] - yref[18:21] = reference['ref_foot_RL'][idx_ref_foot_to_assign[2]] - yref[21:24] = reference['ref_foot_RR'][idx_ref_foot_to_assign[3]] - + yref[0:3] = reference["ref_position"] + yref[3:6] = reference["ref_linear_velocity"] + yref[6:9] = reference["ref_orientation"] + yref[9:12] = reference["ref_angular_velocity"] + yref[12:15] = reference["ref_foot_FL"][idx_ref_foot_to_assign[0]] + yref[15:18] = reference["ref_foot_FR"][idx_ref_foot_to_assign[1]] + yref[18:21] = reference["ref_foot_RL"][idx_ref_foot_to_assign[2]] + yref[21:24] = reference["ref_foot_RR"][idx_ref_foot_to_assign[3]] # Update the idx_ref_foot_to_assign. Every time there is a change in the contact phase # between 1 and 0, it means that the leg go into swing and a new reference is needed!!! - if(j > 1 and j idx_ref_foot_to_assign[0]+1): + if j > 1 and j < self.horizon - 1: + if FL_contact_sequence[j + 1] == 0 and FL_contact_sequence[j] == 1: + if reference["ref_foot_FL"].shape[0] > idx_ref_foot_to_assign[0] + 1: idx_ref_foot_to_assign[0] += 1 - if(FR_contact_sequence[j+1] == 0 and FR_contact_sequence[j] == 1): - if(reference['ref_foot_FR'].shape[0] > idx_ref_foot_to_assign[1]+1): + if FR_contact_sequence[j + 1] == 0 and FR_contact_sequence[j] == 1: + if reference["ref_foot_FR"].shape[0] > idx_ref_foot_to_assign[1] + 1: idx_ref_foot_to_assign[1] += 1 - if(RL_contact_sequence[j+1] == 0 and RL_contact_sequence[j] == 1): - if(reference['ref_foot_RL'].shape[0] > idx_ref_foot_to_assign[2]+1): + if RL_contact_sequence[j + 1] == 0 and RL_contact_sequence[j] == 1: + if reference["ref_foot_RL"].shape[0] > idx_ref_foot_to_assign[2] + 1: idx_ref_foot_to_assign[2] += 1 - if(RR_contact_sequence[j+1] == 0 and RR_contact_sequence[j] == 1): - if(reference['ref_foot_RR'].shape[0] > idx_ref_foot_to_assign[3]+1): - idx_ref_foot_to_assign[3] += 1 - + if RR_contact_sequence[j + 1] == 0 and RR_contact_sequence[j] == 1: + if reference["ref_foot_RR"].shape[0] > idx_ref_foot_to_assign[3] + 1: + idx_ref_foot_to_assign[3] += 1 # Calculate the reference force z for the leg in stance # It's simply mass*acc/number_of_legs_in_stance!! # Force x and y are always 0 - number_of_legs_in_stance = np.array([FL_contact_sequence[j], FR_contact_sequence[j], - RL_contact_sequence [j], RR_contact_sequence[j]]).sum() - if(number_of_legs_in_stance == 0): + number_of_legs_in_stance = np.array( + [FL_contact_sequence[j], FR_contact_sequence[j], RL_contact_sequence[j], RR_contact_sequence[j]] + ).sum() + if number_of_legs_in_stance == 0: reference_force_stance_legs = 0 else: reference_force_stance_legs = (mass * 9.81) / number_of_legs_in_stance - + reference_force_fl_z = reference_force_stance_legs * FL_contact_sequence[j] reference_force_fr_z = reference_force_stance_legs * FR_contact_sequence[j] reference_force_rl_z = reference_force_stance_legs * RL_contact_sequence[j] @@ -1261,30 +1284,36 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, yref[35] = reference_force_fr_z yref[38] = reference_force_rl_z yref[41] = reference_force_rr_z - + # Setting the reference to acados - if(self.use_DDP): - if(j == 0): + if self.use_DDP: + if j == 0: num_l2_penalties = self.ocp.model.cost_y_expr_0.shape[0] - (self.states_dim + self.inputs_dim) else: num_l2_penalties = self.ocp.model.cost_y_expr.shape[0] - (self.states_dim + self.inputs_dim) - - yref_tot = np.concatenate((yref, np.zeros(num_l2_penalties,) )) + + yref_tot = np.concatenate( + ( + yref, + np.zeros( + num_l2_penalties, + ), + ) + ) self.acados_ocp_solver.set(j, "yref", yref_tot) else: self.acados_ocp_solver.set(j, "yref", yref) - # Fill last step horizon reference (self.states_dim - no control action!!) - yref_N = np.zeros(shape=(self.states_dim ,)) - yref_N[0:3] = reference['ref_position'] - yref_N[3:6] = reference['ref_linear_velocity'] - yref_N[6:9] = reference['ref_orientation'] - yref_N[9:12] = reference['ref_angular_velocity'] - yref_N[12:15] = reference['ref_foot_FL'][idx_ref_foot_to_assign[0]] - yref_N[15:18] = reference['ref_foot_FR'][idx_ref_foot_to_assign[1]] - yref_N[18:21] = reference['ref_foot_RL'][idx_ref_foot_to_assign[2]] - yref_N[21:24] = reference['ref_foot_RR'][idx_ref_foot_to_assign[3]] + yref_N = np.zeros(shape=(self.states_dim,)) + yref_N[0:3] = reference["ref_position"] + yref_N[3:6] = reference["ref_linear_velocity"] + yref_N[6:9] = reference["ref_orientation"] + yref_N[9:12] = reference["ref_angular_velocity"] + yref_N[12:15] = reference["ref_foot_FL"][idx_ref_foot_to_assign[0]] + yref_N[15:18] = reference["ref_foot_FR"][idx_ref_foot_to_assign[1]] + yref_N[18:21] = reference["ref_foot_RL"][idx_ref_foot_to_assign[2]] + yref_N[21:24] = reference["ref_foot_RR"][idx_ref_foot_to_assign[3]] yref_N[32] = reference_force_fl_z yref_N[35] = reference_force_fr_z yref_N[38] = reference_force_rl_z @@ -1292,119 +1321,132 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, # Setting the reference to acados self.acados_ocp_solver.set(self.horizon, "yref", yref_N) - - - - # Fill stance param, friction and stance proximity + # Fill stance param, friction and stance proximity # (stance proximity will disable foothold optimization near a stance!!) - mu = config.mpc_params['mu'] + mu = config.mpc_params["mu"] yaw = state["orientation"][2] - + # Stance Proximity ugly routine. Basically we disable foothold optimization # in the proximity of a stance phase (the real foot cannot travel too fast in # a small time!!) - stance_proximity_FL = np.zeros((self.horizon, )) - stance_proximity_FR = np.zeros((self.horizon, )) - stance_proximity_RL = np.zeros((self.horizon, )) - stance_proximity_RR = np.zeros((self.horizon, )) + stance_proximity_FL = np.zeros((self.horizon,)) + stance_proximity_FR = np.zeros((self.horizon,)) + stance_proximity_RL = np.zeros((self.horizon,)) + stance_proximity_RR = np.zeros((self.horizon,)) for j in range(self.horizon): - if(FL_contact_sequence[j] == 0): - if(j+1) < self.horizon: - if(FL_contact_sequence[j+1] == 1): - stance_proximity_FL[j] = 1*0 - if(j+2) < self.horizon: - if(FL_contact_sequence[j+1] == 0 and FL_contact_sequence[j+2] == 1): - stance_proximity_FL[j] = 1*0 - - if(FR_contact_sequence[j] == 0): - if(j+1) < self.horizon: - if(FR_contact_sequence[j+1] == 1): - stance_proximity_FR[j] = 1*0 - if(j+2) < self.horizon: - if(FR_contact_sequence[j+1] == 0 and FR_contact_sequence[j+2] == 1): - stance_proximity_FR[j] = 1*0 - - if(RL_contact_sequence[j] == 0): - if(j+1) < self.horizon: - if(RL_contact_sequence[j+1] == 1): - stance_proximity_RL[j] = 1*0 - if(j+2) < self.horizon: - if(RL_contact_sequence[j+1] == 0 and RL_contact_sequence[j+2] == 1): - stance_proximity_RL[j] = 1*0 - - if(RR_contact_sequence[j] == 0): - if(j+1) < self.horizon: - if(RR_contact_sequence[j+1] == 1): - stance_proximity_RR[j] = 1*0 - if(j+2) < self.horizon: - if(RR_contact_sequence[j+1] == 0 and RR_contact_sequence[j+2] == 1): - stance_proximity_RR[j] = 1*0 - - - + if FL_contact_sequence[j] == 0: + if (j + 1) < self.horizon: + if FL_contact_sequence[j + 1] == 1: + stance_proximity_FL[j] = 1 * 0 + if (j + 2) < self.horizon: + if FL_contact_sequence[j + 1] == 0 and FL_contact_sequence[j + 2] == 1: + stance_proximity_FL[j] = 1 * 0 + + if FR_contact_sequence[j] == 0: + if (j + 1) < self.horizon: + if FR_contact_sequence[j + 1] == 1: + stance_proximity_FR[j] = 1 * 0 + if (j + 2) < self.horizon: + if FR_contact_sequence[j + 1] == 0 and FR_contact_sequence[j + 2] == 1: + stance_proximity_FR[j] = 1 * 0 + + if RL_contact_sequence[j] == 0: + if (j + 1) < self.horizon: + if RL_contact_sequence[j + 1] == 1: + stance_proximity_RL[j] = 1 * 0 + if (j + 2) < self.horizon: + if RL_contact_sequence[j + 1] == 0 and RL_contact_sequence[j + 2] == 1: + stance_proximity_RL[j] = 1 * 0 + + if RR_contact_sequence[j] == 0: + if (j + 1) < self.horizon: + if RR_contact_sequence[j + 1] == 1: + stance_proximity_RR[j] = 1 * 0 + if (j + 2) < self.horizon: + if RR_contact_sequence[j + 1] == 0 and RR_contact_sequence[j + 2] == 1: + stance_proximity_RR[j] = 1 * 0 # Set the parameters to acados for j in range(self.horizon): - # If we have estimated an external wrench, we can compensate it for all steps # or less (maybe the disturbance is not costant along the horizon!) - if(config.mpc_params['external_wrenches_compensation'] and - config.mpc_params['external_wrenches_compensation_num_step'] and - j < config.mpc_params['external_wrenches_compensation_num_step']): + if ( + config.mpc_params["external_wrenches_compensation"] + and config.mpc_params["external_wrenches_compensation_num_step"] + and j < config.mpc_params["external_wrenches_compensation_num_step"] + ): external_wrenches_estimated_param = copy.deepcopy(external_wrenches) - external_wrenches_estimated_param = external_wrenches_estimated_param.reshape((6, )) + external_wrenches_estimated_param = external_wrenches_estimated_param.reshape((6,)) else: external_wrenches_estimated_param = np.zeros((6,)) - - param = np.array([FL_contact_sequence[j], FR_contact_sequence[j], - RL_contact_sequence[j], RR_contact_sequence[j], mu, - stance_proximity_FL[j], - stance_proximity_FR[j], - stance_proximity_RL[j], - stance_proximity_RR[j], - state["position"][0], state["position"][1], - state["position"][2], state["orientation"][2], - external_wrenches_estimated_param[0], external_wrenches_estimated_param[1], - external_wrenches_estimated_param[2], external_wrenches_estimated_param[3], - external_wrenches_estimated_param[4], external_wrenches_estimated_param[5], - inertia[0], inertia[1], inertia[2], inertia[3], inertia[4], inertia[5], - inertia[6], inertia[7], inertia[8], mass]) + param = np.array( + [ + FL_contact_sequence[j], + FR_contact_sequence[j], + RL_contact_sequence[j], + RR_contact_sequence[j], + mu, + stance_proximity_FL[j], + stance_proximity_FR[j], + stance_proximity_RL[j], + stance_proximity_RR[j], + state["position"][0], + state["position"][1], + state["position"][2], + state["orientation"][2], + external_wrenches_estimated_param[0], + external_wrenches_estimated_param[1], + external_wrenches_estimated_param[2], + external_wrenches_estimated_param[3], + external_wrenches_estimated_param[4], + external_wrenches_estimated_param[5], + inertia[0], + inertia[1], + inertia[2], + inertia[3], + inertia[4], + inertia[5], + inertia[6], + inertia[7], + inertia[8], + mass, + ] + ) self.acados_ocp_solver.set(j, "p", copy.deepcopy(param)) - - - - # Set initial state constraint. We teleported the robot foothold # to the previous optimal foothold. This is done to avoid the optimization # of a foothold that is not considered at all at touchdown! In any case, - # the height cames always from the VFA - if(FL_contact_sequence[0] == 0): + # the height cames always from the VFA + if FL_contact_sequence[0] == 0: state["foot_FL"] = reference["ref_foot_FL"][0] - if(FR_contact_sequence[0] == 0): + if FR_contact_sequence[0] == 0: state["foot_FR"] = reference["ref_foot_FR"][0] - if(RL_contact_sequence[0] == 0): + if RL_contact_sequence[0] == 0: state["foot_RL"] = reference["ref_foot_RL"][0] - - if(RR_contact_sequence[0] == 0): - state["foot_RR"] = reference["ref_foot_RR"][0] - + if RR_contact_sequence[0] == 0: + state["foot_RR"] = reference["ref_foot_RR"][0] - if(self.use_integrators): + if self.use_integrators: # Compute error for integral action alpha_integrator = config.mpc_params["alpha_integrator"] - self.integral_errors[0] += (state["position"][2] - reference["ref_position"][2])*alpha_integrator - self.integral_errors[1] += (state["linear_velocity"][0] - reference["ref_linear_velocity"][0])*alpha_integrator - self.integral_errors[2] += (state["linear_velocity"][1] - reference["ref_linear_velocity"][1])*alpha_integrator - self.integral_errors[3] += (state["linear_velocity"][2] - reference["ref_linear_velocity"][2])*alpha_integrator - self.integral_errors[4] += (state["orientation"][0] - reference["ref_orientation"][0])*(alpha_integrator) - self.integral_errors[5] += (state["orientation"][1] - reference["ref_orientation"][1])*alpha_integrator - + self.integral_errors[0] += (state["position"][2] - reference["ref_position"][2]) * alpha_integrator + self.integral_errors[1] += ( + state["linear_velocity"][0] - reference["ref_linear_velocity"][0] + ) * alpha_integrator + self.integral_errors[2] += ( + state["linear_velocity"][1] - reference["ref_linear_velocity"][1] + ) * alpha_integrator + self.integral_errors[3] += ( + state["linear_velocity"][2] - reference["ref_linear_velocity"][2] + ) * alpha_integrator + self.integral_errors[4] += (state["orientation"][0] - reference["ref_orientation"][0]) * (alpha_integrator) + self.integral_errors[5] += (state["orientation"][1] - reference["ref_orientation"][1]) * alpha_integrator cap_integrator_z = config.mpc_params["integrator_cap"][0] cap_integrator_x_dot = config.mpc_params["integrator_cap"][1] @@ -1413,81 +1455,104 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, cap_integrator_roll = config.mpc_params["integrator_cap"][4] cap_integrator_pitch = config.mpc_params["integrator_cap"][5] - self.integral_errors[0] = np.where(np.abs(self.integral_errors[0]) > cap_integrator_z, cap_integrator_z*np.sign(self.integral_errors[0]), self.integral_errors[0]) - self.integral_errors[1] = np.where(np.abs(self.integral_errors[1]) > cap_integrator_x_dot, cap_integrator_x_dot*np.sign(self.integral_errors[1]), self.integral_errors[1]) - self.integral_errors[2] = np.where(np.abs(self.integral_errors[2]) > cap_integrator_y_dot, cap_integrator_y_dot*np.sign(self.integral_errors[2]), self.integral_errors[2]) - self.integral_errors[3] = np.where(np.abs(self.integral_errors[3]) > cap_integrator_z_dot, cap_integrator_z_dot*np.sign(self.integral_errors[3]), self.integral_errors[3]) - self.integral_errors[4] = np.where(np.abs(self.integral_errors[4]) > cap_integrator_roll, cap_integrator_roll*np.sign(self.integral_errors[4]), self.integral_errors[4]) - self.integral_errors[5] = np.where(np.abs(self.integral_errors[5]) > cap_integrator_pitch, cap_integrator_pitch*np.sign(self.integral_errors[5]), self.integral_errors[5]) - - if(self.verbose): + self.integral_errors[0] = np.where( + np.abs(self.integral_errors[0]) > cap_integrator_z, + cap_integrator_z * np.sign(self.integral_errors[0]), + self.integral_errors[0], + ) + self.integral_errors[1] = np.where( + np.abs(self.integral_errors[1]) > cap_integrator_x_dot, + cap_integrator_x_dot * np.sign(self.integral_errors[1]), + self.integral_errors[1], + ) + self.integral_errors[2] = np.where( + np.abs(self.integral_errors[2]) > cap_integrator_y_dot, + cap_integrator_y_dot * np.sign(self.integral_errors[2]), + self.integral_errors[2], + ) + self.integral_errors[3] = np.where( + np.abs(self.integral_errors[3]) > cap_integrator_z_dot, + cap_integrator_z_dot * np.sign(self.integral_errors[3]), + self.integral_errors[3], + ) + self.integral_errors[4] = np.where( + np.abs(self.integral_errors[4]) > cap_integrator_roll, + cap_integrator_roll * np.sign(self.integral_errors[4]), + self.integral_errors[4], + ) + self.integral_errors[5] = np.where( + np.abs(self.integral_errors[5]) > cap_integrator_pitch, + cap_integrator_pitch * np.sign(self.integral_errors[5]), + self.integral_errors[5], + ) + + if self.verbose: print("self.integral_errors\n", self.integral_errors) - - # Set initial state constraint acados, converting first the dictionary to np array - state_acados = np.concatenate((state["position"], state["linear_velocity"], - state["orientation"], state["angular_velocity"], - state["foot_FL"], state["foot_FR"], - state["foot_RL"], state["foot_RR"], - self.integral_errors, - self.force_FL, - self.force_FR, - self.force_RL, - self.force_RR)).reshape((self.states_dim, 1)) + state_acados = np.concatenate( + ( + state["position"], + state["linear_velocity"], + state["orientation"], + state["angular_velocity"], + state["foot_FL"], + state["foot_FR"], + state["foot_RL"], + state["foot_RR"], + self.integral_errors, + self.force_FL, + self.force_FR, + self.force_RL, + self.force_RR, + ) + ).reshape((self.states_dim, 1)) self.acados_ocp_solver.set(0, "lbx", state_acados) self.acados_ocp_solver.set(0, "ubx", state_acados) - - - # Set stage constraint - h_R_w = np.array([np.cos(yaw), np.sin(yaw), - -np.sin(yaw), np.cos(yaw)]) - if(self.use_foothold_constraints or self.use_stability_constraints): - - stance_proximity = np.vstack((stance_proximity_FL, stance_proximity_FR, - stance_proximity_RL, stance_proximity_RR)) + h_R_w = np.array([np.cos(yaw), np.sin(yaw), -np.sin(yaw), np.cos(yaw)]) + if self.use_foothold_constraints or self.use_stability_constraints: + stance_proximity = np.vstack( + (stance_proximity_FL, stance_proximity_FR, stance_proximity_RL, stance_proximity_RR) + ) self.set_stage_constraint(constraint, state, reference, contact_sequence, h_R_w, stance_proximity) - - - # Set Warm start in case... - if(self.use_warm_start): - self.set_warm_start(state_acados, reference, FL_contact_sequence, FR_contact_sequence, RL_contact_sequence, RR_contact_sequence) - - - + if self.use_warm_start: + self.set_warm_start( + state_acados, + reference, + FL_contact_sequence, + FR_contact_sequence, + RL_contact_sequence, + RR_contact_sequence, + ) # Solve ocp via RTI or normal ocp if self.use_RTI: # feedback phase - self.acados_ocp_solver.options_set('rti_phase', 2) + self.acados_ocp_solver.options_set("rti_phase", 2) status = self.acados_ocp_solver.solve() - if(self.verbose): - print("feedback phase time: ", self.acados_ocp_solver.get_stats('time_tot')) - + if self.verbose: + print("feedback phase time: ", self.acados_ocp_solver.get_stats("time_tot")) else: status = self.acados_ocp_solver.solve() - if(self.verbose): - print("ocp time: ", self.acados_ocp_solver.get_stats('time_tot')) - - - + if self.verbose: + print("ocp time: ", self.acados_ocp_solver.get_stats("time_tot")) - if(config.mpc_params['dt'] <= 0.02 or ( - config.mpc_params['use_nonuniform_discretization'] and config.mpc_params['dt_fine_grained'] <= 0.02)): + if config.mpc_params["dt"] <= 0.02 or ( + config.mpc_params["use_nonuniform_discretization"] and config.mpc_params["dt_fine_grained"] <= 0.02 + ): optimal_next_state_index = 2 else: optimal_next_state_index = 1 self.optimal_next_state = self.acados_ocp_solver.get(optimal_next_state_index, "x")[0:24] - - - if(self.use_input_prediction): + + if self.use_input_prediction: optimal_GRF = self.acados_ocp_solver.get(optimal_next_state_index, "x")[30:42] self.force_FL = optimal_GRF[0:3] self.force_FR = optimal_GRF[3:6] @@ -1501,15 +1566,12 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, self.force_RL = optimal_next_GRF[6:9] self.force_RR = optimal_next_GRF[9:12] - - optimal_foothold = np.zeros((4, 3)) - optimal_footholds_assigned = np.zeros((4, ), dtype='bool') - + optimal_footholds_assigned = np.zeros((4,), dtype="bool") # We need to provide the next touchdown foothold position. # We first take the foothold in stance now (they are not optimized!) - # and flag them as True (aka "assigned") + # and flag them as True (aka "assigned") if FL_contact_sequence[0] == 1: optimal_foothold[0] = state["foot_FL"] optimal_footholds_assigned[0] = True @@ -1523,120 +1585,177 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, optimal_foothold[3] = state["foot_RR"] optimal_footholds_assigned[3] = True - - - - # Then we take the foothold at the next touchdown from the one - # that are not flagged as True from before, and saturate them!! + # Then we take the foothold at the next touchdown from the one + # that are not flagged as True from before, and saturate them!! # P.S. The saturation is in the horizontal frame - h_R_w = h_R_w.reshape((2, 2)) + h_R_w = h_R_w.reshape((2, 2)) for j in range(1, self.horizon): - if(FL_contact_sequence[j] != FL_contact_sequence[j-1] and not optimal_footholds_assigned[0]): + if FL_contact_sequence[j] != FL_contact_sequence[j - 1] and not optimal_footholds_assigned[0]: optimal_foothold[0] = self.acados_ocp_solver.get(j, "x")[12:15] optimal_footholds_assigned[0] = True - - if(constraint is None): - first_up_constraint_FL = np.array([reference["ref_foot_FL"][0][0], reference["ref_foot_FL"][0][1], reference["ref_foot_FL"][0][2] + 0.002]) - first_low_constraint_FL = np.array([reference["ref_foot_FL"][0][0], reference["ref_foot_FL"][0][1], reference["ref_foot_FL"][0][2] - 0.002]) - first_up_constraint_FL[0:2] = h_R_w@(first_up_constraint_FL[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_FL[0:2] = h_R_w@(first_low_constraint_FL[0:2] - state["position"][0:2]) - 0.15 + if constraint is None: + first_up_constraint_FL = np.array( + [ + reference["ref_foot_FL"][0][0], + reference["ref_foot_FL"][0][1], + reference["ref_foot_FL"][0][2] + 0.002, + ] + ) + first_low_constraint_FL = np.array( + [ + reference["ref_foot_FL"][0][0], + reference["ref_foot_FL"][0][1], + reference["ref_foot_FL"][0][2] - 0.002, + ] + ) + + first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - state["position"][0:2]) + 0.15 + first_low_constraint_FL[0:2] = ( + h_R_w @ (first_low_constraint_FL[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_FL = np.array([constraint[0][0], constraint[1][0], constraint[2][0] + 0.002]) first_low_constraint_FL = np.array([constraint[9][0], constraint[10][0], constraint[11][0] - 0.002]) - first_up_constraint_FL[0:2] = h_R_w@(first_up_constraint_FL[0:2] - state["position"][0:2]) - first_low_constraint_FL[0:2] = h_R_w@(first_low_constraint_FL[0:2] - state["position"][0:2]) + first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - state["position"][0:2]) + first_low_constraint_FL[0:2] = h_R_w @ (first_low_constraint_FL[0:2] - state["position"][0:2]) - optimal_foothold[0][0:2] = h_R_w@(optimal_foothold[0][0:2] - state["position"][0:2]) - optimal_foothold[0][0:2] = np.clip(optimal_foothold[0][0:2], first_low_constraint_FL[0:2], first_up_constraint_FL[0:2]) - optimal_foothold[0][0:2] = h_R_w.T@optimal_foothold[0][0:2] + state["position"][0:2] + optimal_foothold[0][0:2] = h_R_w @ (optimal_foothold[0][0:2] - state["position"][0:2]) + optimal_foothold[0][0:2] = np.clip( + optimal_foothold[0][0:2], first_low_constraint_FL[0:2], first_up_constraint_FL[0:2] + ) + optimal_foothold[0][0:2] = h_R_w.T @ optimal_foothold[0][0:2] + state["position"][0:2] - - if(FR_contact_sequence[j] != FR_contact_sequence[j-1] and not optimal_footholds_assigned[1]): + if FR_contact_sequence[j] != FR_contact_sequence[j - 1] and not optimal_footholds_assigned[1]: optimal_foothold[1] = self.acados_ocp_solver.get(j, "x")[15:18] optimal_footholds_assigned[1] = True - if(constraint is None): - first_up_constraint_FR = np.array([reference["ref_foot_FR"][0][0], reference["ref_foot_FR"][0][1], reference["ref_foot_FR"][0][2] + 0.002]) - first_low_constraint_FR = np.array([reference["ref_foot_FR"][0][0], reference["ref_foot_FR"][0][1], reference["ref_foot_FR"][0][2] - 0.002]) - - first_up_constraint_FR[0:2] = h_R_w@(first_up_constraint_FR[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_FR[0:2] = h_R_w@(first_low_constraint_FR[0:2] - state["position"][0:2]) - 0.15 + if constraint is None: + first_up_constraint_FR = np.array( + [ + reference["ref_foot_FR"][0][0], + reference["ref_foot_FR"][0][1], + reference["ref_foot_FR"][0][2] + 0.002, + ] + ) + first_low_constraint_FR = np.array( + [ + reference["ref_foot_FR"][0][0], + reference["ref_foot_FR"][0][1], + reference["ref_foot_FR"][0][2] - 0.002, + ] + ) + + first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - state["position"][0:2]) + 0.15 + first_low_constraint_FR[0:2] = ( + h_R_w @ (first_low_constraint_FR[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_FR = np.array([constraint[0][1], constraint[1][1], constraint[2][1] + 0.002]) first_low_constraint_FR = np.array([constraint[9][1], constraint[10][1], constraint[11][1] - 0.002]) - first_up_constraint_FR[0:2] = h_R_w@(first_up_constraint_FR[0:2] - state["position"][0:2]) - first_low_constraint_FR[0:2] = h_R_w@(first_low_constraint_FR[0:2] - state["position"][0:2]) + first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - state["position"][0:2]) + first_low_constraint_FR[0:2] = h_R_w @ (first_low_constraint_FR[0:2] - state["position"][0:2]) - optimal_foothold[1][0:2] = h_R_w@(optimal_foothold[1][0:2] - state["position"][0:2]) - optimal_foothold[1][0:2] = np.clip(optimal_foothold[1][0:2], first_low_constraint_FR[0:2], first_up_constraint_FR[0:2]) - optimal_foothold[1][0:2] = h_R_w.T@optimal_foothold[1][0:2] + state["position"][0:2] + optimal_foothold[1][0:2] = h_R_w @ (optimal_foothold[1][0:2] - state["position"][0:2]) + optimal_foothold[1][0:2] = np.clip( + optimal_foothold[1][0:2], first_low_constraint_FR[0:2], first_up_constraint_FR[0:2] + ) + optimal_foothold[1][0:2] = h_R_w.T @ optimal_foothold[1][0:2] + state["position"][0:2] - - if(RL_contact_sequence[j] != RL_contact_sequence[j-1] and not optimal_footholds_assigned[2]): + if RL_contact_sequence[j] != RL_contact_sequence[j - 1] and not optimal_footholds_assigned[2]: optimal_foothold[2] = self.acados_ocp_solver.get(j, "x")[18:21] optimal_footholds_assigned[2] = True - if(constraint is None): - first_up_constraint_RL = np.array([reference["ref_foot_RL"][0][0], reference["ref_foot_RL"][0][1], reference["ref_foot_RL"][0][2] + 0.002]) - first_low_constraint_RL = np.array([reference["ref_foot_RL"][0][0], reference["ref_foot_RL"][0][1], reference["ref_foot_RL"][0][2] - 0.002]) - - first_up_constraint_RL[0:2] = h_R_w@(first_up_constraint_RL[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_RL[0:2] = h_R_w@(first_low_constraint_RL[0:2] - state["position"][0:2]) - 0.15 + if constraint is None: + first_up_constraint_RL = np.array( + [ + reference["ref_foot_RL"][0][0], + reference["ref_foot_RL"][0][1], + reference["ref_foot_RL"][0][2] + 0.002, + ] + ) + first_low_constraint_RL = np.array( + [ + reference["ref_foot_RL"][0][0], + reference["ref_foot_RL"][0][1], + reference["ref_foot_RL"][0][2] - 0.002, + ] + ) + + first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - state["position"][0:2]) + 0.15 + first_low_constraint_RL[0:2] = ( + h_R_w @ (first_low_constraint_RL[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_RL = np.array([constraint[0][2], constraint[1][2], constraint[2][2] + 0.002]) first_low_constraint_RL = np.array([constraint[9][2], constraint[10][2], constraint[11][2] - 0.002]) - first_up_constraint_RL[0:2] = h_R_w@(first_up_constraint_RL[0:2] - state["position"][0:2]) - first_low_constraint_RL[0:2] = h_R_w@(first_low_constraint_RL[0:2] - state["position"][0:2]) - - optimal_foothold[2][0:2] = h_R_w@(optimal_foothold[2][0:2] - state["position"][0:2]) - optimal_foothold[2][0:2] = np.clip(optimal_foothold[2][0:2], first_low_constraint_RL[0:2], first_up_constraint_RL[0:2]) - optimal_foothold[2][0:2] = h_R_w.T@optimal_foothold[2][0:2] + state["position"][0:2] - - if(RR_contact_sequence[j] != RR_contact_sequence[j-1] and not optimal_footholds_assigned[3]): + first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - state["position"][0:2]) + first_low_constraint_RL[0:2] = h_R_w @ (first_low_constraint_RL[0:2] - state["position"][0:2]) + + optimal_foothold[2][0:2] = h_R_w @ (optimal_foothold[2][0:2] - state["position"][0:2]) + optimal_foothold[2][0:2] = np.clip( + optimal_foothold[2][0:2], first_low_constraint_RL[0:2], first_up_constraint_RL[0:2] + ) + optimal_foothold[2][0:2] = h_R_w.T @ optimal_foothold[2][0:2] + state["position"][0:2] + + if RR_contact_sequence[j] != RR_contact_sequence[j - 1] and not optimal_footholds_assigned[3]: optimal_foothold[3] = self.acados_ocp_solver.get(j, "x")[21:24] optimal_footholds_assigned[3] = True - if(constraint is None): - first_up_constraint_RR = np.array([reference["ref_foot_RR"][0][0], reference["ref_foot_RR"][0][1], reference["ref_foot_RR"][0][2] + 0.002]) - first_low_constraint_RR = np.array([reference["ref_foot_RR"][0][0], reference["ref_foot_RR"][0][1], reference["ref_foot_RR"][0][2] - 0.002]) - - first_up_constraint_RR[0:2] = h_R_w@(first_up_constraint_RR[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_RR[0:2] = h_R_w@(first_low_constraint_RR[0:2] - state["position"][0:2]) - 0.15 + if constraint is None: + first_up_constraint_RR = np.array( + [ + reference["ref_foot_RR"][0][0], + reference["ref_foot_RR"][0][1], + reference["ref_foot_RR"][0][2] + 0.002, + ] + ) + first_low_constraint_RR = np.array( + [ + reference["ref_foot_RR"][0][0], + reference["ref_foot_RR"][0][1], + reference["ref_foot_RR"][0][2] - 0.002, + ] + ) + + first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - state["position"][0:2]) + 0.15 + first_low_constraint_RR[0:2] = ( + h_R_w @ (first_low_constraint_RR[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_RR = np.array([constraint[0][3], constraint[1][3], constraint[2][3] + 0.002]) first_low_constraint_RR = np.array([constraint[9][3], constraint[10][3], constraint[11][3] - 0.002]) - first_up_constraint_RR[0:2] = h_R_w@(first_up_constraint_RR[0:2] - state["position"][0:2]) - first_low_constraint_RR[0:2] = h_R_w@(first_low_constraint_RR[0:2] - state["position"][0:2]) - - optimal_foothold[3][0:2] = h_R_w@(optimal_foothold[3][0:2] - state["position"][0:2]) - optimal_foothold[3][0:2] = np.clip(optimal_foothold[3][0:2], first_low_constraint_RR[0:2], first_up_constraint_RR[0:2]) - optimal_foothold[3][0:2] = h_R_w.T@optimal_foothold[3][0:2] + state["position"][0:2] - - - # If in the prediction horizon, the foot is never in stance, we replicate the reference - #to not confuse the swing controller - if(optimal_footholds_assigned[0] == False): + first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - state["position"][0:2]) + first_low_constraint_RR[0:2] = h_R_w @ (first_low_constraint_RR[0:2] - state["position"][0:2]) + + optimal_foothold[3][0:2] = h_R_w @ (optimal_foothold[3][0:2] - state["position"][0:2]) + optimal_foothold[3][0:2] = np.clip( + optimal_foothold[3][0:2], first_low_constraint_RR[0:2], first_up_constraint_RR[0:2] + ) + optimal_foothold[3][0:2] = h_R_w.T @ optimal_foothold[3][0:2] + state["position"][0:2] + + # If in the prediction horizon, the foot is never in stance, we replicate the reference + # to not confuse the swing controller + if optimal_footholds_assigned[0] == False: optimal_foothold[0] = reference["ref_foot_FL"][0] - if(optimal_footholds_assigned[1] == False): + if optimal_footholds_assigned[1] == False: optimal_foothold[1] = reference["ref_foot_FR"][0] - if(optimal_footholds_assigned[2] == False): + if optimal_footholds_assigned[2] == False: optimal_foothold[2] = reference["ref_foot_RL"][0] - if(optimal_footholds_assigned[3] == False): + if optimal_footholds_assigned[3] == False: optimal_foothold[3] = reference["ref_foot_RR"][0] - if(self.verbose): + if self.verbose: self.acados_ocp_solver.print_statistics() - # Check if QPs converged, if not just use the reference footholds # and a GRF over Z distribuited between the leg in stance - if(status == 1 or status == 4): - if(self.verbose): + if status == 1 or status == 4: + if self.verbose: print("status", status) if FL_contact_sequence[0] == 0: optimal_foothold[0] = reference["ref_foot_FL"][0] @@ -1646,39 +1765,33 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, optimal_foothold[2] = reference["ref_foot_RL"][0] if RR_contact_sequence[0] == 0: optimal_foothold[3] = reference["ref_foot_RR"][0] - - number_of_legs_in_stance = np.array([FL_contact_sequence[0], FR_contact_sequence[0], - RL_contact_sequence [0], RR_contact_sequence[0]]).sum() - if(number_of_legs_in_stance == 0): + + number_of_legs_in_stance = np.array( + [FL_contact_sequence[0], FR_contact_sequence[0], RL_contact_sequence[0], RR_contact_sequence[0]] + ).sum() + if number_of_legs_in_stance == 0: reference_force_stance_legs = 0 else: reference_force_stance_legs = (mass * 9.81) / number_of_legs_in_stance - + reference_force_fl_z = reference_force_stance_legs * FL_contact_sequence[0] reference_force_fr_z = reference_force_stance_legs * FR_contact_sequence[0] reference_force_rl_z = reference_force_stance_legs * RL_contact_sequence[0] reference_force_rr_z = reference_force_stance_legs * RR_contact_sequence[0] - optimal_GRF = np.zeros((12, )) + optimal_GRF = np.zeros((12,)) optimal_GRF[2] = reference_force_fl_z optimal_GRF[5] = reference_force_fr_z optimal_GRF[8] = reference_force_rl_z optimal_GRF[11] = reference_force_rr_z - - optimal_GRF = self.previous_optimal_GRF self.reset() - - # Save the previous optimal GRF, the previous status and the previous contact sequence self.previous_optimal_GRF = optimal_GRF self.previous_status = status self.previous_contact_sequence = contact_sequence - - - # Decenter the optimal foothold and the next state (they were centered around zero at the beginning) optimal_foothold[0] = optimal_foothold[0] + self.initial_base_position optimal_foothold[1] = optimal_foothold[1] + self.initial_base_position @@ -1691,6 +1804,5 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, self.optimal_next_state[18:21] = optimal_foothold[2] self.optimal_next_state[21:24] = optimal_foothold[3] - # Return the optimal GRF, the optimal foothold, the next state and the status of the optimization return optimal_GRF, optimal_foothold, self.optimal_next_state, status diff --git a/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_model.py b/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_model.py index 8c373f5..fc862ee 100644 --- a/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_model.py +++ b/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_model.py @@ -3,19 +3,18 @@ # Authors: Giulio Turrisi - -import time -import unittest -import casadi as cs +import os +import casadi as cs import numpy as np from acados_template import AcadosModel -import os dir_path = os.path.dirname(os.path.realpath(__file__)) import sys + sys.path.append(dir_path) -sys.path.append(dir_path + '/../../') +sys.path.append(dir_path + "/../../") from liecasadi import SO3 @@ -25,39 +24,50 @@ use_fixed_inertia = True use_centroidal_model = True -if(use_adam): +if use_adam: # ADAM import - from adam.casadi import KinDynComputations from adam import Representations + from adam.casadi import KinDynComputations else: # PINOCCHIO import import pinocchio as pin from pinocchio import casadi as cpin - # Class that defines the prediction model of the NMPC class KinoDynamic_Model: - def __init__(self,) -> None: - - - if(config.robot == 'go2'): - urdf_filename = dir_path + '/../../../gym-quadruped/gym_quadruped/robot_model/go2/go2.urdf' - elif(config.robot == 'aliengo'): - #urdf_filename = dir_path + '/../../../gym-quadruped/gym_quadruped/robot_model/aliengo/aliengo.urdf' - urdf_filename = '/home/iit.local/gturrisi/personal_ws_home/gym-quadruped/gym_quadruped/robot_model/aliengo/aliengo.urdf' - elif(config.robot == 'hyqreal'): - urdf_filename = dir_path + '/../../../gym-quadruped/gym_quadruped/robot_model/hyqreal/hyqreal.urdf' - elif(config.robot == 'mini_cheetah'): - urdf_filename = dir_path + '/../../../gym-quadruped/gym_quadruped/robot_model/mini_cheetah/mini_cheetah.urdf' - - - joint_list = ['FL_hip_joint', 'FL_thigh_joint', 'FL_calf_joint', - 'FR_hip_joint', 'FR_thigh_joint', 'FR_calf_joint', - 'RL_hip_joint', 'RL_thigh_joint', 'RL_calf_joint', - 'RR_hip_joint', 'RR_thigh_joint', 'RR_calf_joint'] - - + def __init__( + self, + ) -> None: + if config.robot == "go2": + urdf_filename = dir_path + "/../../../gym-quadruped/gym_quadruped/robot_model/go2/go2.urdf" + elif config.robot == "aliengo": + # urdf_filename = dir_path + '/../../../gym-quadruped/gym_quadruped/robot_model/aliengo/aliengo.urdf' + urdf_filename = ( + "/home/iit.local/gturrisi/personal_ws_home/gym-quadruped/gym_quadruped/robot_model/aliengo/aliengo.urdf" + ) + elif config.robot == "hyqreal": + urdf_filename = dir_path + "/../../../gym-quadruped/gym_quadruped/robot_model/hyqreal/hyqreal.urdf" + elif config.robot == "mini_cheetah": + urdf_filename = ( + dir_path + "/../../../gym-quadruped/gym_quadruped/robot_model/mini_cheetah/mini_cheetah.urdf" + ) + + joint_list = [ + "FL_hip_joint", + "FL_thigh_joint", + "FL_calf_joint", + "FR_hip_joint", + "FR_thigh_joint", + "FR_calf_joint", + "RL_hip_joint", + "RL_thigh_joint", + "RL_calf_joint", + "RR_hip_joint", + "RR_thigh_joint", + "RR_calf_joint", + ] + self.kindyn = KinDynComputations(urdfstring=urdf_filename, joints_name_list=joint_list) self.kindyn.set_frame_velocity_representation(representation=Representations.MIXED_REPRESENTATION) self.mass_mass_fun = self.kindyn.mass_matrix_fun() @@ -65,9 +75,8 @@ def __init__(self,) -> None: self.bias_force_fun = self.kindyn.bias_force_fun() self.gravity_fun = self.kindyn.gravity_term_fun() self.coriolis_fun = self.kindyn.coriolis_term_fun() - - - if(use_adam): + + if use_adam: self.forward_kinematics_FL_fun = self.kindyn.forward_kinematics_fun("FL_foot") self.forward_kinematics_FR_fun = self.kindyn.forward_kinematics_fun("FR_foot") self.forward_kinematics_RL_fun = self.kindyn.forward_kinematics_fun("RL_foot") @@ -77,7 +86,7 @@ def __init__(self,) -> None: self.jacobian_FR_fun = self.kindyn.jacobian_fun("FR_foot") self.jacobian_RL_fun = self.kindyn.jacobian_fun("RL_foot") self.jacobian_RR_fun = self.kindyn.jacobian_fun("RR_foot") - + else: model_pin = pin.buildModelFromUrdf(urdf_filename) data_pin = model_pin.createData() @@ -85,41 +94,53 @@ def __init__(self,) -> None: cmodel_pin = cpin.Model(model_pin) cdata_pin = cmodel_pin.createData() cq_pin = cs.SX.sym("q", model_pin.nq, 1) - + # precompute the forward kinematics graph cpin.framesForwardKinematics(cmodel_pin, cdata_pin, cq_pin) # takes the ID of the feet, and generate a casadi function for a generic forward kinematics FL_foot_id = model_pin.getFrameId("FL_foot_fixed") - #self.FL_foot_id = self.model.getJointId("FL_foot_fixed") - self.forward_kinematics_FL_fun = cs.Function("FL_foot_pos", [cq_pin], [cdata_pin.oMf[FL_foot_id].translation]) + # self.FL_foot_id = self.model.getJointId("FL_foot_fixed") + self.forward_kinematics_FL_fun = cs.Function( + "FL_foot_pos", [cq_pin], [cdata_pin.oMf[FL_foot_id].translation] + ) FR_foot_id = model_pin.getFrameId("FR_foot_fixed") - #self.FR_foot_id = self.model.getJointId("FR_foot_fixed") - self.forward_kinematics_FR_fun= cs.Function("FR_foot_pos", [cq_pin], [cdata_pin.oMf[FR_foot_id].translation]) + # self.FR_foot_id = self.model.getJointId("FR_foot_fixed") + self.forward_kinematics_FR_fun = cs.Function( + "FR_foot_pos", [cq_pin], [cdata_pin.oMf[FR_foot_id].translation] + ) RL_foot_id = model_pin.getFrameId("RL_foot_fixed") - #self.RL_foot_id = self.model.getJointId("RL_foot_fixed") - self.forward_kinematics_RR_fun = cs.Function("RL_foot_pos", [cq_pin], [cdata_pin.oMf[RL_foot_id].translation]) + # self.RL_foot_id = self.model.getJointId("RL_foot_fixed") + self.forward_kinematics_RR_fun = cs.Function( + "RL_foot_pos", [cq_pin], [cdata_pin.oMf[RL_foot_id].translation] + ) RR_foot_id = model_pin.getFrameId("RR_foot_fixed") - #self.RR_foot_id = self.model.getJointId("RR_foot_fixed") - self.forwabase_transformrd_kinematics_RL_fun = cs.Function("RR_foot_pos", [cq_pin], [cdata_pin.oMf[RR_foot_id].translation]) - - J_FL = pin.getFrameJacobian(model_pin, data_pin, FL_foot_id, pin.ReferenceFrame.LOCAL_WORLD_ALIGNED) # in joint frame - J_FR = pin.getFrameJacobian(model_pin, data_pin, FR_foot_id, pin.ReferenceFrame.LOCAL_WORLD_ALIGNED) # in joint frame - J_RL = pin.getFrameJacobian(model_pin, data_pin, RL_foot_id, pin.ReferenceFrame.LOCAL_WORLD_ALIGNED) # in joint frame - J_RR = pin.getFrameJacobian(model_pin, data_pin, RR_foot_id, pin.ReferenceFrame.LOCAL_WORLD_ALIGNED) # in joint frame + # self.RR_foot_id = self.model.getJointId("RR_foot_fixed") + self.forwabase_transformrd_kinematics_RL_fun = cs.Function( + "RR_foot_pos", [cq_pin], [cdata_pin.oMf[RR_foot_id].translation] + ) + + J_FL = pin.getFrameJacobian( + model_pin, data_pin, FL_foot_id, pin.ReferenceFrame.LOCAL_WORLD_ALIGNED + ) # in joint frame + J_FR = pin.getFrameJacobian( + model_pin, data_pin, FR_foot_id, pin.ReferenceFrame.LOCAL_WORLD_ALIGNED + ) # in joint frame + J_RL = pin.getFrameJacobian( + model_pin, data_pin, RL_foot_id, pin.ReferenceFrame.LOCAL_WORLD_ALIGNED + ) # in joint frame + J_RR = pin.getFrameJacobian( + model_pin, data_pin, RR_foot_id, pin.ReferenceFrame.LOCAL_WORLD_ALIGNED + ) # in joint frame self.jacobian_FL_fun = cs.Function("jacobian_FL", [cq_pin], [J_FL]) self.jacobian_FR_fun = cs.Function("jacobian_FR", [cq_pin], [J_FR]) self.jacobian_RL_fun = cs.Function("jacobian_RL", [cq_pin], [J_RL]) self.jacobian_RR_fun = cs.Function("jacobian_RR", [cq_pin], [J_RR]) - - - - # Define state and its casadi variables com_position_x = cs.SX.sym("com_position_x") com_position_y = cs.SX.sym("com_position_y") @@ -128,20 +149,19 @@ def __init__(self,) -> None: com_velocity_x = cs.SX.sym("com_velocity_x") com_velocity_y = cs.SX.sym("com_velocity_y") com_velocity_z = cs.SX.sym("com_velocity_z") - + roll = cs.SX.sym("roll", 1, 1) pitch = cs.SX.sym("pitch", 1, 1) yaw = cs.SX.sym("yaw", 1, 1) omega_x = cs.SX.sym("omega_x", 1, 1) omega_y = cs.SX.sym("omega_y", 1, 1) omega_z = cs.SX.sym("omega_z", 1, 1) - + joint_position_fl = cs.SX.sym("joint_position_fl", 3, 1) joint_position_fr = cs.SX.sym("joint_position_fr", 3, 1) joint_position_rl = cs.SX.sym("joint_position_rl", 3, 1) joint_position_rr = cs.SX.sym("joint_position_rr", 3, 1) - com_position_z_integral = cs.SX.sym("com_position_z_integral") com_velocity_x_integral = cs.SX.sym("com_velocity_x_integral") com_velocity_y_integral = cs.SX.sym("com_velocity_y_integral") @@ -149,47 +169,46 @@ def __init__(self,) -> None: roll_integral = cs.SX.sym("roll_integral") pitch_integral = cs.SX.sym("pitch_integral") - - self.states = cs.vertcat(com_position_x, - com_position_y, - com_position_z, - com_velocity_x, - com_velocity_y, - com_velocity_z, - roll, - pitch, - yaw, - omega_x, - omega_y, - omega_z, - joint_position_fl, - joint_position_fr, - joint_position_rl, - joint_position_rr, - com_position_z_integral, - com_velocity_x_integral, - com_velocity_y_integral, - com_velocity_z_integral, - roll_integral, - pitch_integral) - - - - # Define state dot - self.states_dot = cs.vertcat(cs.SX.sym("linear_com_vel", 3, 1), - cs.SX.sym("linear_com_acc", 3, 1), - cs.SX.sym("euler_rates_base", 3, 1), - cs.SX.sym("angular_acc_base", 3, 1), - cs.SX.sym("joint_vel_FL", 3, 1), - cs.SX.sym("joint_vel_FR", 3, 1), - cs.SX.sym("joint_vel_RL", 3, 1), - cs.SX.sym("joint_vel_RR", 3, 1), - cs.SX.sym("linear_com_vel_z_integral", 1, 1), - cs.SX.sym("linear_com_acc_integral", 3, 1), - cs.SX.sym("euler_rates_roll_integral", 1, 1), - cs.SX.sym("euler_rates_pitch_integral", 1, 1)) - - + self.states = cs.vertcat( + com_position_x, + com_position_y, + com_position_z, + com_velocity_x, + com_velocity_y, + com_velocity_z, + roll, + pitch, + yaw, + omega_x, + omega_y, + omega_z, + joint_position_fl, + joint_position_fr, + joint_position_rl, + joint_position_rr, + com_position_z_integral, + com_velocity_x_integral, + com_velocity_y_integral, + com_velocity_z_integral, + roll_integral, + pitch_integral, + ) + + # Define state dot + self.states_dot = cs.vertcat( + cs.SX.sym("linear_com_vel", 3, 1), + cs.SX.sym("linear_com_acc", 3, 1), + cs.SX.sym("euler_rates_base", 3, 1), + cs.SX.sym("angular_acc_base", 3, 1), + cs.SX.sym("joint_vel_FL", 3, 1), + cs.SX.sym("joint_vel_FR", 3, 1), + cs.SX.sym("joint_vel_RL", 3, 1), + cs.SX.sym("joint_vel_RR", 3, 1), + cs.SX.sym("linear_com_vel_z_integral", 1, 1), + cs.SX.sym("linear_com_acc_integral", 3, 1), + cs.SX.sym("euler_rates_roll_integral", 1, 1), + cs.SX.sym("euler_rates_pitch_integral", 1, 1), + ) # Define input and its casadi variables joint_velocity_fl = cs.SX.sym("joint_velocity_fl", 3, 1) @@ -202,28 +221,27 @@ def __init__(self,) -> None: foot_force_rl = cs.SX.sym("foot_force_rl", 3, 1) foot_force_rr = cs.SX.sym("foot_force_rr", 3, 1) - self.inputs = cs.vertcat(joint_velocity_fl, - joint_velocity_fr, - joint_velocity_rl, - joint_velocity_rr, - foot_force_fl, - foot_force_fr, - foot_force_rl, - foot_force_rr) - + self.inputs = cs.vertcat( + joint_velocity_fl, + joint_velocity_fr, + joint_velocity_rl, + joint_velocity_rr, + foot_force_fl, + foot_force_fr, + foot_force_rl, + foot_force_rr, + ) # Usefull for debug what things goes where in y_ref in the compute_control function, # because there are a lot of variables self.y_ref = cs.vertcat(self.states, self.inputs) - # Define acados parameters that can be changed at runtine self.stanceFL = cs.SX.sym("stanceFL", 1, 1) self.stanceFR = cs.SX.sym("stanceFR", 1, 1) self.stanceRL = cs.SX.sym("stanceRL", 1, 1) self.stanceRR = cs.SX.sym("stanceRR", 1, 1) - self.stance_param = cs.vertcat(self.stanceFL , self.stanceFR , self.stanceRL , self.stanceRR) - + self.stance_param = cs.vertcat(self.stanceFL, self.stanceFR, self.stanceRL, self.stanceRR) self.mu_friction = cs.SX.sym("mu_friction", 1, 1) self.stance_proximity = cs.SX.sym("stanceProximity", 4, 1) @@ -235,48 +253,44 @@ def __init__(self,) -> None: self.inertia = cs.SX.sym("inertia", 9, 1) self.mass = cs.SX.sym("mass", 1, 1) - - - def compute_b_R_w(self, roll: float, pitch: float, yaw: float) -> np.ndarray: - #Z Y X rotations! + # Z Y X rotations! Rx = cs.SX.eye(3) - Rx[0,0] = 1 - Rx[0,1] = 0 - Rx[0,2] = 0 - Rx[1,0] = 0 - Rx[1,1] = cs.cos(roll) - Rx[1,2] = cs.sin(roll) - Rx[2,0] = 0 - Rx[2,1] = -cs.sin(roll) - Rx[2,2] = cs.cos(roll) - + Rx[0, 0] = 1 + Rx[0, 1] = 0 + Rx[0, 2] = 0 + Rx[1, 0] = 0 + Rx[1, 1] = cs.cos(roll) + Rx[1, 2] = cs.sin(roll) + Rx[2, 0] = 0 + Rx[2, 1] = -cs.sin(roll) + Rx[2, 2] = cs.cos(roll) + Ry = cs.SX.eye(3) - Ry[0,0] = cs.cos(pitch) - Ry[0,1] = 0 - Ry[0,2] = -cs.sin(pitch) - Ry[1,0] = 0 - Ry[1,1] = 1 - Ry[1,2] = 0 - Ry[2,0] = cs.sin(pitch) - Ry[2,1] = 0 - Ry[2,2] = cs.cos(pitch) + Ry[0, 0] = cs.cos(pitch) + Ry[0, 1] = 0 + Ry[0, 2] = -cs.sin(pitch) + Ry[1, 0] = 0 + Ry[1, 1] = 1 + Ry[1, 2] = 0 + Ry[2, 0] = cs.sin(pitch) + Ry[2, 1] = 0 + Ry[2, 2] = cs.cos(pitch) Rz = cs.SX.eye(3) - Rz[0,0] = cs.cos(yaw) - Rz[0,1] = cs.sin(yaw) - Rz[0,2] = 0 - Rz[1,0] = -cs.sin(yaw) - Rz[1,1] = cs.cos(yaw) - Rz[1,2] = 0 - Rz[2,0] = 0 - Rz[2,1] = 0 - Rz[2,2] = 1 - - b_R_w = Rx@Ry@Rz + Rz[0, 0] = cs.cos(yaw) + Rz[0, 1] = cs.sin(yaw) + Rz[0, 2] = 0 + Rz[1, 0] = -cs.sin(yaw) + Rz[1, 1] = cs.cos(yaw) + Rz[1, 2] = 0 + Rz[2, 0] = 0 + Rz[2, 1] = 0 + Rz[2, 2] = 1 + + b_R_w = Rx @ Ry @ Rz return b_R_w - def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.ndarray) -> cs.SX: """ This method computes the symbolic forward dynamics of the robot. It is used inside @@ -291,7 +305,7 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda Returns: A CasADi SX object of shape (29,) representing the predicted state of the robot. """ - + # Saving for clarity a bunch of variables joint_velocity_fl = inputs[0:3] joint_velocity_fr = inputs[3:6] @@ -313,34 +327,27 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda stanceRL = param[2] stanceRR = param[3] - external_wrench_linear = param[13:16] external_wrench_angular = param[16:19] - # FINAL linear_com_vel STATE (1) linear_com_vel = states[3:6] - # Start to write the component of euler_rates_base and angular_acc_base STATES w = states[9:12] roll = states[6] pitch = states[7] yaw = states[8] - + conj_euler_rates = cs.SX.eye(3) conj_euler_rates[1, 1] = cs.cos(roll) - conj_euler_rates[2, 2] = cs.cos(pitch)*cs.cos(roll) + conj_euler_rates[2, 2] = cs.cos(pitch) * cs.cos(roll) conj_euler_rates[2, 1] = -cs.sin(roll) conj_euler_rates[0, 2] = -cs.sin(pitch) - conj_euler_rates[1, 2] = cs.cos(pitch)*cs.sin(roll) + conj_euler_rates[1, 2] = cs.cos(pitch) * cs.sin(roll) - - # FINAL euler_rates_base STATE (3) - euler_rates_base = cs.inv(conj_euler_rates)@w - - + euler_rates_base = cs.inv(conj_euler_rates) @ w # Compute the homogeneous transformation matrix # b_R_w = self.compute_b_R_w(roll, pitch, yaw) @@ -350,67 +357,56 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda H[0:3, 0:3] = b_R_w.T H[0:3, 3] = com_position - - # Compute com pos, feet pos and inertia via ADAM joint_position = cs.vertcat(joint_position_fl, joint_position_fr, joint_position_rl, joint_position_rr) joints_velocities = cs.vertcat(joint_velocity_fl, joint_velocity_fr, joint_velocity_rl, joint_velocity_rr) - + self.foot_position_fl = self.forward_kinematics_FL_fun(H, joint_position)[0:3, 3] self.foot_position_fr = self.forward_kinematics_FR_fun(H, joint_position)[0:3, 3] self.foot_position_rl = self.forward_kinematics_RL_fun(H, joint_position)[0:3, 3] self.foot_position_rr = self.forward_kinematics_RR_fun(H, joint_position)[0:3, 3] - inertia = param[19:28] inertia = inertia.reshape((3, 3)) - if(not use_fixed_inertia): + if not use_fixed_inertia: inertia = self.mass_mass_fun(H, joint_position)[3:6, 3:6] - # FINAL linear_com_acc STATE (2) - temp = foot_force_fl@stanceFL - temp += foot_force_fr@stanceFR - temp += foot_force_rl@stanceRL - temp += foot_force_rr@stanceRR + temp = foot_force_fl @ stanceFL + temp += foot_force_fr @ stanceFR + temp += foot_force_rl @ stanceRL + temp += foot_force_rr @ stanceRR temp += external_wrench_linear gravity = np.array([0, 0, -9.81]) mass = self.kindyn.get_total_mass() - linear_com_acc = (1/mass)@temp + gravity - + linear_com_acc = (1 / mass) @ temp + gravity # FINAL euler_rates_base STATE (3) - temp2 = cs.skew(self.foot_position_fl - com_position)@foot_force_fl@stanceFL - temp2 += cs.skew(self.foot_position_fr - com_position)@foot_force_fr@stanceFR - temp2 += cs.skew(self.foot_position_rl - com_position)@foot_force_rl@stanceRL - temp2 += cs.skew(self.foot_position_rr - com_position)@foot_force_rr@stanceRR + temp2 = cs.skew(self.foot_position_fl - com_position) @ foot_force_fl @ stanceFL + temp2 += cs.skew(self.foot_position_fr - com_position) @ foot_force_fr @ stanceFR + temp2 += cs.skew(self.foot_position_rl - com_position) @ foot_force_rl @ stanceRL + temp2 += cs.skew(self.foot_position_rr - com_position) @ foot_force_rr @ stanceRR temp2 = temp2 + external_wrench_angular - angular_acc_base = cs.inv(inertia)@(b_R_w@temp2 - cs.skew(w)@inertia@w) + angular_acc_base = cs.inv(inertia) @ (b_R_w @ temp2 - cs.skew(w) @ inertia @ w) - #breakpoint() + # breakpoint() - if(not use_centroidal_model): - u_wrenches = self.jacobian_FL_fun(H, joint_position)[0:3, :].T@foot_force_fl@stanceFL - u_wrenches += self.jacobian_FR_fun(H, joint_position)[0:3, :].T@foot_force_fr@stanceFR - u_wrenches += self.jacobian_RL_fun(H, joint_position)[0:3, :].T@foot_force_rl@stanceRL - u_wrenches += self.jacobian_RR_fun(H, joint_position)[0:3, :].T@foot_force_rr@stanceRR + if not use_centroidal_model: + u_wrenches = self.jacobian_FL_fun(H, joint_position)[0:3, :].T @ foot_force_fl @ stanceFL + u_wrenches += self.jacobian_FR_fun(H, joint_position)[0:3, :].T @ foot_force_fr @ stanceFR + u_wrenches += self.jacobian_RL_fun(H, joint_position)[0:3, :].T @ foot_force_rl @ stanceRL + u_wrenches += self.jacobian_RR_fun(H, joint_position)[0:3, :].T @ foot_force_rr @ stanceRR u_wrenches = u_wrenches[0:6] - joints_velocities = inputs[0:12] base_velocities = cs.vertcat(linear_com_vel, w) eta = self.bias_force_fun(H, joint_position, base_velocities, joints_velocities)[0:6] inertia_base = self.mass_mass_fun(H, joint_position)[0:6, 0:6] - acc = cs.inv(inertia_base)@(-eta + u_wrenches) - - + acc = cs.inv(inertia_base) @ (-eta + u_wrenches) linear_com_acc = acc[0:3] angular_acc_base = acc[3:6] - - - # FINAL linear_foot_vel STATES (5,6,7,8) linear_joint_vel_FL = joint_velocity_fl linear_joint_vel_FR = joint_velocity_fr @@ -426,23 +422,38 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda integral_states[4] += roll integral_states[5] += pitch - # The order of the return should be equal to the order of the states_dot - return cs.vertcat(linear_com_vel, linear_com_acc, euler_rates_base, angular_acc_base, - linear_joint_vel_FL, linear_joint_vel_FR, linear_joint_vel_RL, linear_joint_vel_RR, integral_states) - - - - - def export_robot_model(self,) -> AcadosModel: + return cs.vertcat( + linear_com_vel, + linear_com_acc, + euler_rates_base, + angular_acc_base, + linear_joint_vel_FL, + linear_joint_vel_FR, + linear_joint_vel_RL, + linear_joint_vel_RR, + integral_states, + ) + + def export_robot_model( + self, + ) -> AcadosModel: """ This method set some general properties of the NMPC, such as the params, prediction mode, etc...! It will be called in centroidal_nmpc.py """ - + # dynamics - self.param = cs.vertcat(self.stance_param, self.mu_friction, self.stance_proximity, self.base_position, - self.base_yaw, self.external_wrench, self.inertia, self.mass) + self.param = cs.vertcat( + self.stance_param, + self.mu_friction, + self.stance_proximity, + self.base_position, + self.base_yaw, + self.external_wrench, + self.inertia, + self.mass, + ) f_expl = self.forward_dynamics(self.states, self.inputs, self.param) f_impl = self.states_dot - f_expl @@ -455,6 +466,4 @@ def export_robot_model(self,) -> AcadosModel: acados_model.p = self.param acados_model.name = "kinodynamic_model" - return acados_model - diff --git a/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_nmpc.py b/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_nmpc.py index dfbc5eb..7df705f 100644 --- a/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_nmpc.py +++ b/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_nmpc.py @@ -1,62 +1,57 @@ # Description: This file contains the class for the NMPC controller -# Authors: Giulio Turrisi - +# Authors: Giulio Turrisi - -from acados_template import AcadosOcp, AcadosOcpSolver -from .kinodynamic_model import KinoDynamic_Model +import copy +import os + +import casadi as cs import numpy as np import scipy.linalg -import casadi as cs -import copy -import math +from acados_template import AcadosOcp, AcadosOcpSolver -import time +from .kinodynamic_model import KinoDynamic_Model -import os dir_path = os.path.dirname(os.path.realpath(__file__)) import sys + sys.path.append(dir_path) -sys.path.append(dir_path + '/../../') +sys.path.append(dir_path + "/../../") + + +from quadruped_pympc import config -from quadruped_pympc import config -from liecasadi import SO3 # Class for the Acados NMPC, the model is in another file! class Acados_NMPC_KinoDynamic: def __init__(self): - - self.horizon = config.mpc_params['horizon'] # Define the number of discretization steps - self.dt = config.mpc_params['dt'] - self.T_horizon = self.horizon*self.dt - self.use_RTI = config.mpc_params['use_RTI'] - self.use_integrators = config.mpc_params['use_integrators'] - self.use_warm_start = config.mpc_params['use_warm_start'] - self.use_foothold_constraints = config.mpc_params['use_foothold_constraints'] - - - self.use_static_stability = config.mpc_params['use_static_stability'] - self.use_zmp_stability = config.mpc_params['use_zmp_stability'] + self.horizon = config.mpc_params["horizon"] # Define the number of discretization steps + self.dt = config.mpc_params["dt"] + self.T_horizon = self.horizon * self.dt + self.use_RTI = config.mpc_params["use_RTI"] + self.use_integrators = config.mpc_params["use_integrators"] + self.use_warm_start = config.mpc_params["use_warm_start"] + self.use_foothold_constraints = config.mpc_params["use_foothold_constraints"] + + self.use_static_stability = config.mpc_params["use_static_stability"] + self.use_zmp_stability = config.mpc_params["use_zmp_stability"] self.use_stability_constraints = self.use_static_stability or self.use_zmp_stability - self.use_DDP = config.mpc_params['use_DDP'] + self.use_DDP = config.mpc_params["use_DDP"] - self.previous_status = -1 self.previous_contact_sequence = np.zeros((4, self.horizon)) self.optimal_next_state = np.zeros((24,)) self.previous_optimal_GRF = np.zeros((12,)) - self.integral_errors = np.zeros((6,)) - # For centering the variable around 0, 0, 0 (World frame) self.initial_base_position = np.array([0, 0, 0]) self.previous_yaw = None - # Create the class of the centroidal model and instantiate the acados model self.centroidal_model = KinoDynamic_Model() acados_model = self.centroidal_model.export_robot_model() @@ -65,10 +60,7 @@ def __init__(self): # Create the acados ocp solver self.ocp = self.create_ocp_solver_description(acados_model) - self.acados_ocp_solver = AcadosOcpSolver( - self.ocp, json_file="nmpc_kynodinamic" + ".json" - ) - + self.acados_ocp_solver = AcadosOcpSolver(self.ocp, json_file="nmpc_kynodinamic" + ".json") # Initialize solver for stage in range(self.horizon + 1): @@ -76,17 +68,12 @@ def __init__(self): for stage in range(self.horizon): self.acados_ocp_solver.set(stage, "u", np.zeros((self.inputs_dim,))) - if(self.use_RTI): + if self.use_RTI: # first preparation phase - self.acados_ocp_solver.options_set('rti_phase', 1) - status = self.acados_ocp_solver.solve() - - - - - + self.acados_ocp_solver.options_set("rti_phase", 1) + status = self.acados_ocp_solver.solve() - # Set cost, constraints and options + # Set cost, constraints and options def create_ocp_solver_description(self, acados_model) -> AcadosOcp: # Create ocp object to formulate the OCP ocp = AcadosOcp() @@ -103,23 +90,32 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: ocp.cost.cost_type = "NONLINEAR_LS" ocp.cost.cost_type_e = "NONLINEAR_LS" - - - ocp.model.cost_y_expr = cs.vertcat(ocp.model.x[0:12], self.centroidal_model.foot_position_fl, - self.centroidal_model.foot_position_fr, self.centroidal_model.foot_position_rl, - self.centroidal_model.foot_position_rr, ocp.model.x[24:], ocp.model.u, - ocp.model.x[12:24]) - ocp.model.cost_y_expr_e = cs.vertcat(ocp.model.x[0:12], self.centroidal_model.foot_position_fl, - self.centroidal_model.foot_position_fr, self.centroidal_model.foot_position_rl, - self.centroidal_model.foot_position_rr, ocp.model.x[24:], - ocp.model.x[12:24]) + ocp.model.cost_y_expr = cs.vertcat( + ocp.model.x[0:12], + self.centroidal_model.foot_position_fl, + self.centroidal_model.foot_position_fr, + self.centroidal_model.foot_position_rl, + self.centroidal_model.foot_position_rr, + ocp.model.x[24:], + ocp.model.u, + ocp.model.x[12:24], + ) + ocp.model.cost_y_expr_e = cs.vertcat( + ocp.model.x[0:12], + self.centroidal_model.foot_position_fl, + self.centroidal_model.foot_position_fr, + self.centroidal_model.foot_position_rl, + self.centroidal_model.foot_position_rr, + ocp.model.x[24:], + ocp.model.x[12:24], + ) ny = nx + 12 + nu ny_e = nx + 12 ocp.cost.W_e = Q_mat ocp.cost.W = scipy.linalg.block_diag(Q_mat, R_mat) - + ocp.cost.Vx = np.zeros((ny, nx)) ocp.cost.Vx[:nx, :nx] = np.eye(nx) @@ -133,55 +129,46 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: ocp.cost.yref_e = np.zeros((ny_e,)) # Set friction and foothold constraints - expr_h_friction, \ - self.constr_uh_friction, \ - self.constr_lh_friction = self.create_friction_cone_constraints() - + expr_h_friction, self.constr_uh_friction, self.constr_lh_friction = self.create_friction_cone_constraints() + ocp.model.con_h_expr = expr_h_friction ocp.constraints.uh = self.constr_uh_friction ocp.constraints.lh = self.constr_lh_friction - #ocp.model.con_h_expr_0 = expr_h_friction - #ocp.constraints.uh_0 = self.constr_uh_friction - #ocp.constraints.lh_0 = self.constr_lh_friction - nsh = expr_h_friction.shape[0] + # ocp.model.con_h_expr_0 = expr_h_friction + # ocp.constraints.uh_0 = self.constr_uh_friction + # ocp.constraints.lh_0 = self.constr_lh_friction + nsh = expr_h_friction.shape[0] nsh_state_constraint_start = copy.copy(nsh) - - if(self.use_foothold_constraints): - expr_h_foot, \ - self.constr_uh_foot, \ - self.constr_lh_foot = self.create_foothold_constraints() + if self.use_foothold_constraints: + expr_h_foot, self.constr_uh_foot, self.constr_lh_foot = self.create_foothold_constraints() - ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_foot) + ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_foot) ocp.constraints.uh = np.concatenate((ocp.constraints.uh, self.constr_uh_foot)) ocp.constraints.lh = np.concatenate((ocp.constraints.lh, self.constr_lh_foot)) nsh += expr_h_foot.shape[0] - - # Set stability constraints - if(self.use_stability_constraints): + # Set stability constraints + if self.use_stability_constraints: self.nsh_stability_start = copy.copy(nsh) - expr_h_support_polygon, \ - self.constr_uh_support_polygon, \ - self.constr_lh_support_polygon = self.create_stability_constraints() + expr_h_support_polygon, self.constr_uh_support_polygon, self.constr_lh_support_polygon = ( + self.create_stability_constraints() + ) - ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_support_polygon) + ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_support_polygon) ocp.constraints.uh = np.concatenate((ocp.constraints.uh, self.constr_uh_support_polygon)) ocp.constraints.lh = np.concatenate((ocp.constraints.lh, self.constr_lh_support_polygon)) nsh += expr_h_support_polygon.shape[0] self.nsh_stability_end = copy.copy(nsh) - self.nsh_foot_vel_start = copy.copy(nsh) - - expr_h_foot_vel, \ - self.constr_uh_foot_vel, \ - self.constr_lh_foot_vel = self.create_foot_vel_constraints() - ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_foot_vel) + + expr_h_foot_vel, self.constr_uh_foot_vel, self.constr_lh_foot_vel = self.create_foot_vel_constraints() + ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_foot_vel) ocp.constraints.uh = np.concatenate((ocp.constraints.uh, self.constr_uh_foot_vel)) ocp.constraints.lh = np.concatenate((ocp.constraints.lh, self.constr_lh_foot_vel)) - ocp.model.con_h_expr_0 = cs.vertcat(expr_h_friction,expr_h_foot_vel) + ocp.model.con_h_expr_0 = cs.vertcat(expr_h_friction, expr_h_foot_vel) ocp.constraints.uh_0 = np.concatenate((self.constr_uh_friction, self.constr_uh_foot_vel)) ocp.constraints.lh_0 = np.concatenate((self.constr_lh_friction, self.constr_lh_foot_vel)) @@ -189,19 +176,23 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: nsh_state_constraint_end = copy.copy(nsh) - # Set slack variable configuration: num_state_cstr = nsh_state_constraint_end - nsh_state_constraint_start - if(num_state_cstr > 0): - ocp.constraints.lsh = np.zeros(num_state_cstr) # Lower bounds on slacks corresponding to soft lower bounds for nonlinear constraints - ocp.constraints.ush = np.zeros(num_state_cstr) # Lower bounds on slacks corresponding to soft upper bounds for nonlinear constraints - ocp.constraints.idxsh = np.array(range(nsh_state_constraint_start, nsh_state_constraint_end)) # Jsh + if num_state_cstr > 0: + ocp.constraints.lsh = np.zeros( + num_state_cstr + ) # Lower bounds on slacks corresponding to soft lower bounds for nonlinear constraints + ocp.constraints.ush = np.zeros( + num_state_cstr + ) # Lower bounds on slacks corresponding to soft upper bounds for nonlinear constraints + ocp.constraints.idxsh = np.array(range(nsh_state_constraint_start, nsh_state_constraint_end)) # Jsh ns = num_state_cstr - ocp.cost.zl = 1000 * np.ones((ns,)) # gradient wrt lower slack at intermediate shooting nodes (1 to N-1) - ocp.cost.Zl = 1 * np.ones((ns,)) # diagonal of Hessian wrt lower slack at intermediate shooting nodes (1 to N-1) - ocp.cost.zu = 1000 * np.ones((ns,)) - ocp.cost.Zu = 1 * np.ones((ns,)) - + ocp.cost.zl = 1000 * np.ones((ns,)) # gradient wrt lower slack at intermediate shooting nodes (1 to N-1) + ocp.cost.Zl = 1 * np.ones( + (ns,) + ) # diagonal of Hessian wrt lower slack at intermediate shooting nodes (1 to N-1) + ocp.cost.zu = 1000 * np.ones((ns,)) + ocp.cost.Zu = 1 * np.ones((ns,)) # Variables to save the upper and lower bound of the constraints applied list_upper_bound = [] @@ -212,13 +203,12 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: self.upper_bound = np.array(list_upper_bound, dtype=object) self.lower_bound = np.array(list_lower_bound, dtype=object) - # Set initial state constraint X0 = np.zeros(shape=(nx,)) ocp.constraints.x0 = X0 # Set initialize parameters - init_contact_status = np.array([1., 1., 1., 1.]) + init_contact_status = np.array([1.0, 1.0, 1.0, 1.0]) init_mu = np.array([0.5]) init_stance_proximity = np.array([0, 0, 0, 0]) init_base_position = np.array([0, 0, 0]) @@ -226,91 +216,98 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: init_external_wrench = np.array([0, 0, 0, 0, 0, 0]) init_inertia = np.ones((9,)) init_mass = np.array([config.mass]) - - ocp.parameter_values = np.concatenate((init_contact_status, init_mu, init_stance_proximity, - init_base_position, init_base_yaw, init_external_wrench, - init_inertia, init_mass)) + ocp.parameter_values = np.concatenate( + ( + init_contact_status, + init_mu, + init_stance_proximity, + init_base_position, + init_base_yaw, + init_external_wrench, + init_inertia, + init_mass, + ) + ) - # Set options - ocp.solver_options.qp_solver = "PARTIAL_CONDENSING_HPIPM" # FULL_CONDENSING_QPOASES PARTIAL_CONDENSING_OSQP PARTIAL_CONDENSING_HPIPM + ocp.solver_options.qp_solver = ( + "PARTIAL_CONDENSING_HPIPM" # FULL_CONDENSING_QPOASES PARTIAL_CONDENSING_OSQP PARTIAL_CONDENSING_HPIPM + ) ocp.solver_options.hessian_approx = "GAUSS_NEWTON" # 'GAUSS_NEWTON', 'EXACT' - ocp.solver_options.integrator_type = "ERK" #ERK IRK GNSF DISCRETE - if(self.use_DDP): - ocp.solver_options.nlp_solver_type = 'DDP' - ocp.solver_options.nlp_solver_max_iter = config.mpc_params['num_qp_iterations'] - #ocp.solver_options.globalization = 'MERIT_BACKTRACKING' + ocp.solver_options.integrator_type = "ERK" # ERK IRK GNSF DISCRETE + if self.use_DDP: + ocp.solver_options.nlp_solver_type = "DDP" + ocp.solver_options.nlp_solver_max_iter = config.mpc_params["num_qp_iterations"] + # ocp.solver_options.globalization = 'MERIT_BACKTRACKING' ocp.solver_options.with_adaptive_levenberg_marquardt = True - ocp.translate_to_feasibility_problem(keep_x0=True, keep_cost=True) - elif(self.use_RTI): - ocp.solver_options.nlp_solver_type = "SQP_RTI" + elif self.use_RTI: + ocp.solver_options.nlp_solver_type = "SQP_RTI" ocp.solver_options.nlp_solver_max_iter = 1 - # Set the RTI type for the advanced RTI method + # Set the RTI type for the advanced RTI method # (see https://arxiv.org/pdf/2403.07101.pdf) - if(config.mpc_params['as_rti_type'] == "AS-RTI-A"): + if config.mpc_params["as_rti_type"] == "AS-RTI-A": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 0 - elif(config.mpc_params['as_rti_type'] == "AS-RTI-B"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-B": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 1 - elif(config.mpc_params['as_rti_type'] == "AS-RTI-C"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-C": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 2 - elif(config.mpc_params['as_rti_type'] == "AS-RTI-D"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-D": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 3 else: - ocp.solver_options.nlp_solver_type = "SQP" - ocp.solver_options.nlp_solver_max_iter = config.mpc_params['num_qp_iterations'] - #ocp.solver_options.globalization = "MERIT_BACKTRACKING" # FIXED_STEP, MERIT_BACKTRACKING - - if(config.mpc_params['solver_mode'] == "balance"): + ocp.solver_options.nlp_solver_type = "SQP" + ocp.solver_options.nlp_solver_max_iter = config.mpc_params["num_qp_iterations"] + # ocp.solver_options.globalization = "MERIT_BACKTRACKING" # FIXED_STEP, MERIT_BACKTRACKING + + if config.mpc_params["solver_mode"] == "balance": ocp.solver_options.hpipm_mode = "BALANCE" - elif(config.mpc_params['solver_mode'] == "robust"): + elif config.mpc_params["solver_mode"] == "robust": ocp.solver_options.hpipm_mode = "ROBUST" - elif(config.mpc_params['solver_mode'] == "speed"): + elif config.mpc_params["solver_mode"] == "speed": ocp.solver_options.qp_solver_iter_max = 10 ocp.solver_options.hpipm_mode = "SPEED" - elif(config.mpc_params['solver_mode'] == "crazy_speed"): + elif config.mpc_params["solver_mode"] == "crazy_speed": ocp.solver_options.qp_solver_iter_max = 5 ocp.solver_options.hpipm_mode = "SPEED_ABS" ocp.solver_options.line_search_use_sufficient_descent = 1 - #ocp.solver_options.regularize_method = "PROJECT_REDUC_HESS" - #ocp.solver_options.nlp_solver_ext_qp_res = 1 + # ocp.solver_options.regularize_method = "PROJECT_REDUC_HESS" + # ocp.solver_options.nlp_solver_ext_qp_res = 1 ocp.solver_options.levenberg_marquardt = 1e-2 - # Set prediction horizon ocp.solver_options.tf = self.T_horizon - # Nonuniform discretization - if(config.mpc_params['use_nonuniform_discretization']): - time_steps_fine_grained = np.tile(config.mpc_params['dt_fine_grained'], config.mpc_params['horizon_fine_grained']) - time_steps = np.concatenate((time_steps_fine_grained, np.tile(self.dt, self.horizon-config.mpc_params['horizon_fine_grained']))) - shooting_nodes = np.zeros((self.horizon+1,)) + if config.mpc_params["use_nonuniform_discretization"]: + time_steps_fine_grained = np.tile( + config.mpc_params["dt_fine_grained"], config.mpc_params["horizon_fine_grained"] + ) + time_steps = np.concatenate( + (time_steps_fine_grained, np.tile(self.dt, self.horizon - config.mpc_params["horizon_fine_grained"])) + ) + shooting_nodes = np.zeros((self.horizon + 1,)) for i in range(len(time_steps)): - shooting_nodes[i+1] = shooting_nodes[i] + time_steps[i] + shooting_nodes[i + 1] = shooting_nodes[i] + time_steps[i] ocp.solver_options.shooting_nodes = shooting_nodes - - - return ocp - - def create_foot_vel_constraints(self,) -> None: - + def create_foot_vel_constraints( + self, + ) -> None: qvel_joints_FL = self.centroidal_model.inputs[0:3] qvel_joints_FR = self.centroidal_model.inputs[3:6] qvel_joints_RL = self.centroidal_model.inputs[6:9] - qvel_joints_RR = self.centroidal_model.inputs[9:12] + qvel_joints_RR = self.centroidal_model.inputs[9:12] joint_position = self.centroidal_model.states[12:24] com_position = self.centroidal_model.states[0:3] @@ -324,72 +321,68 @@ def create_foot_vel_constraints(self,) -> None: H[0:3, 0:3] = b_R_w.T H[0:3, 3] = com_position - - qvel = cs.vertcat(com_velocity, com_angular_velocity, qvel_joints_FL, qvel_joints_FR, qvel_joints_RL, qvel_joints_RR) - - #breakpoint() - - foot_vel_FL = self.centroidal_model.jacobian_FL_fun(H, joint_position)[0:3,:]@qvel #qvel_FL - foot_vel_FR = self.centroidal_model.jacobian_FR_fun(H, joint_position)[0:3,:]@qvel #qvel_FR - foot_vel_RL = self.centroidal_model.jacobian_RL_fun(H, joint_position)[0:3,:]@qvel #qvel_RL - foot_vel_RR = self.centroidal_model.jacobian_RR_fun(H, joint_position)[0:3,:]@qvel #qvel_RR + qvel = cs.vertcat( + com_velocity, com_angular_velocity, qvel_joints_FL, qvel_joints_FR, qvel_joints_RL, qvel_joints_RR + ) + # breakpoint() - ub = np.ones(12)*1000 - lb = -np.ones(12)*1000 + foot_vel_FL = self.centroidal_model.jacobian_FL_fun(H, joint_position)[0:3, :] @ qvel # qvel_FL + foot_vel_FR = self.centroidal_model.jacobian_FR_fun(H, joint_position)[0:3, :] @ qvel # qvel_FR + foot_vel_RL = self.centroidal_model.jacobian_RL_fun(H, joint_position)[0:3, :] @ qvel # qvel_RL + foot_vel_RR = self.centroidal_model.jacobian_RR_fun(H, joint_position)[0:3, :] @ qvel # qvel_RR + ub = np.ones(12) * 1000 + lb = -np.ones(12) * 1000 Jb = cs.vertcat(foot_vel_FL, foot_vel_FR, foot_vel_RL, foot_vel_RR) - + return Jb, ub, lb - # Create a constraint for stability (COM, ZMP or CP inside support polygon) - def create_stability_constraints(self,) -> None: - + def create_stability_constraints( + self, + ) -> None: base_w = self.centroidal_model.states[0:3] base_vel_w = self.centroidal_model.states[3:6] - FL = self.centroidal_model.foot_position_fl FR = self.centroidal_model.foot_position_fr RL = self.centroidal_model.foot_position_rl RR = self.centroidal_model.foot_position_rr - #yaw = self.centroidal_model.base_yaw[0] + # yaw = self.centroidal_model.base_yaw[0] yaw = self.centroidal_model.states[8] h_R_w = cs.SX.zeros(2, 2) - h_R_w[0,0] = cs.cos(yaw) - h_R_w[0,1] = cs.sin(yaw) - h_R_w[1,0] = -cs.sin(yaw) - h_R_w[1,1] = cs.cos(yaw) - - FL[0:2] = h_R_w@(FL[0:2] - base_w[0:2]) - FR[0:2] = h_R_w@(FR[0:2] - base_w[0:2]) - RL[0:2] = h_R_w@(RL[0:2] - base_w[0:2]) - RR[0:2] = h_R_w@(RR[0:2] - base_w[0:2]) - - - if(self.use_static_stability): + h_R_w[0, 0] = cs.cos(yaw) + h_R_w[0, 1] = cs.sin(yaw) + h_R_w[1, 0] = -cs.sin(yaw) + h_R_w[1, 1] = cs.cos(yaw) + + FL[0:2] = h_R_w @ (FL[0:2] - base_w[0:2]) + FR[0:2] = h_R_w @ (FR[0:2] - base_w[0:2]) + RL[0:2] = h_R_w @ (RL[0:2] - base_w[0:2]) + RR[0:2] = h_R_w @ (RR[0:2] - base_w[0:2]) + + if self.use_static_stability: x = 0.0 y = 0.0 else: # Compute the ZMP robotHeight = base_w[2] - foot_force_fl = self.centroidal_model.inputs[12:15]#@self.centroidal_model.param[0] - foot_force_fr = self.centroidal_model.inputs[15:18]#@self.centroidal_model.param[1] - foot_force_rl = self.centroidal_model.inputs[18:21]#@self.centroidal_model.param[2] - foot_force_rr = self.centroidal_model.inputs[21:24]#@self.centroidal_model.param[3] + foot_force_fl = self.centroidal_model.inputs[12:15] # @self.centroidal_model.param[0] + foot_force_fr = self.centroidal_model.inputs[15:18] # @self.centroidal_model.param[1] + foot_force_rl = self.centroidal_model.inputs[18:21] # @self.centroidal_model.param[2] + foot_force_rr = self.centroidal_model.inputs[21:24] # @self.centroidal_model.param[3] temp = foot_force_fl + foot_force_fr + foot_force_rl + foot_force_rr gravity = np.array([0, 0, -9.81]) - linear_com_acc = (1/self.centroidal_model.mass)@temp + gravity - zmp = base_w[0:2] - linear_com_acc[0:2]*(robotHeight/(-gravity[2])) - #zmp = base_w[0:2] - base_vel_w[0:2]*(robotHeight/gravity[2]) - zmp = h_R_w@(zmp - base_w[0:2]) + linear_com_acc = (1 / self.centroidal_model.mass) @ temp + gravity + zmp = base_w[0:2] - linear_com_acc[0:2] * (robotHeight / (-gravity[2])) + # zmp = base_w[0:2] - base_vel_w[0:2]*(robotHeight/gravity[2]) + zmp = h_R_w @ (zmp - base_w[0:2]) x = zmp[0] y = zmp[1] - y_FL = FL[1] y_FR = FR[1] y_RL = RL[1] @@ -400,67 +393,65 @@ def create_stability_constraints(self,) -> None: x_RL = RL[0] x_RR = RR[0] - #LF - RF : x < (x2 - x1) (y - y1) / (y2 - y1) + x1 - #RF - RH: y > (y2 - y1) (x - x1) / (x2 - x1) + y1 - #RH - LH : x > (x2 - x1) (y - y1) / (y2 - y1) + x1 - #LH - LF: y < (y2 - y1) (x - x1) / (x2 - x1) + y1 - - #FL and FR cannot stay at the same x! #constrint should be less than zero - constraint_FL_FR = x - (x_FR - x_FL)*(y - y_FL) / (y_FR - y_FL + 0.001) - x_FL - - #FR and RR cannot stay at the same y! #constraint should be bigger than zero - constraint_FR_RR = y - (y_RR - y_FR)*(x - x_FR) / (x_RR - x_FR + 0.001) - y_FR - - #RL and RR cannot stay at the same x! #constraint should be bigger than zero - constraint_RR_RL = x - (x_RL - x_RR)*(y - y_RR) / (y_RL - y_RR + 0.001) - x_RR - - #FL and RL cannot stay at the same y! #constraint should be less than zero - constraint_RL_FL = y - (y_FL - y_RL)*(x - x_RL) / (x_FL - x_RL + 0.001) - y_RL - + # LF - RF : x < (x2 - x1) (y - y1) / (y2 - y1) + x1 + # RF - RH: y > (y2 - y1) (x - x1) / (x2 - x1) + y1 + # RH - LH : x > (x2 - x1) (y - y1) / (y2 - y1) + x1 + # LH - LF: y < (y2 - y1) (x - x1) / (x2 - x1) + y1 + + # FL and FR cannot stay at the same x! #constrint should be less than zero + constraint_FL_FR = x - (x_FR - x_FL) * (y - y_FL) / (y_FR - y_FL + 0.001) - x_FL + + # FR and RR cannot stay at the same y! #constraint should be bigger than zero + constraint_FR_RR = y - (y_RR - y_FR) * (x - x_FR) / (x_RR - x_FR + 0.001) - y_FR + + # RL and RR cannot stay at the same x! #constraint should be bigger than zero + constraint_RR_RL = x - (x_RL - x_RR) * (y - y_RR) / (y_RL - y_RR + 0.001) - x_RR + + # FL and RL cannot stay at the same y! #constraint should be less than zero + constraint_RL_FL = y - (y_FL - y_RL) * (x - x_RL) / (x_FL - x_RL + 0.001) - y_RL + # the diagonal stuff can be at any point... - constraint_FL_RR = y - (y_RR - y_FL)*(x - x_FL) / (x_RR - x_FL + 0.001) - y_FL #bigger than zero - constraint_FR_RL = y - (y_RL - y_FR)*(x - x_FR) / (x_RL - x_FR + 0.001) - y_FR #bigger than zero + constraint_FL_RR = y - (y_RR - y_FL) * (x - x_FL) / (x_RR - x_FL + 0.001) - y_FL # bigger than zero + constraint_FR_RL = y - (y_RL - y_FR) * (x - x_FR) / (x_RL - x_FR + 0.001) - y_FR # bigger than zero # upper and lower bound - ub = np.ones(6)*1000 - lb = -np.ones(6)*1000 + ub = np.ones(6) * 1000 + lb = -np.ones(6) * 1000 - #constraint_FL_FR + # constraint_FL_FR ub[0] = 0 lb[0] = -1000 - - #constraint_FR_RR + + # constraint_FR_RR ub[1] = 1000 lb[1] = 0 - #constraint_RR_RL + # constraint_RR_RL ub[2] = 1000 lb[2] = 0 - #constraint_RL_FL + # constraint_RL_FL ub[3] = 0 lb[3] = -1000 - - #constraint_FL_RR + + # constraint_FL_RR ub[4] = 0 lb[4] = -1000 - - #constraint_FR_RL + + # constraint_FR_RL ub[5] = 0 lb[5] = -1000 - - - Jb = cs.vertcat(constraint_FL_FR, constraint_FR_RR, - constraint_RR_RL, constraint_RL_FL, - constraint_FL_RR, constraint_FR_RL) + Jb = cs.vertcat( + constraint_FL_FR, constraint_FR_RR, constraint_RR_RL, constraint_RL_FL, constraint_FL_RR, constraint_FR_RL + ) return Jb, ub, lb - - # Create a standard foothold box constraint - def create_foothold_constraints(self,): + def create_foothold_constraints( + self, + ): """ This function calculates the symbolic foothold constraints for the centroidal NMPC problem. @@ -470,52 +461,45 @@ def create_foothold_constraints(self,): lbu: A numpy array of shape (8,) representing the lower bounds of the foothold constraints. """ - ubu = np.ones(12)*1000 - lbu = -np.ones(12)*1000 - + ubu = np.ones(12) * 1000 + lbu = -np.ones(12) * 1000 - # The foothold constraint in acados are in the horizontal frame, + # The foothold constraint in acados are in the horizontal frame, # but they arrive to us in the world frame. We need to rotate them # using the robot yaw yaw = self.centroidal_model.base_yaw[0] h_R_w = cs.SX.zeros(2, 2) - h_R_w[0,0] = cs.cos(yaw) - h_R_w[0,1] = cs.sin(yaw) - h_R_w[1,0] = -cs.sin(yaw) - h_R_w[1,1] = cs.cos(yaw) - + h_R_w[0, 0] = cs.cos(yaw) + h_R_w[0, 1] = cs.sin(yaw) + h_R_w[1, 0] = -cs.sin(yaw) + h_R_w[1, 1] = cs.cos(yaw) # and translate them using the robot base position base = self.centroidal_model.base_position[0:3] - - - foot_position_fl = cs.SX.zeros(3,1) - foot_position_fl[0:2] = h_R_w@cs.vertcat(self.centroidal_model.foot_position_fl - base[0:2]) + foot_position_fl = cs.SX.zeros(3, 1) + foot_position_fl[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.foot_position_fl - base[0:2]) foot_position_fl[2] = self.centroidal_model.states[14] - - foot_position_fr = cs.SX.zeros(3,1) - foot_position_fr[0:2] = h_R_w@cs.vertcat(self.centroidal_model.foot_position_fr - base[0:2]) + + foot_position_fr = cs.SX.zeros(3, 1) + foot_position_fr[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.foot_position_fr - base[0:2]) foot_position_fr[2] = self.centroidal_model.states[17] - - foot_position_rl = cs.SX.zeros(3,1) - foot_position_rl[0:2] = h_R_w@cs.vertcat(self.centroidal_model.foot_position_rl - base[0:2]) + + foot_position_rl = cs.SX.zeros(3, 1) + foot_position_rl[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.foot_position_rl - base[0:2]) foot_position_rl[2] = self.centroidal_model.states[20] - - foot_position_rr = cs.SX.zeros(3,1) - foot_position_rr[0:2] = h_R_w@cs.vertcat(self.centroidal_model.foot_position_rr - base[0:2]) + + foot_position_rr = cs.SX.zeros(3, 1) + foot_position_rr[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.foot_position_rr - base[0:2]) foot_position_rr[2] = self.centroidal_model.states[23] - - - Jbu = cs.vertcat(foot_position_fl, foot_position_fr, - foot_position_rl, foot_position_rr) + Jbu = cs.vertcat(foot_position_fl, foot_position_fr, foot_position_rl, foot_position_rr) return Jbu, ubu, lbu - - # Create the friction cone constraint - def create_friction_cone_constraints(self,) ->None: + def create_friction_cone_constraints( + self, + ) -> None: """ This function calculates the symbolic friction cone constraints for the centroidal NMPC problem. @@ -528,47 +512,45 @@ def create_friction_cone_constraints(self,) ->None: t = np.array([1, 0, 0]) b = np.array([0, 1, 0]) mu = self.centroidal_model.mu_friction - f_max = config.mpc_params['grf_max'] - f_min = config.mpc_params['grf_min'] + f_max = config.mpc_params["grf_max"] + f_min = config.mpc_params["grf_min"] # Derivation can be found in the paper # "High-slope terrain locomotion for torque-controlled quadruped robots", # M Focchi, A Del Prete, I Havoutis, R Featherstone, DG Caldwell, C Semini Jbu = cs.SX.zeros(20, 12) - Jbu[0, :3] = -n*mu + t - Jbu[1, :3] = -n*mu + b - Jbu[2, :3] = n*mu + b - Jbu[3, :3] = n*mu + t + Jbu[0, :3] = -n * mu + t + Jbu[1, :3] = -n * mu + b + Jbu[2, :3] = n * mu + b + Jbu[3, :3] = n * mu + t Jbu[4, :3] = n - Jbu[5, 3:6] = -n*mu + t - Jbu[6, 3:6] = -n*mu + b - Jbu[7, 3:6] = n*mu + b - Jbu[8, 3:6] = n*mu + t + Jbu[5, 3:6] = -n * mu + t + Jbu[6, 3:6] = -n * mu + b + Jbu[7, 3:6] = n * mu + b + Jbu[8, 3:6] = n * mu + t Jbu[9, 3:6] = n - - Jbu[10, 6:9] = -n*mu + t - Jbu[11, 6:9] = -n*mu + b - Jbu[12, 6:9] = n*mu + b - Jbu[13, 6:9] = n*mu + t + + Jbu[10, 6:9] = -n * mu + t + Jbu[11, 6:9] = -n * mu + b + Jbu[12, 6:9] = n * mu + b + Jbu[13, 6:9] = n * mu + t Jbu[14, 6:9] = n - - Jbu[15, 9:12] = -n*mu + t - Jbu[16, 9:12] = -n*mu + b - Jbu[17, 9:12] = n*mu + b - Jbu[18, 9:12] = n*mu + t + + Jbu[15, 9:12] = -n * mu + t + Jbu[16, 9:12] = -n * mu + b + Jbu[17, 9:12] = n * mu + b + Jbu[18, 9:12] = n * mu + t Jbu[19, 9:12] = n # C matrix - Jbu = Jbu@cs.vertcat(self.centroidal_model.inputs[12:24]) - - + Jbu = Jbu @ cs.vertcat(self.centroidal_model.inputs[12:24]) # lower bound (-1000 is "almost" -inf) lbu = np.zeros(20) lbu[0] = -10000000 lbu[1] = -10000000 - lbu[2] = 0 + lbu[2] = 0 lbu[3] = 0 lbu[4] = f_min lbu[5:10] = lbu[0:5] @@ -577,7 +559,7 @@ def create_friction_cone_constraints(self,) ->None: # upper bound (1000 is "almost" inf) ubu = np.zeros(20) - ubu[0] = 0 + ubu[0] = 0 ubu[1] = 0 ubu[2] = 10000000 ubu[3] = 10000000 @@ -588,61 +570,87 @@ def create_friction_cone_constraints(self,) ->None: return Jbu, ubu, lbu - - def set_weight(self, nx, nu): # Define the weight matrices for the cost function - Q_position = np.array([0, 0, 1500]) #x, y, z - Q_velocity = np.array([200, 200, 200]) #x_vel, y_vel, z_vel - Q_base_angle = np.array([500, 500, 0]) #roll, pitch, yaw - Q_base_angle_rates = np.array([20, 20, 50]) #roll_rate, pitch_rate, yaw_rate - - Q_foot_pos = np.array([300, 300, 30000]) #f_x, f_y, f_z (should be 4 times this, once per foot) - Q_com_position_z_integral = np.array([50]) #integral of z_com - Q_com_velocity_x_integral = np.array([10]) #integral of x_com - Q_com_velocity_y_integral = np.array([10]) #integral of y_com - Q_com_velocity_z_integral = np.array([10]) #integral of z_com_vel - Q_roll_integral_integral = np.array([10]) #integral of roll - Q_pitch_integral_integral = np.array([10]) #integral of pitch + Q_position = np.array([0, 0, 1500]) # x, y, z + Q_velocity = np.array([200, 200, 200]) # x_vel, y_vel, z_vel + Q_base_angle = np.array([500, 500, 0]) # roll, pitch, yaw + Q_base_angle_rates = np.array([20, 20, 50]) # roll_rate, pitch_rate, yaw_rate - Q_joint_angle = np.array([0.1, 0.1, 0.1]) + Q_foot_pos = np.array([300, 300, 30000]) # f_x, f_y, f_z (should be 4 times this, once per foot) + Q_com_position_z_integral = np.array([50]) # integral of z_com + Q_com_velocity_x_integral = np.array([10]) # integral of x_com + Q_com_velocity_y_integral = np.array([10]) # integral of y_com + Q_com_velocity_z_integral = np.array([10]) # integral of z_com_vel + Q_roll_integral_integral = np.array([10]) # integral of roll + Q_pitch_integral_integral = np.array([10]) # integral of pitch + Q_joint_angle = np.array([0.1, 0.1, 0.1]) - R_joint_vel = np.array([0.0001, 0.0001, 0.0001]) - if(config.robot == "hyqreal"): - R_foot_force = np.array([0.00001, 0.00001, 0.00001])#f_x, f_y, f_z (should be 4 times this, once per foot) + R_joint_vel = np.array([0.0001, 0.0001, 0.0001]) + if config.robot == "hyqreal": + R_foot_force = np.array( + [0.00001, 0.00001, 0.00001] + ) # f_x, f_y, f_z (should be 4 times this, once per foot) else: R_foot_force = np.array([0.001, 0.001, 0.001]) - Q_mat = np.diag(np.concatenate((Q_position, Q_velocity, - Q_base_angle, Q_base_angle_rates, - Q_foot_pos, Q_foot_pos, Q_foot_pos, Q_foot_pos, - Q_com_position_z_integral, Q_com_velocity_x_integral, - Q_com_velocity_y_integral, Q_com_velocity_z_integral, - Q_roll_integral_integral, Q_pitch_integral_integral, - Q_joint_angle, Q_joint_angle, Q_joint_angle, Q_joint_angle))) - - R_mat = np.diag(np.concatenate((R_joint_vel, R_joint_vel, R_joint_vel, R_joint_vel, - R_foot_force, R_foot_force, R_foot_force, R_foot_force))) + Q_mat = np.diag( + np.concatenate( + ( + Q_position, + Q_velocity, + Q_base_angle, + Q_base_angle_rates, + Q_foot_pos, + Q_foot_pos, + Q_foot_pos, + Q_foot_pos, + Q_com_position_z_integral, + Q_com_velocity_x_integral, + Q_com_velocity_y_integral, + Q_com_velocity_z_integral, + Q_roll_integral_integral, + Q_pitch_integral_integral, + Q_joint_angle, + Q_joint_angle, + Q_joint_angle, + Q_joint_angle, + ) + ) + ) - return Q_mat, R_mat - + R_mat = np.diag( + np.concatenate( + ( + R_joint_vel, + R_joint_vel, + R_joint_vel, + R_joint_vel, + R_foot_force, + R_foot_force, + R_foot_force, + R_foot_force, + ) + ) + ) + return Q_mat, R_mat def reset(self): self.acados_ocp_solver.reset() - self.acados_ocp_solver = AcadosOcpSolver(self.ocp, - json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json", - build=False, generate=False) - - - + self.acados_ocp_solver = AcadosOcpSolver( + self.ocp, + json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json", + build=False, + generate=False, + ) def set_stage_constraint(self, constraint, state, reference, contact_sequence, h_R_w, stance_proximity=None): """ Set the stage constraint for the centroidal NMPC problem. We only consider the stance constraint, and the swing - constraint up to 2 maximum references. + constraint up to 2 maximum references. Args: constraint (numpy.ndarray or None): Constraint passed from outside (e.g. vision). If None, nominal foothold is used. @@ -655,10 +663,9 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h Returns: None """ - #try: + # try: - - # Take the array of the contact sequence and split + # Take the array of the contact sequence and split # it in 4 arrays for clarity FL_contact_sequence = contact_sequence[0] FR_contact_sequence = contact_sequence[1] @@ -676,74 +683,67 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h RL_reference_foot = reference["ref_foot_RL"] RR_reference_foot = reference["ref_foot_RR"] - - # Take the base position and the yaw rotation matrix. This is needed to + # Take the base position and the yaw rotation matrix. This is needed to # express the foothold constraint in the horizontal frame base = state["position"] - - h_R_w = h_R_w.reshape((2,2)) - + h_R_w = h_R_w.reshape((2, 2)) # Divide the constraint in upper and lower bound. The constraint are - # represented by 4 vertex, but we only use box constraint hence + # represented by 4 vertex, but we only use box constraint hence # we need only 2 vertex for each constraint (first an last) # For the leg in stance now, we simply enlarge the actual position as a box constraint # maybe this is not needed, but we cannot disable constraint!! - + # FL Stance Constraint stance_up_constraint_FL = np.array([FL_actual_foot[0], FL_actual_foot[1], FL_actual_foot[2] + 0.002]) - stance_up_constraint_FL[0:2] = h_R_w@(stance_up_constraint_FL[0:2] - base[0:2]) + stance_up_constraint_FL[0:2] = h_R_w @ (stance_up_constraint_FL[0:2] - base[0:2]) stance_up_constraint_FL[0:2] = stance_up_constraint_FL[0:2] + 0.1 stance_up_constraint_FL[2] = stance_up_constraint_FL[2] + 0.01 - + stance_low_constraint_FL = np.array([FL_actual_foot[0], FL_actual_foot[1], FL_actual_foot[2] - 0.002]) - stance_low_constraint_FL[0:2] = h_R_w@(stance_low_constraint_FL[0:2] - base[0:2]) + stance_low_constraint_FL[0:2] = h_R_w @ (stance_low_constraint_FL[0:2] - base[0:2]) stance_low_constraint_FL[0:2] = stance_low_constraint_FL[0:2] - 0.1 stance_low_constraint_FL[2] = stance_low_constraint_FL[2] - 0.01 - # FR Stance Constraint stance_up_constraint_FR = np.array([FR_actual_foot[0], FR_actual_foot[1], FR_actual_foot[2] + 0.002]) - stance_up_constraint_FR[0:2] = h_R_w@(stance_up_constraint_FR[0:2] - base[0:2]) + stance_up_constraint_FR[0:2] = h_R_w @ (stance_up_constraint_FR[0:2] - base[0:2]) stance_up_constraint_FR[0:2] = stance_up_constraint_FR[0:2] + 0.1 stance_up_constraint_FR[2] = stance_up_constraint_FR[2] + 0.01 stance_low_constraint_FR = np.array([FR_actual_foot[0], FR_actual_foot[1], FR_actual_foot[2] - 0.002]) - stance_low_constraint_FR[0:2] = h_R_w@(stance_low_constraint_FR[0:2] - base[0:2]) + stance_low_constraint_FR[0:2] = h_R_w @ (stance_low_constraint_FR[0:2] - base[0:2]) stance_low_constraint_FR[0:2] = stance_low_constraint_FR[0:2] - 0.1 stance_low_constraint_FR[2] = stance_low_constraint_FR[2] - 0.01 - # RL Stance Constraint stance_up_constraint_RL = np.array([RL_actual_foot[0], RL_actual_foot[1], RL_actual_foot[2] + 0.002]) - stance_up_constraint_RL[0:2] = h_R_w@(stance_up_constraint_RL[0:2]- base[0:2]) + stance_up_constraint_RL[0:2] = h_R_w @ (stance_up_constraint_RL[0:2] - base[0:2]) stance_up_constraint_RL[0:2] = stance_up_constraint_RL[0:2] + 0.1 stance_up_constraint_RL[2] = stance_up_constraint_RL[2] + 0.01 - + stance_low_constraint_RL = np.array([RL_actual_foot[0], RL_actual_foot[1], RL_actual_foot[2] - 0.002]) - stance_low_constraint_RL[0:2] = h_R_w@(stance_low_constraint_RL[0:2]- base[0:2]) + stance_low_constraint_RL[0:2] = h_R_w @ (stance_low_constraint_RL[0:2] - base[0:2]) stance_low_constraint_RL[0:2] = stance_low_constraint_RL[0:2] - 0.1 stance_low_constraint_RL[2] = stance_low_constraint_RL[2] - 0.01 - # RR Stance Constraint stance_up_constraint_RR = np.array([RR_actual_foot[0], RR_actual_foot[1], RR_actual_foot[2] + 0.002]) - stance_up_constraint_RR[0:2] = h_R_w@(stance_up_constraint_RR[0:2]- base[0:2]) + stance_up_constraint_RR[0:2] = h_R_w @ (stance_up_constraint_RR[0:2] - base[0:2]) stance_up_constraint_RR[0:2] = stance_up_constraint_RR[0:2] + 0.1 stance_up_constraint_RR[2] = stance_up_constraint_RR[2] + 0.01 - + stance_low_constraint_RR = np.array([RR_actual_foot[0], RR_actual_foot[1], RR_actual_foot[2] - 0.002]) - stance_low_constraint_RR[0:2] = h_R_w@(stance_low_constraint_RR[0:2]- base[0:2]) + stance_low_constraint_RR[0:2] = h_R_w @ (stance_low_constraint_RR[0:2] - base[0:2]) stance_low_constraint_RR[0:2] = stance_low_constraint_RR[0:2] - 0.1 stance_low_constraint_RR[2] = stance_low_constraint_RR[2] - 0.01 - # Constraint for the first footholds at the next touchdown. If constraint == True # we have some constraint passed from outside (e.g. vision), otherwise we use the - # nominal foothold enlarged as we do previously - if(constraint is not None): + # nominal foothold enlarged as we do previously + if constraint is not None: # From the VFA first_up_constraint_FL = np.array([constraint[0][0], constraint[1][0], constraint[2][0] + 0.002]) first_up_constraint_FR = np.array([constraint[0][1], constraint[1][1], constraint[2][1] + 0.002]) @@ -754,75 +754,85 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h first_low_constraint_FR = np.array([constraint[9][1], constraint[10][1], constraint[11][1] - 0.002]) first_low_constraint_RL = np.array([constraint[9][2], constraint[10][2], constraint[11][2] - 0.002]) first_low_constraint_RR = np.array([constraint[9][3], constraint[10][3], constraint[11][3] - 0.002]) - + # Rotate the constraint in the horizontal frame - first_up_constraint_FL[0:2] = h_R_w@(first_up_constraint_FL[0:2] - base[0:2]) + first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - base[0:2]) first_up_constraint_FL = first_up_constraint_FL + 0.005 - #first_up_constraint_FL[2] = first_up_constraint_FL[2] + 0.5 - - first_up_constraint_FR[0:2] = h_R_w@(first_up_constraint_FR[0:2] - base[0:2]) + # first_up_constraint_FL[2] = first_up_constraint_FL[2] + 0.5 + + first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - base[0:2]) first_up_constraint_FR = first_up_constraint_FR + 0.005 - #first_up_constraint_FR[2] = first_up_constraint_FR[2] + 0.5 + # first_up_constraint_FR[2] = first_up_constraint_FR[2] + 0.5 - first_up_constraint_RL[0:2] = h_R_w@(first_up_constraint_RL[0:2] - base[0:2]) + first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - base[0:2]) first_up_constraint_RL = first_up_constraint_RL + 0.005 - #first_up_constraint_RL[2] = first_up_constraint_RL[2] + 0.5 + # first_up_constraint_RL[2] = first_up_constraint_RL[2] + 0.5 - first_up_constraint_RR[0:2] = h_R_w@(first_up_constraint_RR[0:2] - base[0:2]) + first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - base[0:2]) first_up_constraint_RR = first_up_constraint_RR + 0.005 - #first_up_constraint_RR[2] = first_up_constraint_RR[2] + 0.5 + # first_up_constraint_RR[2] = first_up_constraint_RR[2] + 0.5 - first_low_constraint_FL[0:2] = h_R_w@(first_low_constraint_FL[0:2] - base[0:2]) + first_low_constraint_FL[0:2] = h_R_w @ (first_low_constraint_FL[0:2] - base[0:2]) first_low_constraint_FL = first_low_constraint_FL - 0.005 - #first_low_constraint_FL[2] = first_low_constraint_FL[2] - 0.5 + # first_low_constraint_FL[2] = first_low_constraint_FL[2] - 0.5 - first_low_constraint_FR[0:2] = h_R_w@(first_low_constraint_FR[0:2] - base[0:2]) + first_low_constraint_FR[0:2] = h_R_w @ (first_low_constraint_FR[0:2] - base[0:2]) first_low_constraint_FR = first_low_constraint_FR - 0.005 - #first_low_constraint_FR[2] = first_low_constraint_FR[2] - 0.5 - - first_low_constraint_RL[0:2] = h_R_w@(first_low_constraint_RL[0:2] - base[0:2]) + # first_low_constraint_FR[2] = first_low_constraint_FR[2] - 0.5 + + first_low_constraint_RL[0:2] = h_R_w @ (first_low_constraint_RL[0:2] - base[0:2]) first_low_constraint_RL = first_low_constraint_RL - 0.005 - #first_low_constraint_RL[2] = first_low_constraint_RL[2] - 0.5 + # first_low_constraint_RL[2] = first_low_constraint_RL[2] - 0.5 - first_low_constraint_RR[0:2] = h_R_w@(first_low_constraint_RR[0:2] - base[0:2]) + first_low_constraint_RR[0:2] = h_R_w @ (first_low_constraint_RR[0:2] - base[0:2]) first_low_constraint_RR = first_low_constraint_RR - 0.005 - #first_low_constraint_RR[2] = first_low_constraint_RR[2] - 0.5 + # first_low_constraint_RR[2] = first_low_constraint_RR[2] - 0.5 else: # Constrain taken from the nominal foothold (NO VISION) # FL first touchdown constraint - first_up_constraint_FL = np.array([FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] + 0.002]) - first_up_constraint_FL[0:2] = h_R_w@(first_up_constraint_FL[0:2] - base[0:2]) + 0.15 - - first_low_constraint_FL = np.array([FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] - 0.002]) - first_low_constraint_FL[0:2] = h_R_w@(first_low_constraint_FL[0:2] - base[0:2]) - 0.15 - + first_up_constraint_FL = np.array( + [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] + 0.002] + ) + first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - base[0:2]) + 0.15 - # FR first touchdown constraint - first_up_constraint_FR = np.array([FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] + 0.002]) - first_up_constraint_FR[0:2] = h_R_w@(first_up_constraint_FR[0:2] - base[0:2]) + 0.15 + first_low_constraint_FL = np.array( + [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] - 0.002] + ) + first_low_constraint_FL[0:2] = h_R_w @ (first_low_constraint_FL[0:2] - base[0:2]) - 0.15 - first_low_constraint_FR = np.array([FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] - 0.002]) - first_low_constraint_FR[0:2] = h_R_w@(first_low_constraint_FR[0:2] - base[0:2]) - 0.15 + # FR first touchdown constraint + first_up_constraint_FR = np.array( + [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] + 0.002] + ) + first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - base[0:2]) + 0.15 + first_low_constraint_FR = np.array( + [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] - 0.002] + ) + first_low_constraint_FR[0:2] = h_R_w @ (first_low_constraint_FR[0:2] - base[0:2]) - 0.15 # RL first touchdown constraint - first_up_constraint_RL = np.array([RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] + 0.002]) - first_up_constraint_RL[0:2] = h_R_w@(first_up_constraint_RL[0:2] - base[0:2]) + 0.15 - - first_low_constraint_RL = np.array([RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] - 0.002]) - first_low_constraint_RL[0:2] = h_R_w@(first_low_constraint_RL[0:2] - base[0:2]) - 0.15 + first_up_constraint_RL = np.array( + [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] + 0.002] + ) + first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - base[0:2]) + 0.15 + first_low_constraint_RL = np.array( + [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] - 0.002] + ) + first_low_constraint_RL[0:2] = h_R_w @ (first_low_constraint_RL[0:2] - base[0:2]) - 0.15 # RR first touchdown constraint - first_up_constraint_RR = np.array([RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] + 0.002]) - first_up_constraint_RR[0:2] = h_R_w@(first_up_constraint_RR[0:2] - base[0:2]) + 0.15 - - first_low_constraint_RR = np.array([RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] - 0.002]) - first_low_constraint_RR[0:2] = h_R_w@(first_low_constraint_RR[0:2] - base[0:2]) - 0.15 - + first_up_constraint_RR = np.array( + [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] + 0.002] + ) + first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - base[0:2]) + 0.15 - + first_low_constraint_RR = np.array( + [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] - 0.002] + ) + first_low_constraint_RR[0:2] = h_R_w @ (first_low_constraint_RR[0:2] - base[0:2]) - 0.15 # we stack all the constraint we have for now up_constraint_FL = np.vstack((stance_up_constraint_FL, first_up_constraint_FL)) @@ -835,115 +845,112 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h low_constraint_RL = np.vstack((stance_low_constraint_RL, first_low_constraint_RL)) low_constraint_RR = np.vstack((stance_low_constraint_RR, first_low_constraint_RR)) - - # If there are more than 1 nominal foothold per leg, we create additional constraints # We do not expect more than two reference footholds... # FL second touchdown constraint - if(FL_reference_foot.shape[0] == 2): - second_up_constraint_FL = np.array([FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] + 0.002]) - second_up_constraint_FL[0:2] = h_R_w@(second_up_constraint_FL[0:2] - base[0:2]) + 0.15 - - second_low_constraint_FL = np.array([FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] - 0.002]) - second_low_constraint_FL[0:2] = h_R_w@(second_low_constraint_FL[0:2] - base[0:2]) - 0.15 + if FL_reference_foot.shape[0] == 2: + second_up_constraint_FL = np.array( + [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] + 0.002] + ) + second_up_constraint_FL[0:2] = h_R_w @ (second_up_constraint_FL[0:2] - base[0:2]) + 0.15 + second_low_constraint_FL = np.array( + [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] - 0.002] + ) + second_low_constraint_FL[0:2] = h_R_w @ (second_low_constraint_FL[0:2] - base[0:2]) - 0.15 up_constraint_FL = np.vstack((up_constraint_FL, second_up_constraint_FL)) low_constraint_FL = np.vstack((low_constraint_FL, second_low_constraint_FL)) # FR second touchdown constraint - if(FR_reference_foot.shape[0] == 2): - second_up_constraint_FR = np.array([FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] + 0.002]) - second_up_constraint_FR[0:2] = h_R_w@(second_up_constraint_FR[0:2] - base[0:2]) + 0.15 - - - second_low_constraint_FR = np.array([FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] - 0.002]) - second_low_constraint_FR[0:2] = h_R_w@(second_low_constraint_FR[0:2]- base[0:2]) - 0.15 + if FR_reference_foot.shape[0] == 2: + second_up_constraint_FR = np.array( + [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] + 0.002] + ) + second_up_constraint_FR[0:2] = h_R_w @ (second_up_constraint_FR[0:2] - base[0:2]) + 0.15 + second_low_constraint_FR = np.array( + [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] - 0.002] + ) + second_low_constraint_FR[0:2] = h_R_w @ (second_low_constraint_FR[0:2] - base[0:2]) - 0.15 up_constraint_FR = np.vstack((up_constraint_FR, second_up_constraint_FR)) low_constraint_FR = np.vstack((low_constraint_FR, second_low_constraint_FR)) - - # RL second touchdown constraint - if(RL_reference_foot.shape[0] == 2): - second_up_constraint_RL = np.array([RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] + 0.002]) - second_up_constraint_RL[0:2] = h_R_w@(second_up_constraint_RL[0:2]- base[0:2]) + 0.15 - - second_low_constraint_RL = np.array([RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] - 0.002]) - second_low_constraint_RL[0:2] = h_R_w@(second_low_constraint_RL[0:2] - base[0:2]) - 0.15 + # RL second touchdown constraint + if RL_reference_foot.shape[0] == 2: + second_up_constraint_RL = np.array( + [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] + 0.002] + ) + second_up_constraint_RL[0:2] = h_R_w @ (second_up_constraint_RL[0:2] - base[0:2]) + 0.15 + second_low_constraint_RL = np.array( + [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] - 0.002] + ) + second_low_constraint_RL[0:2] = h_R_w @ (second_low_constraint_RL[0:2] - base[0:2]) - 0.15 up_constraint_RL = np.vstack((up_constraint_RL, second_up_constraint_RL)) low_constraint_RL = np.vstack((low_constraint_RL, second_low_constraint_RL)) - - # RR second touchdown constraint - if(RR_reference_foot.shape[0] == 2): - second_up_constraint_RR = np.array([RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] + 0.002]) - second_up_constraint_RR[0:2] = h_R_w@(second_up_constraint_RR[0:2] - base[0:2]) + 0.15 - - second_low_constraint_RR = np.array([RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] - 0.002]) - second_low_constraint_RR[0:2] = h_R_w@(second_low_constraint_RR[0:2] - base[0:2]) - 0.15 + # RR second touchdown constraint + if RR_reference_foot.shape[0] == 2: + second_up_constraint_RR = np.array( + [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] + 0.002] + ) + second_up_constraint_RR[0:2] = h_R_w @ (second_up_constraint_RR[0:2] - base[0:2]) + 0.15 + second_low_constraint_RR = np.array( + [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] - 0.002] + ) + second_low_constraint_RR[0:2] = h_R_w @ (second_low_constraint_RR[0:2] - base[0:2]) - 0.15 up_constraint_RR = np.vstack((up_constraint_RR, second_up_constraint_RR)) low_constraint_RR = np.vstack((low_constraint_RR, second_low_constraint_RR)) - - # We pass all the constraints (foothold and friction conte) to acados. # Then one regarding friction are precomputed ub_friction = self.constr_uh_friction lb_friction = self.constr_lh_friction - - # If the foothold is in swing, the idx of the constraint start from 1 idx_constraint = np.array([0, 0, 0, 0]) - if(FL_contact_sequence[0] == 0): + if FL_contact_sequence[0] == 0: idx_constraint[0] = 1 - if(FR_contact_sequence[0] == 0): + if FR_contact_sequence[0] == 0: idx_constraint[1] = 1 - if(RL_contact_sequence[0] == 0): + if RL_contact_sequence[0] == 0: idx_constraint[2] = 1 - if(RR_contact_sequence[0] == 0): + if RR_contact_sequence[0] == 0: idx_constraint[3] = 1 - - for j in range(0, self.horizon): - # take the constraint for the current timestep - ub_foot_FL = up_constraint_FL[idx_constraint[0]] - lb_foot_FL = low_constraint_FL[idx_constraint[0]] + for j in range(0, self.horizon): + # take the constraint for the current timestep + ub_foot_FL = up_constraint_FL[idx_constraint[0]] + lb_foot_FL = low_constraint_FL[idx_constraint[0]] - ub_foot_FR = up_constraint_FR[idx_constraint[1]] + ub_foot_FR = up_constraint_FR[idx_constraint[1]] lb_foot_FR = low_constraint_FR[idx_constraint[1]] - ub_foot_RL = up_constraint_RL[idx_constraint[2]] + ub_foot_RL = up_constraint_RL[idx_constraint[2]] lb_foot_RL = low_constraint_RL[idx_constraint[2]] - ub_foot_RR = up_constraint_RR[idx_constraint[3]] + ub_foot_RR = up_constraint_RR[idx_constraint[3]] lb_foot_RR = low_constraint_RR[idx_constraint[3]] - - - # Concatenate the friction and foothold constraint - ub_foot = copy.deepcopy(np.concatenate((ub_foot_FL, ub_foot_FR, - ub_foot_RL, ub_foot_RR))) - lb_foot = copy.deepcopy(np.concatenate((lb_foot_FL, lb_foot_FR, - lb_foot_RL, lb_foot_RR))) - if(self.use_foothold_constraints): + # Concatenate the friction and foothold constraint + ub_foot = copy.deepcopy(np.concatenate((ub_foot_FL, ub_foot_FR, ub_foot_RL, ub_foot_RR))) + lb_foot = copy.deepcopy(np.concatenate((lb_foot_FL, lb_foot_FR, lb_foot_RL, lb_foot_RR))) + if self.use_foothold_constraints: ub_total = np.concatenate((ub_friction, ub_foot)) lb_total = np.concatenate((lb_friction, lb_foot)) else: ub_total = ub_friction lb_total = lb_friction - - #Constraints for the support polygon depending on the leg in stance + # Constraints for the support polygon depending on the leg in stance # all disabled at the beginning!! - if(self.use_stability_constraints): + if self.use_stability_constraints: ub_support_FL_FR = 1000 lb_support_FL_FR = -1000 @@ -955,20 +962,21 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h ub_support_RL_FL = 1000 lb_support_RL_FL = -1000 - + ub_support_FL_RR = 1000 lb_support_FL_RR = -1000 ub_support_FR_RL = 1000 - lb_support_FR_RL = -1000 - - + lb_support_FR_RL = -1000 + # We have 4 cases for the stability constraint: trot, pace, crawl, full stance - if(FL_contact_sequence[j] == 1 and - FR_contact_sequence[j] == 1 and - RL_contact_sequence[j] == 1 and - RR_contact_sequence[j] == 1): - #FULL STANCE TODO + if ( + FL_contact_sequence[j] == 1 + and FR_contact_sequence[j] == 1 + and RL_contact_sequence[j] == 1 + and RR_contact_sequence[j] == 1 + ): + # FULL STANCE TODO ub_support_FL_FR = 1000 lb_support_FL_FR = -1000 @@ -980,194 +988,190 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h ub_support_RL_FL = 1000 lb_support_RL_FL = -1000 - + ub_support_FL_RR = 1000 lb_support_FL_RR = -1000 ub_support_FR_RL = 1000 lb_support_FR_RL = -1000 - elif(np.array_equal(FL_contact_sequence, RR_contact_sequence) - and np.array_equal(FR_contact_sequence, RL_contact_sequence)): - #TROT - stability_margin = config.mpc_params['trot_stability_margin'] - if(FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0): + elif np.array_equal(FL_contact_sequence, RR_contact_sequence) and np.array_equal( + FR_contact_sequence, RL_contact_sequence + ): + # TROT + stability_margin = config.mpc_params["trot_stability_margin"] + if FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0: ub_support_FL_RR = 0 + stability_margin lb_support_FL_RR = 0 - stability_margin - - if(FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0): + + if FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0: ub_support_FR_RL = 0 + stability_margin lb_support_FR_RL = 0 - stability_margin - elif(np.array_equal(FL_contact_sequence, RL_contact_sequence) - and np.array_equal(FR_contact_sequence, RR_contact_sequence)): - #PACE - stability_margin = config.mpc_params['pace_stability_margin'] - if(FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0): + elif np.array_equal(FL_contact_sequence, RL_contact_sequence) and np.array_equal( + FR_contact_sequence, RR_contact_sequence + ): + # PACE + stability_margin = config.mpc_params["pace_stability_margin"] + if FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0: ub_support_RL_FL = 0 + stability_margin lb_support_RL_FL = 0 - stability_margin - - if(FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0): + if FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0: ub_support_FR_RR = 0 + stability_margin lb_support_FR_RR = 0 - stability_margin - + else: - #CRAWL BACKDIAGONALCRAWL ONLY - stability_margin = config.mpc_params['crawl_stability_margin'] + # CRAWL BACKDIAGONALCRAWL ONLY + stability_margin = config.mpc_params["crawl_stability_margin"] - if(FL_contact_sequence[j] == 1): - if(FR_contact_sequence[j] == 1): - ub_support_FL_FR = -0.0 - stability_margin + if FL_contact_sequence[j] == 1: + if FR_contact_sequence[j] == 1: + ub_support_FL_FR = -0.0 - stability_margin lb_support_FL_FR = -1000 else: ub_support_FL_RR = 1000 - lb_support_FL_RR = 0.0 + stability_margin - + lb_support_FL_RR = 0.0 + stability_margin - if(FR_contact_sequence[j] == 1): - if(RR_contact_sequence[j] == 1): + if FR_contact_sequence[j] == 1: + if RR_contact_sequence[j] == 1: ub_support_FR_RR = 1000 lb_support_FR_RR = 0.0 + stability_margin else: ub_support_FR_RL = 1000 lb_support_FR_RL = 0.0 + stability_margin - - if(RR_contact_sequence[j] == 1): - if(RL_contact_sequence[j] == 1): + if RR_contact_sequence[j] == 1: + if RL_contact_sequence[j] == 1: ub_support_RR_RL = 1000 - lb_support_RR_RL = 0.0 + stability_margin + lb_support_RR_RL = 0.0 + stability_margin else: - ub_support_FL_RR = -0.0 - stability_margin + ub_support_FL_RR = -0.0 - stability_margin lb_support_FL_RR = -1000 - - if(RL_contact_sequence[j] == 1): - if(FL_contact_sequence[j] == 1): - ub_support_RL_FL = -0.0 - stability_margin + if RL_contact_sequence[j] == 1: + if FL_contact_sequence[j] == 1: + ub_support_RL_FL = -0.0 - stability_margin lb_support_RL_FL = -1000 else: ub_support_FR_RL = -0.0 - stability_margin lb_support_FR_RL = -1000 + ub_support = np.array( + [ + ub_support_FL_FR, + ub_support_FR_RR, + ub_support_RR_RL, + ub_support_RL_FL, + ub_support_FL_RR, + ub_support_FR_RL, + ] + ) + lb_support = np.array( + [ + lb_support_FL_FR, + lb_support_FR_RR, + lb_support_RR_RL, + lb_support_RL_FL, + lb_support_FL_RR, + lb_support_FR_RL, + ] + ) - - - ub_support = np.array([ub_support_FL_FR, ub_support_FR_RR, ub_support_RR_RL, ub_support_RL_FL, - ub_support_FL_RR, ub_support_FR_RL]) - lb_support = np.array([lb_support_FL_FR, lb_support_FR_RR, lb_support_RR_RL, lb_support_RL_FL, - lb_support_FL_RR, lb_support_FR_RL]) - ub_total = np.concatenate((ub_total, ub_support)) lb_total = np.concatenate((lb_total, lb_support)) - # No friction constraint at the end! we don't have u_N - if(j == self.horizon): - if(self.use_foothold_constraints): - if(self.use_stability_constraints): + if j == self.horizon: + if self.use_foothold_constraints: + if self.use_stability_constraints: ub_total = np.concatenate((ub_foot, ub_support)) lb_total = np.concatenate((lb_foot, lb_support)) else: ub_total = ub_foot lb_total = lb_foot else: - if(self.use_stability_constraints): + if self.use_stability_constraints: ub_total = ub_support lb_total = lb_support else: continue # No slip constraint - ub_no_slip_foot_vel = self.constr_uh_foot_vel - ub_no_slip_foot_vel[0:3] = ub_no_slip_foot_vel[0:3]*(1-contact_sequence[0,j]) - ub_no_slip_foot_vel[3:6] = ub_no_slip_foot_vel[3:6]*(1-contact_sequence[1,j]) - ub_no_slip_foot_vel[6:9] = ub_no_slip_foot_vel[6:9]*(1-contact_sequence[2,j]) - ub_no_slip_foot_vel[9:12] = ub_no_slip_foot_vel[9:12]*(1-contact_sequence[3,j]) - + ub_no_slip_foot_vel = self.constr_uh_foot_vel + ub_no_slip_foot_vel[0:3] = ub_no_slip_foot_vel[0:3] * (1 - contact_sequence[0, j]) + ub_no_slip_foot_vel[3:6] = ub_no_slip_foot_vel[3:6] * (1 - contact_sequence[1, j]) + ub_no_slip_foot_vel[6:9] = ub_no_slip_foot_vel[6:9] * (1 - contact_sequence[2, j]) + ub_no_slip_foot_vel[9:12] = ub_no_slip_foot_vel[9:12] * (1 - contact_sequence[3, j]) + lb_no_slip_foot_vel = self.constr_lh_foot_vel - lb_no_slip_foot_vel[0:3] = lb_no_slip_foot_vel[0:3]*(1-contact_sequence[0,j]) - lb_no_slip_foot_vel[3:6] = lb_no_slip_foot_vel[3:6]*(1-contact_sequence[1,j]) - lb_no_slip_foot_vel[6:9] = lb_no_slip_foot_vel[6:9]*(1-contact_sequence[2,j]) - lb_no_slip_foot_vel[9:12] = lb_no_slip_foot_vel[9:12]*(1-contact_sequence[3,j]) + lb_no_slip_foot_vel[0:3] = lb_no_slip_foot_vel[0:3] * (1 - contact_sequence[0, j]) + lb_no_slip_foot_vel[3:6] = lb_no_slip_foot_vel[3:6] * (1 - contact_sequence[1, j]) + lb_no_slip_foot_vel[6:9] = lb_no_slip_foot_vel[6:9] * (1 - contact_sequence[2, j]) + lb_no_slip_foot_vel[9:12] = lb_no_slip_foot_vel[9:12] * (1 - contact_sequence[3, j]) ub_total = np.concatenate((ub_friction, ub_no_slip_foot_vel)) lb_total = np.concatenate((lb_friction, lb_no_slip_foot_vel)) - #print("j", j) - #self.acados_ocp_solver.constraints_set(j, "uh", ub_total) - #self.acados_ocp_solver.constraints_set(j, "lh", lb_total) - + # print("j", j) + # self.acados_ocp_solver.constraints_set(j, "uh", ub_total) + # self.acados_ocp_solver.constraints_set(j, "lh", lb_total) # Only friction costraints at the beginning - #if(j == 0): + # if(j == 0): # self.acados_ocp_solver.constraints_set(j, "uh", ub_friction) # self.acados_ocp_solver.constraints_set(j, "lh", lb_friction) - #if(j > 0): + # if(j > 0): # self.acados_ocp_solver.constraints_set(j, "uh", ub_total) # self.acados_ocp_solver.constraints_set(j, "lh", lb_total) - - - #save the constraint for logging + # save the constraint for logging self.upper_bound[j] = ub_total.tolist() self.lower_bound[j] = lb_total.tolist() - - # ugly procedure to update the idx of the constraint - if(j>=1): - if(FL_contact_sequence[j] == 0 and FL_contact_sequence[j-1] == 1): - if(idx_constraint[0] < up_constraint_FL.shape[0] - 1): + if j >= 1: + if FL_contact_sequence[j] == 0 and FL_contact_sequence[j - 1] == 1: + if idx_constraint[0] < up_constraint_FL.shape[0] - 1: idx_constraint[0] += 1 - - if(FR_contact_sequence[j] == 0 and FR_contact_sequence[j-1] == 1): - if(idx_constraint[1] < up_constraint_FR.shape[0] - 1): + + if FR_contact_sequence[j] == 0 and FR_contact_sequence[j - 1] == 1: + if idx_constraint[1] < up_constraint_FR.shape[0] - 1: idx_constraint[1] += 1 - - if(RL_contact_sequence[j] == 0 and RL_contact_sequence[j-1] == 1): - if(idx_constraint[2] < up_constraint_RL.shape[0] - 1): + + if RL_contact_sequence[j] == 0 and RL_contact_sequence[j - 1] == 1: + if idx_constraint[2] < up_constraint_RL.shape[0] - 1: idx_constraint[2] += 1 - - if(RR_contact_sequence[j] == 0 and RR_contact_sequence[j-1] == 1): - if(idx_constraint[3] < up_constraint_RR.shape[0] - 1): + + if RR_contact_sequence[j] == 0 and RR_contact_sequence[j - 1] == 1: + if idx_constraint[3] < up_constraint_RR.shape[0] - 1: idx_constraint[3] += 1 - - #except: + + # except: # print("###WARNING: error in setting the constraints") return - - - - - # Method to perform the centering of the states and the reference around (0, 0, 0) - def perform_scaling(self, state, reference, constraint = None): - - + def perform_scaling(self, state, reference, constraint=None): self.initial_base_position = copy.deepcopy(state["position"]) reference = copy.deepcopy(reference) state = copy.deepcopy(state) - + reference["ref_position"] = reference["ref_position"] - state["position"] reference["ref_foot_FL"] = reference["ref_foot_FL"] - state["position"] reference["ref_foot_FR"] = reference["ref_foot_FR"] - state["position"] reference["ref_foot_RL"] = reference["ref_foot_RL"] - state["position"] reference["ref_foot_RR"] = reference["ref_foot_RR"] - state["position"] - state["foot_FL"] = state["foot_FL"] - state["position"] state["foot_FR"] = state["foot_FR"] - state["position"] state["foot_RL"] = state["foot_RL"] - state["position"] state["foot_RR"] = state["foot_RR"] - state["position"] state["position"] = np.array([0, 0, 0]) - # Perform a wrapping of the yaw angle - if(self.previous_yaw is None): + if self.previous_yaw is None: self.previous_yaw = state["orientation"][2] # Calculate the difference between the current yaw and the previous yaw @@ -1179,60 +1183,56 @@ def perform_scaling(self, state, reference, constraint = None): # Update the current yaw based on the normalized difference state["orientation"][2] = self.previous_yaw + yaw_diff - # Update the previous yaw self.previous_yaw = state["orientation"][2] return state, reference, constraint - - # Main loop for computing the control - def compute_control(self, state, reference, contact_sequence, constraint = None, external_wrenches = np.zeros((6,)), - inertia = config.inertia.reshape((9,)), mass = config.mass): - - - # Take the array of the contact sequence and split it in 4 arrays, + def compute_control( + self, + state, + reference, + contact_sequence, + constraint=None, + external_wrenches=np.zeros((6,)), + inertia=config.inertia.reshape((9,)), + mass=config.mass, + ): + # Take the array of the contact sequence and split it in 4 arrays, # one for each leg FL_contact_sequence = contact_sequence[0] FR_contact_sequence = contact_sequence[1] RL_contact_sequence = contact_sequence[2] RR_contact_sequence = contact_sequence[3] - - # Perform the scaling of the states and the reference - state, \ - reference, \ - constraint = self.perform_scaling(state, reference, constraint) - + state, reference, constraint = self.perform_scaling(state, reference, constraint) # Fill reference (self.states_dim+self.inputs_dim) - #idx_ref_foot_to_assign = np.array([0, 0, 0, 0]) + # idx_ref_foot_to_assign = np.array([0, 0, 0, 0]) for j in range(self.horizon): - yref = np.zeros(shape=(self.states_dim + self.inputs_dim + 12,)) - yref[0:3] = reference['ref_position'] - yref[3:6] = reference['ref_linear_velocity'] - yref[6:9] = reference['ref_orientation'] - yref[9:12] = reference['ref_angular_velocity'] - yref[12:15] = reference['ref_foot_FL'][j] - yref[15:18] = reference['ref_foot_FR'][j] - yref[18:21] = reference['ref_foot_RL'][j] - yref[21:24] = reference['ref_foot_RR'][j] - - + yref[0:3] = reference["ref_position"] + yref[3:6] = reference["ref_linear_velocity"] + yref[6:9] = reference["ref_orientation"] + yref[9:12] = reference["ref_angular_velocity"] + yref[12:15] = reference["ref_foot_FL"][j] + yref[15:18] = reference["ref_foot_FR"][j] + yref[18:21] = reference["ref_foot_RL"][j] + yref[21:24] = reference["ref_foot_RR"][j] # Calculate the reference force z for the leg in stance # It's simply mass*acc/number_of_legs_in_stance!! # Force x and y are always 0 - number_of_legs_in_stance = np.array([FL_contact_sequence[j], FR_contact_sequence[j], - RL_contact_sequence [j], RR_contact_sequence[j]]).sum() - if(number_of_legs_in_stance == 0): + number_of_legs_in_stance = np.array( + [FL_contact_sequence[j], FR_contact_sequence[j], RL_contact_sequence[j], RR_contact_sequence[j]] + ).sum() + if number_of_legs_in_stance == 0: reference_force_stance_legs = 0 else: reference_force_stance_legs = (mass * 9.81) / number_of_legs_in_stance - + reference_force_fl_z = reference_force_stance_legs * FL_contact_sequence[j] reference_force_fr_z = reference_force_stance_legs * FR_contact_sequence[j] reference_force_rl_z = reference_force_stance_legs * RL_contact_sequence[j] @@ -1244,112 +1244,130 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, yref[54:] = reference["ref_joints"] - # Setting the reference to acados - if(self.use_DDP): - if(j == 0): + if self.use_DDP: + if j == 0: num_l2_penalties = self.ocp.model.cost_y_expr_0.shape[0] - (self.states_dim + self.inputs_dim) else: num_l2_penalties = self.ocp.model.cost_y_expr.shape[0] - (self.states_dim + self.inputs_dim) - - yref_tot = np.concatenate((yref, np.zeros(num_l2_penalties,) )) + + yref_tot = np.concatenate( + ( + yref, + np.zeros( + num_l2_penalties, + ), + ) + ) self.acados_ocp_solver.set(j, "yref", yref_tot) else: self.acados_ocp_solver.set(j, "yref", yref) - - - # Fill last step horizon reference (self.states_dim - no control action!!) yref_N = np.zeros(shape=(self.states_dim + 12,)) - yref_N[0:3] = reference['ref_position'] - yref_N[3:6] = reference['ref_linear_velocity'] - yref_N[6:9] = reference['ref_orientation'] - yref_N[9:12] = reference['ref_angular_velocity'] - yref_N[12:15] = reference['ref_foot_FL'][-1] - yref_N[15:18] = reference['ref_foot_FR'][-1] - yref_N[18:21] = reference['ref_foot_RL'][-1] - yref_N[21:24] = reference['ref_foot_RR'][-1] - yref_N[self.states_dim:] = reference["ref_joints"] + yref_N[0:3] = reference["ref_position"] + yref_N[3:6] = reference["ref_linear_velocity"] + yref_N[6:9] = reference["ref_orientation"] + yref_N[9:12] = reference["ref_angular_velocity"] + yref_N[12:15] = reference["ref_foot_FL"][-1] + yref_N[15:18] = reference["ref_foot_FR"][-1] + yref_N[18:21] = reference["ref_foot_RL"][-1] + yref_N[21:24] = reference["ref_foot_RR"][-1] + yref_N[self.states_dim :] = reference["ref_joints"] # Setting the reference to acados self.acados_ocp_solver.set(self.horizon, "yref", yref_N) - - - - # Fill stance param, friction and stance proximity + # Fill stance param, friction and stance proximity # (stance proximity will disable foothold optimization near a stance!!) - mu = config.mpc_params['mu'] + mu = config.mpc_params["mu"] yaw = state["orientation"][2] - - # Stance Proximity - to be removed - stance_proximity_FL = np.zeros((self.horizon, )) - stance_proximity_FR = np.zeros((self.horizon, )) - stance_proximity_RL = np.zeros((self.horizon, )) - stance_proximity_RR = np.zeros((self.horizon, )) - - + # Stance Proximity - to be removed + stance_proximity_FL = np.zeros((self.horizon,)) + stance_proximity_FR = np.zeros((self.horizon,)) + stance_proximity_RL = np.zeros((self.horizon,)) + stance_proximity_RR = np.zeros((self.horizon,)) # Set the parameters to acados for j in range(self.horizon): # If we have estimated an external wrench, we can compensate it for all steps # or less (maybe the disturbance is not costant along the horizon!) - if(config.mpc_params['external_wrenches_compensation'] and - config.mpc_params['external_wrenches_compensation_num_step'] and - j < config.mpc_params['external_wrenches_compensation_num_step']): + if ( + config.mpc_params["external_wrenches_compensation"] + and config.mpc_params["external_wrenches_compensation_num_step"] + and j < config.mpc_params["external_wrenches_compensation_num_step"] + ): external_wrenches_estimated_param = copy.deepcopy(external_wrenches) - external_wrenches_estimated_param = external_wrenches_estimated_param.reshape((6, )) + external_wrenches_estimated_param = external_wrenches_estimated_param.reshape((6,)) else: external_wrenches_estimated_param = np.zeros((6,)) - - - param = np.array([FL_contact_sequence[j], FR_contact_sequence[j], - RL_contact_sequence[j], RR_contact_sequence[j], mu, - stance_proximity_FL[j], - stance_proximity_FR[j], - stance_proximity_RL[j], - stance_proximity_RR[j], - state["position"][0], state["position"][1], - state["position"][2], state["orientation"][2], - external_wrenches_estimated_param[0], external_wrenches_estimated_param[1], - external_wrenches_estimated_param[2], external_wrenches_estimated_param[3], - external_wrenches_estimated_param[4], external_wrenches_estimated_param[5], - inertia[0], inertia[1], inertia[2], inertia[3], inertia[4], inertia[5], - inertia[6], inertia[7], inertia[8], mass]) - self.acados_ocp_solver.set(j, "p", copy.deepcopy(param)) - + param = np.array( + [ + FL_contact_sequence[j], + FR_contact_sequence[j], + RL_contact_sequence[j], + RR_contact_sequence[j], + mu, + stance_proximity_FL[j], + stance_proximity_FR[j], + stance_proximity_RL[j], + stance_proximity_RR[j], + state["position"][0], + state["position"][1], + state["position"][2], + state["orientation"][2], + external_wrenches_estimated_param[0], + external_wrenches_estimated_param[1], + external_wrenches_estimated_param[2], + external_wrenches_estimated_param[3], + external_wrenches_estimated_param[4], + external_wrenches_estimated_param[5], + inertia[0], + inertia[1], + inertia[2], + inertia[3], + inertia[4], + inertia[5], + inertia[6], + inertia[7], + inertia[8], + mass, + ] + ) + self.acados_ocp_solver.set(j, "p", copy.deepcopy(param)) # Set initial state constraint. We teleported the robot foothold # to the previous optimal foothold. This is done to avoid the optimization # of a foothold that is not considered at all at touchdown! In any case, - # the height cames always from the VFA - if(FL_contact_sequence[0] == 0): + # the height cames always from the VFA + if FL_contact_sequence[0] == 0: state["foot_FL"] = reference["ref_foot_FL"][0] - if(FR_contact_sequence[0] == 0): + if FR_contact_sequence[0] == 0: state["foot_FR"] = reference["ref_foot_FR"][0] - if(RL_contact_sequence[0] == 0): + if RL_contact_sequence[0] == 0: state["foot_RL"] = reference["ref_foot_RL"][0] - - if(RR_contact_sequence[0] == 0): - state["foot_RR"] = reference["ref_foot_RR"][0] - - + if RR_contact_sequence[0] == 0: + state["foot_RR"] = reference["ref_foot_RR"][0] - if(self.use_integrators): + if self.use_integrators: # Compute error for integral action alpha_integrator = config.mpc_params["alpha_integrator"] - self.integral_errors[0] += (state["position"][2] - reference["ref_position"][2])*alpha_integrator - self.integral_errors[1] += (state["linear_velocity"][0] - reference["ref_linear_velocity"][0])*alpha_integrator - self.integral_errors[2] += (state["linear_velocity"][1] - reference["ref_linear_velocity"][1])*alpha_integrator - self.integral_errors[3] += (state["linear_velocity"][2] - reference["ref_linear_velocity"][2])*alpha_integrator - self.integral_errors[4] += (state["orientation"][0] - reference["ref_orientation"][0])*(alpha_integrator) - self.integral_errors[5] += (state["orientation"][1] - reference["ref_orientation"][1])*alpha_integrator - + self.integral_errors[0] += (state["position"][2] - reference["ref_position"][2]) * alpha_integrator + self.integral_errors[1] += ( + state["linear_velocity"][0] - reference["ref_linear_velocity"][0] + ) * alpha_integrator + self.integral_errors[2] += ( + state["linear_velocity"][1] - reference["ref_linear_velocity"][1] + ) * alpha_integrator + self.integral_errors[3] += ( + state["linear_velocity"][2] - reference["ref_linear_velocity"][2] + ) * alpha_integrator + self.integral_errors[4] += (state["orientation"][0] - reference["ref_orientation"][0]) * (alpha_integrator) + self.integral_errors[5] += (state["orientation"][1] - reference["ref_orientation"][1]) * alpha_integrator cap_integrator_z = config.mpc_params["integrator_cap"][0] cap_integrator_x_dot = config.mpc_params["integrator_cap"][1] @@ -1358,66 +1376,91 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, cap_integrator_roll = config.mpc_params["integrator_cap"][4] cap_integrator_pitch = config.mpc_params["integrator_cap"][5] - self.integral_errors[0] = np.where(np.abs(self.integral_errors[0]) > cap_integrator_z, cap_integrator_z*np.sign(self.integral_errors[0]), self.integral_errors[0]) - self.integral_errors[1] = np.where(np.abs(self.integral_errors[1]) > cap_integrator_x_dot, cap_integrator_x_dot*np.sign(self.integral_errors[1]), self.integral_errors[1]) - self.integral_errors[2] = np.where(np.abs(self.integral_errors[2]) > cap_integrator_y_dot, cap_integrator_y_dot*np.sign(self.integral_errors[2]), self.integral_errors[2]) - self.integral_errors[3] = np.where(np.abs(self.integral_errors[3]) > cap_integrator_z_dot, cap_integrator_z_dot*np.sign(self.integral_errors[3]), self.integral_errors[3]) - self.integral_errors[4] = np.where(np.abs(self.integral_errors[4]) > cap_integrator_roll, cap_integrator_roll*np.sign(self.integral_errors[4]), self.integral_errors[4]) - self.integral_errors[5] = np.where(np.abs(self.integral_errors[5]) > cap_integrator_pitch, cap_integrator_pitch*np.sign(self.integral_errors[5]), self.integral_errors[5]) + self.integral_errors[0] = np.where( + np.abs(self.integral_errors[0]) > cap_integrator_z, + cap_integrator_z * np.sign(self.integral_errors[0]), + self.integral_errors[0], + ) + self.integral_errors[1] = np.where( + np.abs(self.integral_errors[1]) > cap_integrator_x_dot, + cap_integrator_x_dot * np.sign(self.integral_errors[1]), + self.integral_errors[1], + ) + self.integral_errors[2] = np.where( + np.abs(self.integral_errors[2]) > cap_integrator_y_dot, + cap_integrator_y_dot * np.sign(self.integral_errors[2]), + self.integral_errors[2], + ) + self.integral_errors[3] = np.where( + np.abs(self.integral_errors[3]) > cap_integrator_z_dot, + cap_integrator_z_dot * np.sign(self.integral_errors[3]), + self.integral_errors[3], + ) + self.integral_errors[4] = np.where( + np.abs(self.integral_errors[4]) > cap_integrator_roll, + cap_integrator_roll * np.sign(self.integral_errors[4]), + self.integral_errors[4], + ) + self.integral_errors[5] = np.where( + np.abs(self.integral_errors[5]) > cap_integrator_pitch, + cap_integrator_pitch * np.sign(self.integral_errors[5]), + self.integral_errors[5], + ) print("self.integral_errors: ", self.integral_errors) - # Set initial state constraint acados, converting first the dictionary to np array - state_acados = np.concatenate((state["position"], state["linear_velocity"], - state["orientation"], state["angular_velocity"], - state["joint_FL"], state["joint_FR"], - state["joint_RL"], state["joint_RR"], - self.integral_errors)).reshape((self.states_dim, 1)) + state_acados = np.concatenate( + ( + state["position"], + state["linear_velocity"], + state["orientation"], + state["angular_velocity"], + state["joint_FL"], + state["joint_FR"], + state["joint_RL"], + state["joint_RR"], + self.integral_errors, + ) + ).reshape((self.states_dim, 1)) self.acados_ocp_solver.set(0, "lbx", state_acados) self.acados_ocp_solver.set(0, "ubx", state_acados) - - - # Set stage constraint - h_R_w = np.array([np.cos(yaw), np.sin(yaw), - -np.sin(yaw), np.cos(yaw)]) + h_R_w = np.array([np.cos(yaw), np.sin(yaw), -np.sin(yaw), np.cos(yaw)]) self.set_stage_constraint(constraint, state, reference, contact_sequence, h_R_w) - # Solve ocp via RTI or normal ocp if self.use_RTI: # feedback phase - self.acados_ocp_solver.options_set('rti_phase', 2) + self.acados_ocp_solver.options_set("rti_phase", 2) status = self.acados_ocp_solver.solve() - print("feedback phase time: ", self.acados_ocp_solver.get_stats('time_tot')) + print("feedback phase time: ", self.acados_ocp_solver.get_stats("time_tot")) else: status = self.acados_ocp_solver.solve() - print("ocp time: ", self.acados_ocp_solver.get_stats('time_tot')) - - + print("ocp time: ", self.acados_ocp_solver.get_stats("time_tot")) - # Take the solution + # Take the solution control = self.acados_ocp_solver.get(0, "u") optimal_joint_velocities = control[0:12] - - if(config.mpc_params['use_nonuniform_discretization']): - optimal_joint_acceleration = (optimal_joint_velocities - self.acados_ocp_solver.get(1, "u")[0:12])/config.mpc_params['dt_fine_grained'] + + if config.mpc_params["use_nonuniform_discretization"]: + optimal_joint_acceleration = ( + optimal_joint_velocities - self.acados_ocp_solver.get(1, "u")[0:12] + ) / config.mpc_params["dt_fine_grained"] else: - optimal_joint_acceleration = (optimal_joint_velocities - self.acados_ocp_solver.get(1, "u")[0:12])/config.mpc_params['dt'] - - optimal_GRF = control[12:] - + optimal_joint_acceleration = ( + optimal_joint_velocities - self.acados_ocp_solver.get(1, "u")[0:12] + ) / config.mpc_params["dt"] - + optimal_GRF = control[12:] # We need to provide the next touchdown foothold position. # We first take the foothold in stance now (they are not optimized!) - # and flag them as True (aka "assigned") + # and flag them as True (aka "assigned") optimal_foothold = np.zeros((4, 3)) - optimal_footholds_assigned = np.zeros((4, ), dtype='bool') + optimal_footholds_assigned = np.zeros((4,), dtype="bool") if FL_contact_sequence[0] == 1: optimal_foothold[0] = state["foot_FL"] optimal_footholds_assigned[0] = True @@ -1431,15 +1474,11 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, optimal_foothold[3] = state["foot_RR"] optimal_footholds_assigned[3] = True - - - - # Then we take the foothold at the next touchdown from the one - # that are not flagged as True from before, and saturate them!! - # P.S. The saturation is in the horizontal frame + # Then we take the foothold at the next touchdown from the one + # that are not flagged as True from before, and saturate them!! + # P.S. The saturation is in the horizontal frame for j in range(1, self.horizon): - if(FL_contact_sequence[j] != FL_contact_sequence[j-1] and not optimal_footholds_assigned[0]): - + if FL_contact_sequence[j] != FL_contact_sequence[j - 1] and not optimal_footholds_assigned[0]: base_predicted = self.acados_ocp_solver.get(j, "x")[0:3] roll_predicted, pitch_predicted, yaw_predicted = self.acados_ocp_solver.get(j, "x")[6:9] b_R_w_predicted = self.centroidal_model.compute_b_R_w(roll_predicted, pitch_predicted, yaw_predicted) @@ -1447,82 +1486,69 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, H[0:3, 0:3] = b_R_w_predicted.T H[0:3, 3] = base_predicted - #predicted_joint_position = self.acados_ocp_solver.get(j, "x")[12:24] - #optimal_foothold[0] = np.array(self.centroidal_model.forward_kinematics_FL_fun(H, predicted_joint_position)[0:3, 3]) + # predicted_joint_position = self.acados_ocp_solver.get(j, "x")[12:24] + # optimal_foothold[0] = np.array(self.centroidal_model.forward_kinematics_FL_fun(H, predicted_joint_position)[0:3, 3]) optimal_foothold[0] = state["foot_FL"] optimal_footholds_assigned[0] = True - - - if(FR_contact_sequence[j] != FR_contact_sequence[j-1] and not optimal_footholds_assigned[1]): - + + if FR_contact_sequence[j] != FR_contact_sequence[j - 1] and not optimal_footholds_assigned[1]: base_predicted = self.acados_ocp_solver.get(j, "x")[0:3] roll_predicted, pitch_predicted, yaw_predicted = self.acados_ocp_solver.get(j, "x")[6:9] b_R_w_predicted = self.centroidal_model.compute_b_R_w(roll_predicted, pitch_predicted, yaw_predicted) H = cs.SX.eye(4) H[0:3, 0:3] = b_R_w_predicted.T H[0:3, 3] = base_predicted - - #predicted_joint_position = self.acados_ocp_solver.get(j, "x")[12:24] - #optimal_foothold[1] = np.array(self.centroidal_model.forward_kinematics_FR_fun(H, predicted_joint_position)[0:3, 3]) + + # predicted_joint_position = self.acados_ocp_solver.get(j, "x")[12:24] + # optimal_foothold[1] = np.array(self.centroidal_model.forward_kinematics_FR_fun(H, predicted_joint_position)[0:3, 3]) optimal_foothold[1] = state["foot_FR"] optimal_footholds_assigned[1] = True - - if(RL_contact_sequence[j] != RL_contact_sequence[j-1] and not optimal_footholds_assigned[2]): - + if RL_contact_sequence[j] != RL_contact_sequence[j - 1] and not optimal_footholds_assigned[2]: base_predicted = self.acados_ocp_solver.get(j, "x")[0:3] roll_predicted, pitch_predicted, yaw_predicted = self.acados_ocp_solver.get(j, "x")[6:9] b_R_w_predicted = self.centroidal_model.compute_b_R_w(roll_predicted, pitch_predicted, yaw_predicted) H = cs.SX.eye(4) H[0:3, 0:3] = b_R_w_predicted.T H[0:3, 3] = base_predicted - - #predicted_joint_position = self.acados_ocp_solver.get(j, "x")[12:24] - #optimal_foothold[2] = np.array(self.centroidal_model.forward_kinematics_RL_fun(H, predicted_joint_position)[0:3, 3]) + + # predicted_joint_position = self.acados_ocp_solver.get(j, "x")[12:24] + # optimal_foothold[2] = np.array(self.centroidal_model.forward_kinematics_RL_fun(H, predicted_joint_position)[0:3, 3]) optimal_foothold[2] = state["foot_RL"] optimal_footholds_assigned[2] = True - - if(RR_contact_sequence[j] != RR_contact_sequence[j-1] and not optimal_footholds_assigned[3]): - + if RR_contact_sequence[j] != RR_contact_sequence[j - 1] and not optimal_footholds_assigned[3]: base_predicted = self.acados_ocp_solver.get(j, "x")[0:3] roll_predicted, pitch_predicted, yaw_predicted = self.acados_ocp_solver.get(j, "x")[6:9] b_R_w_predicted = self.centroidal_model.compute_b_R_w(roll_predicted, pitch_predicted, yaw_predicted) H = cs.SX.eye(4) H[0:3, 0:3] = b_R_w_predicted.T H[0:3, 3] = base_predicted - - #predicted_joint_position = self.acados_ocp_solver.get(j, "x")[12:24] - #optimal_foothold[3] = np.array(self.centroidal_model.forward_kinematics_RR_fun(H, predicted_joint_position)[0:3, 3]) + + # predicted_joint_position = self.acados_ocp_solver.get(j, "x")[12:24] + # optimal_foothold[3] = np.array(self.centroidal_model.forward_kinematics_RR_fun(H, predicted_joint_position)[0:3, 3]) optimal_foothold[3] = state["foot_RR"] optimal_footholds_assigned[3] = True - - - - # If in the prediction horizon, the foot is never in stance, we replicate the reference + # If in the prediction horizon, the foot is never in stance, we replicate the reference # to not confuse the swing controller - if(optimal_footholds_assigned[0] == False): + if optimal_footholds_assigned[0] == False: optimal_foothold[0] = reference["ref_foot_FL"][0] - if(optimal_footholds_assigned[1] == False): + if optimal_footholds_assigned[1] == False: optimal_foothold[1] = reference["ref_foot_FR"][0] - if(optimal_footholds_assigned[2] == False): + if optimal_footholds_assigned[2] == False: optimal_foothold[2] = reference["ref_foot_RL"][0] - if(optimal_footholds_assigned[3] == False): + if optimal_footholds_assigned[3] == False: optimal_foothold[3] = reference["ref_foot_RR"][0] - - - optimal_next_state_index = 1 optimal_next_state = self.acados_ocp_solver.get(optimal_next_state_index, "x")[0:24] self.optimal_next_state = optimal_next_state self.acados_ocp_solver.print_statistics() - # Check if QPs converged, if not just use the reference footholds # and a GRF over Z distribuited between the leg in stance - if(status == 1 or status == 4): + if status == 1 or status == 4: print("status", status) if FL_contact_sequence[0] == 0: optimal_foothold[0] = reference["ref_foot_FL"][0] @@ -1532,37 +1558,33 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, optimal_foothold[2] = reference["ref_foot_RL"][0] if RR_contact_sequence[0] == 0: optimal_foothold[3] = reference["ref_foot_RR"][0] - - number_of_legs_in_stance = np.array([FL_contact_sequence[0], FR_contact_sequence[0], - RL_contact_sequence [0], RR_contact_sequence[0]]).sum() - if(number_of_legs_in_stance == 0): + + number_of_legs_in_stance = np.array( + [FL_contact_sequence[0], FR_contact_sequence[0], RL_contact_sequence[0], RR_contact_sequence[0]] + ).sum() + if number_of_legs_in_stance == 0: reference_force_stance_legs = 0 else: reference_force_stance_legs = (mass * 9.81) / number_of_legs_in_stance - + reference_force_fl_z = reference_force_stance_legs * FL_contact_sequence[0] reference_force_fr_z = reference_force_stance_legs * FR_contact_sequence[0] reference_force_rl_z = reference_force_stance_legs * RL_contact_sequence[0] reference_force_rr_z = reference_force_stance_legs * RR_contact_sequence[0] - optimal_GRF = np.zeros((12, )) + optimal_GRF = np.zeros((12,)) optimal_GRF[2] = reference_force_fl_z optimal_GRF[5] = reference_force_fr_z optimal_GRF[8] = reference_force_rl_z optimal_GRF[11] = reference_force_rr_z - - optimal_GRF = self.previous_optimal_GRF self.acados_ocp_solver.reset() - - # Save the previous optimal GRF, the previous status and the previous contact sequence self.previous_optimal_GRF = optimal_GRF self.previous_status = status self.previous_contact_sequence = contact_sequence - # Decenter the optimal foothold and the next state (they were centered around zero at the beginning) optimal_foothold[0] = optimal_foothold[0] + self.initial_base_position optimal_foothold[1] = optimal_foothold[1] + self.initial_base_position @@ -1572,8 +1594,13 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, optimal_next_state[0:3] = optimal_next_state[0:3] + self.initial_base_position optimal_joint_positions = optimal_next_state[12:24] - - - # Return the optimal GRF, the optimal foothold, optimal_joint_velocities, the next state and the status of the optimization - return optimal_GRF, optimal_foothold, optimal_joint_positions, optimal_joint_velocities, optimal_joint_acceleration, optimal_next_state, status + return ( + optimal_GRF, + optimal_foothold, + optimal_joint_positions, + optimal_joint_velocities, + optimal_joint_acceleration, + optimal_next_state, + status, + ) diff --git a/quadruped_pympc/controllers/gradient/lyapunov/centroidal_model_lyapunov.py b/quadruped_pympc/controllers/gradient/lyapunov/centroidal_model_lyapunov.py index 1fe0907..1938cdb 100644 --- a/quadruped_pympc/controllers/gradient/lyapunov/centroidal_model_lyapunov.py +++ b/quadruped_pympc/controllers/gradient/lyapunov/centroidal_model_lyapunov.py @@ -3,33 +3,32 @@ # Authors: Giulio Turrisi - -import time -import unittest -import casadi as cs +import os +import casadi as cs import numpy as np from acados_template import AcadosModel -import os dir_path = os.path.dirname(os.path.realpath(__file__)) import sys + sys.path.append(dir_path) -sys.path.append(dir_path + '/../../../') +sys.path.append(dir_path + "/../../../") import config + # Class that defines the prediction model of the NMPC class Centroidal_Model_Lyapunov: - def __init__(self,) -> None: + def __init__( + self, + ) -> None: """ This method initializes the Centroidal_Model, which creates the prediction model of the NMPC. """ - - - # Define state and its casadi variables com_position_x = cs.SX.sym("com_position_x") com_position_y = cs.SX.sym("com_position_y") @@ -51,7 +50,6 @@ def __init__(self,) -> None: foot_position_rl = cs.SX.sym("foot_position_rl", 3, 1) foot_position_rr = cs.SX.sym("foot_position_rr", 3, 1) - com_position_z_integral = cs.SX.sym("com_position_z_integral") com_velocity_x_integral = cs.SX.sym("com_velocity_x_integral") com_velocity_y_integral = cs.SX.sym("com_velocity_y_integral") @@ -67,55 +65,54 @@ def __init__(self,) -> None: eta = cs.SX.sym("eta", 6, 1) phi = cs.SX.sym("theta", 3, 1) - - self.states = cs.vertcat(com_position_x, - com_position_y, - com_position_z, - com_velocity_x, - com_velocity_y, - com_velocity_z, - roll, - pitch, - yaw, - omega_x, - omega_y, - omega_z, - foot_position_fl, - foot_position_fr, - foot_position_rl, - foot_position_rr, - com_position_z_integral, - com_velocity_x_integral, - com_velocity_y_integral, - com_velocity_z_integral, - roll_integral, - pitch_integral, - z1, - z2, - eta, - phi) - - + self.states = cs.vertcat( + com_position_x, + com_position_y, + com_position_z, + com_velocity_x, + com_velocity_y, + com_velocity_z, + roll, + pitch, + yaw, + omega_x, + omega_y, + omega_z, + foot_position_fl, + foot_position_fr, + foot_position_rl, + foot_position_rr, + com_position_z_integral, + com_velocity_x_integral, + com_velocity_y_integral, + com_velocity_z_integral, + roll_integral, + pitch_integral, + z1, + z2, + eta, + phi, + ) # Define state dot - self.states_dot = cs.vertcat(cs.SX.sym("linear_com_vel", 3, 1), - cs.SX.sym("linear_com_acc", 3, 1), - cs.SX.sym("euler_rates_base", 3, 1), - cs.SX.sym("angular_acc_base", 3, 1), - cs.SX.sym("linear_vel_foot_FL", 3, 1), - cs.SX.sym("linear_vel_foot_FR", 3, 1), - cs.SX.sym("linear_vel_foot_RL", 3, 1), - cs.SX.sym("linear_vel_foot_RR", 3, 1), - cs.SX.sym("linear_com_vel_z_integral", 1, 1), - cs.SX.sym("linear_com_acc_integral", 3, 1), - cs.SX.sym("euler_rates_roll_integral", 1, 1), - cs.SX.sym("euler_rates_pitch_integral", 1, 1), - cs.SX.sym("z1_dot", 3, 1), - cs.SX.sym("z2_dot", 3, 1), - cs.SX.sym("eta_dot", 6, 1), - cs.SX.sym("phi_dot", 3, 1)) - - + self.states_dot = cs.vertcat( + cs.SX.sym("linear_com_vel", 3, 1), + cs.SX.sym("linear_com_acc", 3, 1), + cs.SX.sym("euler_rates_base", 3, 1), + cs.SX.sym("angular_acc_base", 3, 1), + cs.SX.sym("linear_vel_foot_FL", 3, 1), + cs.SX.sym("linear_vel_foot_FR", 3, 1), + cs.SX.sym("linear_vel_foot_RL", 3, 1), + cs.SX.sym("linear_vel_foot_RR", 3, 1), + cs.SX.sym("linear_com_vel_z_integral", 1, 1), + cs.SX.sym("linear_com_acc_integral", 3, 1), + cs.SX.sym("euler_rates_roll_integral", 1, 1), + cs.SX.sym("euler_rates_pitch_integral", 1, 1), + cs.SX.sym("z1_dot", 3, 1), + cs.SX.sym("z2_dot", 3, 1), + cs.SX.sym("eta_dot", 6, 1), + cs.SX.sym("phi_dot", 3, 1), + ) # Define input and its casadi variables foot_velocity_fl = cs.SX.sym("foot_velocity_fl", 3, 1) @@ -128,28 +125,27 @@ def __init__(self,) -> None: foot_force_rl = cs.SX.sym("foot_force_rl", 3, 1) foot_force_rr = cs.SX.sym("foot_force_rr", 3, 1) - self.inputs = cs.vertcat(foot_velocity_fl, - foot_velocity_fr, - foot_velocity_rl, - foot_velocity_rr, - foot_force_fl, - foot_force_fr, - foot_force_rl, - foot_force_rr) - + self.inputs = cs.vertcat( + foot_velocity_fl, + foot_velocity_fr, + foot_velocity_rl, + foot_velocity_rr, + foot_force_fl, + foot_force_fr, + foot_force_rl, + foot_force_rr, + ) # Usefull for debug what things goes where in y_ref in the compute_control function, # because there are a lot of variables self.y_ref = cs.vertcat(self.states, self.inputs) - # Define acados parameters that can be changed at runtine self.stanceFL = cs.SX.sym("stanceFL", 1, 1) self.stanceFR = cs.SX.sym("stanceFR", 1, 1) self.stanceRL = cs.SX.sym("stanceRL", 1, 1) self.stanceRR = cs.SX.sym("stanceRR", 1, 1) - self.stance_param = cs.vertcat(self.stanceFL , self.stanceFR , self.stanceRL , self.stanceRR) - + self.stance_param = cs.vertcat(self.stanceFL, self.stanceFR, self.stanceRL, self.stanceRR) self.mu_friction = cs.SX.sym("mu_friction", 1, 1) self.stance_proximity = cs.SX.sym("stanceProximity", 4, 1) @@ -161,16 +157,19 @@ def __init__(self,) -> None: self.inertia = cs.SX.sym("inertia", 9, 1) self.mass = cs.SX.sym("mass", 1, 1) - # Not so useful, i can instantiate a casadi function for the fd - param = cs.vertcat(self.stance_param, self.mu_friction, self.stance_proximity, - self.base_position, self.base_yaw, self.external_wrench, self.inertia, self.mass) + param = cs.vertcat( + self.stance_param, + self.mu_friction, + self.stance_proximity, + self.base_position, + self.base_yaw, + self.external_wrench, + self.inertia, + self.mass, + ) fd = self.forward_dynamics(self.states, self.inputs, param) - self.fun_forward_dynamics = cs.Function('fun_forward_dynamics', [self.states, self.inputs, param], [fd]) - - - - + self.fun_forward_dynamics = cs.Function("fun_forward_dynamics", [self.states, self.inputs, param], [fd]) def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.ndarray) -> cs.SX: """ @@ -221,7 +220,6 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda gravity = np.array([0, 0, -9.81]) - # Lyapunov states (ETA will be computed at the end of this function) K_z1 = config.mpc_params["K_z1"] K_z2 = config.mpc_params["K_z2"] @@ -234,54 +232,44 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda phi_dot = -z2 # Calculate the adaptive law - squared_K_z1 = np.array([K_z1[0]*K_z1[0], K_z1[1]*K_z1[1], K_z1[2]*K_z1[2]]) - F_star = mass*(-(K_z1 + K_z2)*z2 + squared_K_z1*z1 - gravity + p_ddot) + squared_K_z1 = np.array([K_z1[0] * K_z1[0], K_z1[1] * K_z1[1], K_z1[2] * K_z1[2]]) + F_star = mass * (-(K_z1 + K_z2) * z2 + squared_K_z1 * z1 - gravity + p_ddot) F_star -= phi + z1_dot = -K_z1 * z1 + z2 - - z1_dot = -K_z1*z1 + z2 - - F_delta = foot_force_fl@stanceFL - F_delta += foot_force_fr@stanceFR - F_delta += foot_force_rl@stanceRL - F_delta += foot_force_rr@stanceRR - z2_dot = -K_z2*z2 + 1/mass@(F_delta + F_star) + gravity - p_ddot + phi - - + F_delta = foot_force_fl @ stanceFL + F_delta += foot_force_fr @ stanceFR + F_delta += foot_force_rl @ stanceRL + F_delta += foot_force_rr @ stanceRR + z2_dot = -K_z2 * z2 + 1 / mass @ (F_delta + F_star) + gravity - p_ddot + phi # Redistribute the lyapunov forces to the leg num_leg_contact = stanceFL + stanceFR + stanceRL + stanceRR - foot_force_fl_final = foot_force_fl@stanceFL + (F_star/num_leg_contact)@stanceFL - foot_force_fr_final = foot_force_fr@stanceFR + (F_star/num_leg_contact)@stanceFR - foot_force_rl_final = foot_force_rl@stanceRL + (F_star/num_leg_contact)@stanceRL - foot_force_rr_final = foot_force_rr@stanceRR + (F_star/num_leg_contact)@stanceRR + foot_force_fl_final = foot_force_fl @ stanceFL + (F_star / num_leg_contact) @ stanceFL + foot_force_fr_final = foot_force_fr @ stanceFR + (F_star / num_leg_contact) @ stanceFR + foot_force_rl_final = foot_force_rl @ stanceRL + (F_star / num_leg_contact) @ stanceRL + foot_force_rr_final = foot_force_rr @ stanceRR + (F_star / num_leg_contact) @ stanceRR # linear wrench exterted by the robot (considering mpc + adaptive law) - temp = foot_force_fl_final@stanceFL - temp += foot_force_fr_final@stanceFR - temp += foot_force_rl_final@stanceRL - temp += foot_force_rr_final@stanceRR + temp = foot_force_fl_final @ stanceFL + temp += foot_force_fr_final @ stanceFR + temp += foot_force_rl_final @ stanceRL + temp += foot_force_rr_final @ stanceRR temp += external_wrench_linear - - # angular wrench exerted by the robot (considering mpc + adaptive law) - temp2 = cs.skew(foot_position_fl - com_position)@foot_force_fl_final@stanceFL - temp2 += cs.skew(foot_position_fr - com_position)@foot_force_fr_final@stanceFR - temp2 += cs.skew(foot_position_rl - com_position)@foot_force_rl_final@stanceRL - temp2 += cs.skew(foot_position_rr - com_position)@foot_force_rr_final@stanceRR + temp2 = cs.skew(foot_position_fl - com_position) @ foot_force_fl_final @ stanceFL + temp2 += cs.skew(foot_position_fr - com_position) @ foot_force_fr_final @ stanceFR + temp2 += cs.skew(foot_position_rl - com_position) @ foot_force_rl_final @ stanceRL + temp2 += cs.skew(foot_position_rr - com_position) @ foot_force_rr_final @ stanceRR temp2 = temp2 + external_wrench_angular - - # FINAL linear_com_vel STATE (1) linear_com_vel = states[3:6] - # FINAL linear_com_acc STATE (2) - linear_com_acc = (1/mass)@temp + gravity - phi - + linear_com_acc = (1 / mass) @ temp + gravity - phi # Start to write the component of euler_rates_base and angular_acc_base STATES w = states[9:12] @@ -291,74 +279,66 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda conj_euler_rates = cs.SX.eye(3) conj_euler_rates[1, 1] = cs.cos(roll) - conj_euler_rates[2, 2] = cs.cos(pitch)*cs.cos(roll) + conj_euler_rates[2, 2] = cs.cos(pitch) * cs.cos(roll) conj_euler_rates[2, 1] = -cs.sin(roll) conj_euler_rates[0, 2] = -cs.sin(pitch) - conj_euler_rates[1, 2] = cs.cos(pitch)*cs.sin(roll) - - - + conj_euler_rates[1, 2] = cs.cos(pitch) * cs.sin(roll) # FINAL euler_rates_base STATE (3) - euler_rates_base = cs.inv(conj_euler_rates)@w - + euler_rates_base = cs.inv(conj_euler_rates) @ w # FINAL angular_acc_base STATE (4) - #Z Y X rotations! + # Z Y X rotations! Rx = cs.SX.eye(3) - Rx[0,0] = 1 - Rx[0,1] = 0 - Rx[0,2] = 0 - Rx[1,0] = 0 - Rx[1,1] = cs.cos(roll) - Rx[1,2] = cs.sin(roll) - Rx[2,0] = 0 - Rx[2,1] = -cs.sin(roll) - Rx[2,2] = cs.cos(roll) - + Rx[0, 0] = 1 + Rx[0, 1] = 0 + Rx[0, 2] = 0 + Rx[1, 0] = 0 + Rx[1, 1] = cs.cos(roll) + Rx[1, 2] = cs.sin(roll) + Rx[2, 0] = 0 + Rx[2, 1] = -cs.sin(roll) + Rx[2, 2] = cs.cos(roll) Ry = cs.SX.eye(3) - Ry[0,0] = cs.cos(pitch) - Ry[0,1] = 0 - Ry[0,2] = -cs.sin(pitch) - Ry[1,0] = 0 - Ry[1,1] = 1 - Ry[1,2] = 0 - Ry[2,0] = cs.sin(pitch) - Ry[2,1] = 0 - Ry[2,2] = cs.cos(pitch) + Ry[0, 0] = cs.cos(pitch) + Ry[0, 1] = 0 + Ry[0, 2] = -cs.sin(pitch) + Ry[1, 0] = 0 + Ry[1, 1] = 1 + Ry[1, 2] = 0 + Ry[2, 0] = cs.sin(pitch) + Ry[2, 1] = 0 + Ry[2, 2] = cs.cos(pitch) Rz = cs.SX.eye(3) - Rz[0,0] = cs.cos(yaw) - Rz[0,1] = cs.sin(yaw) - Rz[0,2] = 0 - Rz[1,0] = -cs.sin(yaw) - Rz[1,1] = cs.cos(yaw) - Rz[1,2] = 0 - Rz[2,0] = 0 - Rz[2,1] = 0 - Rz[2,2] = 1 - - b_R_w = Rx@Ry@Rz - + Rz[0, 0] = cs.cos(yaw) + Rz[0, 1] = cs.sin(yaw) + Rz[0, 2] = 0 + Rz[1, 0] = -cs.sin(yaw) + Rz[1, 1] = cs.cos(yaw) + Rz[1, 2] = 0 + Rz[2, 0] = 0 + Rz[2, 1] = 0 + Rz[2, 2] = 1 + + b_R_w = Rx @ Ry @ Rz pos_phi = np.array([0, 0, 0.2]) - pos_phi_base = b_R_w.T@pos_phi - temp2 -= cs.skew(pos_phi_base)@phi - angular_acc_base = cs.inv(inertia)@(b_R_w@temp2 - cs.skew(w)@inertia@w) - - + pos_phi_base = b_R_w.T @ pos_phi + temp2 -= cs.skew(pos_phi_base) @ phi + angular_acc_base = cs.inv(inertia) @ (b_R_w @ temp2 - cs.skew(w) @ inertia @ w) # FINAL linear_foot_vel STATES (5,6,7,8) - if(not config.mpc_params["use_foothold_optimization"]): - foot_velocity_fl = foot_velocity_fl@0.0 - foot_velocity_fr = foot_velocity_fr@0.0 - foot_velocity_rl = foot_velocity_rl@0.0 - foot_velocity_rr = foot_velocity_rr@0.0 - linear_foot_vel_FL = foot_velocity_fl@(1-stanceFL)@(1-stance_proximity_FL) - linear_foot_vel_FR = foot_velocity_fr@(1-stanceFR)@(1-stance_proximity_FR) - linear_foot_vel_RL = foot_velocity_rl@(1-stanceRL)@(1-stance_proximity_RL) - linear_foot_vel_RR = foot_velocity_rr@(1-stanceRR)@(1-stance_proximity_RR) + if not config.mpc_params["use_foothold_optimization"]: + foot_velocity_fl = foot_velocity_fl @ 0.0 + foot_velocity_fr = foot_velocity_fr @ 0.0 + foot_velocity_rl = foot_velocity_rl @ 0.0 + foot_velocity_rr = foot_velocity_rr @ 0.0 + linear_foot_vel_FL = foot_velocity_fl @ (1 - stanceFL) @ (1 - stance_proximity_FL) + linear_foot_vel_FR = foot_velocity_fr @ (1 - stanceFR) @ (1 - stance_proximity_FR) + linear_foot_vel_RL = foot_velocity_rl @ (1 - stanceRL) @ (1 - stance_proximity_RL) + linear_foot_vel_RR = foot_velocity_rr @ (1 - stanceRR) @ (1 - stance_proximity_RR) # Integral states integral_states = states[24:30] @@ -369,29 +349,46 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda integral_states[4] += roll integral_states[5] += pitch - # STILL LYAPUNOV FOR ETA - eta_dot_temp = cs.inv(inertia)@(b_R_w@temp2 - cs.skew(w)@inertia@w) + eta_dot_temp = cs.inv(inertia) @ (b_R_w @ temp2 - cs.skew(w) @ inertia @ w) eta_dot = cs.vertcat(w, eta_dot_temp) - # The order of the return should be equal to the order of the states_dot - return cs.vertcat(linear_com_vel, linear_com_acc, euler_rates_base, angular_acc_base, - linear_foot_vel_FL,linear_foot_vel_FR, linear_foot_vel_RL, linear_foot_vel_RR, integral_states, - z1_dot, z2_dot, eta_dot, phi_dot) - - - - - def export_robot_model(self,) -> AcadosModel: + return cs.vertcat( + linear_com_vel, + linear_com_acc, + euler_rates_base, + angular_acc_base, + linear_foot_vel_FL, + linear_foot_vel_FR, + linear_foot_vel_RL, + linear_foot_vel_RR, + integral_states, + z1_dot, + z2_dot, + eta_dot, + phi_dot, + ) + + def export_robot_model( + self, + ) -> AcadosModel: """ This method set some general properties of the NMPC, such as the params, prediction mode, etc...! It will be called in centroidal_nmpc.py """ # dynamics - self.param = cs.vertcat(self.stance_param, self.mu_friction, self.stance_proximity, self.base_position, - self.base_yaw, self.external_wrench, self.inertia, self.mass) + self.param = cs.vertcat( + self.stance_param, + self.mu_friction, + self.stance_proximity, + self.base_position, + self.base_yaw, + self.external_wrench, + self.inertia, + self.mass, + ) f_expl = self.forward_dynamics(self.states, self.inputs, self.param) f_impl = self.states_dot - f_expl @@ -404,6 +401,4 @@ def export_robot_model(self,) -> AcadosModel: acados_model.p = self.param acados_model.name = "centroidal_model" - return acados_model - diff --git a/quadruped_pympc/controllers/gradient/lyapunov/centroidal_nmpc_lyapunov.py b/quadruped_pympc/controllers/gradient/lyapunov/centroidal_nmpc_lyapunov.py index 3f1e170..5dad4f5 100644 --- a/quadruped_pympc/controllers/gradient/lyapunov/centroidal_nmpc_lyapunov.py +++ b/quadruped_pympc/controllers/gradient/lyapunov/centroidal_nmpc_lyapunov.py @@ -2,43 +2,37 @@ import pathlib # Authors: Giulio Turrisi - - from acados_template import AcadosOcp, AcadosOcpSolver + ACADOS_INFTY = 1000 -import numpy as np -import scipy.linalg -import casadi as cs import copy - +import casadi as cs +import config +import numpy as np +import scipy.linalg from .centroidal_model_lyapunov import Centroidal_Model_Lyapunov -import config - - # Class for the Acados NMPC, the model is in another file! class Acados_NMPC_Lyapunov: def __init__(self): - - self.horizon = config.mpc_params['horizon'] # Define the number of discretization steps - self.dt = config.mpc_params['dt'] - self.T_horizon = self.horizon*self.dt - self.use_RTI = config.mpc_params['use_RTI'] - self.use_integrators = config.mpc_params['use_integrators'] - self.use_warm_start = config.mpc_params['use_warm_start'] - self.use_foothold_constraints = config.mpc_params['use_foothold_constraints'] - - - self.use_static_stability = config.mpc_params['use_static_stability'] - self.use_zmp_stability = config.mpc_params['use_zmp_stability'] + self.horizon = config.mpc_params["horizon"] # Define the number of discretization steps + self.dt = config.mpc_params["dt"] + self.T_horizon = self.horizon * self.dt + self.use_RTI = config.mpc_params["use_RTI"] + self.use_integrators = config.mpc_params["use_integrators"] + self.use_warm_start = config.mpc_params["use_warm_start"] + self.use_foothold_constraints = config.mpc_params["use_foothold_constraints"] + + self.use_static_stability = config.mpc_params["use_static_stability"] + self.use_zmp_stability = config.mpc_params["use_zmp_stability"] self.use_stability_constraints = self.use_static_stability or self.use_zmp_stability - self.use_DDP = config.mpc_params['use_DDP'] - - self.verbose = config.mpc_params['verbose'] + self.use_DDP = config.mpc_params["use_DDP"] + self.verbose = config.mpc_params["verbose"] self.previous_status = -1 self.previous_contact_sequence = np.zeros((4, self.horizon)) @@ -47,14 +41,11 @@ def __init__(self): self.phi_predicted = np.zeros((3,)) - self.integral_errors = np.zeros((6,)) - # For centering the variable around 0, 0, 0 (World frame) self.initial_base_position = np.array([0, 0, 0]) - # Create the class of the centroidal model and instantiate the acados model self.centroidal_model = Centroidal_Model_Lyapunov() acados_model = self.centroidal_model.export_robot_model() @@ -66,13 +57,12 @@ def __init__(self): # Create the acados ocp solver self.ocp = self.create_ocp_solver_description(acados_model) - code_export_dir = pathlib.Path(__file__).parent / 'c_generated_code' + code_export_dir = pathlib.Path(__file__).parent / "c_generated_code" self.ocp.code_export_directory = str(code_export_dir) - self.acados_ocp_solver = AcadosOcpSolver(self.ocp, - json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + - ".json") - + self.acados_ocp_solver = AcadosOcpSolver( + self.ocp, json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json" + ) # Initialize solver for stage in range(self.horizon + 1): @@ -80,18 +70,11 @@ def __init__(self): for stage in range(self.horizon): self.acados_ocp_solver.set(stage, "u", np.zeros((self.inputs_dim,))) - if(self.use_RTI): + if self.use_RTI: # first preparation phase - self.acados_ocp_solver.options_set('rti_phase', 1) + self.acados_ocp_solver.options_set("rti_phase", 1) status = self.acados_ocp_solver.solve() - - - - - - - # Set cost, constraints and options def create_ocp_solver_description(self, acados_model) -> AcadosOcp: # Create ocp object to formulate the OCP @@ -128,9 +111,7 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: ocp.cost.yref_e = np.zeros((ny_e,)) # Set friction and foothold constraints - expr_h_friction, \ - self.constr_uh_friction, \ - self.constr_lh_friction = self.create_friction_cone_constraints() + expr_h_friction, self.constr_uh_friction, self.constr_lh_friction = self.create_friction_cone_constraints() ocp.model.con_h_expr = expr_h_friction ocp.constraints.uh = self.constr_uh_friction @@ -142,23 +123,20 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: nsh_state_constraint_start = copy.copy(nsh) # Set foothold constraints - if(self.use_foothold_constraints): - expr_h_foot, \ - self.constr_uh_foot, \ - self.constr_lh_foot = self.create_foothold_constraints() + if self.use_foothold_constraints: + expr_h_foot, self.constr_uh_foot, self.constr_lh_foot = self.create_foothold_constraints() ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_foot) ocp.constraints.uh = np.concatenate((ocp.constraints.uh, self.constr_uh_foot)) ocp.constraints.lh = np.concatenate((ocp.constraints.lh, self.constr_lh_foot)) nsh += expr_h_foot.shape[0] - # Set stability constraints - if(self.use_stability_constraints): + if self.use_stability_constraints: self.nsh_stability_start = copy.copy(nsh) - expr_h_support_polygon, \ - self.constr_uh_support_polygon, \ - self.constr_lh_support_polygon = self.create_stability_constraints() + expr_h_support_polygon, self.constr_uh_support_polygon, self.constr_lh_support_polygon = ( + self.create_stability_constraints() + ) ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_support_polygon) ocp.constraints.uh = np.concatenate((ocp.constraints.uh, self.constr_uh_support_polygon)) @@ -166,44 +144,42 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: nsh += expr_h_support_polygon.shape[0] self.nsh_stability_end = copy.copy(nsh) - # Set lyapunov constraints - expr_h_support_lyapunov, \ - self.constr_uh_lyapunov, \ - self.constr_lh_lyapunov = self.create_lyapunov_constraints() + expr_h_support_lyapunov, self.constr_uh_lyapunov, self.constr_lh_lyapunov = self.create_lyapunov_constraints() ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_support_lyapunov) ocp.constraints.uh = np.concatenate((ocp.constraints.uh, self.constr_uh_lyapunov)) ocp.constraints.lh = np.concatenate((ocp.constraints.lh, self.constr_lh_lyapunov)) nsh += expr_h_support_lyapunov.shape[0] - # Set residual dynamics constraints - expr_h_residual_dynamics, \ - self.constr_uh_residual_dynamics, \ - self.constr_lh_residual_dynamics = self.create_residual_dynamics_constraint() + expr_h_residual_dynamics, self.constr_uh_residual_dynamics, self.constr_lh_residual_dynamics = ( + self.create_residual_dynamics_constraint() + ) ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_residual_dynamics) ocp.constraints.uh = np.concatenate((ocp.constraints.uh, self.constr_uh_residual_dynamics)) ocp.constraints.lh = np.concatenate((ocp.constraints.lh, self.constr_lh_residual_dynamics)) nsh += expr_h_residual_dynamics.shape[0] - nsh_state_constraint_end = copy.copy(nsh) - - # Set slack variable configuration: num_state_cstr = nsh_state_constraint_end - nsh_state_constraint_start - if(num_state_cstr > 0): - ocp.constraints.lsh = np.zeros(num_state_cstr) # Lower bounds on slacks corresponding to soft lower bounds for nonlinear constraints - ocp.constraints.ush = np.zeros(num_state_cstr) # Lower bounds on slacks corresponding to soft upper bounds for nonlinear constraints - ocp.constraints.idxsh = np.array(range(nsh_state_constraint_start, nsh_state_constraint_end)) # Jsh + if num_state_cstr > 0: + ocp.constraints.lsh = np.zeros( + num_state_cstr + ) # Lower bounds on slacks corresponding to soft lower bounds for nonlinear constraints + ocp.constraints.ush = np.zeros( + num_state_cstr + ) # Lower bounds on slacks corresponding to soft upper bounds for nonlinear constraints + ocp.constraints.idxsh = np.array(range(nsh_state_constraint_start, nsh_state_constraint_end)) # Jsh ns = num_state_cstr - ocp.cost.zl = 1000 * np.ones((ns,)) # gradient wrt lower slack at intermediate shooting nodes (1 to N-1) - ocp.cost.Zl = 1 * np.ones((ns,)) # diagonal of Hessian wrt lower slack at intermediate shooting nodes (1 to N-1) + ocp.cost.zl = 1000 * np.ones((ns,)) # gradient wrt lower slack at intermediate shooting nodes (1 to N-1) + ocp.cost.Zl = 1 * np.ones( + (ns,) + ) # diagonal of Hessian wrt lower slack at intermediate shooting nodes (1 to N-1) ocp.cost.zu = 1000 * np.ones((ns,)) ocp.cost.Zu = 1 * np.ones((ns,)) - # Variables to save the upper and lower bound of the constraints applied list_upper_bound = [] list_lower_bound = [] @@ -213,13 +189,12 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: self.upper_bound = np.array(list_upper_bound, dtype=object) self.lower_bound = np.array(list_lower_bound, dtype=object) - # Set initial state constraint X0 = np.zeros(shape=(nx,)) ocp.constraints.x0 = X0 # Set initialize parameters - init_contact_status = np.array([1., 1., 1., 1.]) + init_contact_status = np.array([1.0, 1.0, 1.0, 1.0]) init_mu = np.array([0.5]) init_stance_proximity = np.array([0, 0, 0, 0]) init_base_position = np.array([0, 0, 0]) @@ -228,60 +203,69 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: init_inertia = config.inertia.reshape((9,)) init_mass = np.array([config.mass]) - ocp.parameter_values = np.concatenate((init_contact_status, init_mu, init_stance_proximity, - init_base_position, init_base_yaw, init_external_wrench, - init_inertia, init_mass)) - - + ocp.parameter_values = np.concatenate( + ( + init_contact_status, + init_mu, + init_stance_proximity, + init_base_position, + init_base_yaw, + init_external_wrench, + init_inertia, + init_mass, + ) + ) # Set options - ocp.solver_options.qp_solver = "PARTIAL_CONDENSING_HPIPM" # FULL_CONDENSING_QPOASES PARTIAL_CONDENSING_OSQP PARTIAL_CONDENSING_HPIPM + ocp.solver_options.qp_solver = ( + "PARTIAL_CONDENSING_HPIPM" # FULL_CONDENSING_QPOASES PARTIAL_CONDENSING_OSQP PARTIAL_CONDENSING_HPIPM + ) ocp.solver_options.hessian_approx = "GAUSS_NEWTON" # 'GAUSS_NEWTON', 'EXACT' - ocp.solver_options.integrator_type = "ERK" #ERK IRK GNSF DISCRETE - if(self.use_DDP): - ocp.solver_options.nlp_solver_type = 'DDP' - ocp.solver_options.nlp_solver_max_iter = config.mpc_params['num_qp_iterations'] - #ocp.solver_options.globalization = 'MERIT_BACKTRACKING' + ocp.solver_options.integrator_type = "ERK" # ERK IRK GNSF DISCRETE + if self.use_DDP: + ocp.solver_options.nlp_solver_type = "DDP" + ocp.solver_options.nlp_solver_max_iter = config.mpc_params["num_qp_iterations"] + # ocp.solver_options.globalization = 'MERIT_BACKTRACKING' ocp.solver_options.with_adaptive_levenberg_marquardt = True - ocp.cost.cost_type = 'NONLINEAR_LS' - ocp.cost.cost_type_e = 'NONLINEAR_LS' + ocp.cost.cost_type = "NONLINEAR_LS" + ocp.cost.cost_type_e = "NONLINEAR_LS" ocp.model.cost_y_expr = cs.vertcat(ocp.model.x, ocp.model.u) ocp.model.cost_y_expr_e = ocp.model.x ocp.translate_to_feasibility_problem(keep_x0=True, keep_cost=True) - elif(self.use_RTI): + elif self.use_RTI: ocp.solver_options.nlp_solver_type = "SQP_RTI" ocp.solver_options.nlp_solver_max_iter = 1 # Set the RTI type for the advanced RTI method # (see https://arxiv.org/pdf/2403.07101.pdf) - if(config.mpc_params['as_rti_type'] == "AS-RTI-A"): + if config.mpc_params["as_rti_type"] == "AS-RTI-A": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 0 - elif(config.mpc_params['as_rti_type'] == "AS-RTI-B"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-B": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 1 - elif(config.mpc_params['as_rti_type'] == "AS-RTI-C"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-C": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 2 - elif(config.mpc_params['as_rti_type'] == "AS-RTI-D"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-D": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 3 else: ocp.solver_options.nlp_solver_type = "SQP" - ocp.solver_options.nlp_solver_max_iter = config.mpc_params['num_qp_iterations'] - #ocp.solver_options.globalization = "MERIT_BACKTRACKING" # FIXED_STEP, MERIT_BACKTRACKING + ocp.solver_options.nlp_solver_max_iter = config.mpc_params["num_qp_iterations"] + # ocp.solver_options.globalization = "MERIT_BACKTRACKING" # FIXED_STEP, MERIT_BACKTRACKING - if (config.mpc_params['solver_mode'] == "balance"): + if config.mpc_params["solver_mode"] == "balance": ocp.solver_options.hpipm_mode = "BALANCE" - elif (config.mpc_params['solver_mode'] == "robust"): + elif config.mpc_params["solver_mode"] == "robust": ocp.solver_options.hpipm_mode = "ROBUST" - elif (config.mpc_params['solver_mode'] == "fast"): + elif config.mpc_params["solver_mode"] == "fast": ocp.solver_options.qp_solver_iter_max = 10 ocp.solver_options.hpipm_mode = "SPEED" - elif (config.mpc_params['solver_mode'] == "crazy_speed"): + elif config.mpc_params["solver_mode"] == "crazy_speed": ocp.solver_options.qp_solver_iter_max = 5 ocp.solver_options.hpipm_mode = "SPEED_ABS" @@ -290,29 +274,28 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: # ocp.solver_options.nlp_solver_ext_qp_res = 1 ocp.solver_options.levenberg_marquardt = 1e-3 - # Set prediction horizon ocp.solver_options.tf = self.T_horizon - # Nonuniform discretization - if(config.mpc_params['use_nonuniform_discretization']): - time_steps_fine_grained = np.tile(config.mpc_params['dt_fine_grained'], config.mpc_params['horizon_fine_grained']) - time_steps = np.concatenate((time_steps_fine_grained, np.tile(self.dt, self.horizon-config.mpc_params['horizon_fine_grained']))) - shooting_nodes = np.zeros((self.horizon+1,)) + if config.mpc_params["use_nonuniform_discretization"]: + time_steps_fine_grained = np.tile( + config.mpc_params["dt_fine_grained"], config.mpc_params["horizon_fine_grained"] + ) + time_steps = np.concatenate( + (time_steps_fine_grained, np.tile(self.dt, self.horizon - config.mpc_params["horizon_fine_grained"])) + ) + shooting_nodes = np.zeros((self.horizon + 1,)) for i in range(len(time_steps)): - shooting_nodes[i+1] = shooting_nodes[i] + time_steps[i] + shooting_nodes[i + 1] = shooting_nodes[i] + time_steps[i] ocp.solver_options.shooting_nodes = shooting_nodes - - - return ocp - # Create Lyapunov constraints - def create_lyapunov_constraints(self,): - + def create_lyapunov_constraints( + self, + ): z1 = self.centroidal_model.states[30:33] z2 = self.centroidal_model.states[33:36] K_z1 = config.mpc_params["K_z1"] @@ -320,94 +303,87 @@ def create_lyapunov_constraints(self,): mass = self.centroidal_model.param[28] - stanceFL = self.centroidal_model.param[0] stanceFR = self.centroidal_model.param[1] stanceRL = self.centroidal_model.param[2] stanceRR = self.centroidal_model.param[3] - - foot_force_fl = self.centroidal_model.inputs[12:15] foot_force_fr = self.centroidal_model.inputs[15:18] foot_force_rl = self.centroidal_model.inputs[18:21] foot_force_rr = self.centroidal_model.inputs[21:24] - F_delta = foot_force_fl@stanceFL - F_delta += foot_force_fr@stanceFR - F_delta += foot_force_rl@stanceRL - F_delta += foot_force_rr@stanceRR - + F_delta = foot_force_fl @ stanceFL + F_delta += foot_force_fr @ stanceFR + F_delta += foot_force_rl @ stanceRL + F_delta += foot_force_rr @ stanceRR - V_dot = -z1.T@np.diag(K_z1)@z1 - z2.T@np.diag(K_z2)@z2 + z1.T@z2 + z2.T@F_delta/mass - #V_dot += self.centroidal_model.states[6:12].T@self.acados_model.f_expl_expr[6:12] + V_dot = -z1.T @ np.diag(K_z1) @ z1 - z2.T @ np.diag(K_z2) @ z2 + z1.T @ z2 + z2.T @ F_delta / mass + # V_dot += self.centroidal_model.states[6:12].T@self.acados_model.f_expl_expr[6:12] - ub = np.ones(1)*0 - lb = np.ones(1)*-1000 + ub = np.ones(1) * 0 + lb = np.ones(1) * -1000 return V_dot, ub, lb - # Create Residual Dynamics constraint - def create_residual_dynamics_constraint(self,): + def create_residual_dynamics_constraint( + self, + ): w = self.centroidal_model.states[9:12] angles = self.centroidal_model.states[6:9] eta = cs.vertcat(angles, w) - - residual = eta.T@eta - ub = np.array([config.mpc_params['residual_dynamics_upper_bound']]) + residual = eta.T @ eta + ub = np.array([config.mpc_params["residual_dynamics_upper_bound"]]) lb = np.zeros(1) return residual, ub, lb - # Create a constraint for stability (COM, ZMP or CP inside support polygon) - def create_stability_constraints(self,) -> None: - + def create_stability_constraints( + self, + ) -> None: base_w = self.centroidal_model.states[0:3] base_vel_w = self.centroidal_model.states[3:6] - FL = self.centroidal_model.states[12:15] FR = self.centroidal_model.states[15:18] RL = self.centroidal_model.states[18:21] RR = self.centroidal_model.states[21:24] - #yaw = self.centroidal_model.base_yaw[0] + # yaw = self.centroidal_model.base_yaw[0] yaw = self.centroidal_model.states[8] h_R_w = cs.SX.zeros(2, 2) - h_R_w[0,0] = cs.cos(yaw) - h_R_w[0,1] = cs.sin(yaw) - h_R_w[1,0] = -cs.sin(yaw) - h_R_w[1,1] = cs.cos(yaw) + h_R_w[0, 0] = cs.cos(yaw) + h_R_w[0, 1] = cs.sin(yaw) + h_R_w[1, 0] = -cs.sin(yaw) + h_R_w[1, 1] = cs.cos(yaw) - FL[0:2] = h_R_w@(FL[0:2] - base_w[0:2]) - FR[0:2] = h_R_w@(FR[0:2] - base_w[0:2]) - RL[0:2] = h_R_w@(RL[0:2] - base_w[0:2]) - RR[0:2] = h_R_w@(RR[0:2] - base_w[0:2]) + FL[0:2] = h_R_w @ (FL[0:2] - base_w[0:2]) + FR[0:2] = h_R_w @ (FR[0:2] - base_w[0:2]) + RL[0:2] = h_R_w @ (RL[0:2] - base_w[0:2]) + RR[0:2] = h_R_w @ (RR[0:2] - base_w[0:2]) - - if(self.use_static_stability): + if self.use_static_stability: x = 0.0 y = 0.0 else: # Compute the ZMP robotHeight = base_w[2] - foot_force_fl = self.centroidal_model.inputs[12:15]#@self.centroidal_model.param[0] - foot_force_fr = self.centroidal_model.inputs[15:18]#@self.centroidal_model.param[1] - foot_force_rl = self.centroidal_model.inputs[18:21]#@self.centroidal_model.param[2] - foot_force_rr = self.centroidal_model.inputs[21:24]#@self.centroidal_model.param[3] + foot_force_fl = self.centroidal_model.inputs[12:15] # @self.centroidal_model.param[0] + foot_force_fr = self.centroidal_model.inputs[15:18] # @self.centroidal_model.param[1] + foot_force_rl = self.centroidal_model.inputs[18:21] # @self.centroidal_model.param[2] + foot_force_rr = self.centroidal_model.inputs[21:24] # @self.centroidal_model.param[3] temp = foot_force_fl + foot_force_fr + foot_force_rl + foot_force_rr gravity = np.array([0, 0, -9.81]) - linear_com_acc = (1/self.centroidal_model.mass)@temp + gravity - zmp = base_w[0:2] - linear_com_acc[0:2]*(robotHeight/(-gravity[2])) - #zmp = base_w[0:2] - base_vel_w[0:2]*(robotHeight/gravity[2]) - zmp = h_R_w@(zmp - base_w[0:2]) + linear_com_acc = (1 / self.centroidal_model.mass) @ temp + gravity + zmp = base_w[0:2] - linear_com_acc[0:2] * (robotHeight / (-gravity[2])) + # zmp = base_w[0:2] - base_vel_w[0:2]*(robotHeight/gravity[2]) + zmp = h_R_w @ (zmp - base_w[0:2]) x = zmp[0] y = zmp[1] - y_FL = FL[1] y_FR = FR[1] y_RL = RL[1] @@ -418,32 +394,32 @@ def create_stability_constraints(self,) -> None: x_RL = RL[0] x_RR = RR[0] - #LF - RF : x < (x2 - x1) (y - y1) / (y2 - y1) + x1 - #RF - RH: y > (y2 - y1) (x - x1) / (x2 - x1) + y1 - #RH - LH : x > (x2 - x1) (y - y1) / (y2 - y1) + x1 - #LH - LF: y < (y2 - y1) (x - x1) / (x2 - x1) + y1 + # LF - RF : x < (x2 - x1) (y - y1) / (y2 - y1) + x1 + # RF - RH: y > (y2 - y1) (x - x1) / (x2 - x1) + y1 + # RH - LH : x > (x2 - x1) (y - y1) / (y2 - y1) + x1 + # LH - LF: y < (y2 - y1) (x - x1) / (x2 - x1) + y1 - #FL and FR cannot stay at the same x! #constrint should be less than zero - constraint_FL_FR = x - (x_FR - x_FL)*(y - y_FL) / (y_FR - y_FL + 0.001) - x_FL + # FL and FR cannot stay at the same x! #constrint should be less than zero + constraint_FL_FR = x - (x_FR - x_FL) * (y - y_FL) / (y_FR - y_FL + 0.001) - x_FL - #FR and RR cannot stay at the same y! #constraint should be bigger than zero - constraint_FR_RR = y - (y_RR - y_FR)*(x - x_FR) / (x_RR - x_FR + 0.001) - y_FR + # FR and RR cannot stay at the same y! #constraint should be bigger than zero + constraint_FR_RR = y - (y_RR - y_FR) * (x - x_FR) / (x_RR - x_FR + 0.001) - y_FR - #RL and RR cannot stay at the same x! #constraint should be bigger than zero - constraint_RR_RL = x - (x_RL - x_RR)*(y - y_RR) / (y_RL - y_RR + 0.001) - x_RR + # RL and RR cannot stay at the same x! #constraint should be bigger than zero + constraint_RR_RL = x - (x_RL - x_RR) * (y - y_RR) / (y_RL - y_RR + 0.001) - x_RR - #FL and RL cannot stay at the same y! #constraint should be less than zero - constraint_RL_FL = y - (y_FL - y_RL)*(x - x_RL) / (x_FL - x_RL + 0.001) - y_RL + # FL and RL cannot stay at the same y! #constraint should be less than zero + constraint_RL_FL = y - (y_FL - y_RL) * (x - x_RL) / (x_FL - x_RL + 0.001) - y_RL # the diagonal stuff can be at any point... - constraint_FL_RR = y - (y_RR - y_FL)*(x - x_FL) / (x_RR - x_FL + 0.001) - y_FL #bigger than zero - constraint_FR_RL = y - (y_RL - y_FR)*(x - x_FR) / (x_RL - x_FR + 0.001) - y_FR #bigger than zero + constraint_FL_RR = y - (y_RR - y_FL) * (x - x_FL) / (x_RR - x_FL + 0.001) - y_FL # bigger than zero + constraint_FR_RL = y - (y_RL - y_FR) * (x - x_FR) / (x_RL - x_FR + 0.001) - y_FR # bigger than zero # upper and lower bound ub = np.ones(6) * ACADOS_INFTY lb = -np.ones(6) * ACADOS_INFTY - #constraint_FL_FR + # constraint_FL_FR ub[0] = 0 lb[0] = -ACADOS_INFTY @@ -455,29 +431,28 @@ def create_stability_constraints(self,) -> None: ub[2] = ACADOS_INFTY lb[2] = 0 - #constraint_RL_FL + # constraint_RL_FL ub[3] = 0 lb[3] = -ACADOS_INFTY - #constraint_FL_RR + # constraint_FL_RR ub[4] = 0 lb[4] = -ACADOS_INFTY - #constraint_FR_RL + # constraint_FR_RL ub[5] = 0 lb[5] = -ACADOS_INFTY - Jb = cs.vertcat(constraint_FL_FR, constraint_FR_RR, - constraint_RR_RL, constraint_RL_FL, - constraint_FL_RR, constraint_FR_RL) - + Jb = cs.vertcat( + constraint_FL_FR, constraint_FR_RR, constraint_RR_RL, constraint_RL_FL, constraint_FL_RR, constraint_FR_RL + ) return Jb, ub, lb - - # Create a standard foothold box constraint - def create_foothold_constraints(self,): + def create_foothold_constraints( + self, + ): """ This function calculates the symbolic foothold constraints for the centroidal NMPC problem. @@ -495,43 +470,37 @@ def create_foothold_constraints(self,): # using the robot yaw yaw = self.centroidal_model.base_yaw[0] h_R_w = cs.SX.zeros(2, 2) - h_R_w[0,0] = cs.cos(yaw) - h_R_w[0,1] = cs.sin(yaw) - h_R_w[1,0] = -cs.sin(yaw) - h_R_w[1,1] = cs.cos(yaw) - + h_R_w[0, 0] = cs.cos(yaw) + h_R_w[0, 1] = cs.sin(yaw) + h_R_w[1, 0] = -cs.sin(yaw) + h_R_w[1, 1] = cs.cos(yaw) # and translate them using the robot base position base = self.centroidal_model.base_position[0:3] - - - foot_position_fl = cs.SX.zeros(3,1) - foot_position_fl[0:2] = h_R_w@cs.vertcat(self.centroidal_model.states[12:14] - base[0:2]) + foot_position_fl = cs.SX.zeros(3, 1) + foot_position_fl[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.states[12:14] - base[0:2]) foot_position_fl[2] = self.centroidal_model.states[14] - foot_position_fr = cs.SX.zeros(3,1) - foot_position_fr[0:2] = h_R_w@cs.vertcat(self.centroidal_model.states[15:17] - base[0:2]) + foot_position_fr = cs.SX.zeros(3, 1) + foot_position_fr[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.states[15:17] - base[0:2]) foot_position_fr[2] = self.centroidal_model.states[17] - foot_position_rl = cs.SX.zeros(3,1) - foot_position_rl[0:2] = h_R_w@cs.vertcat(self.centroidal_model.states[18:20] - base[0:2]) + foot_position_rl = cs.SX.zeros(3, 1) + foot_position_rl[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.states[18:20] - base[0:2]) foot_position_rl[2] = self.centroidal_model.states[20] - foot_position_rr = cs.SX.zeros(3,1) - foot_position_rr[0:2] = h_R_w@cs.vertcat(self.centroidal_model.states[21:23] - base[0:2]) + foot_position_rr = cs.SX.zeros(3, 1) + foot_position_rr[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.states[21:23] - base[0:2]) foot_position_rr[2] = self.centroidal_model.states[23] - - - Jbu = cs.vertcat(foot_position_fl, foot_position_fr, - foot_position_rl, foot_position_rr) + Jbu = cs.vertcat(foot_position_fl, foot_position_fr, foot_position_rl, foot_position_rr) return Jbu, ubu, lbu - - # Create the friction cone constraint - def create_friction_cone_constraints(self,) ->None: + def create_friction_cone_constraints( + self, + ) -> None: """ This function calculates the symbolic friction cone constraints for the centroidal NMPC problem. @@ -544,86 +513,79 @@ def create_friction_cone_constraints(self,) ->None: t = np.array([1, 0, 0]) b = np.array([0, 1, 0]) mu = self.centroidal_model.mu_friction - f_max = config.mpc_params['grf_max'] - f_min = config.mpc_params['grf_min'] + f_max = config.mpc_params["grf_max"] + f_min = config.mpc_params["grf_min"] # Derivation can be found in the paper # "High-slope terrain locomotion for torque-controlled quadruped robots", # M Focchi, A Del Prete, I Havoutis, R Featherstone, DG Caldwell, C Semini Jbu = cs.SX.zeros(20, 12) - Jbu[0, :3] = -n*mu + t - Jbu[1, :3] = -n*mu + b - Jbu[2, :3] = n*mu + b - Jbu[3, :3] = n*mu + t + Jbu[0, :3] = -n * mu + t + Jbu[1, :3] = -n * mu + b + Jbu[2, :3] = n * mu + b + Jbu[3, :3] = n * mu + t Jbu[4, :3] = n - Jbu[5, 3:6] = -n*mu + t - Jbu[6, 3:6] = -n*mu + b - Jbu[7, 3:6] = n*mu + b - Jbu[8, 3:6] = n*mu + t + Jbu[5, 3:6] = -n * mu + t + Jbu[6, 3:6] = -n * mu + b + Jbu[7, 3:6] = n * mu + b + Jbu[8, 3:6] = n * mu + t Jbu[9, 3:6] = n - Jbu[10, 6:9] = -n*mu + t - Jbu[11, 6:9] = -n*mu + b - Jbu[12, 6:9] = n*mu + b - Jbu[13, 6:9] = n*mu + t + Jbu[10, 6:9] = -n * mu + t + Jbu[11, 6:9] = -n * mu + b + Jbu[12, 6:9] = n * mu + b + Jbu[13, 6:9] = n * mu + t Jbu[14, 6:9] = n - Jbu[15, 9:12] = -n*mu + t - Jbu[16, 9:12] = -n*mu + b - Jbu[17, 9:12] = n*mu + b - Jbu[18, 9:12] = n*mu + t + Jbu[15, 9:12] = -n * mu + t + Jbu[16, 9:12] = -n * mu + b + Jbu[17, 9:12] = n * mu + b + Jbu[18, 9:12] = n * mu + t Jbu[19, 9:12] = n - - # Calculate the adaptive law + # Calculate the adaptive law K_z1 = config.mpc_params["K_z1"] K_z2 = config.mpc_params["K_z2"] - p_ddot = np.array([0, 0, 0]) #?? + p_ddot = np.array([0, 0, 0]) # ?? gravity = np.array([0, 0, -9.81]) mass = self.centroidal_model.param[28] z1 = self.centroidal_model.states[30:33] z2 = self.centroidal_model.states[33:36] phi = self.centroidal_model.states[42:45] - squared_K_z1 = np.array([K_z1[0]*K_z1[0], K_z1[1]*K_z1[1], K_z1[2]*K_z1[2]]) - F_star = mass*(-(K_z1 + K_z2)*z2 + squared_K_z1*z1 - gravity + p_ddot) + squared_K_z1 = np.array([K_z1[0] * K_z1[0], K_z1[1] * K_z1[1], K_z1[2] * K_z1[2]]) + F_star = mass * (-(K_z1 + K_z2) * z2 + squared_K_z1 * z1 - gravity + p_ddot) F_star -= phi - stanceFL = self.centroidal_model.param[0] stanceFR = self.centroidal_model.param[1] stanceRL = self.centroidal_model.param[2] stanceRR = self.centroidal_model.param[3] - - foot_force_fl = self.centroidal_model.inputs[12:15] foot_force_fr = self.centroidal_model.inputs[15:18] foot_force_rl = self.centroidal_model.inputs[18:21] foot_force_rr = self.centroidal_model.inputs[21:24] - - # Redistribute the lyapunov forces to the leg num_leg_contact = stanceFL + stanceFR + stanceRL + stanceRR - #foot_force_fl_final = foot_force_fl@stanceFL + (F_star/num_leg_contact)@stanceFL - #foot_force_fr_final = foot_force_fr@stanceFR + (F_star/num_leg_contact)@stanceFR - #foot_force_rl_final = foot_force_rl@stanceRL + (F_star/num_leg_contact)@stanceRL - #foot_force_rr_final = foot_force_rr@stanceRR + (F_star/num_leg_contact)@stanceRR + # foot_force_fl_final = foot_force_fl@stanceFL + (F_star/num_leg_contact)@stanceFL + # foot_force_fr_final = foot_force_fr@stanceFR + (F_star/num_leg_contact)@stanceFR + # foot_force_rl_final = foot_force_rl@stanceRL + (F_star/num_leg_contact)@stanceRL + # foot_force_rr_final = foot_force_rr@stanceRR + (F_star/num_leg_contact)@stanceRR - foot_force_fl_final = foot_force_fl + (F_star/num_leg_contact) - foot_force_fr_final = foot_force_fr + (F_star/num_leg_contact) - foot_force_rl_final = foot_force_rl + (F_star/num_leg_contact) - foot_force_rr_final = foot_force_rr + (F_star/num_leg_contact) - - feet_forces_final = cs.vertcat(foot_force_fl_final, foot_force_fr_final, foot_force_rl_final, foot_force_rr_final) + foot_force_fl_final = foot_force_fl + (F_star / num_leg_contact) + foot_force_fr_final = foot_force_fr + (F_star / num_leg_contact) + foot_force_rl_final = foot_force_rl + (F_star / num_leg_contact) + foot_force_rr_final = foot_force_rr + (F_star / num_leg_contact) + feet_forces_final = cs.vertcat( + foot_force_fl_final, foot_force_fr_final, foot_force_rl_final, foot_force_rr_final + ) # C matrix - Jbu = Jbu@feet_forces_final - - + Jbu = Jbu @ feet_forces_final # lower bound (-1000 is "almost" -inf) lbu = np.zeros(20) @@ -649,53 +611,69 @@ def create_friction_cone_constraints(self,) ->None: return Jbu, ubu, lbu - - def set_weight(self, nx, nu): # Define the weight matrices for the cost function - Q_position = np.array([0, 0, 1500]) #x, y, z - Q_velocity = np.array([200, 200, 200]) #x_vel, y_vel, z_vel - Q_base_angle = np.array([500, 500, 0]) #roll, pitch, yaw - Q_base_angle_rates = np.array([20, 20, 50]) #roll_rate, pitch_rate, yaw_rate - Q_foot_pos = np.array([300, 300, 300]) #f_x, f_y, f_z (should be 4 times this, once per foot) - Q_com_position_z_integral = np.array([50]) #integral of z_com - Q_com_velocity_x_integral = np.array([10]) #integral of x_com - Q_com_velocity_y_integral = np.array([10]) #integral of y_com - Q_com_velocity_z_integral = np.array([10]) #integral of z_com_vel - Q_roll_integral_integral = np.array([10]) #integral of roll - Q_pitch_integral_integral = np.array([10]) #integral of pitch - - Q_lyapunov_states = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) #z1, z2, eta, phi - - R_foot_vel = np.array([0.0001, 0.0001, 0.00001]) #v_x, v_y, v_z (should be 4 times this, once per foot) - R_foot_force = np.array([0.001, 0.001, 0.001])#f_x, f_y, f_z (should be 4 times this, once per foot) - - Q_mat = np.diag(np.concatenate((Q_position, Q_velocity, - Q_base_angle, Q_base_angle_rates, - Q_foot_pos, Q_foot_pos, Q_foot_pos, Q_foot_pos, - Q_com_position_z_integral, Q_com_velocity_x_integral, - Q_com_velocity_y_integral, Q_com_velocity_z_integral, - Q_roll_integral_integral, Q_pitch_integral_integral, - Q_lyapunov_states))) - - R_mat = np.diag(np.concatenate((R_foot_vel, R_foot_vel, R_foot_vel, R_foot_vel, - R_foot_force, R_foot_force, R_foot_force, R_foot_force))) + Q_position = np.array([0, 0, 1500]) # x, y, z + Q_velocity = np.array([200, 200, 200]) # x_vel, y_vel, z_vel + Q_base_angle = np.array([500, 500, 0]) # roll, pitch, yaw + Q_base_angle_rates = np.array([20, 20, 50]) # roll_rate, pitch_rate, yaw_rate + Q_foot_pos = np.array([300, 300, 300]) # f_x, f_y, f_z (should be 4 times this, once per foot) + Q_com_position_z_integral = np.array([50]) # integral of z_com + Q_com_velocity_x_integral = np.array([10]) # integral of x_com + Q_com_velocity_y_integral = np.array([10]) # integral of y_com + Q_com_velocity_z_integral = np.array([10]) # integral of z_com_vel + Q_roll_integral_integral = np.array([10]) # integral of roll + Q_pitch_integral_integral = np.array([10]) # integral of pitch + + Q_lyapunov_states = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) # z1, z2, eta, phi + + R_foot_vel = np.array([0.0001, 0.0001, 0.00001]) # v_x, v_y, v_z (should be 4 times this, once per foot) + R_foot_force = np.array([0.001, 0.001, 0.001]) # f_x, f_y, f_z (should be 4 times this, once per foot) + + Q_mat = np.diag( + np.concatenate( + ( + Q_position, + Q_velocity, + Q_base_angle, + Q_base_angle_rates, + Q_foot_pos, + Q_foot_pos, + Q_foot_pos, + Q_foot_pos, + Q_com_position_z_integral, + Q_com_velocity_x_integral, + Q_com_velocity_y_integral, + Q_com_velocity_z_integral, + Q_roll_integral_integral, + Q_pitch_integral_integral, + Q_lyapunov_states, + ) + ) + ) + + R_mat = np.diag( + np.concatenate( + (R_foot_vel, R_foot_vel, R_foot_vel, R_foot_vel, R_foot_force, R_foot_force, R_foot_force, R_foot_force) + ) + ) return Q_mat, R_mat - - def reset(self): self.acados_ocp_solver.reset() - self.acados_ocp_solver = AcadosOcpSolver(self.ocp, - json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json", - build=False, generate=False) - self.phi_predicted = self.phi_predicted*0.0 - - - - def set_residual_dynamics_constraint(self,): + self.acados_ocp_solver = AcadosOcpSolver( + self.ocp, + json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json", + build=False, + generate=False, + ) + self.phi_predicted = self.phi_predicted * 0.0 + + def set_residual_dynamics_constraint( + self, + ): try: ub_friction = self.constr_uh_friction lb_friction = self.constr_lh_friction @@ -704,30 +682,24 @@ def set_residual_dynamics_constraint(self,): lb_lyapunov = self.constr_lh_lyapunov for j in range(0, self.horizon): - ub_residual = config.mpc_params['residual_dynamics_upper_bound']*(1/(j+1)) + ub_residual = config.mpc_params["residual_dynamics_upper_bound"] * (1 / (j + 1)) lb_residual = 0.0 # Only friction costraints at the beginning - if(j == 0): + if j == 0: self.acados_ocp_solver.constraints_set(j, "uh", ub_friction) self.acados_ocp_solver.constraints_set(j, "lh", lb_friction) - if(j > 0): + if j > 0: ub_total = np.concatenate((ub_friction, ub_lyapunov, np.array([ub_residual]))) lb_total = np.concatenate((lb_friction, lb_lyapunov, np.array([lb_residual]))) self.acados_ocp_solver.constraints_set(j, "uh", ub_total) self.acados_ocp_solver.constraints_set(j, "lh", lb_total) - except: - if(self.verbose): + if self.verbose: print("Error in setting the residual constraint") pass - - - - - def set_stage_constraint(self, constraint, state, reference, contact_sequence, h_R_w, stance_proximity): """ Set the stage constraint for the centroidal NMPC problem. We only consider the stance constraint, and the swing @@ -745,8 +717,6 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h None """ try: - - # Take the array of the contact sequence and split # it in 4 arrays for clarity FL_contact_sequence = contact_sequence[0] @@ -765,14 +735,11 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h RL_reference_foot = reference["ref_foot_RL"] RR_reference_foot = reference["ref_foot_RR"] - # Take the base position and the yaw rotation matrix. This is needed to # express the foothold constraint in the horizontal frame base = state["position"] - h_R_w = h_R_w.reshape((2,2)) - - + h_R_w = h_R_w.reshape((2, 2)) # Divide the constraint in upper and lower bound. The constraint are # represented by 4 vertex, but we only use box constraint hence @@ -783,56 +750,52 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h # FL Stance Constraint stance_up_constraint_FL = np.array([FL_actual_foot[0], FL_actual_foot[1], FL_actual_foot[2] + 0.002]) - stance_up_constraint_FL[0:2] = h_R_w@(stance_up_constraint_FL[0:2] - base[0:2]) + stance_up_constraint_FL[0:2] = h_R_w @ (stance_up_constraint_FL[0:2] - base[0:2]) stance_up_constraint_FL[0:2] = stance_up_constraint_FL[0:2] + 0.1 stance_up_constraint_FL[2] = stance_up_constraint_FL[2] + 0.01 stance_low_constraint_FL = np.array([FL_actual_foot[0], FL_actual_foot[1], FL_actual_foot[2] - 0.002]) - stance_low_constraint_FL[0:2] = h_R_w@(stance_low_constraint_FL[0:2] - base[0:2]) + stance_low_constraint_FL[0:2] = h_R_w @ (stance_low_constraint_FL[0:2] - base[0:2]) stance_low_constraint_FL[0:2] = stance_low_constraint_FL[0:2] - 0.1 stance_low_constraint_FL[2] = stance_low_constraint_FL[2] - 0.01 - # FR Stance Constraint stance_up_constraint_FR = np.array([FR_actual_foot[0], FR_actual_foot[1], FR_actual_foot[2] + 0.002]) - stance_up_constraint_FR[0:2] = h_R_w@(stance_up_constraint_FR[0:2] - base[0:2]) + stance_up_constraint_FR[0:2] = h_R_w @ (stance_up_constraint_FR[0:2] - base[0:2]) stance_up_constraint_FR[0:2] = stance_up_constraint_FR[0:2] + 0.1 stance_up_constraint_FR[2] = stance_up_constraint_FR[2] + 0.01 stance_low_constraint_FR = np.array([FR_actual_foot[0], FR_actual_foot[1], FR_actual_foot[2] - 0.002]) - stance_low_constraint_FR[0:2] = h_R_w@(stance_low_constraint_FR[0:2] - base[0:2]) + stance_low_constraint_FR[0:2] = h_R_w @ (stance_low_constraint_FR[0:2] - base[0:2]) stance_low_constraint_FR[0:2] = stance_low_constraint_FR[0:2] - 0.1 stance_low_constraint_FR[2] = stance_low_constraint_FR[2] - 0.01 - # RL Stance Constraint stance_up_constraint_RL = np.array([RL_actual_foot[0], RL_actual_foot[1], RL_actual_foot[2] + 0.002]) - stance_up_constraint_RL[0:2] = h_R_w@(stance_up_constraint_RL[0:2]- base[0:2]) + stance_up_constraint_RL[0:2] = h_R_w @ (stance_up_constraint_RL[0:2] - base[0:2]) stance_up_constraint_RL[0:2] = stance_up_constraint_RL[0:2] + 0.1 stance_up_constraint_RL[2] = stance_up_constraint_RL[2] + 0.01 stance_low_constraint_RL = np.array([RL_actual_foot[0], RL_actual_foot[1], RL_actual_foot[2] - 0.002]) - stance_low_constraint_RL[0:2] = h_R_w@(stance_low_constraint_RL[0:2]- base[0:2]) + stance_low_constraint_RL[0:2] = h_R_w @ (stance_low_constraint_RL[0:2] - base[0:2]) stance_low_constraint_RL[0:2] = stance_low_constraint_RL[0:2] - 0.1 stance_low_constraint_RL[2] = stance_low_constraint_RL[2] - 0.01 - # RR Stance Constraint stance_up_constraint_RR = np.array([RR_actual_foot[0], RR_actual_foot[1], RR_actual_foot[2] + 0.002]) - stance_up_constraint_RR[0:2] = h_R_w@(stance_up_constraint_RR[0:2]- base[0:2]) + stance_up_constraint_RR[0:2] = h_R_w @ (stance_up_constraint_RR[0:2] - base[0:2]) stance_up_constraint_RR[0:2] = stance_up_constraint_RR[0:2] + 0.1 stance_up_constraint_RR[2] = stance_up_constraint_RR[2] + 0.01 stance_low_constraint_RR = np.array([RR_actual_foot[0], RR_actual_foot[1], RR_actual_foot[2] - 0.002]) - stance_low_constraint_RR[0:2] = h_R_w@(stance_low_constraint_RR[0:2]- base[0:2]) + stance_low_constraint_RR[0:2] = h_R_w @ (stance_low_constraint_RR[0:2] - base[0:2]) stance_low_constraint_RR[0:2] = stance_low_constraint_RR[0:2] - 0.1 stance_low_constraint_RR[2] = stance_low_constraint_RR[2] - 0.01 - # Constraint for the first footholds at the next touchdown. If constraint == True # we have some constraint passed from outside (e.g. vision), otherwise we use the # nominal foothold enlarged as we do previously - if(constraint is not None): + if constraint is not None: # From the VFA first_up_constraint_FL = np.array([constraint[0][0], constraint[1][0], constraint[2][0] + 0.002]) first_up_constraint_FR = np.array([constraint[0][1], constraint[1][1], constraint[2][1] + 0.002]) @@ -845,73 +808,83 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h first_low_constraint_RR = np.array([constraint[9][3], constraint[10][3], constraint[11][3] - 0.002]) # Rotate the constraint in the horizontal frame - first_up_constraint_FL[0:2] = h_R_w@(first_up_constraint_FL[0:2] - base[0:2]) + first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - base[0:2]) first_up_constraint_FL = first_up_constraint_FL + 0.005 - #first_up_constraint_FL[2] = first_up_constraint_FL[2] + 0.5 + # first_up_constraint_FL[2] = first_up_constraint_FL[2] + 0.5 - first_up_constraint_FR[0:2] = h_R_w@(first_up_constraint_FR[0:2] - base[0:2]) + first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - base[0:2]) first_up_constraint_FR = first_up_constraint_FR + 0.005 - #first_up_constraint_FR[2] = first_up_constraint_FR[2] + 0.5 + # first_up_constraint_FR[2] = first_up_constraint_FR[2] + 0.5 - first_up_constraint_RL[0:2] = h_R_w@(first_up_constraint_RL[0:2] - base[0:2]) + first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - base[0:2]) first_up_constraint_RL = first_up_constraint_RL + 0.005 - #first_up_constraint_RL[2] = first_up_constraint_RL[2] + 0.5 + # first_up_constraint_RL[2] = first_up_constraint_RL[2] + 0.5 - first_up_constraint_RR[0:2] = h_R_w@(first_up_constraint_RR[0:2] - base[0:2]) + first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - base[0:2]) first_up_constraint_RR = first_up_constraint_RR + 0.005 - #first_up_constraint_RR[2] = first_up_constraint_RR[2] + 0.5 + # first_up_constraint_RR[2] = first_up_constraint_RR[2] + 0.5 - first_low_constraint_FL[0:2] = h_R_w@(first_low_constraint_FL[0:2] - base[0:2]) + first_low_constraint_FL[0:2] = h_R_w @ (first_low_constraint_FL[0:2] - base[0:2]) first_low_constraint_FL = first_low_constraint_FL - 0.005 - #first_low_constraint_FL[2] = first_low_constraint_FL[2] - 0.5 + # first_low_constraint_FL[2] = first_low_constraint_FL[2] - 0.5 - first_low_constraint_FR[0:2] = h_R_w@(first_low_constraint_FR[0:2] - base[0:2]) + first_low_constraint_FR[0:2] = h_R_w @ (first_low_constraint_FR[0:2] - base[0:2]) first_low_constraint_FR = first_low_constraint_FR - 0.005 - #first_low_constraint_FR[2] = first_low_constraint_FR[2] - 0.5 + # first_low_constraint_FR[2] = first_low_constraint_FR[2] - 0.5 - first_low_constraint_RL[0:2] = h_R_w@(first_low_constraint_RL[0:2] - base[0:2]) + first_low_constraint_RL[0:2] = h_R_w @ (first_low_constraint_RL[0:2] - base[0:2]) first_low_constraint_RL = first_low_constraint_RL - 0.005 - #first_low_constraint_RL[2] = first_low_constraint_RL[2] - 0.5 + # first_low_constraint_RL[2] = first_low_constraint_RL[2] - 0.5 - first_low_constraint_RR[0:2] = h_R_w@(first_low_constraint_RR[0:2] - base[0:2]) + first_low_constraint_RR[0:2] = h_R_w @ (first_low_constraint_RR[0:2] - base[0:2]) first_low_constraint_RR = first_low_constraint_RR - 0.005 - #first_low_constraint_RR[2] = first_low_constraint_RR[2] - 0.5 + # first_low_constraint_RR[2] = first_low_constraint_RR[2] - 0.5 else: # Constrain taken from the nominal foothold (NO VISION) # FL first touchdown constraint - first_up_constraint_FL = np.array([FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] + 0.002]) - first_up_constraint_FL[0:2] = h_R_w@(first_up_constraint_FL[0:2] - base[0:2]) + 0.15 - - first_low_constraint_FL = np.array([FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] - 0.002]) - first_low_constraint_FL[0:2] = h_R_w@(first_low_constraint_FL[0:2] - base[0:2]) - 0.15 + first_up_constraint_FL = np.array( + [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] + 0.002] + ) + first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - base[0:2]) + 0.15 + first_low_constraint_FL = np.array( + [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] - 0.002] + ) + first_low_constraint_FL[0:2] = h_R_w @ (first_low_constraint_FL[0:2] - base[0:2]) - 0.15 # FR first touchdown constraint - first_up_constraint_FR = np.array([FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] + 0.002]) - first_up_constraint_FR[0:2] = h_R_w@(first_up_constraint_FR[0:2] - base[0:2]) + 0.15 - - first_low_constraint_FR = np.array([FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] - 0.002]) - first_low_constraint_FR[0:2] = h_R_w@(first_low_constraint_FR[0:2] - base[0:2]) - 0.15 + first_up_constraint_FR = np.array( + [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] + 0.002] + ) + first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - base[0:2]) + 0.15 + first_low_constraint_FR = np.array( + [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] - 0.002] + ) + first_low_constraint_FR[0:2] = h_R_w @ (first_low_constraint_FR[0:2] - base[0:2]) - 0.15 # RL first touchdown constraint - first_up_constraint_RL = np.array([RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] + 0.002]) - first_up_constraint_RL[0:2] = h_R_w@(first_up_constraint_RL[0:2] - base[0:2]) + 0.15 - - first_low_constraint_RL = np.array([RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] - 0.002]) - first_low_constraint_RL[0:2] = h_R_w@(first_low_constraint_RL[0:2] - base[0:2]) - 0.15 + first_up_constraint_RL = np.array( + [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] + 0.002] + ) + first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - base[0:2]) + 0.15 + first_low_constraint_RL = np.array( + [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] - 0.002] + ) + first_low_constraint_RL[0:2] = h_R_w @ (first_low_constraint_RL[0:2] - base[0:2]) - 0.15 # RR first touchdown constraint - first_up_constraint_RR = np.array([RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] + 0.002]) - first_up_constraint_RR[0:2] = h_R_w@(first_up_constraint_RR[0:2] - base[0:2]) + 0.15 - - first_low_constraint_RR = np.array([RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] - 0.002]) - first_low_constraint_RR[0:2] = h_R_w@(first_low_constraint_RR[0:2] - base[0:2]) - 0.15 - - + first_up_constraint_RR = np.array( + [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] + 0.002] + ) + first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - base[0:2]) + 0.15 + first_low_constraint_RR = np.array( + [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] - 0.002] + ) + first_low_constraint_RR[0:2] = h_R_w @ (first_low_constraint_RR[0:2] - base[0:2]) - 0.15 # we stack all the constraint we have for now up_constraint_FL = np.vstack((stance_up_constraint_FL, first_up_constraint_FL)) @@ -924,64 +897,69 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h low_constraint_RL = np.vstack((stance_low_constraint_RL, first_low_constraint_RL)) low_constraint_RR = np.vstack((stance_low_constraint_RR, first_low_constraint_RR)) - - # If there are more than 1 nominal foothold per leg, we create additional constraints # We do not expect more than two reference footholds... # FL second touchdown constraint - if(FL_reference_foot.shape[0] == 2): - second_up_constraint_FL = np.array([FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] + 0.002]) - second_up_constraint_FL[0:2] = h_R_w@(second_up_constraint_FL[0:2] - base[0:2]) + 0.15 - - second_low_constraint_FL = np.array([FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] - 0.002]) - second_low_constraint_FL[0:2] = h_R_w@(second_low_constraint_FL[0:2] - base[0:2]) - 0.15 + if FL_reference_foot.shape[0] == 2: + second_up_constraint_FL = np.array( + [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] + 0.002] + ) + second_up_constraint_FL[0:2] = h_R_w @ (second_up_constraint_FL[0:2] - base[0:2]) + 0.15 + second_low_constraint_FL = np.array( + [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] - 0.002] + ) + second_low_constraint_FL[0:2] = h_R_w @ (second_low_constraint_FL[0:2] - base[0:2]) - 0.15 up_constraint_FL = np.vstack((up_constraint_FL, second_up_constraint_FL)) low_constraint_FL = np.vstack((low_constraint_FL, second_low_constraint_FL)) # FR second touchdown constraint - if(FR_reference_foot.shape[0] == 2): - second_up_constraint_FR = np.array([FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] + 0.002]) - second_up_constraint_FR[0:2] = h_R_w@(second_up_constraint_FR[0:2] - base[0:2]) + 0.15 - - - second_low_constraint_FR = np.array([FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] - 0.002]) - second_low_constraint_FR[0:2] = h_R_w@(second_low_constraint_FR[0:2]- base[0:2]) - 0.15 + if FR_reference_foot.shape[0] == 2: + second_up_constraint_FR = np.array( + [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] + 0.002] + ) + second_up_constraint_FR[0:2] = h_R_w @ (second_up_constraint_FR[0:2] - base[0:2]) + 0.15 + second_low_constraint_FR = np.array( + [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] - 0.002] + ) + second_low_constraint_FR[0:2] = h_R_w @ (second_low_constraint_FR[0:2] - base[0:2]) - 0.15 up_constraint_FR = np.vstack((up_constraint_FR, second_up_constraint_FR)) low_constraint_FR = np.vstack((low_constraint_FR, second_low_constraint_FR)) # RL second touchdown constraint - if(RL_reference_foot.shape[0] == 2): - second_up_constraint_RL = np.array([RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] + 0.002]) - second_up_constraint_RL[0:2] = h_R_w@(second_up_constraint_RL[0:2]- base[0:2]) + 0.15 - - - second_low_constraint_RL = np.array([RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] - 0.002]) - second_low_constraint_RL[0:2] = h_R_w@(second_low_constraint_RL[0:2] - base[0:2]) - 0.15 + if RL_reference_foot.shape[0] == 2: + second_up_constraint_RL = np.array( + [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] + 0.002] + ) + second_up_constraint_RL[0:2] = h_R_w @ (second_up_constraint_RL[0:2] - base[0:2]) + 0.15 + second_low_constraint_RL = np.array( + [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] - 0.002] + ) + second_low_constraint_RL[0:2] = h_R_w @ (second_low_constraint_RL[0:2] - base[0:2]) - 0.15 up_constraint_RL = np.vstack((up_constraint_RL, second_up_constraint_RL)) low_constraint_RL = np.vstack((low_constraint_RL, second_low_constraint_RL)) # RR second touchdown constraint - if(RR_reference_foot.shape[0] == 2): - second_up_constraint_RR = np.array([RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] + 0.002]) - second_up_constraint_RR[0:2] = h_R_w@(second_up_constraint_RR[0:2] - base[0:2]) + 0.15 - - - second_low_constraint_RR = np.array([RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] - 0.002]) - second_low_constraint_RR[0:2] = h_R_w@(second_low_constraint_RR[0:2] - base[0:2]) - 0.15 + if RR_reference_foot.shape[0] == 2: + second_up_constraint_RR = np.array( + [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] + 0.002] + ) + second_up_constraint_RR[0:2] = h_R_w @ (second_up_constraint_RR[0:2] - base[0:2]) + 0.15 + second_low_constraint_RR = np.array( + [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] - 0.002] + ) + second_low_constraint_RR[0:2] = h_R_w @ (second_low_constraint_RR[0:2] - base[0:2]) - 0.15 up_constraint_RR = np.vstack((up_constraint_RR, second_up_constraint_RR)) low_constraint_RR = np.vstack((low_constraint_RR, second_low_constraint_RR)) - - # We pass all the constraints (foothold and friction conte) to acados. # Then one regarding friction are precomputed ub_friction = self.constr_uh_friction @@ -989,17 +967,15 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h # If the foothold is in swing, the idx of the constraint start from 1 idx_constraint = np.array([0, 0, 0, 0]) - if(FL_contact_sequence[0] == 0): + if FL_contact_sequence[0] == 0: idx_constraint[0] = 1 - if(FR_contact_sequence[0] == 0): + if FR_contact_sequence[0] == 0: idx_constraint[1] = 1 - if(RL_contact_sequence[0] == 0): + if RL_contact_sequence[0] == 0: idx_constraint[2] = 1 - if(RR_contact_sequence[0] == 0): + if RR_contact_sequence[0] == 0: idx_constraint[3] = 1 - - for j in range(0, self.horizon): # take the constraint for the current timestep ub_foot_FL = up_constraint_FL[idx_constraint[0]] @@ -1014,24 +990,19 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h ub_foot_RR = up_constraint_RR[idx_constraint[3]] lb_foot_RR = low_constraint_RR[idx_constraint[3]] - - # Concatenate the friction and foothold constraint - ub_foot = copy.deepcopy(np.concatenate((ub_foot_FL, ub_foot_FR, - ub_foot_RL, ub_foot_RR))) - lb_foot = copy.deepcopy(np.concatenate((lb_foot_FL, lb_foot_FR, - lb_foot_RL, lb_foot_RR))) - if(self.use_foothold_constraints): + ub_foot = copy.deepcopy(np.concatenate((ub_foot_FL, ub_foot_FR, ub_foot_RL, ub_foot_RR))) + lb_foot = copy.deepcopy(np.concatenate((lb_foot_FL, lb_foot_FR, lb_foot_RL, lb_foot_RR))) + if self.use_foothold_constraints: ub_total = np.concatenate((ub_friction, ub_foot)) lb_total = np.concatenate((lb_friction, lb_foot)) else: ub_total = ub_friction lb_total = lb_friction - - #Constraints for the support polygon depending on the leg in stance + # Constraints for the support polygon depending on the leg in stance # all disabled at the beginning!! - if(self.use_stability_constraints): + if self.use_stability_constraints: ub_support_FL_FR = 1000 lb_support_FL_FR = -1000 @@ -1050,13 +1021,14 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h ub_support_FR_RL = 1000 lb_support_FR_RL = -1000 - # We have 4 cases for the stability constraint: trot, pace, crawl, full stance - if(FL_contact_sequence[j] == 1 and - FR_contact_sequence[j] == 1 and - RL_contact_sequence[j] == 1 and - RR_contact_sequence[j] == 1): - #FULL STANCE TODO + if ( + FL_contact_sequence[j] == 1 + and FR_contact_sequence[j] == 1 + and RL_contact_sequence[j] == 1 + and RR_contact_sequence[j] == 1 + ): + # FULL STANCE TODO ub_support_FL_FR = 1000 lb_support_FL_FR = -1000 @@ -1075,141 +1047,153 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h ub_support_FR_RL = 1000 lb_support_FR_RL = -1000 - elif(np.array_equal(FL_contact_sequence, RR_contact_sequence) - and np.array_equal(FR_contact_sequence, RL_contact_sequence)): - #TROT - stability_margin = config.mpc_params['trot_stability_margin'] - if(FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0): + elif np.array_equal(FL_contact_sequence, RR_contact_sequence) and np.array_equal( + FR_contact_sequence, RL_contact_sequence + ): + # TROT + stability_margin = config.mpc_params["trot_stability_margin"] + if FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0: ub_support_FL_RR = 0 + stability_margin lb_support_FL_RR = 0 - stability_margin - if(FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0): + if FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0: ub_support_FR_RL = 0 + stability_margin lb_support_FR_RL = 0 - stability_margin - elif(np.array_equal(FL_contact_sequence, RL_contact_sequence) - and np.array_equal(FR_contact_sequence, RR_contact_sequence)): - #PACE - stability_margin = config.mpc_params['pace_stability_margin'] - if(FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0): + elif np.array_equal(FL_contact_sequence, RL_contact_sequence) and np.array_equal( + FR_contact_sequence, RR_contact_sequence + ): + # PACE + stability_margin = config.mpc_params["pace_stability_margin"] + if FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0: ub_support_RL_FL = 0 + stability_margin lb_support_RL_FL = 0 - stability_margin - - if(FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0): + if FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0: ub_support_FR_RR = 0 + stability_margin lb_support_FR_RR = 0 - stability_margin else: - #CRAWL BACKDIAGONALCRAWL ONLY - stability_margin = config.mpc_params['crawl_stability_margin'] + # CRAWL BACKDIAGONALCRAWL ONLY + stability_margin = config.mpc_params["crawl_stability_margin"] - if(FL_contact_sequence[j] == 1): - if(FR_contact_sequence[j] == 1): + if FL_contact_sequence[j] == 1: + if FR_contact_sequence[j] == 1: ub_support_FL_FR = -0.0 - stability_margin lb_support_FL_FR = -1000 else: ub_support_FL_RR = 1000 lb_support_FL_RR = 0.0 + stability_margin - - if(FR_contact_sequence[j] == 1): - if(RR_contact_sequence[j] == 1): + if FR_contact_sequence[j] == 1: + if RR_contact_sequence[j] == 1: ub_support_FR_RR = 1000 lb_support_FR_RR = 0.0 + stability_margin else: ub_support_FR_RL = 1000 lb_support_FR_RL = 0.0 + stability_margin - - if(RR_contact_sequence[j] == 1): - if(RL_contact_sequence[j] == 1): + if RR_contact_sequence[j] == 1: + if RL_contact_sequence[j] == 1: ub_support_RR_RL = 1000 lb_support_RR_RL = 0.0 + stability_margin else: ub_support_FL_RR = -0.0 - stability_margin lb_support_FL_RR = -1000 - - if(RL_contact_sequence[j] == 1): - if(FL_contact_sequence[j] == 1): + if RL_contact_sequence[j] == 1: + if FL_contact_sequence[j] == 1: ub_support_RL_FL = -0.0 - stability_margin lb_support_RL_FL = -1000 else: ub_support_FR_RL = -0.0 - stability_margin lb_support_FR_RL = -1000 - - - - ub_support = np.array([ub_support_FL_FR, ub_support_FR_RR, ub_support_RR_RL, ub_support_RL_FL, - ub_support_FL_RR, ub_support_FR_RL]) - lb_support = np.array([lb_support_FL_FR, lb_support_FR_RR, lb_support_RR_RL, lb_support_RL_FL, - lb_support_FL_RR, lb_support_FR_RL]) + ub_support = np.array( + [ + ub_support_FL_FR, + ub_support_FR_RR, + ub_support_RR_RL, + ub_support_RL_FL, + ub_support_FL_RR, + ub_support_FR_RL, + ] + ) + lb_support = np.array( + [ + lb_support_FL_FR, + lb_support_FR_RR, + lb_support_RR_RL, + lb_support_RL_FL, + lb_support_FL_RR, + lb_support_FR_RL, + ] + ) ub_total = np.concatenate((ub_total, ub_support)) lb_total = np.concatenate((lb_total, lb_support)) - # No friction constraint at the end! we don't have u_N - if(j == self.horizon): - if(self.use_foothold_constraints): - if(self.use_stability_constraints): + if j == self.horizon: + if self.use_foothold_constraints: + if self.use_stability_constraints: ub_total = np.concatenate((ub_foot, ub_support)) lb_total = np.concatenate((lb_foot, lb_support)) else: ub_total = ub_foot lb_total = lb_foot else: - if(self.use_stability_constraints): + if self.use_stability_constraints: ub_total = ub_support lb_total = lb_support else: continue # Only friction costraints at the beginning - if(j == 0): + if j == 0: self.acados_ocp_solver.constraints_set(j, "uh", ub_friction) self.acados_ocp_solver.constraints_set(j, "lh", lb_friction) - if(j > 0): + if j > 0: self.acados_ocp_solver.constraints_set(j, "uh", ub_total) self.acados_ocp_solver.constraints_set(j, "lh", lb_total) - - - #save the constraint for logging + # save the constraint for logging self.upper_bound[j] = ub_total.tolist() self.lower_bound[j] = lb_total.tolist() - - # ugly procedure to update the idx of the constraint - if(j>=1): - if(FL_contact_sequence[j] == 0 and FL_contact_sequence[j-1] == 1): - if(idx_constraint[0] < up_constraint_FL.shape[0] - 1): + if j >= 1: + if FL_contact_sequence[j] == 0 and FL_contact_sequence[j - 1] == 1: + if idx_constraint[0] < up_constraint_FL.shape[0] - 1: idx_constraint[0] += 1 - if(FR_contact_sequence[j] == 0 and FR_contact_sequence[j-1] == 1): - if(idx_constraint[1] < up_constraint_FR.shape[0] - 1): + if FR_contact_sequence[j] == 0 and FR_contact_sequence[j - 1] == 1: + if idx_constraint[1] < up_constraint_FR.shape[0] - 1: idx_constraint[1] += 1 - if(RL_contact_sequence[j] == 0 and RL_contact_sequence[j-1] == 1): - if(idx_constraint[2] < up_constraint_RL.shape[0] - 1): + if RL_contact_sequence[j] == 0 and RL_contact_sequence[j - 1] == 1: + if idx_constraint[2] < up_constraint_RL.shape[0] - 1: idx_constraint[2] += 1 - if(RR_contact_sequence[j] == 0 and RR_contact_sequence[j-1] == 1): - if(idx_constraint[3] < up_constraint_RR.shape[0] - 1): + if RR_contact_sequence[j] == 0 and RR_contact_sequence[j - 1] == 1: + if idx_constraint[3] < up_constraint_RR.shape[0] - 1: idx_constraint[3] += 1 except: - if(self.verbose): + if self.verbose: print("###WARNING: error in setting the constraints") return - - - def set_warm_start(self, state_acados, reference, FL_contact_sequence, FR_contact_sequence, RL_contact_sequence, RR_contact_sequence): + def set_warm_start( + self, + state_acados, + reference, + FL_contact_sequence, + FR_contact_sequence, + RL_contact_sequence, + RR_contact_sequence, + ): """ Sets the warm start for the optimization problem. This could be useful in the case some old guess is outside some new constraints.. Maybe we could have some @@ -1231,50 +1215,45 @@ def set_warm_start(self, state_acados, reference, FL_contact_sequence, FR_contac # modify onlya subset of it warm_start = copy.deepcopy(self.acados_ocp_solver.get(j, "x")) - # Update the idx_ref_foot_to_assign. Every time there is a change in the contact phase # between 1 and 0, it means that the leg go into swing and a new reference is needed!!! - if(j > 1 and j idx_ref_foot_to_assign[0]+1): + if j > 1 and j < self.horizon - 1: + if FL_contact_sequence[j] == 0 and FL_contact_sequence[j - 1] == 1: + if reference["ref_foot_FL"].shape[0] > idx_ref_foot_to_assign[0] + 1: idx_ref_foot_to_assign[0] += 1 - if(FR_contact_sequence[j] == 0 and FR_contact_sequence[j-1] == 1): - if(reference['ref_foot_FR'].shape[0] > idx_ref_foot_to_assign[1]+1): + if FR_contact_sequence[j] == 0 and FR_contact_sequence[j - 1] == 1: + if reference["ref_foot_FR"].shape[0] > idx_ref_foot_to_assign[1] + 1: idx_ref_foot_to_assign[1] += 1 - if(RL_contact_sequence[j] == 0 and RL_contact_sequence[j-1] == 1): - if(reference['ref_foot_RL'].shape[0] > idx_ref_foot_to_assign[2]+1): + if RL_contact_sequence[j] == 0 and RL_contact_sequence[j - 1] == 1: + if reference["ref_foot_RL"].shape[0] > idx_ref_foot_to_assign[2] + 1: idx_ref_foot_to_assign[2] += 1 - if(RR_contact_sequence[j] == 0 and RR_contact_sequence[j-1] == 1): - if(reference['ref_foot_RR'].shape[0] > idx_ref_foot_to_assign[3]+1): + if RR_contact_sequence[j] == 0 and RR_contact_sequence[j - 1] == 1: + if reference["ref_foot_RR"].shape[0] > idx_ref_foot_to_assign[3] + 1: idx_ref_foot_to_assign[3] += 1 warm_start[8] = state_acados[8] - if(idx_ref_foot_to_assign[0] == 0): + if idx_ref_foot_to_assign[0] == 0: warm_start[12:15] = state_acados[12:15].reshape((3,)) else: warm_start[12:15] = reference["ref_foot_FL"][0] - if(idx_ref_foot_to_assign[1] == 0): + if idx_ref_foot_to_assign[1] == 0: warm_start[15:18] = state_acados[15:18].reshape((3,)) else: warm_start[15:18] = reference["ref_foot_FR"][0] - if(idx_ref_foot_to_assign[2] == 0): + if idx_ref_foot_to_assign[2] == 0: warm_start[18:21] = state_acados[18:21].reshape((3,)) else: warm_start[18:21] = reference["ref_foot_RL"][0] - if(idx_ref_foot_to_assign[3] == 0): + if idx_ref_foot_to_assign[3] == 0: warm_start[21:24] = state_acados[21:24].reshape((3,)) else: warm_start[21:24] = reference["ref_foot_RR"][0] - self.acados_ocp_solver.set(j, "x", warm_start) - # Method to perform the centering of the states and the reference around (0, 0, 0) - def perform_scaling(self, state, reference, constraint = None): - - + def perform_scaling(self, state, reference, constraint=None): self.initial_base_position = copy.deepcopy(state["position"]) reference = copy.deepcopy(reference) state = copy.deepcopy(state) @@ -1285,23 +1264,26 @@ def perform_scaling(self, state, reference, constraint = None): reference["ref_foot_RL"] = reference["ref_foot_RL"] - state["position"] reference["ref_foot_RR"] = reference["ref_foot_RR"] - state["position"] - state["foot_FL"] = state["foot_FL"] - state["position"] state["foot_FR"] = state["foot_FR"] - state["position"] state["foot_RL"] = state["foot_RL"] - state["position"] state["foot_RR"] = state["foot_RR"] - state["position"] state["position"] = np.array([0, 0, 0]) - reference["ref_position"][0:2] = 0 #we do not close the loop in position + reference["ref_position"][0:2] = 0 # we do not close the loop in position return state, reference, constraint - - # Main loop for computing the control - def compute_control(self, state, reference, contact_sequence, constraint = None, external_wrenches = np.zeros((6,)), - inertia = config.inertia.reshape((9,)), mass = config.mass): - - + def compute_control( + self, + state, + reference, + contact_sequence, + constraint=None, + external_wrenches=np.zeros((6,)), + inertia=config.inertia.reshape((9,)), + mass=config.mass, + ): # Take the array of the contact sequence and split it in 4 arrays, # one for each leg FL_contact_sequence = contact_sequence[0] @@ -1309,191 +1291,203 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, RL_contact_sequence = contact_sequence[2] RR_contact_sequence = contact_sequence[3] - - - # Perform the scaling of the states and the reference - state, \ - reference, \ - constraint = self.perform_scaling(state, reference, constraint) - - + state, reference, constraint = self.perform_scaling(state, reference, constraint) start_reference_position = copy.deepcopy(reference["ref_position"]) # Fill reference (self.states_dim+self.inputs_dim) idx_ref_foot_to_assign = np.array([0, 0, 0, 0]) for j in range(self.horizon): - # Update the reference position for the next step - start_reference_position = start_reference_position + reference["ref_linear_velocity"]*config.mpc_params['dt'] - + start_reference_position = ( + start_reference_position + reference["ref_linear_velocity"] * config.mpc_params["dt"] + ) yref = np.zeros(shape=(self.states_dim + self.inputs_dim,)) yref[0:3] = start_reference_position - yref[3:6] = reference['ref_linear_velocity'] - yref[6:9] = reference['ref_orientation'] - yref[9:12] = reference['ref_angular_velocity'] - yref[12:15] = reference['ref_foot_FL'][idx_ref_foot_to_assign[0]] - yref[15:18] = reference['ref_foot_FR'][idx_ref_foot_to_assign[1]] - yref[18:21] = reference['ref_foot_RL'][idx_ref_foot_to_assign[2]] - yref[21:24] = reference['ref_foot_RR'][idx_ref_foot_to_assign[3]] - + yref[3:6] = reference["ref_linear_velocity"] + yref[6:9] = reference["ref_orientation"] + yref[9:12] = reference["ref_angular_velocity"] + yref[12:15] = reference["ref_foot_FL"][idx_ref_foot_to_assign[0]] + yref[15:18] = reference["ref_foot_FR"][idx_ref_foot_to_assign[1]] + yref[18:21] = reference["ref_foot_RL"][idx_ref_foot_to_assign[2]] + yref[21:24] = reference["ref_foot_RR"][idx_ref_foot_to_assign[3]] # Update the idx_ref_foot_to_assign. Every time there is a change in the contact phase # between 1 and 0, it means that the leg go into swing and a new reference is needed!!! - if(j > 1 and j idx_ref_foot_to_assign[0]+1): + if j > 1 and j < self.horizon - 1: + if FL_contact_sequence[j + 1] == 0 and FL_contact_sequence[j] == 1: + if reference["ref_foot_FL"].shape[0] > idx_ref_foot_to_assign[0] + 1: idx_ref_foot_to_assign[0] += 1 - if(FR_contact_sequence[j+1] == 0 and FR_contact_sequence[j] == 1): - if(reference['ref_foot_FR'].shape[0] > idx_ref_foot_to_assign[1]+1): + if FR_contact_sequence[j + 1] == 0 and FR_contact_sequence[j] == 1: + if reference["ref_foot_FR"].shape[0] > idx_ref_foot_to_assign[1] + 1: idx_ref_foot_to_assign[1] += 1 - if(RL_contact_sequence[j+1] == 0 and RL_contact_sequence[j] == 1): - if(reference['ref_foot_RL'].shape[0] > idx_ref_foot_to_assign[2]+1): + if RL_contact_sequence[j + 1] == 0 and RL_contact_sequence[j] == 1: + if reference["ref_foot_RL"].shape[0] > idx_ref_foot_to_assign[2] + 1: idx_ref_foot_to_assign[2] += 1 - if(RR_contact_sequence[j+1] == 0 and RR_contact_sequence[j] == 1): - if(reference['ref_foot_RR'].shape[0] > idx_ref_foot_to_assign[3]+1): + if RR_contact_sequence[j + 1] == 0 and RR_contact_sequence[j] == 1: + if reference["ref_foot_RR"].shape[0] > idx_ref_foot_to_assign[3] + 1: idx_ref_foot_to_assign[3] += 1 - - - # Setting the reference to acados - if(self.use_DDP): - if(j == 0): + if self.use_DDP: + if j == 0: num_l2_penalties = self.ocp.model.cost_y_expr_0.shape[0] - (self.states_dim + self.inputs_dim) else: num_l2_penalties = self.ocp.model.cost_y_expr.shape[0] - (self.states_dim + self.inputs_dim) - yref_tot = np.concatenate((yref, np.zeros(num_l2_penalties,) )) + yref_tot = np.concatenate( + ( + yref, + np.zeros( + num_l2_penalties, + ), + ) + ) self.acados_ocp_solver.set(j, "yref", yref_tot) else: self.acados_ocp_solver.set(j, "yref", yref) - - # Fill last step horizon reference (self.states_dim - no control action!!) - yref_N = np.zeros(shape=(self.states_dim ,)) + yref_N = np.zeros(shape=(self.states_dim,)) yref_N[0:3] = start_reference_position - yref_N[3:6] = reference['ref_linear_velocity'] - yref_N[6:9] = reference['ref_orientation'] - yref_N[9:12] = reference['ref_angular_velocity'] - yref_N[12:15] = reference['ref_foot_FL'][idx_ref_foot_to_assign[0]] - yref_N[15:18] = reference['ref_foot_FR'][idx_ref_foot_to_assign[1]] - yref_N[18:21] = reference['ref_foot_RL'][idx_ref_foot_to_assign[2]] - yref_N[21:24] = reference['ref_foot_RR'][idx_ref_foot_to_assign[3]] + yref_N[3:6] = reference["ref_linear_velocity"] + yref_N[6:9] = reference["ref_orientation"] + yref_N[9:12] = reference["ref_angular_velocity"] + yref_N[12:15] = reference["ref_foot_FL"][idx_ref_foot_to_assign[0]] + yref_N[15:18] = reference["ref_foot_FR"][idx_ref_foot_to_assign[1]] + yref_N[18:21] = reference["ref_foot_RL"][idx_ref_foot_to_assign[2]] + yref_N[21:24] = reference["ref_foot_RR"][idx_ref_foot_to_assign[3]] # Setting the reference to acados self.acados_ocp_solver.set(self.horizon, "yref", yref_N) - - - # Fill stance param, friction and stance proximity # (stance proximity will disable foothold optimization near a stance!!) - mu = config.mpc_params['mu'] + mu = config.mpc_params["mu"] yaw = state["orientation"][2] # Stance Proximity ugly routine. Basically we disable foothold optimization # in the proximity of a stance phase (the real foot cannot travel too fast in # a small time!!) - stance_proximity_FL = np.zeros((self.horizon, )) - stance_proximity_FR = np.zeros((self.horizon, )) - stance_proximity_RL = np.zeros((self.horizon, )) - stance_proximity_RR = np.zeros((self.horizon, )) + stance_proximity_FL = np.zeros((self.horizon,)) + stance_proximity_FR = np.zeros((self.horizon,)) + stance_proximity_RL = np.zeros((self.horizon,)) + stance_proximity_RR = np.zeros((self.horizon,)) for j in range(self.horizon): - if(FL_contact_sequence[j] == 0): - if(j+1) < self.horizon: - if(FL_contact_sequence[j+1] == 1): - stance_proximity_FL[j] = 1*0 - if(j+2) < self.horizon: - if(FL_contact_sequence[j+1] == 0 and FL_contact_sequence[j+2] == 1): - stance_proximity_FL[j] = 1*0 - - if(FR_contact_sequence[j] == 0): - if(j+1) < self.horizon: - if(FR_contact_sequence[j+1] == 1): - stance_proximity_FR[j] = 1*0 - if(j+2) < self.horizon: - if(FR_contact_sequence[j+1] == 0 and FR_contact_sequence[j+2] == 1): - stance_proximity_FR[j] = 1*0 - - if(RL_contact_sequence[j] == 0): - if(j+1) < self.horizon: - if(RL_contact_sequence[j+1] == 1): - stance_proximity_RL[j] = 1*0 - if(j+2) < self.horizon: - if(RL_contact_sequence[j+1] == 0 and RL_contact_sequence[j+2] == 1): - stance_proximity_RL[j] = 1*0 - - if(RR_contact_sequence[j] == 0): - if(j+1) < self.horizon: - if(RR_contact_sequence[j+1] == 1): - stance_proximity_RR[j] = 1*0 - if(j+2) < self.horizon: - if(RR_contact_sequence[j+1] == 0 and RR_contact_sequence[j+2] == 1): - stance_proximity_RR[j] = 1*0 - - + if FL_contact_sequence[j] == 0: + if (j + 1) < self.horizon: + if FL_contact_sequence[j + 1] == 1: + stance_proximity_FL[j] = 1 * 0 + if (j + 2) < self.horizon: + if FL_contact_sequence[j + 1] == 0 and FL_contact_sequence[j + 2] == 1: + stance_proximity_FL[j] = 1 * 0 + + if FR_contact_sequence[j] == 0: + if (j + 1) < self.horizon: + if FR_contact_sequence[j + 1] == 1: + stance_proximity_FR[j] = 1 * 0 + if (j + 2) < self.horizon: + if FR_contact_sequence[j + 1] == 0 and FR_contact_sequence[j + 2] == 1: + stance_proximity_FR[j] = 1 * 0 + + if RL_contact_sequence[j] == 0: + if (j + 1) < self.horizon: + if RL_contact_sequence[j + 1] == 1: + stance_proximity_RL[j] = 1 * 0 + if (j + 2) < self.horizon: + if RL_contact_sequence[j + 1] == 0 and RL_contact_sequence[j + 2] == 1: + stance_proximity_RL[j] = 1 * 0 + + if RR_contact_sequence[j] == 0: + if (j + 1) < self.horizon: + if RR_contact_sequence[j + 1] == 1: + stance_proximity_RR[j] = 1 * 0 + if (j + 2) < self.horizon: + if RR_contact_sequence[j + 1] == 0 and RR_contact_sequence[j + 2] == 1: + stance_proximity_RR[j] = 1 * 0 # Set the parameters to acados for j in range(self.horizon): # If we have estimated an external wrench, we can compensate it for all steps # or less (maybe the disturbance is not costant along the horizon!) - if(config.mpc_params['external_wrenches_compensation'] and - config.mpc_params['external_wrenches_compensation_num_step'] and - j < config.mpc_params['external_wrenches_compensation_num_step']): + if ( + config.mpc_params["external_wrenches_compensation"] + and config.mpc_params["external_wrenches_compensation_num_step"] + and j < config.mpc_params["external_wrenches_compensation_num_step"] + ): external_wrenches_estimated_param = copy.deepcopy(external_wrenches) - external_wrenches_estimated_param = external_wrenches_estimated_param.reshape((6, )) + external_wrenches_estimated_param = external_wrenches_estimated_param.reshape((6,)) else: external_wrenches_estimated_param = np.zeros((6,)) - - param = np.array([FL_contact_sequence[j], FR_contact_sequence[j], - RL_contact_sequence[j], RR_contact_sequence[j], mu, - stance_proximity_FL[j], - stance_proximity_FR[j], - stance_proximity_RL[j], - stance_proximity_RR[j], - state["position"][0], state["position"][1], - state["position"][2], state["orientation"][2], - external_wrenches_estimated_param[0], external_wrenches_estimated_param[1], - external_wrenches_estimated_param[2], external_wrenches_estimated_param[3], - external_wrenches_estimated_param[4], external_wrenches_estimated_param[5], - inertia[0], inertia[1], inertia[2], inertia[3], inertia[4], inertia[5], - inertia[6], inertia[7], inertia[8], mass]) + param = np.array( + [ + FL_contact_sequence[j], + FR_contact_sequence[j], + RL_contact_sequence[j], + RR_contact_sequence[j], + mu, + stance_proximity_FL[j], + stance_proximity_FR[j], + stance_proximity_RL[j], + stance_proximity_RR[j], + state["position"][0], + state["position"][1], + state["position"][2], + state["orientation"][2], + external_wrenches_estimated_param[0], + external_wrenches_estimated_param[1], + external_wrenches_estimated_param[2], + external_wrenches_estimated_param[3], + external_wrenches_estimated_param[4], + external_wrenches_estimated_param[5], + inertia[0], + inertia[1], + inertia[2], + inertia[3], + inertia[4], + inertia[5], + inertia[6], + inertia[7], + inertia[8], + mass, + ] + ) self.acados_ocp_solver.set(j, "p", copy.deepcopy(param)) - - # Set initial state constraint. We teleported the robot foothold # to the previous optimal foothold. This is done to avoid the optimization # of a foothold that is not considered at all at touchdown! In any case, # the height cames always from the VFA - if(FL_contact_sequence[0] == 0): + if FL_contact_sequence[0] == 0: state["foot_FL"] = reference["ref_foot_FL"][0] - if(FR_contact_sequence[0] == 0): + if FR_contact_sequence[0] == 0: state["foot_FR"] = reference["ref_foot_FR"][0] - if(RL_contact_sequence[0] == 0): + if RL_contact_sequence[0] == 0: state["foot_RL"] = reference["ref_foot_RL"][0] - if(RR_contact_sequence[0] == 0): + if RR_contact_sequence[0] == 0: state["foot_RR"] = reference["ref_foot_RR"][0] - - if(self.use_integrators): + if self.use_integrators: # Compute error for integral action alpha_integrator = config.mpc_params["alpha_integrator"] - self.integral_errors[0] += (state["position"][2] - reference["ref_position"][2])*alpha_integrator - self.integral_errors[1] += (state["linear_velocity"][0] - reference["ref_linear_velocity"][0])*alpha_integrator - self.integral_errors[2] += (state["linear_velocity"][1] - reference["ref_linear_velocity"][1])*alpha_integrator - self.integral_errors[3] += (state["linear_velocity"][2] - reference["ref_linear_velocity"][2])*alpha_integrator - self.integral_errors[4] += (state["orientation"][0] - reference["ref_orientation"][0])*(alpha_integrator) - self.integral_errors[5] += (state["orientation"][1] - reference["ref_orientation"][1])*alpha_integrator - + self.integral_errors[0] += (state["position"][2] - reference["ref_position"][2]) * alpha_integrator + self.integral_errors[1] += ( + state["linear_velocity"][0] - reference["ref_linear_velocity"][0] + ) * alpha_integrator + self.integral_errors[2] += ( + state["linear_velocity"][1] - reference["ref_linear_velocity"][1] + ) * alpha_integrator + self.integral_errors[3] += ( + state["linear_velocity"][2] - reference["ref_linear_velocity"][2] + ) * alpha_integrator + self.integral_errors[4] += (state["orientation"][0] - reference["ref_orientation"][0]) * (alpha_integrator) + self.integral_errors[5] += (state["orientation"][1] - reference["ref_orientation"][1]) * alpha_integrator cap_integrator_z = config.mpc_params["integrator_cap"][0] cap_integrator_x_dot = config.mpc_params["integrator_cap"][1] @@ -1502,97 +1496,128 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, cap_integrator_roll = config.mpc_params["integrator_cap"][4] cap_integrator_pitch = config.mpc_params["integrator_cap"][5] - self.integral_errors[0] = np.where(np.abs(self.integral_errors[0]) > cap_integrator_z, cap_integrator_z*np.sign(self.integral_errors[0]), self.integral_errors[0]) - self.integral_errors[1] = np.where(np.abs(self.integral_errors[1]) > cap_integrator_x_dot, cap_integrator_x_dot*np.sign(self.integral_errors[1]), self.integral_errors[1]) - self.integral_errors[2] = np.where(np.abs(self.integral_errors[2]) > cap_integrator_y_dot, cap_integrator_y_dot*np.sign(self.integral_errors[2]), self.integral_errors[2]) - self.integral_errors[3] = np.where(np.abs(self.integral_errors[3]) > cap_integrator_z_dot, cap_integrator_z_dot*np.sign(self.integral_errors[3]), self.integral_errors[3]) - self.integral_errors[4] = np.where(np.abs(self.integral_errors[4]) > cap_integrator_roll, cap_integrator_roll*np.sign(self.integral_errors[4]), self.integral_errors[4]) - self.integral_errors[5] = np.where(np.abs(self.integral_errors[5]) > cap_integrator_pitch, cap_integrator_pitch*np.sign(self.integral_errors[5]), self.integral_errors[5]) - - if(self.verbose): + self.integral_errors[0] = np.where( + np.abs(self.integral_errors[0]) > cap_integrator_z, + cap_integrator_z * np.sign(self.integral_errors[0]), + self.integral_errors[0], + ) + self.integral_errors[1] = np.where( + np.abs(self.integral_errors[1]) > cap_integrator_x_dot, + cap_integrator_x_dot * np.sign(self.integral_errors[1]), + self.integral_errors[1], + ) + self.integral_errors[2] = np.where( + np.abs(self.integral_errors[2]) > cap_integrator_y_dot, + cap_integrator_y_dot * np.sign(self.integral_errors[2]), + self.integral_errors[2], + ) + self.integral_errors[3] = np.where( + np.abs(self.integral_errors[3]) > cap_integrator_z_dot, + cap_integrator_z_dot * np.sign(self.integral_errors[3]), + self.integral_errors[3], + ) + self.integral_errors[4] = np.where( + np.abs(self.integral_errors[4]) > cap_integrator_roll, + cap_integrator_roll * np.sign(self.integral_errors[4]), + self.integral_errors[4], + ) + self.integral_errors[5] = np.where( + np.abs(self.integral_errors[5]) > cap_integrator_pitch, + cap_integrator_pitch * np.sign(self.integral_errors[5]), + self.integral_errors[5], + ) + + if self.verbose: print("self.integral_errors: ", self.integral_errors) - # Calculate state of the lyapunov function K_z1 = config.mpc_params["K_z1"] K_z2 = config.mpc_params["K_z2"] - z1 = state["position"] - reference['ref_position'] - z2 = state["linear_velocity"] - reference['ref_linear_velocity'] + K_z1*z1 + z1 = state["position"] - reference["ref_position"] + z2 = state["linear_velocity"] - reference["ref_linear_velocity"] + K_z1 * z1 eta = np.concatenate((state["orientation"], state["angular_velocity"])) phi = self.phi_predicted - - # Set initial state constraint acados, converting first the dictionary to np array - state_acados = np.concatenate((state["position"], state["linear_velocity"], - state["orientation"], state["angular_velocity"], - state["foot_FL"], state["foot_FR"], - state["foot_RL"], state["foot_RR"], - self.integral_errors, - z1, z2, eta, phi)).reshape((self.states_dim, 1)) + state_acados = np.concatenate( + ( + state["position"], + state["linear_velocity"], + state["orientation"], + state["angular_velocity"], + state["foot_FL"], + state["foot_FR"], + state["foot_RL"], + state["foot_RR"], + self.integral_errors, + z1, + z2, + eta, + phi, + ) + ).reshape((self.states_dim, 1)) self.acados_ocp_solver.set(0, "lbx", state_acados) self.acados_ocp_solver.set(0, "ubx", state_acados) - # Set Warm start in case... - if(self.use_warm_start): - self.set_warm_start(state_acados, reference, FL_contact_sequence, FR_contact_sequence, RL_contact_sequence, RR_contact_sequence) - + if self.use_warm_start: + self.set_warm_start( + state_acados, + reference, + FL_contact_sequence, + FR_contact_sequence, + RL_contact_sequence, + RR_contact_sequence, + ) # Set stage constraint - h_R_w = np.array([np.cos(yaw), np.sin(yaw), - -np.sin(yaw), np.cos(yaw)]) - if(self.use_foothold_constraints or self.use_stability_constraints): - - stance_proximity = np.vstack((stance_proximity_FL, stance_proximity_FR, - stance_proximity_RL, stance_proximity_RR)) + h_R_w = np.array([np.cos(yaw), np.sin(yaw), -np.sin(yaw), np.cos(yaw)]) + if self.use_foothold_constraints or self.use_stability_constraints: + stance_proximity = np.vstack( + (stance_proximity_FL, stance_proximity_FR, stance_proximity_RL, stance_proximity_RR) + ) self.set_stage_constraint(constraint, state, reference, contact_sequence, h_R_w, stance_proximity) - if(config.mpc_params["use_residual_dynamics_decay"]): + if config.mpc_params["use_residual_dynamics_decay"]: self.set_residual_dynamics_constraint() - # Solve ocp via RTI or normal ocp if self.use_RTI: # feedback phase - self.acados_ocp_solver.options_set('rti_phase', 2) + self.acados_ocp_solver.options_set("rti_phase", 2) status = self.acados_ocp_solver.solve() - if(self.verbose): - print("feedback phase time: ", self.acados_ocp_solver.get_stats('time_tot')) + if self.verbose: + print("feedback phase time: ", self.acados_ocp_solver.get_stats("time_tot")) else: status = self.acados_ocp_solver.solve() - if(self.verbose): - print("ocp time: ", self.acados_ocp_solver.get_stats('time_tot')) - + if self.verbose: + print("ocp time: ", self.acados_ocp_solver.get_stats("time_tot")) # Take the solution control = self.acados_ocp_solver.get(0, "u") optimal_GRF = control[12:] optimal_foothold = np.zeros((4, 3)) - optimal_footholds_assigned = np.zeros((4, ), dtype='bool') + optimal_footholds_assigned = np.zeros((4,), dtype="bool") self.phi_predicted = self.acados_ocp_solver.get(1, "x")[42:45] # Now i add the component from the adaptive control law p_ddot = np.array([0, 0, 0]) gravity = np.array([0, 0, -9.81]) - squared_K_z1 = np.array([K_z1[0]*K_z1[0], K_z1[1]*K_z1[1], K_z1[2]*K_z1[2]]) - F_star = mass*(-(K_z1 + K_z2)*z2 + squared_K_z1*z1 - gravity + p_ddot) + squared_K_z1 = np.array([K_z1[0] * K_z1[0], K_z1[1] * K_z1[1], K_z1[2] * K_z1[2]]) + F_star = mass * (-(K_z1 + K_z2) * z2 + squared_K_z1 * z1 - gravity + p_ddot) F_star -= phi - num_leg_contact = FL_contact_sequence[0] + FR_contact_sequence[0] num_leg_contact += RL_contact_sequence[0] + RR_contact_sequence[0] - GRF_star = (F_star/num_leg_contact) - - - - optimal_GRF[0:3] = copy.deepcopy(optimal_GRF[0:3]*FL_contact_sequence[0] + GRF_star*FL_contact_sequence[0]) - optimal_GRF[3:6] = copy.deepcopy(optimal_GRF[3:6]*FR_contact_sequence[0] + GRF_star*FR_contact_sequence[0]) - optimal_GRF[6:9] = copy.deepcopy(optimal_GRF[6:9]*RL_contact_sequence[0] + GRF_star*RL_contact_sequence[0]) - optimal_GRF[9:12] = copy.deepcopy(optimal_GRF[9:12]*RR_contact_sequence[0] + GRF_star*RR_contact_sequence[0]) - + GRF_star = F_star / num_leg_contact + optimal_GRF[0:3] = copy.deepcopy(optimal_GRF[0:3] * FL_contact_sequence[0] + GRF_star * FL_contact_sequence[0]) + optimal_GRF[3:6] = copy.deepcopy(optimal_GRF[3:6] * FR_contact_sequence[0] + GRF_star * FR_contact_sequence[0]) + optimal_GRF[6:9] = copy.deepcopy(optimal_GRF[6:9] * RL_contact_sequence[0] + GRF_star * RL_contact_sequence[0]) + optimal_GRF[9:12] = copy.deepcopy( + optimal_GRF[9:12] * RR_contact_sequence[0] + GRF_star * RR_contact_sequence[0] + ) # We need to provide the next touchdown foothold position. # We first take the foothold in stance now (they are not optimized!) @@ -1610,132 +1635,186 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, optimal_foothold[3] = state["foot_RR"] optimal_footholds_assigned[3] = True - - # Then we take the foothold at the next touchdown from the one # that are not flagged as True from before, and saturate them!! # P.S. The saturation is in the horizontal frame h_R_w = h_R_w.reshape((2, 2)) for j in range(1, self.horizon): - if(FL_contact_sequence[j] != FL_contact_sequence[j-1] and not optimal_footholds_assigned[0]): + if FL_contact_sequence[j] != FL_contact_sequence[j - 1] and not optimal_footholds_assigned[0]: optimal_foothold[0] = self.acados_ocp_solver.get(j, "x")[12:15] optimal_footholds_assigned[0] = True - if(constraint is None): - first_up_constraint_FL = np.array([reference["ref_foot_FL"][0][0], reference["ref_foot_FL"][0][1], reference["ref_foot_FL"][0][2] + 0.002]) - first_low_constraint_FL = np.array([reference["ref_foot_FL"][0][0], reference["ref_foot_FL"][0][1], reference["ref_foot_FL"][0][2] - 0.002]) - - first_up_constraint_FL[0:2] = h_R_w@(first_up_constraint_FL[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_FL[0:2] = h_R_w@(first_low_constraint_FL[0:2] - state["position"][0:2]) - 0.15 + if constraint is None: + first_up_constraint_FL = np.array( + [ + reference["ref_foot_FL"][0][0], + reference["ref_foot_FL"][0][1], + reference["ref_foot_FL"][0][2] + 0.002, + ] + ) + first_low_constraint_FL = np.array( + [ + reference["ref_foot_FL"][0][0], + reference["ref_foot_FL"][0][1], + reference["ref_foot_FL"][0][2] - 0.002, + ] + ) + + first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - state["position"][0:2]) + 0.15 + first_low_constraint_FL[0:2] = ( + h_R_w @ (first_low_constraint_FL[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_FL = np.array([constraint[0][0], constraint[1][0], constraint[2][0] + 0.002]) first_low_constraint_FL = np.array([constraint[9][0], constraint[10][0], constraint[11][0] - 0.002]) - first_up_constraint_FL[0:2] = h_R_w@(first_up_constraint_FL[0:2] - state["position"][0:2]) - first_low_constraint_FL[0:2] = h_R_w@(first_low_constraint_FL[0:2] - state["position"][0:2]) - - optimal_foothold[0][0:2] = h_R_w@(optimal_foothold[0][0:2] - state["position"][0:2]) - optimal_foothold[0][0:2] = np.clip(optimal_foothold[0][0:2], first_low_constraint_FL[0:2], first_up_constraint_FL[0:2]) - optimal_foothold[0][0:2] = h_R_w.T@optimal_foothold[0][0:2] + state["position"][0:2] + first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - state["position"][0:2]) + first_low_constraint_FL[0:2] = h_R_w @ (first_low_constraint_FL[0:2] - state["position"][0:2]) + optimal_foothold[0][0:2] = h_R_w @ (optimal_foothold[0][0:2] - state["position"][0:2]) + optimal_foothold[0][0:2] = np.clip( + optimal_foothold[0][0:2], first_low_constraint_FL[0:2], first_up_constraint_FL[0:2] + ) + optimal_foothold[0][0:2] = h_R_w.T @ optimal_foothold[0][0:2] + state["position"][0:2] - if(FR_contact_sequence[j] != FR_contact_sequence[j-1] and not optimal_footholds_assigned[1]): + if FR_contact_sequence[j] != FR_contact_sequence[j - 1] and not optimal_footholds_assigned[1]: optimal_foothold[1] = self.acados_ocp_solver.get(j, "x")[15:18] optimal_footholds_assigned[1] = True - if(constraint is None): - first_up_constraint_FR = np.array([reference["ref_foot_FR"][0][0], reference["ref_foot_FR"][0][1], reference["ref_foot_FR"][0][2] + 0.002]) - first_low_constraint_FR = np.array([reference["ref_foot_FR"][0][0], reference["ref_foot_FR"][0][1], reference["ref_foot_FR"][0][2] - 0.002]) - - first_up_constraint_FR[0:2] = h_R_w@(first_up_constraint_FR[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_FR[0:2] = h_R_w@(first_low_constraint_FR[0:2] - state["position"][0:2]) - 0.15 + if constraint is None: + first_up_constraint_FR = np.array( + [ + reference["ref_foot_FR"][0][0], + reference["ref_foot_FR"][0][1], + reference["ref_foot_FR"][0][2] + 0.002, + ] + ) + first_low_constraint_FR = np.array( + [ + reference["ref_foot_FR"][0][0], + reference["ref_foot_FR"][0][1], + reference["ref_foot_FR"][0][2] - 0.002, + ] + ) + + first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - state["position"][0:2]) + 0.15 + first_low_constraint_FR[0:2] = ( + h_R_w @ (first_low_constraint_FR[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_FR = np.array([constraint[0][1], constraint[1][1], constraint[2][1] + 0.002]) first_low_constraint_FR = np.array([constraint[9][1], constraint[10][1], constraint[11][1] - 0.002]) - first_up_constraint_FR[0:2] = h_R_w@(first_up_constraint_FR[0:2] - state["position"][0:2]) - first_low_constraint_FR[0:2] = h_R_w@(first_low_constraint_FR[0:2] - state["position"][0:2]) + first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - state["position"][0:2]) + first_low_constraint_FR[0:2] = h_R_w @ (first_low_constraint_FR[0:2] - state["position"][0:2]) - optimal_foothold[1][0:2] = h_R_w@(optimal_foothold[1][0:2] - state["position"][0:2]) - optimal_foothold[1][0:2] = np.clip(optimal_foothold[1][0:2], first_low_constraint_FR[0:2], first_up_constraint_FR[0:2]) - optimal_foothold[1][0:2] = h_R_w.T@optimal_foothold[1][0:2] + state["position"][0:2] + optimal_foothold[1][0:2] = h_R_w @ (optimal_foothold[1][0:2] - state["position"][0:2]) + optimal_foothold[1][0:2] = np.clip( + optimal_foothold[1][0:2], first_low_constraint_FR[0:2], first_up_constraint_FR[0:2] + ) + optimal_foothold[1][0:2] = h_R_w.T @ optimal_foothold[1][0:2] + state["position"][0:2] - - if(RL_contact_sequence[j] != RL_contact_sequence[j-1] and not optimal_footholds_assigned[2]): + if RL_contact_sequence[j] != RL_contact_sequence[j - 1] and not optimal_footholds_assigned[2]: optimal_foothold[2] = self.acados_ocp_solver.get(j, "x")[18:21] optimal_footholds_assigned[2] = True - if(constraint is None): - first_up_constraint_RL = np.array([reference["ref_foot_RL"][0][0], reference["ref_foot_RL"][0][1], reference["ref_foot_RL"][0][2] + 0.002]) - first_low_constraint_RL = np.array([reference["ref_foot_RL"][0][0], reference["ref_foot_RL"][0][1], reference["ref_foot_RL"][0][2] - 0.002]) - - first_up_constraint_RL[0:2] = h_R_w@(first_up_constraint_RL[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_RL[0:2] = h_R_w@(first_low_constraint_RL[0:2] - state["position"][0:2]) - 0.15 + if constraint is None: + first_up_constraint_RL = np.array( + [ + reference["ref_foot_RL"][0][0], + reference["ref_foot_RL"][0][1], + reference["ref_foot_RL"][0][2] + 0.002, + ] + ) + first_low_constraint_RL = np.array( + [ + reference["ref_foot_RL"][0][0], + reference["ref_foot_RL"][0][1], + reference["ref_foot_RL"][0][2] - 0.002, + ] + ) + + first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - state["position"][0:2]) + 0.15 + first_low_constraint_RL[0:2] = ( + h_R_w @ (first_low_constraint_RL[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_RL = np.array([constraint[0][2], constraint[1][2], constraint[2][2] + 0.002]) first_low_constraint_RL = np.array([constraint[9][2], constraint[10][2], constraint[11][2] - 0.002]) - first_up_constraint_RL[0:2] = h_R_w@(first_up_constraint_RL[0:2] - state["position"][0:2]) - first_low_constraint_RL[0:2] = h_R_w@(first_low_constraint_RL[0:2] - state["position"][0:2]) + first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - state["position"][0:2]) + first_low_constraint_RL[0:2] = h_R_w @ (first_low_constraint_RL[0:2] - state["position"][0:2]) - optimal_foothold[2][0:2] = h_R_w@(optimal_foothold[2][0:2] - state["position"][0:2]) - optimal_foothold[2][0:2] = np.clip(optimal_foothold[2][0:2], first_low_constraint_RL[0:2], first_up_constraint_RL[0:2]) - optimal_foothold[2][0:2] = h_R_w.T@optimal_foothold[2][0:2] + state["position"][0:2] + optimal_foothold[2][0:2] = h_R_w @ (optimal_foothold[2][0:2] - state["position"][0:2]) + optimal_foothold[2][0:2] = np.clip( + optimal_foothold[2][0:2], first_low_constraint_RL[0:2], first_up_constraint_RL[0:2] + ) + optimal_foothold[2][0:2] = h_R_w.T @ optimal_foothold[2][0:2] + state["position"][0:2] - if(RR_contact_sequence[j] != RR_contact_sequence[j-1] and not optimal_footholds_assigned[3]): + if RR_contact_sequence[j] != RR_contact_sequence[j - 1] and not optimal_footholds_assigned[3]: optimal_foothold[3] = self.acados_ocp_solver.get(j, "x")[21:24] optimal_footholds_assigned[3] = True - if(constraint is None): - first_up_constraint_RR = np.array([reference["ref_foot_RR"][0][0], reference["ref_foot_RR"][0][1], reference["ref_foot_RR"][0][2] + 0.002]) - first_low_constraint_RR = np.array([reference["ref_foot_RR"][0][0], reference["ref_foot_RR"][0][1], reference["ref_foot_RR"][0][2] - 0.002]) - - first_up_constraint_RR[0:2] = h_R_w@(first_up_constraint_RR[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_RR[0:2] = h_R_w@(first_low_constraint_RR[0:2] - state["position"][0:2]) - 0.15 + if constraint is None: + first_up_constraint_RR = np.array( + [ + reference["ref_foot_RR"][0][0], + reference["ref_foot_RR"][0][1], + reference["ref_foot_RR"][0][2] + 0.002, + ] + ) + first_low_constraint_RR = np.array( + [ + reference["ref_foot_RR"][0][0], + reference["ref_foot_RR"][0][1], + reference["ref_foot_RR"][0][2] - 0.002, + ] + ) + + first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - state["position"][0:2]) + 0.15 + first_low_constraint_RR[0:2] = ( + h_R_w @ (first_low_constraint_RR[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_RR = np.array([constraint[0][3], constraint[1][3], constraint[2][3] + 0.002]) first_low_constraint_RR = np.array([constraint[9][3], constraint[10][3], constraint[11][3] - 0.002]) - first_up_constraint_RR[0:2] = h_R_w@(first_up_constraint_RR[0:2] - state["position"][0:2]) - first_low_constraint_RR[0:2] = h_R_w@(first_low_constraint_RR[0:2] - state["position"][0:2]) - - optimal_foothold[3][0:2] = h_R_w@(optimal_foothold[3][0:2] - state["position"][0:2]) - optimal_foothold[3][0:2] = np.clip(optimal_foothold[3][0:2], first_low_constraint_RR[0:2], first_up_constraint_RR[0:2]) - optimal_foothold[3][0:2] = h_R_w.T@optimal_foothold[3][0:2] + state["position"][0:2] - - - + first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - state["position"][0:2]) + first_low_constraint_RR[0:2] = h_R_w @ (first_low_constraint_RR[0:2] - state["position"][0:2]) + optimal_foothold[3][0:2] = h_R_w @ (optimal_foothold[3][0:2] - state["position"][0:2]) + optimal_foothold[3][0:2] = np.clip( + optimal_foothold[3][0:2], first_low_constraint_RR[0:2], first_up_constraint_RR[0:2] + ) + optimal_foothold[3][0:2] = h_R_w.T @ optimal_foothold[3][0:2] + state["position"][0:2] # If in the prediction horizon, the foot is never in stance, we replicate the reference - #to not confuse the swing controller - if(optimal_footholds_assigned[0] == False): + # to not confuse the swing controller + if optimal_footholds_assigned[0] == False: optimal_foothold[0] = reference["ref_foot_FL"][0] - if(optimal_footholds_assigned[1] == False): + if optimal_footholds_assigned[1] == False: optimal_foothold[1] = reference["ref_foot_FR"][0] - if(optimal_footholds_assigned[2] == False): + if optimal_footholds_assigned[2] == False: optimal_foothold[2] = reference["ref_foot_RL"][0] - if(optimal_footholds_assigned[3] == False): + if optimal_footholds_assigned[3] == False: optimal_foothold[3] = reference["ref_foot_RR"][0] - - - - if(config.mpc_params['dt'] <= 0.02 or (config.mpc_params['use_nonuniform_discretization'] and config.mpc_params['dt_fine_grained'] <= 0.02)): + if config.mpc_params["dt"] <= 0.02 or ( + config.mpc_params["use_nonuniform_discretization"] and config.mpc_params["dt_fine_grained"] <= 0.02 + ): optimal_next_state_index = 2 else: optimal_next_state_index = 1 optimal_next_state = self.acados_ocp_solver.get(optimal_next_state_index, "x")[0:24] self.optimal_next_state = optimal_next_state - if(self.verbose): + if self.verbose: self.acados_ocp_solver.print_statistics() - # Check if QPs converged, if not just use the reference footholds # and a GRF over Z distribuited between the leg in stance - if(status == 1 or status == 4): - if(self.verbose): + if status == 1 or status == 4: + if self.verbose: print("status", status) if FL_contact_sequence[0] == 0: optimal_foothold[0] = reference["ref_foot_FL"][0] @@ -1746,35 +1825,29 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, if RR_contact_sequence[0] == 0: optimal_foothold[3] = reference["ref_foot_RR"][0] - number_of_legs_in_stance = np.array([FL_contact_sequence[0], FR_contact_sequence[0], - RL_contact_sequence [0], RR_contact_sequence[0]]).sum() + number_of_legs_in_stance = np.array( + [FL_contact_sequence[0], FR_contact_sequence[0], RL_contact_sequence[0], RR_contact_sequence[0]] + ).sum() reference_force_stance_legs = (mass * 9.81) / number_of_legs_in_stance reference_force_fl_z = reference_force_stance_legs * FL_contact_sequence[0] reference_force_fr_z = reference_force_stance_legs * FR_contact_sequence[0] reference_force_rl_z = reference_force_stance_legs * RL_contact_sequence[0] reference_force_rr_z = reference_force_stance_legs * RR_contact_sequence[0] - optimal_GRF = np.zeros((12, )) + optimal_GRF = np.zeros((12,)) optimal_GRF[2] = reference_force_fl_z optimal_GRF[5] = reference_force_fr_z optimal_GRF[8] = reference_force_rl_z optimal_GRF[11] = reference_force_rr_z - - optimal_GRF = self.previous_optimal_GRF self.acados_ocp_solver.reset() - - - - # Save the previous optimal GRF, the previous status and the previous contact sequence self.previous_optimal_GRF = optimal_GRF self.previous_status = status self.previous_contact_sequence = contact_sequence - # Decenter the optimal foothold and the next state (they were centered around zero at the beginning) optimal_foothold[0] = optimal_foothold[0] + self.initial_base_position optimal_foothold[1] = optimal_foothold[1] + self.initial_base_position @@ -1787,6 +1860,5 @@ def compute_control(self, state, reference, contact_sequence, constraint = None, optimal_next_state[18:21] = optimal_foothold[2] optimal_next_state[21:24] = optimal_foothold[3] - # Return the optimal GRF, the optimal foothold, the next state and the status of the optimization and the phi predicted return optimal_GRF, optimal_foothold, optimal_next_state, status diff --git a/quadruped_pympc/controllers/gradient/nominal/centroidal_model_nominal.py b/quadruped_pympc/controllers/gradient/nominal/centroidal_model_nominal.py index b66aa13..fa89272 100644 --- a/quadruped_pympc/controllers/gradient/nominal/centroidal_model_nominal.py +++ b/quadruped_pympc/controllers/gradient/nominal/centroidal_model_nominal.py @@ -3,34 +3,32 @@ # Authors: Giulio Turrisi - -import time -import unittest -import casadi as cs +import os +import casadi as cs import numpy as np from acados_template import AcadosModel -import os dir_path = os.path.dirname(os.path.realpath(__file__)) import sys + sys.path.append(dir_path) -sys.path.append(dir_path + '/../../') +sys.path.append(dir_path + "/../../") from quadruped_pympc import config # Class that defines the prediction model of the NMPC class Centroidal_Model_Nominal: - def __init__(self,) -> None: + def __init__( + self, + ) -> None: """ This method initializes the foothold generator Centroidal_Model, which creates the prediction model of the NMPC. """ - - - # Define state and its casadi variables com_position_x = cs.SX.sym("com_position_x") com_position_y = cs.SX.sym("com_position_y") @@ -39,20 +37,19 @@ def __init__(self,) -> None: com_velocity_x = cs.SX.sym("com_velocity_x") com_velocity_y = cs.SX.sym("com_velocity_y") com_velocity_z = cs.SX.sym("com_velocity_z") - + roll = cs.SX.sym("roll", 1, 1) pitch = cs.SX.sym("pitch", 1, 1) yaw = cs.SX.sym("yaw", 1, 1) omega_x = cs.SX.sym("omega_x", 1, 1) omega_y = cs.SX.sym("omega_y", 1, 1) omega_z = cs.SX.sym("omega_z", 1, 1) - + foot_position_fl = cs.SX.sym("foot_position_fl", 3, 1) foot_position_fr = cs.SX.sym("foot_position_fr", 3, 1) foot_position_rl = cs.SX.sym("foot_position_rl", 3, 1) foot_position_rr = cs.SX.sym("foot_position_rr", 3, 1) - com_position_z_integral = cs.SX.sym("com_position_z_integral") com_velocity_x_integral = cs.SX.sym("com_velocity_x_integral") com_velocity_y_integral = cs.SX.sym("com_velocity_y_integral") @@ -63,47 +60,46 @@ def __init__(self,) -> None: omega_y_integral = cs.SX.sym("omega_y_integral") omega_z_integral = cs.SX.sym("omega_z_integral") - - self.states = cs.vertcat(com_position_x, - com_position_y, - com_position_z, - com_velocity_x, - com_velocity_y, - com_velocity_z, - roll, - pitch, - yaw, - omega_x, - omega_y, - omega_z, - foot_position_fl, - foot_position_fr, - foot_position_rl, - foot_position_rr, - com_position_z_integral, - com_velocity_x_integral, - com_velocity_y_integral, - com_velocity_z_integral, - roll_integral, - pitch_integral) - - - - # Define state dot - self.states_dot = cs.vertcat(cs.SX.sym("linear_com_vel", 3, 1), - cs.SX.sym("linear_com_acc", 3, 1), - cs.SX.sym("euler_rates_base", 3, 1), - cs.SX.sym("angular_acc_base", 3, 1), - cs.SX.sym("linear_vel_foot_FL", 3, 1), - cs.SX.sym("linear_vel_foot_FR", 3, 1), - cs.SX.sym("linear_vel_foot_RL", 3, 1), - cs.SX.sym("linear_vel_foot_RR", 3, 1), - cs.SX.sym("linear_com_vel_z_integral", 1, 1), - cs.SX.sym("linear_com_acc_integral", 3, 1), - cs.SX.sym("euler_rates_roll_integral", 1, 1), - cs.SX.sym("euler_rates_pitch_integral", 1, 1)) - - + self.states = cs.vertcat( + com_position_x, + com_position_y, + com_position_z, + com_velocity_x, + com_velocity_y, + com_velocity_z, + roll, + pitch, + yaw, + omega_x, + omega_y, + omega_z, + foot_position_fl, + foot_position_fr, + foot_position_rl, + foot_position_rr, + com_position_z_integral, + com_velocity_x_integral, + com_velocity_y_integral, + com_velocity_z_integral, + roll_integral, + pitch_integral, + ) + + # Define state dot + self.states_dot = cs.vertcat( + cs.SX.sym("linear_com_vel", 3, 1), + cs.SX.sym("linear_com_acc", 3, 1), + cs.SX.sym("euler_rates_base", 3, 1), + cs.SX.sym("angular_acc_base", 3, 1), + cs.SX.sym("linear_vel_foot_FL", 3, 1), + cs.SX.sym("linear_vel_foot_FR", 3, 1), + cs.SX.sym("linear_vel_foot_RL", 3, 1), + cs.SX.sym("linear_vel_foot_RR", 3, 1), + cs.SX.sym("linear_com_vel_z_integral", 1, 1), + cs.SX.sym("linear_com_acc_integral", 3, 1), + cs.SX.sym("euler_rates_roll_integral", 1, 1), + cs.SX.sym("euler_rates_pitch_integral", 1, 1), + ) # Define input and its casadi variables foot_velocity_fl = cs.SX.sym("foot_velocity_fl", 3, 1) @@ -116,28 +112,27 @@ def __init__(self,) -> None: foot_force_rl = cs.SX.sym("foot_force_rl", 3, 1) foot_force_rr = cs.SX.sym("foot_force_rr", 3, 1) - self.inputs = cs.vertcat(foot_velocity_fl, - foot_velocity_fr, - foot_velocity_rl, - foot_velocity_rr, - foot_force_fl, - foot_force_fr, - foot_force_rl, - foot_force_rr) - + self.inputs = cs.vertcat( + foot_velocity_fl, + foot_velocity_fr, + foot_velocity_rl, + foot_velocity_rr, + foot_force_fl, + foot_force_fr, + foot_force_rl, + foot_force_rr, + ) # Usefull for debug what things goes where in y_ref in the compute_control function, # because there are a lot of variables self.y_ref = cs.vertcat(self.states, self.inputs) - # Define acados parameters that can be changed at runtine self.stanceFL = cs.SX.sym("stanceFL", 1, 1) self.stanceFR = cs.SX.sym("stanceFR", 1, 1) self.stanceRL = cs.SX.sym("stanceRL", 1, 1) self.stanceRR = cs.SX.sym("stanceRR", 1, 1) - self.stance_param = cs.vertcat(self.stanceFL , self.stanceFR , self.stanceRL , self.stanceRR) - + self.stance_param = cs.vertcat(self.stanceFL, self.stanceFR, self.stanceRL, self.stanceRR) self.mu_friction = cs.SX.sym("mu_friction", 1, 1) self.stance_proximity = cs.SX.sym("stanceProximity", 4, 1) @@ -149,16 +144,19 @@ def __init__(self,) -> None: self.inertia = cs.SX.sym("inertia", 9, 1) self.mass = cs.SX.sym("mass", 1, 1) - # Not so useful, i can instantiate a casadi function for the fd - param = cs.vertcat(self.stance_param, self.mu_friction, self.stance_proximity, - self.base_position, self.base_yaw, self.external_wrench, self.inertia, self.mass) + param = cs.vertcat( + self.stance_param, + self.mu_friction, + self.stance_proximity, + self.base_position, + self.base_yaw, + self.external_wrench, + self.inertia, + self.mass, + ) fd = self.forward_dynamics(self.states, self.inputs, param) - self.fun_forward_dynamics = cs.Function('fun_forward_dynamics', [self.states, self.inputs, param], [fd]) - - - - + self.fun_forward_dynamics = cs.Function("fun_forward_dynamics", [self.states, self.inputs, param], [fd]) def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.ndarray) -> cs.SX: """ @@ -174,7 +172,7 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda Returns: A CasADi SX object of shape (29,) representing the predicted state of the robot. """ - + # Saving for clarity a bunch of variables foot_velocity_fl = inputs[0:3] foot_velocity_fr = inputs[3:6] @@ -207,103 +205,92 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda inertia = inertia.reshape((3, 3)) mass = param[28] - - - # FINAL linear_com_vel STATE (1) linear_com_vel = states[3:6] - # FINAL linear_com_acc STATE (2) - temp = foot_force_fl@stanceFL - temp += foot_force_fr@stanceFR - temp += foot_force_rl@stanceRL - temp += foot_force_rr@stanceRR + temp = foot_force_fl @ stanceFL + temp += foot_force_fr @ stanceFR + temp += foot_force_rl @ stanceRL + temp += foot_force_rr @ stanceRR temp += external_wrench_linear gravity = np.array([0, 0, -9.81]) - linear_com_acc = (1/mass)@temp + gravity - - + linear_com_acc = (1 / mass) @ temp + gravity + # Start to write the component of euler_rates_base and angular_acc_base STATES w = states[9:12] roll = states[6] pitch = states[7] yaw = states[8] - + conj_euler_rates = cs.SX.eye(3) conj_euler_rates[1, 1] = cs.cos(roll) - conj_euler_rates[2, 2] = cs.cos(pitch)*cs.cos(roll) + conj_euler_rates[2, 2] = cs.cos(pitch) * cs.cos(roll) conj_euler_rates[2, 1] = -cs.sin(roll) conj_euler_rates[0, 2] = -cs.sin(pitch) - conj_euler_rates[1, 2] = cs.cos(pitch)*cs.sin(roll) - + conj_euler_rates[1, 2] = cs.cos(pitch) * cs.sin(roll) - temp2 = cs.skew(foot_position_fl - com_position)@foot_force_fl@stanceFL - temp2 += cs.skew(foot_position_fr - com_position)@foot_force_fr@stanceFR - temp2 += cs.skew(foot_position_rl - com_position)@foot_force_rl@stanceRL - temp2 += cs.skew(foot_position_rr - com_position)@foot_force_rr@stanceRR + temp2 = cs.skew(foot_position_fl - com_position) @ foot_force_fl @ stanceFL + temp2 += cs.skew(foot_position_fr - com_position) @ foot_force_fr @ stanceFR + temp2 += cs.skew(foot_position_rl - com_position) @ foot_force_rl @ stanceRL + temp2 += cs.skew(foot_position_rr - com_position) @ foot_force_rr @ stanceRR - # FINAL euler_rates_base STATE (3) - euler_rates_base = cs.inv(conj_euler_rates)@w - + euler_rates_base = cs.inv(conj_euler_rates) @ w # FINAL angular_acc_base STATE (4) - #Z Y X rotations! + # Z Y X rotations! Rx = cs.SX.eye(3) - Rx[0,0] = 1 - Rx[0,1] = 0 - Rx[0,2] = 0 - Rx[1,0] = 0 - Rx[1,1] = cs.cos(roll) - Rx[1,2] = cs.sin(roll) - Rx[2,0] = 0 - Rx[2,1] = -cs.sin(roll) - Rx[2,2] = cs.cos(roll) - + Rx[0, 0] = 1 + Rx[0, 1] = 0 + Rx[0, 2] = 0 + Rx[1, 0] = 0 + Rx[1, 1] = cs.cos(roll) + Rx[1, 2] = cs.sin(roll) + Rx[2, 0] = 0 + Rx[2, 1] = -cs.sin(roll) + Rx[2, 2] = cs.cos(roll) Ry = cs.SX.eye(3) - Ry[0,0] = cs.cos(pitch) - Ry[0,1] = 0 - Ry[0,2] = -cs.sin(pitch) - Ry[1,0] = 0 - Ry[1,1] = 1 - Ry[1,2] = 0 - Ry[2,0] = cs.sin(pitch) - Ry[2,1] = 0 - Ry[2,2] = cs.cos(pitch) + Ry[0, 0] = cs.cos(pitch) + Ry[0, 1] = 0 + Ry[0, 2] = -cs.sin(pitch) + Ry[1, 0] = 0 + Ry[1, 1] = 1 + Ry[1, 2] = 0 + Ry[2, 0] = cs.sin(pitch) + Ry[2, 1] = 0 + Ry[2, 2] = cs.cos(pitch) Rz = cs.SX.eye(3) - Rz[0,0] = cs.cos(yaw) - Rz[0,1] = cs.sin(yaw) - Rz[0,2] = 0 - Rz[1,0] = -cs.sin(yaw) - Rz[1,1] = cs.cos(yaw) - Rz[1,2] = 0 - Rz[2,0] = 0 - Rz[2,1] = 0 - Rz[2,2] = 1 - - b_R_w = Rx@Ry@Rz + Rz[0, 0] = cs.cos(yaw) + Rz[0, 1] = cs.sin(yaw) + Rz[0, 2] = 0 + Rz[1, 0] = -cs.sin(yaw) + Rz[1, 1] = cs.cos(yaw) + Rz[1, 2] = 0 + Rz[2, 0] = 0 + Rz[2, 1] = 0 + Rz[2, 2] = 1 + + b_R_w = Rx @ Ry @ Rz temp2 = temp2 + external_wrench_angular - angular_acc_base = cs.inv(inertia)@(b_R_w@temp2 - cs.skew(w)@inertia@w) - - #angular_acc_base = -cs.inv(self.inertia)@cs.skew(w)@self.inertia@w + cs.inv(self.inertia)@b_R_w@temp2 - #angular_acc_base = -cs.inv(self.inertia)@cs.skew(w)@self.inertia@w + cs.inv(self.inertia)@b_R_w@temp2 + external_wrench_angular - + angular_acc_base = cs.inv(inertia) @ (b_R_w @ temp2 - cs.skew(w) @ inertia @ w) + # angular_acc_base = -cs.inv(self.inertia)@cs.skew(w)@self.inertia@w + cs.inv(self.inertia)@b_R_w@temp2 + # angular_acc_base = -cs.inv(self.inertia)@cs.skew(w)@self.inertia@w + cs.inv(self.inertia)@b_R_w@temp2 + external_wrench_angular # FINAL linear_foot_vel STATES (5,6,7,8) - if(not config.mpc_params["use_foothold_optimization"]): - foot_velocity_fl = foot_velocity_fl@0.0 - foot_velocity_fr = foot_velocity_fr@0.0 - foot_velocity_rl = foot_velocity_rl@0.0 - foot_velocity_rr = foot_velocity_rr@0.0 - linear_foot_vel_FL = foot_velocity_fl@(1-stanceFL)@(1-stance_proximity_FL) - linear_foot_vel_FR = foot_velocity_fr@(1-stanceFR)@(1-stance_proximity_FR) - linear_foot_vel_RL = foot_velocity_rl@(1-stanceRL)@(1-stance_proximity_RL) - linear_foot_vel_RR = foot_velocity_rr@(1-stanceRR)@(1-stance_proximity_RR) + if not config.mpc_params["use_foothold_optimization"]: + foot_velocity_fl = foot_velocity_fl @ 0.0 + foot_velocity_fr = foot_velocity_fr @ 0.0 + foot_velocity_rl = foot_velocity_rl @ 0.0 + foot_velocity_rr = foot_velocity_rr @ 0.0 + linear_foot_vel_FL = foot_velocity_fl @ (1 - stanceFL) @ (1 - stance_proximity_FL) + linear_foot_vel_FR = foot_velocity_fr @ (1 - stanceFR) @ (1 - stance_proximity_FR) + linear_foot_vel_RL = foot_velocity_rl @ (1 - stanceRL) @ (1 - stance_proximity_RL) + linear_foot_vel_RR = foot_velocity_rr @ (1 - stanceRR) @ (1 - stance_proximity_RR) # Integral states integral_states = states[24:] @@ -314,23 +301,38 @@ def forward_dynamics(self, states: np.ndarray, inputs: np.ndarray, param: np.nda integral_states[4] += roll integral_states[5] += pitch - # The order of the return should be equal to the order of the states_dot - return cs.vertcat(linear_com_vel, linear_com_acc, euler_rates_base, angular_acc_base, - linear_foot_vel_FL,linear_foot_vel_FR, linear_foot_vel_RL, linear_foot_vel_RR, integral_states) - - - - - def export_robot_model(self,) -> AcadosModel: + return cs.vertcat( + linear_com_vel, + linear_com_acc, + euler_rates_base, + angular_acc_base, + linear_foot_vel_FL, + linear_foot_vel_FR, + linear_foot_vel_RL, + linear_foot_vel_RR, + integral_states, + ) + + def export_robot_model( + self, + ) -> AcadosModel: """ This method set some general properties of the NMPC, such as the params, prediction mode, etc...! It will be called in centroidal_nmpc.py """ - + # dynamics - self.param = cs.vertcat(self.stance_param, self.mu_friction, self.stance_proximity, self.base_position, - self.base_yaw, self.external_wrench, self.inertia, self.mass) + self.param = cs.vertcat( + self.stance_param, + self.mu_friction, + self.stance_proximity, + self.base_position, + self.base_yaw, + self.external_wrench, + self.inertia, + self.mass, + ) f_expl = self.forward_dynamics(self.states, self.inputs, self.param) f_impl = self.states_dot - f_expl @@ -343,6 +345,4 @@ def export_robot_model(self,) -> AcadosModel: acados_model.p = self.param acados_model.name = "centroidal_model" - return acados_model - diff --git a/quadruped_pympc/controllers/gradient/nominal/centroidal_nmpc_gait_adaptive.py b/quadruped_pympc/controllers/gradient/nominal/centroidal_nmpc_gait_adaptive.py index 5c4d03d..30e865f 100644 --- a/quadruped_pympc/controllers/gradient/nominal/centroidal_nmpc_gait_adaptive.py +++ b/quadruped_pympc/controllers/gradient/nominal/centroidal_nmpc_gait_adaptive.py @@ -1,38 +1,39 @@ # Description: This file contains the class for the NMPC controller -# Authors: Giulio Turrisi - +# Authors: Giulio Turrisi - -import numpy as np -import scipy.linalg +import copy import os import time -import copy import casadi as cs +import numpy as np +import scipy.linalg from acados_template import AcadosOcp, AcadosOcpBatchSolver + ACADOS_INFTY = 1000 from quadruped_pympc import config + from .centroidal_model_nominal import Centroidal_Model_Nominal # Class for the Acados NMPC, the model is in another file! -class Acados_NMPC_GaitAdaptive(): +class Acados_NMPC_GaitAdaptive: def __init__(self): - - self.horizon = config.mpc_params['horizon'] # Define the number of discretization steps - self.dt = config.mpc_params['dt'] + self.horizon = config.mpc_params["horizon"] # Define the number of discretization steps + self.dt = config.mpc_params["dt"] self.T_horizon = self.horizon * self.dt - self.use_RTI = config.mpc_params['use_RTI'] - self.use_integrators = config.mpc_params['use_integrators'] - self.use_warm_start = config.mpc_params['use_warm_start'] - self.use_foothold_constraints = config.mpc_params['use_foothold_constraints'] + self.use_RTI = config.mpc_params["use_RTI"] + self.use_integrators = config.mpc_params["use_integrators"] + self.use_warm_start = config.mpc_params["use_warm_start"] + self.use_foothold_constraints = config.mpc_params["use_foothold_constraints"] - self.use_static_stability = config.mpc_params['use_static_stability'] - self.use_zmp_stability = config.mpc_params['use_zmp_stability'] + self.use_static_stability = config.mpc_params["use_static_stability"] + self.use_zmp_stability = config.mpc_params["use_zmp_stability"] self.use_stability_constraints = self.use_static_stability or self.use_zmp_stability - self.use_DDP = config.mpc_params['use_DDP'] + self.use_DDP = config.mpc_params["use_DDP"] self.previous_status = -1 self.previous_contact_sequence = np.zeros((4, self.horizon)) @@ -51,20 +52,24 @@ def __init__(self): self.inputs_dim = acados_model.u.size()[0] # Create the acados ocp solver - self.ocp = self.create_ocp_solver_description(acados_model, num_threads_in_batch_solve=len( - config.mpc_params['step_freq_available'])) + self.ocp = self.create_ocp_solver_description( + acados_model, num_threads_in_batch_solve=len(config.mpc_params["step_freq_available"]) + ) # Batch solver - use_batch_solver = config.mpc_params['optimize_step_freq'] - num_batch = len(config.mpc_params['step_freq_available']) + use_batch_solver = config.mpc_params["optimize_step_freq"] + num_batch = len(config.mpc_params["step_freq_available"]) self.batch = num_batch # batch_ocp = self.create_ocp_solver_description(acados_model, num_threads_in_batch_solve) batch_ocp = self.ocp dir_path = os.path.dirname(os.path.realpath(__file__)) - batch_ocp.code_export_directory = dir_path + '/c_generated_code_gait_adaptive' - self.batch_solver = AcadosOcpBatchSolver(batch_ocp, self.batch, verbose=False, - json_file=self.ocp.code_export_directory + "/centroidal_nmpc_batch" - + ".json") + batch_ocp.code_export_directory = dir_path + "/c_generated_code_gait_adaptive" + self.batch_solver = AcadosOcpBatchSolver( + batch_ocp, + self.batch, + verbose=False, + json_file=self.ocp.code_export_directory + "/centroidal_nmpc_batch" + ".json", + ) # Initialize solvers for stage in range(self.horizon + 1): @@ -102,7 +107,7 @@ def create_ocp_solver_description(self, acados_model, num_threads_in_batch_solve ocp.cost.Vx[:nx, :nx] = np.eye(nx) Vu = np.zeros((ny, nu)) - Vu[nx: nx + nu, 0:nu] = np.eye(nu) + Vu[nx : nx + nu, 0:nu] = np.eye(nu) ocp.cost.Vu = Vu ocp.cost.Vx_e = np.eye(nx) @@ -111,9 +116,7 @@ def create_ocp_solver_description(self, acados_model, num_threads_in_batch_solve ocp.cost.yref_e = np.zeros((ny_e,)) # Set friction and foothold constraints - expr_h_friction, \ - self.constr_uh_friction, \ - self.constr_lh_friction = self.create_friction_cone_constraints() + expr_h_friction, self.constr_uh_friction, self.constr_lh_friction = self.create_friction_cone_constraints() ocp.model.con_h_expr = expr_h_friction ocp.constraints.uh = self.constr_uh_friction @@ -124,10 +127,8 @@ def create_ocp_solver_description(self, acados_model, num_threads_in_batch_solve nsh = expr_h_friction.shape[0] nsh_state_constraint_start = copy.copy(nsh) - if (self.use_foothold_constraints): - expr_h_foot, \ - self.constr_uh_foot, \ - self.constr_lh_foot = self.create_foothold_constraints() + if self.use_foothold_constraints: + expr_h_foot, self.constr_uh_foot, self.constr_lh_foot = self.create_foothold_constraints() ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_foot) ocp.constraints.uh = np.concatenate((ocp.constraints.uh, self.constr_uh_foot)) @@ -135,11 +136,11 @@ def create_ocp_solver_description(self, acados_model, num_threads_in_batch_solve nsh += expr_h_foot.shape[0] # Set stability constraints - if (self.use_stability_constraints): + if self.use_stability_constraints: self.nsh_stability_start = copy.copy(nsh) - expr_h_support_polygon, \ - self.constr_uh_support_polygon, \ - self.constr_lh_support_polygon = self.create_stability_constraints() + expr_h_support_polygon, self.constr_uh_support_polygon, self.constr_lh_support_polygon = ( + self.create_stability_constraints() + ) ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_support_polygon) ocp.constraints.uh = np.concatenate((ocp.constraints.uh, self.constr_uh_support_polygon)) @@ -151,16 +152,19 @@ def create_ocp_solver_description(self, acados_model, num_threads_in_batch_solve # Set slack variable configuration: num_state_cstr = nsh_state_constraint_end - nsh_state_constraint_start - if (num_state_cstr > 0): + if num_state_cstr > 0: ocp.constraints.lsh = np.zeros( - num_state_cstr) # Lower bounds on slacks corresponding to soft lower bounds for nonlinear constraints + num_state_cstr + ) # Lower bounds on slacks corresponding to soft lower bounds for nonlinear constraints ocp.constraints.ush = np.zeros( - num_state_cstr) # Lower bounds on slacks corresponding to soft upper bounds for nonlinear constraints + num_state_cstr + ) # Lower bounds on slacks corresponding to soft upper bounds for nonlinear constraints ocp.constraints.idxsh = np.array(range(nsh_state_constraint_start, nsh_state_constraint_end)) # Jsh ns = num_state_cstr ocp.cost.zl = 1000 * np.ones((ns,)) # gradient wrt lower slack at intermediate shooting nodes (1 to N-1) ocp.cost.Zl = 1 * np.ones( - (ns,)) # diagonal of Hessian wrt lower slack at intermediate shooting nodes (1 to N-1) + (ns,) + ) # diagonal of Hessian wrt lower slack at intermediate shooting nodes (1 to N-1) ocp.cost.zu = 1000 * np.ones((ns,)) ocp.cost.Zu = 1 * np.ones((ns,)) @@ -178,7 +182,7 @@ def create_ocp_solver_description(self, acados_model, num_threads_in_batch_solve ocp.constraints.x0 = X0 # Set initialize parameters - init_contact_status = np.array([1., 1., 1., 1.]) + init_contact_status = np.array([1.0, 1.0, 1.0, 1.0]) init_mu = np.array([0.5]) init_stance_proximity = np.array([0, 0, 0, 0]) init_base_position = np.array([0, 0, 0]) @@ -187,59 +191,68 @@ def create_ocp_solver_description(self, acados_model, num_threads_in_batch_solve init_inertia = config.inertia.reshape((9,)) init_mass = np.array([config.mass]) - ocp.parameter_values = np.concatenate((init_contact_status, init_mu, init_stance_proximity, - init_base_position, init_base_yaw, init_external_wrench, - init_inertia, init_mass)) + ocp.parameter_values = np.concatenate( + ( + init_contact_status, + init_mu, + init_stance_proximity, + init_base_position, + init_base_yaw, + init_external_wrench, + init_inertia, + init_mass, + ) + ) # Set options ocp.solver_options.qp_solver = "PARTIAL_CONDENSING_HPIPM" # FULL_CONDENSING_QPOASES PARTIAL_CONDENSING_OSQP # PARTIAL_CONDENSING_HPIPM ocp.solver_options.hessian_approx = "GAUSS_NEWTON" # 'GAUSS_NEWTON', 'EXACT' ocp.solver_options.integrator_type = "ERK" # ERK IRK GNSF DISCRETE - if (self.use_DDP): - ocp.solver_options.nlp_solver_type = 'DDP' - ocp.solver_options.nlp_solver_max_iter = config.mpc_params['num_qp_iterations'] + if self.use_DDP: + ocp.solver_options.nlp_solver_type = "DDP" + ocp.solver_options.nlp_solver_max_iter = config.mpc_params["num_qp_iterations"] # ocp.solver_options.globalization = 'MERIT_BACKTRACKING' ocp.solver_options.with_adaptive_levenberg_marquardt = True - ocp.cost.cost_type = 'NONLINEAR_LS' - ocp.cost.cost_type_e = 'NONLINEAR_LS' + ocp.cost.cost_type = "NONLINEAR_LS" + ocp.cost.cost_type_e = "NONLINEAR_LS" ocp.model.cost_y_expr = cs.vertcat(ocp.model.x, ocp.model.u) ocp.model.cost_y_expr_e = ocp.model.x ocp.translate_to_feasibility_problem(keep_x0=True, keep_cost=True) - elif (self.use_RTI): + elif self.use_RTI: ocp.solver_options.nlp_solver_type = "SQP_RTI" ocp.solver_options.nlp_solver_max_iter = 1 - # Set the RTI type for the advanced RTI method + # Set the RTI type for the advanced RTI method # (see https://arxiv.org/pdf/2403.07101.pdf) - if (config.mpc_params['as_rti_type'] == "AS-RTI-A"): + if config.mpc_params["as_rti_type"] == "AS-RTI-A": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 0 - elif (config.mpc_params['as_rti_type'] == "AS-RTI-B"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-B": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 1 - elif (config.mpc_params['as_rti_type'] == "AS-RTI-C"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-C": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 2 - elif (config.mpc_params['as_rti_type'] == "AS-RTI-D"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-D": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 3 else: ocp.solver_options.nlp_solver_type = "SQP" - ocp.solver_options.nlp_solver_max_iter = config.mpc_params['num_qp_iterations'] + ocp.solver_options.nlp_solver_max_iter = config.mpc_params["num_qp_iterations"] # ocp.solver_options.globalization = "MERIT_BACKTRACKING" # FIXED_STEP, MERIT_BACKTRACKING - if (config.mpc_params['solver_mode'] == "balance"): + if config.mpc_params["solver_mode"] == "balance": ocp.solver_options.hpipm_mode = "BALANCE" - elif (config.mpc_params['solver_mode'] == "robust"): + elif config.mpc_params["solver_mode"] == "robust": ocp.solver_options.hpipm_mode = "ROBUST" - elif (config.mpc_params['solver_mode'] == "fast"): + elif config.mpc_params["solver_mode"] == "fast": ocp.solver_options.qp_solver_iter_max = 10 ocp.solver_options.hpipm_mode = "SPEED" - elif (config.mpc_params['solver_mode'] == "crazy_speed"): + elif config.mpc_params["solver_mode"] == "crazy_speed": ocp.solver_options.qp_solver_iter_max = 5 ocp.solver_options.hpipm_mode = "SPEED_ABS" @@ -253,11 +266,13 @@ def create_ocp_solver_description(self, acados_model, num_threads_in_batch_solve ocp.solver_options.tf = self.T_horizon # Nonuniform discretization - if (config.mpc_params['use_nonuniform_discretization']): - time_steps_fine_grained = np.tile(config.mpc_params['dt_fine_grained'], - config.mpc_params['horizon_fine_grained']) + if config.mpc_params["use_nonuniform_discretization"]: + time_steps_fine_grained = np.tile( + config.mpc_params["dt_fine_grained"], config.mpc_params["horizon_fine_grained"] + ) time_steps = np.concatenate( - (time_steps_fine_grained, np.tile(self.dt, self.horizon - config.mpc_params['horizon_fine_grained']))) + (time_steps_fine_grained, np.tile(self.dt, self.horizon - config.mpc_params["horizon_fine_grained"])) + ) shooting_nodes = np.zeros((self.horizon + 1,)) for i in range(len(time_steps)): shooting_nodes[i + 1] = shooting_nodes[i] + time_steps[i] @@ -266,8 +281,9 @@ def create_ocp_solver_description(self, acados_model, num_threads_in_batch_solve return ocp # Create a constraint for stability (COM, ZMP or CP inside support polygon) - def create_stability_constraints(self, ) -> None: - + def create_stability_constraints( + self, + ) -> None: base_w = self.centroidal_model.states[0:3] base_vel_w = self.centroidal_model.states[3:6] @@ -289,7 +305,7 @@ def create_stability_constraints(self, ) -> None: RL[0:2] = h_R_w @ (RL[0:2] - base_w[0:2]) RR[0:2] = h_R_w @ (RR[0:2] - base_w[0:2]) - if (self.use_static_stability): + if self.use_static_stability: x = 0.0 y = 0.0 else: @@ -367,15 +383,16 @@ def create_stability_constraints(self, ) -> None: ub[5] = 0 lb[5] = -ACADOS_INFTY - Jb = cs.vertcat(constraint_FL_FR, constraint_FR_RR, - constraint_RR_RL, constraint_RL_FL, - constraint_FL_RR, constraint_FR_RL) - + Jb = cs.vertcat( + constraint_FL_FR, constraint_FR_RR, constraint_RR_RL, constraint_RL_FL, constraint_FL_RR, constraint_FR_RL + ) return Jb, ub, lb # Create a standard foothold box constraint - def create_foothold_constraints(self, ): + def create_foothold_constraints( + self, + ): """ This function calculates the symbolic foothold constraints for the centroidal NMPC problem. @@ -417,12 +434,13 @@ def create_foothold_constraints(self, ): foot_position_rr[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.states[21:23] - base[0:2]) foot_position_rr[2] = self.centroidal_model.states[23] - Jbu = cs.vertcat(foot_position_fl, foot_position_fr, - foot_position_rl, foot_position_rr) + Jbu = cs.vertcat(foot_position_fl, foot_position_fr, foot_position_rl, foot_position_rr) return Jbu, ubu, lbu # Create the friction cone constraint - def create_friction_cone_constraints(self, ) -> None: + def create_friction_cone_constraints( + self, + ) -> None: """ This function calculates the symbolic friction cone constraints for the centroidal NMPC problem. @@ -435,8 +453,8 @@ def create_friction_cone_constraints(self, ) -> None: t = np.array([1, 0, 0]) b = np.array([0, 1, 0]) mu = self.centroidal_model.mu_friction - f_max = config.mpc_params['grf_max'] - f_min = config.mpc_params['grf_min'] + f_max = config.mpc_params["grf_max"] + f_min = config.mpc_params["grf_min"] # Derivation can be found in the paper # "High-slope terrain locomotion for torque-controlled quadruped robots", @@ -509,21 +527,39 @@ def set_weight(self, nx, nu): Q_pitch_integral_integral = np.array([10]) # integral of pitch R_foot_vel = np.array([0.0001, 0.0001, 0.00001]) # v_x, v_y, v_z (should be 4 times this, once per foot) - if (config.robot == "hyqreal"): + if config.robot == "hyqreal": R_foot_force = np.array( - [0.00001, 0.00001, 0.00001]) # f_x, f_y, f_z (should be 4 times this, once per foot) + [0.00001, 0.00001, 0.00001] + ) # f_x, f_y, f_z (should be 4 times this, once per foot) else: R_foot_force = np.array([0.001, 0.001, 0.001]) - Q_mat = np.diag(np.concatenate((Q_position, Q_velocity, - Q_base_angle, Q_base_angle_rates, - Q_foot_pos, Q_foot_pos, Q_foot_pos, Q_foot_pos, - Q_com_position_z_integral, Q_com_velocity_x_integral, - Q_com_velocity_y_integral, Q_com_velocity_z_integral, - Q_roll_integral_integral, Q_pitch_integral_integral))) - - R_mat = np.diag(np.concatenate((R_foot_vel, R_foot_vel, R_foot_vel, R_foot_vel, - R_foot_force, R_foot_force, R_foot_force, R_foot_force))) + Q_mat = np.diag( + np.concatenate( + ( + Q_position, + Q_velocity, + Q_base_angle, + Q_base_angle_rates, + Q_foot_pos, + Q_foot_pos, + Q_foot_pos, + Q_foot_pos, + Q_com_position_z_integral, + Q_com_velocity_x_integral, + Q_com_velocity_y_integral, + Q_com_velocity_z_integral, + Q_roll_integral_integral, + Q_pitch_integral_integral, + ) + ) + ) + + R_mat = np.diag( + np.concatenate( + (R_foot_vel, R_foot_vel, R_foot_vel, R_foot_vel, R_foot_force, R_foot_force, R_foot_force, R_foot_force) + ) + ) return Q_mat, R_mat @@ -533,7 +569,7 @@ def reset(self): def set_stage_constraint(self, constraint, state, reference, contact_sequence, h_R_w, stance_proximity): """ Set the stage constraint for the centroidal NMPC problem. We only consider the stance constraint, and the swing - constraint up to 2 maximum references. + constraint up to 2 maximum references. Args: constraint (numpy.ndarray or None): Constraint passed from outside (e.g. vision). If None, @@ -548,7 +584,6 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h None """ try: - # Take the array of the contact sequence and split # it in 4 arrays for clarity FL_contact_sequence = contact_sequence[0] @@ -574,7 +609,7 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h h_R_w = h_R_w.reshape((2, 2)) # Divide the constraint in upper and lower bound. The constraint are - # represented by 4 vertex, but we only use box constraint hence + # represented by 4 vertex, but we only use box constraint hence # we need only 2 vertex for each constraint (first an last) # For the leg in stance now, we simply enlarge the actual position as a box constraint @@ -626,8 +661,8 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h # Constraint for the first footholds at the next touchdown. If constraint == True # we have some constraint passed from outside (e.g. vision), otherwise we use the - # nominal foothold enlarged as we do previously - if (constraint is not None): + # nominal foothold enlarged as we do previously + if constraint is not None: # From the VFA first_up_constraint_FL = np.array([constraint[0][0], constraint[1][0], constraint[2][0] + 0.002]) first_up_constraint_FR = np.array([constraint[0][1], constraint[1][1], constraint[2][1] + 0.002]) @@ -676,38 +711,46 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h # FL first touchdown constraint first_up_constraint_FL = np.array( - [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] + 0.002]) + [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] + 0.002] + ) first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - base[0:2]) + 0.15 first_low_constraint_FL = np.array( - [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] - 0.002]) + [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] - 0.002] + ) first_low_constraint_FL[0:2] = h_R_w @ (first_low_constraint_FL[0:2] - base[0:2]) - 0.15 # FR first touchdown constraint first_up_constraint_FR = np.array( - [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] + 0.002]) + [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] + 0.002] + ) first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - base[0:2]) + 0.15 first_low_constraint_FR = np.array( - [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] - 0.002]) + [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] - 0.002] + ) first_low_constraint_FR[0:2] = h_R_w @ (first_low_constraint_FR[0:2] - base[0:2]) - 0.15 # RL first touchdown constraint first_up_constraint_RL = np.array( - [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] + 0.002]) + [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] + 0.002] + ) first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - base[0:2]) + 0.15 first_low_constraint_RL = np.array( - [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] - 0.002]) + [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] - 0.002] + ) first_low_constraint_RL[0:2] = h_R_w @ (first_low_constraint_RL[0:2] - base[0:2]) - 0.15 # RR first touchdown constraint first_up_constraint_RR = np.array( - [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] + 0.002]) + [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] + 0.002] + ) first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - base[0:2]) + 0.15 first_low_constraint_RR = np.array( - [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] - 0.002]) + [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] - 0.002] + ) first_low_constraint_RR[0:2] = h_R_w @ (first_low_constraint_RR[0:2] - base[0:2]) - 0.15 # we stack all the constraint we have for now @@ -725,52 +768,60 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h # We do not expect more than two reference footholds... # FL second touchdown constraint - if (FL_reference_foot.shape[0] == 2): + if FL_reference_foot.shape[0] == 2: second_up_constraint_FL = np.array( - [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] + 0.002]) + [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] + 0.002] + ) second_up_constraint_FL[0:2] = h_R_w @ (second_up_constraint_FL[0:2] - base[0:2]) + 0.15 second_low_constraint_FL = np.array( - [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] - 0.002]) + [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] - 0.002] + ) second_low_constraint_FL[0:2] = h_R_w @ (second_low_constraint_FL[0:2] - base[0:2]) - 0.15 up_constraint_FL = np.vstack((up_constraint_FL, second_up_constraint_FL)) low_constraint_FL = np.vstack((low_constraint_FL, second_low_constraint_FL)) # FR second touchdown constraint - if (FR_reference_foot.shape[0] == 2): + if FR_reference_foot.shape[0] == 2: second_up_constraint_FR = np.array( - [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] + 0.002]) + [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] + 0.002] + ) second_up_constraint_FR[0:2] = h_R_w @ (second_up_constraint_FR[0:2] - base[0:2]) + 0.15 second_low_constraint_FR = np.array( - [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] - 0.002]) + [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] - 0.002] + ) second_low_constraint_FR[0:2] = h_R_w @ (second_low_constraint_FR[0:2] - base[0:2]) - 0.15 up_constraint_FR = np.vstack((up_constraint_FR, second_up_constraint_FR)) low_constraint_FR = np.vstack((low_constraint_FR, second_low_constraint_FR)) # RL second touchdown constraint - if (RL_reference_foot.shape[0] == 2): + if RL_reference_foot.shape[0] == 2: second_up_constraint_RL = np.array( - [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] + 0.002]) + [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] + 0.002] + ) second_up_constraint_RL[0:2] = h_R_w @ (second_up_constraint_RL[0:2] - base[0:2]) + 0.15 second_low_constraint_RL = np.array( - [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] - 0.002]) + [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] - 0.002] + ) second_low_constraint_RL[0:2] = h_R_w @ (second_low_constraint_RL[0:2] - base[0:2]) - 0.15 up_constraint_RL = np.vstack((up_constraint_RL, second_up_constraint_RL)) low_constraint_RL = np.vstack((low_constraint_RL, second_low_constraint_RL)) # RR second touchdown constraint - if (RR_reference_foot.shape[0] == 2): + if RR_reference_foot.shape[0] == 2: second_up_constraint_RR = np.array( - [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] + 0.002]) + [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] + 0.002] + ) second_up_constraint_RR[0:2] = h_R_w @ (second_up_constraint_RR[0:2] - base[0:2]) + 0.15 second_low_constraint_RR = np.array( - [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] - 0.002]) + [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] - 0.002] + ) second_low_constraint_RR[0:2] = h_R_w @ (second_low_constraint_RR[0:2] - base[0:2]) - 0.15 up_constraint_RR = np.vstack((up_constraint_RR, second_up_constraint_RR)) @@ -783,17 +834,17 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h # If the foothold is in swing, the idx of the constraint start from 1 idx_constraint = np.array([0, 0, 0, 0]) - if (FL_contact_sequence[0] == 0): + if FL_contact_sequence[0] == 0: idx_constraint[0] = 1 - if (FR_contact_sequence[0] == 0): + if FR_contact_sequence[0] == 0: idx_constraint[1] = 1 - if (RL_contact_sequence[0] == 0): + if RL_contact_sequence[0] == 0: idx_constraint[2] = 1 - if (RR_contact_sequence[0] == 0): + if RR_contact_sequence[0] == 0: idx_constraint[3] = 1 for j in range(0, self.horizon): - # take the constraint for the current timestep + # take the constraint for the current timestep ub_foot_FL = up_constraint_FL[idx_constraint[0]] lb_foot_FL = low_constraint_FL[idx_constraint[0]] @@ -807,11 +858,9 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h lb_foot_RR = low_constraint_RR[idx_constraint[3]] # Concatenate the friction and foothold constraint - ub_foot = copy.deepcopy(np.concatenate((ub_foot_FL, ub_foot_FR, - ub_foot_RL, ub_foot_RR))) - lb_foot = copy.deepcopy(np.concatenate((lb_foot_FL, lb_foot_FR, - lb_foot_RL, lb_foot_RR))) - if (self.use_foothold_constraints): + ub_foot = copy.deepcopy(np.concatenate((ub_foot_FL, ub_foot_FR, ub_foot_RL, ub_foot_RR))) + lb_foot = copy.deepcopy(np.concatenate((lb_foot_FL, lb_foot_FR, lb_foot_RL, lb_foot_RR))) + if self.use_foothold_constraints: ub_total = np.concatenate((ub_friction, ub_foot)) lb_total = np.concatenate((lb_friction, lb_foot)) else: @@ -820,7 +869,7 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h # Constraints for the support polygon depending on the leg in stance # all disabled at the beginning!! - if (self.use_stability_constraints): + if self.use_stability_constraints: ub_support_FL_FR = ACADOS_INFTY lb_support_FL_FR = -ACADOS_INFTY @@ -840,10 +889,12 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h lb_support_FR_RL = -ACADOS_INFTY # We have 4 cases for the stability constraint: trot, pace, crawl, full stance - if (FL_contact_sequence[j] == 1 and - FR_contact_sequence[j] == 1 and - RL_contact_sequence[j] == 1 and - RR_contact_sequence[j] == 1): + if ( + FL_contact_sequence[j] == 1 + and FR_contact_sequence[j] == 1 + and RL_contact_sequence[j] == 1 + and RR_contact_sequence[j] == 1 + ): # FULL STANCE TODO ub_support_FL_FR = ACADOS_INFTY lb_support_FL_FR = -ACADOS_INFTY @@ -863,95 +914,113 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h ub_support_FR_RL = ACADOS_INFTY lb_support_FR_RL = -ACADOS_INFTY - elif (np.array_equal(FL_contact_sequence, RR_contact_sequence) - and np.array_equal(FR_contact_sequence, RL_contact_sequence)): + elif np.array_equal(FL_contact_sequence, RR_contact_sequence) and np.array_equal( + FR_contact_sequence, RL_contact_sequence + ): # TROT - stability_margin = config.mpc_params['trot_stability_margin'] - if (FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0): + stability_margin = config.mpc_params["trot_stability_margin"] + if FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0: ub_support_FL_RR = 0 + stability_margin lb_support_FL_RR = 0 - stability_margin - if (FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0): + if FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0: ub_support_FR_RL = 0 + stability_margin lb_support_FR_RL = 0 - stability_margin - elif (np.array_equal(FL_contact_sequence, RL_contact_sequence) - and np.array_equal(FR_contact_sequence, RR_contact_sequence)): + elif np.array_equal(FL_contact_sequence, RL_contact_sequence) and np.array_equal( + FR_contact_sequence, RR_contact_sequence + ): # PACE - stability_margin = config.mpc_params['pace_stability_margin'] - if (FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0): + stability_margin = config.mpc_params["pace_stability_margin"] + if FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0: ub_support_RL_FL = 0 + stability_margin lb_support_RL_FL = 0 - stability_margin - if (FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0): + if FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0: ub_support_FR_RR = 0 + stability_margin lb_support_FR_RR = 0 - stability_margin else: # CRAWL BACKDIAGONALCRAWL ONLY - stability_margin = config.mpc_params['crawl_stability_margin'] + stability_margin = config.mpc_params["crawl_stability_margin"] - if (FL_contact_sequence[j] == 1): - if (FR_contact_sequence[j] == 1): + if FL_contact_sequence[j] == 1: + if FR_contact_sequence[j] == 1: ub_support_FL_FR = -0.0 - stability_margin lb_support_FL_FR = -ACADOS_INFTY else: ub_support_FL_RR = ACADOS_INFTY lb_support_FL_RR = 0.0 + stability_margin - if (FR_contact_sequence[j] == 1): - if (RR_contact_sequence[j] == 1): + if FR_contact_sequence[j] == 1: + if RR_contact_sequence[j] == 1: ub_support_FR_RR = ACADOS_INFTY lb_support_FR_RR = 0.0 + stability_margin else: ub_support_FR_RL = ACADOS_INFTY lb_support_FR_RL = 0.0 + stability_margin - if (RR_contact_sequence[j] == 1): - if (RL_contact_sequence[j] == 1): + if RR_contact_sequence[j] == 1: + if RL_contact_sequence[j] == 1: ub_support_RR_RL = ACADOS_INFTY lb_support_RR_RL = 0.0 + stability_margin else: ub_support_FL_RR = -0.0 - stability_margin lb_support_FL_RR = -ACADOS_INFTY - if (RL_contact_sequence[j] == 1): - if (FL_contact_sequence[j] == 1): + if RL_contact_sequence[j] == 1: + if FL_contact_sequence[j] == 1: ub_support_RL_FL = -0.0 - stability_margin lb_support_RL_FL = -ACADOS_INFTY else: ub_support_FR_RL = -0.0 - stability_margin lb_support_FR_RL = -ACADOS_INFTY - ub_support = np.array([ub_support_FL_FR, ub_support_FR_RR, ub_support_RR_RL, ub_support_RL_FL, - ub_support_FL_RR, ub_support_FR_RL]) - lb_support = np.array([lb_support_FL_FR, lb_support_FR_RR, lb_support_RR_RL, lb_support_RL_FL, - lb_support_FL_RR, lb_support_FR_RL]) + ub_support = np.array( + [ + ub_support_FL_FR, + ub_support_FR_RR, + ub_support_RR_RL, + ub_support_RL_FL, + ub_support_FL_RR, + ub_support_FR_RL, + ] + ) + lb_support = np.array( + [ + lb_support_FL_FR, + lb_support_FR_RR, + lb_support_RR_RL, + lb_support_RL_FL, + lb_support_FL_RR, + lb_support_FR_RL, + ] + ) ub_total = np.concatenate((ub_total, ub_support)) lb_total = np.concatenate((lb_total, lb_support)) # No friction constraint at the end! we don't have u_N - if (j == self.horizon): - if (self.use_foothold_constraints): - if (self.use_stability_constraints): + if j == self.horizon: + if self.use_foothold_constraints: + if self.use_stability_constraints: ub_total = np.concatenate((ub_foot, ub_support)) lb_total = np.concatenate((lb_foot, lb_support)) else: ub_total = ub_foot lb_total = lb_foot else: - if (self.use_stability_constraints): + if self.use_stability_constraints: ub_total = ub_support lb_total = lb_support else: continue # Only friction costraints at the beginning - if (j == 0): + if j == 0: self.acados_ocp_solver.constraints_set(j, "uh", ub_friction) self.acados_ocp_solver.constraints_set(j, "lh", lb_friction) - if (j > 0): + if j > 0: self.acados_ocp_solver.constraints_set(j, "uh", ub_total) self.acados_ocp_solver.constraints_set(j, "lh", lb_total) @@ -960,21 +1029,21 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h self.lower_bound[j] = lb_total.tolist() # ugly procedure to update the idx of the constraint - if (j >= 1): - if (FL_contact_sequence[j] == 0 and FL_contact_sequence[j - 1] == 1): - if (idx_constraint[0] < up_constraint_FL.shape[0] - 1): + if j >= 1: + if FL_contact_sequence[j] == 0 and FL_contact_sequence[j - 1] == 1: + if idx_constraint[0] < up_constraint_FL.shape[0] - 1: idx_constraint[0] += 1 - if (FR_contact_sequence[j] == 0 and FR_contact_sequence[j - 1] == 1): - if (idx_constraint[1] < up_constraint_FR.shape[0] - 1): + if FR_contact_sequence[j] == 0 and FR_contact_sequence[j - 1] == 1: + if idx_constraint[1] < up_constraint_FR.shape[0] - 1: idx_constraint[1] += 1 - if (RL_contact_sequence[j] == 0 and RL_contact_sequence[j - 1] == 1): - if (idx_constraint[2] < up_constraint_RL.shape[0] - 1): + if RL_contact_sequence[j] == 0 and RL_contact_sequence[j - 1] == 1: + if idx_constraint[2] < up_constraint_RL.shape[0] - 1: idx_constraint[2] += 1 - if (RR_contact_sequence[j] == 0 and RR_contact_sequence[j - 1] == 1): - if (idx_constraint[3] < up_constraint_RR.shape[0] - 1): + if RR_contact_sequence[j] == 0 and RR_contact_sequence[j - 1] == 1: + if idx_constraint[3] < up_constraint_RR.shape[0] - 1: idx_constraint[3] += 1 except: @@ -984,7 +1053,6 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h # Method to perform the centering of the states and the reference around (0, 0, 0) def perform_scaling(self, state, reference, constraint=None): - self.initial_base_position = copy.deepcopy(state["position"]) reference = copy.deepcopy(reference) state = copy.deepcopy(state) @@ -1004,68 +1072,81 @@ def perform_scaling(self, state, reference, constraint=None): return state, reference, constraint # Main loop for computing batched control - def compute_batch_control(self, state, reference, contact_sequence, constraint=None, - external_wrenches=np.zeros((6,)), - inertia=config.inertia.reshape((9,)), mass=config.mass): + def compute_batch_control( + self, + state, + reference, + contact_sequence, + constraint=None, + external_wrenches=np.zeros((6,)), + inertia=config.inertia.reshape((9,)), + mass=config.mass, + ): start = time.time() costs = [] # Perform the scaling of the states and the reference - state, \ - reference, \ - constraint = self.perform_scaling(state, reference, constraint) + state, reference, constraint = self.perform_scaling(state, reference, constraint) # Fill reference (self.states_dim+self.inputs_dim) idx_ref_foot_to_assign = np.array([0, 0, 0, 0]) yref = np.zeros(shape=(self.states_dim + self.inputs_dim,)) - yref[0:3] = reference['ref_position'] - yref[3:6] = reference['ref_linear_velocity'] - yref[6:9] = reference['ref_orientation'] - yref[9:12] = reference['ref_angular_velocity'] - yref[12:15] = reference['ref_foot_FL'][idx_ref_foot_to_assign[0]] - yref[15:18] = reference['ref_foot_FR'][idx_ref_foot_to_assign[1]] - yref[18:21] = reference['ref_foot_RL'][idx_ref_foot_to_assign[2]] - yref[21:24] = reference['ref_foot_RR'][idx_ref_foot_to_assign[3]] + yref[0:3] = reference["ref_position"] + yref[3:6] = reference["ref_linear_velocity"] + yref[6:9] = reference["ref_orientation"] + yref[9:12] = reference["ref_angular_velocity"] + yref[12:15] = reference["ref_foot_FL"][idx_ref_foot_to_assign[0]] + yref[15:18] = reference["ref_foot_FR"][idx_ref_foot_to_assign[1]] + yref[18:21] = reference["ref_foot_RL"][idx_ref_foot_to_assign[2]] + yref[21:24] = reference["ref_foot_RR"][idx_ref_foot_to_assign[3]] yref_N = np.zeros(shape=(self.states_dim,)) - yref_N[0:3] = reference['ref_position'] - yref_N[3:6] = reference['ref_linear_velocity'] - yref_N[6:9] = reference['ref_orientation'] - yref_N[9:12] = reference['ref_angular_velocity'] + yref_N[0:3] = reference["ref_position"] + yref_N[3:6] = reference["ref_linear_velocity"] + yref_N[6:9] = reference["ref_orientation"] + yref_N[9:12] = reference["ref_angular_velocity"] # Set initial state constraint. We teleported the robot foothold # to the previous optimal foothold. This is done to avoid the optimization # of a foothold that is not considered at all at touchdown! In any case, - # the height cames always from the VFA + # the height cames always from the VFA FL_contact_sequence = contact_sequence[0][0] FR_contact_sequence = contact_sequence[0][1] RL_contact_sequence = contact_sequence[0][2] RR_contact_sequence = contact_sequence[0][3] - if (FL_contact_sequence[0] == 0): + if FL_contact_sequence[0] == 0: state["foot_FL"] = reference["ref_foot_FL"][0] - if (FR_contact_sequence[0] == 0): + if FR_contact_sequence[0] == 0: state["foot_FR"] = reference["ref_foot_FR"][0] - if (RL_contact_sequence[0] == 0): + if RL_contact_sequence[0] == 0: state["foot_RL"] = reference["ref_foot_RL"][0] - if (RR_contact_sequence[0] == 0): + if RR_contact_sequence[0] == 0: state["foot_RR"] = reference["ref_foot_RR"][0] # Fill stance param, friction and stance proximity # (stance proximity will disable foothold optimization near a stance!!) - mu = config.mpc_params['mu'] - - state_acados = np.concatenate((state["position"], state["linear_velocity"], - state["orientation"], state["angular_velocity"], - state["foot_FL"], state["foot_FR"], - state["foot_RL"], state["foot_RR"], - self.integral_errors)) # .reshape((self.states_dim, 1)) + mu = config.mpc_params["mu"] + + state_acados = np.concatenate( + ( + state["position"], + state["linear_velocity"], + state["orientation"], + state["angular_velocity"], + state["foot_FL"], + state["foot_FR"], + state["foot_RL"], + state["foot_RR"], + self.integral_errors, + ) + ) # .reshape((self.states_dim, 1)) for n in range(self.batch): - # Take the array of the contact sequence and split it in 4 arrays, + # Take the array of the contact sequence and split it in 4 arrays, # one for each leg FL_contact_sequence = contact_sequence[n][0] FR_contact_sequence = contact_sequence[n][1] @@ -1073,72 +1154,91 @@ def compute_batch_control(self, state, reference, contact_sequence, constraint=N RR_contact_sequence = contact_sequence[n][3] for j in range(self.horizon): - # Update the idx_ref_foot_to_assign. Every time there is a change in the contact phase # between 1 and 0, it means that the leg go into swing and a new reference is needed!!! - if (j > 1 and j < self.horizon - 1): - if (FL_contact_sequence[j + 1] == 0 and FL_contact_sequence[j] == 1): - if (reference['ref_foot_FL'].shape[0] > idx_ref_foot_to_assign[0] + 1): + if j > 1 and j < self.horizon - 1: + if FL_contact_sequence[j + 1] == 0 and FL_contact_sequence[j] == 1: + if reference["ref_foot_FL"].shape[0] > idx_ref_foot_to_assign[0] + 1: idx_ref_foot_to_assign[0] += 1 - if (FR_contact_sequence[j + 1] == 0 and FR_contact_sequence[j] == 1): - if (reference['ref_foot_FR'].shape[0] > idx_ref_foot_to_assign[1] + 1): + if FR_contact_sequence[j + 1] == 0 and FR_contact_sequence[j] == 1: + if reference["ref_foot_FR"].shape[0] > idx_ref_foot_to_assign[1] + 1: idx_ref_foot_to_assign[1] += 1 - if (RL_contact_sequence[j + 1] == 0 and RL_contact_sequence[j] == 1): - if (reference['ref_foot_RL'].shape[0] > idx_ref_foot_to_assign[2] + 1): + if RL_contact_sequence[j + 1] == 0 and RL_contact_sequence[j] == 1: + if reference["ref_foot_RL"].shape[0] > idx_ref_foot_to_assign[2] + 1: idx_ref_foot_to_assign[2] += 1 - if (RR_contact_sequence[j + 1] == 0 and RR_contact_sequence[j] == 1): - if (reference['ref_foot_RR'].shape[0] > idx_ref_foot_to_assign[3] + 1): + if RR_contact_sequence[j + 1] == 0 and RR_contact_sequence[j] == 1: + if reference["ref_foot_RR"].shape[0] > idx_ref_foot_to_assign[3] + 1: idx_ref_foot_to_assign[3] += 1 # Setting the reference to acados self.batch_solver.ocp_solvers[n].set(j, "yref", yref) # Fill last step horizon reference (self.states_dim - no control action!!) - yref_N[12:15] = reference['ref_foot_FL'][idx_ref_foot_to_assign[0]] - yref_N[15:18] = reference['ref_foot_FR'][idx_ref_foot_to_assign[1]] - yref_N[18:21] = reference['ref_foot_RL'][idx_ref_foot_to_assign[2]] - yref_N[21:24] = reference['ref_foot_RR'][idx_ref_foot_to_assign[3]] + yref_N[12:15] = reference["ref_foot_FL"][idx_ref_foot_to_assign[0]] + yref_N[15:18] = reference["ref_foot_FR"][idx_ref_foot_to_assign[1]] + yref_N[18:21] = reference["ref_foot_RL"][idx_ref_foot_to_assign[2]] + yref_N[21:24] = reference["ref_foot_RR"][idx_ref_foot_to_assign[3]] # Setting the reference to acados self.batch_solver.ocp_solvers[n].set(self.horizon, "yref", yref_N) # Set the parameters to acados for j in range(self.horizon): - param = np.array([FL_contact_sequence[j], FR_contact_sequence[j], - RL_contact_sequence[j], RR_contact_sequence[j], mu, - 0.0, - 0.0, - 0.0, - 0.0, - state["position"][0], state["position"][1], - state["position"][2], state["orientation"][2], - 0.0, 0.0, - 0.0, 0.0, - 0.0, 0.0, - inertia[0], inertia[1], inertia[2], inertia[3], inertia[4], inertia[5], - inertia[6], inertia[7], inertia[8], mass]) + param = np.array( + [ + FL_contact_sequence[j], + FR_contact_sequence[j], + RL_contact_sequence[j], + RR_contact_sequence[j], + mu, + 0.0, + 0.0, + 0.0, + 0.0, + state["position"][0], + state["position"][1], + state["position"][2], + state["orientation"][2], + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + inertia[0], + inertia[1], + inertia[2], + inertia[3], + inertia[4], + inertia[5], + inertia[6], + inertia[7], + inertia[8], + mass, + ] + ) self.batch_solver.ocp_solvers[n].set(j, "p", param) # Set initial state constraint acados, converting first the dictionary to np array self.batch_solver.ocp_solvers[n].set(0, "lbx", state_acados) self.batch_solver.ocp_solvers[n].set(0, "ubx", state_acados) - t_elapsed2 = (time.time() - start) + t_elapsed2 = time.time() - start # Solve the batched ocp t0 = time.time() self.batch_solver.solve() - t_elapsed = (time.time() - t0) + t_elapsed = time.time() - t0 print("time_python: ", t_elapsed2) print("time_solver: ", t_elapsed) for n in range(self.batch): cost_single_qp = self.batch_solver.ocp_solvers[n].get_cost() - if (n != 0): - cost_single_qp += 3 * ( - config.mpc_params["step_freq_available"][n] - config.mpc_params["step_freq_available"][ - 0]) ** 2 + if n != 0: + cost_single_qp += ( + 3 * (config.mpc_params["step_freq_available"][n] - config.mpc_params["step_freq_available"][0]) ** 2 + ) costs.append(cost_single_qp) best_freq_index = np.argmin(costs) diff --git a/quadruped_pympc/controllers/gradient/nominal/centroidal_nmpc_nominal.py b/quadruped_pympc/controllers/gradient/nominal/centroidal_nmpc_nominal.py index fbcca99..b995162 100644 --- a/quadruped_pympc/controllers/gradient/nominal/centroidal_nmpc_nominal.py +++ b/quadruped_pympc/controllers/gradient/nominal/centroidal_nmpc_nominal.py @@ -1,38 +1,39 @@ # Description: This file contains the class for the NMPC controller import pathlib -# Authors: Giulio Turrisi - - +# Authors: Giulio Turrisi - from acados_template import AcadosOcp, AcadosOcpSolver + ACADOS_INFTY = 1000 -from .centroidal_model_nominal import Centroidal_Model_Nominal +import copy + +import casadi as cs import numpy as np import scipy.linalg -import casadi as cs -import copy import quadruped_pympc.config as config +from .centroidal_model_nominal import Centroidal_Model_Nominal + # Class for the Acados NMPC, the model is in another file! class Acados_NMPC_Nominal: def __init__(self): - - self.horizon = config.mpc_params['horizon'] # Define the number of discretization steps - self.dt = config.mpc_params['dt'] + self.horizon = config.mpc_params["horizon"] # Define the number of discretization steps + self.dt = config.mpc_params["dt"] self.T_horizon = self.horizon * self.dt - self.use_RTI = config.mpc_params['use_RTI'] - self.use_integrators = config.mpc_params['use_integrators'] - self.use_warm_start = config.mpc_params['use_warm_start'] - self.use_foothold_constraints = config.mpc_params['use_foothold_constraints'] + self.use_RTI = config.mpc_params["use_RTI"] + self.use_integrators = config.mpc_params["use_integrators"] + self.use_warm_start = config.mpc_params["use_warm_start"] + self.use_foothold_constraints = config.mpc_params["use_foothold_constraints"] - self.use_static_stability = config.mpc_params['use_static_stability'] - self.use_zmp_stability = config.mpc_params['use_zmp_stability'] + self.use_static_stability = config.mpc_params["use_static_stability"] + self.use_zmp_stability = config.mpc_params["use_zmp_stability"] self.use_stability_constraints = self.use_static_stability or self.use_zmp_stability - self.use_DDP = config.mpc_params['use_DDP'] + self.use_DDP = config.mpc_params["use_DDP"] - self.verbose = config.mpc_params['verbose'] + self.verbose = config.mpc_params["verbose"] self.previous_status = -1 self.previous_contact_sequence = np.zeros((4, self.horizon)) @@ -53,12 +54,12 @@ def __init__(self): # Create the acados ocp solver self.ocp = self.create_ocp_solver_description(acados_model) - code_export_dir = pathlib.Path(__file__).parent / 'c_generated_code' + code_export_dir = pathlib.Path(__file__).parent / "c_generated_code" self.ocp.code_export_directory = str(code_export_dir) - self.acados_ocp_solver = AcadosOcpSolver(self.ocp, - json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + - ".json") + self.acados_ocp_solver = AcadosOcpSolver( + self.ocp, json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json" + ) # Initialize solver for stage in range(self.horizon + 1): @@ -66,9 +67,9 @@ def __init__(self): for stage in range(self.horizon): self.acados_ocp_solver.set(stage, "u", np.zeros((self.inputs_dim,))) - if (self.use_RTI): + if self.use_RTI: # first preparation phase - self.acados_ocp_solver.options_set('rti_phase', 1) + self.acados_ocp_solver.options_set("rti_phase", 1) status = self.acados_ocp_solver.solve() # Set cost, constraints and options @@ -99,7 +100,7 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: ocp.cost.Vx[:nx, :nx] = np.eye(nx) Vu = np.zeros((ny, nu)) - Vu[nx: nx + nu, 0:nu] = np.eye(nu) + Vu[nx : nx + nu, 0:nu] = np.eye(nu) ocp.cost.Vu = Vu ocp.cost.Vx_e = np.eye(nx) @@ -108,9 +109,7 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: ocp.cost.yref_e = np.zeros((ny_e,)) # Set friction and foothold constraints - expr_h_friction, \ - self.constr_uh_friction, \ - self.constr_lh_friction = self.create_friction_cone_constraints() + expr_h_friction, self.constr_uh_friction, self.constr_lh_friction = self.create_friction_cone_constraints() ocp.model.con_h_expr = expr_h_friction ocp.constraints.uh = self.constr_uh_friction @@ -121,10 +120,8 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: nsh = expr_h_friction.shape[0] nsh_state_constraint_start = copy.copy(nsh) - if (self.use_foothold_constraints): - expr_h_foot, \ - self.constr_uh_foot, \ - self.constr_lh_foot = self.create_foothold_constraints() + if self.use_foothold_constraints: + expr_h_foot, self.constr_uh_foot, self.constr_lh_foot = self.create_foothold_constraints() ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_foot) ocp.constraints.uh = np.concatenate((ocp.constraints.uh, self.constr_uh_foot)) @@ -132,11 +129,11 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: nsh += expr_h_foot.shape[0] # Set stability constraints - if (self.use_stability_constraints): + if self.use_stability_constraints: self.nsh_stability_start = copy.copy(nsh) - expr_h_support_polygon, \ - self.constr_uh_support_polygon, \ - self.constr_lh_support_polygon = self.create_stability_constraints() + expr_h_support_polygon, self.constr_uh_support_polygon, self.constr_lh_support_polygon = ( + self.create_stability_constraints() + ) ocp.model.con_h_expr = cs.vertcat(ocp.model.con_h_expr, expr_h_support_polygon) ocp.constraints.uh = np.concatenate((ocp.constraints.uh, self.constr_uh_support_polygon)) @@ -148,16 +145,19 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: # Set slack variable configuration: num_state_cstr = nsh_state_constraint_end - nsh_state_constraint_start - if (num_state_cstr > 0): + if num_state_cstr > 0: ocp.constraints.lsh = np.zeros( - num_state_cstr) # Lower bounds on slacks corresponding to soft lower bounds for nonlinear constraints + num_state_cstr + ) # Lower bounds on slacks corresponding to soft lower bounds for nonlinear constraints ocp.constraints.ush = np.zeros( - num_state_cstr) # Lower bounds on slacks corresponding to soft upper bounds for nonlinear constraints + num_state_cstr + ) # Lower bounds on slacks corresponding to soft upper bounds for nonlinear constraints ocp.constraints.idxsh = np.array(range(nsh_state_constraint_start, nsh_state_constraint_end)) # Jsh ns = num_state_cstr ocp.cost.zl = 1000 * np.ones((ns,)) # gradient wrt lower slack at intermediate shooting nodes (1 to N-1) ocp.cost.Zl = 1 * np.ones( - (ns,)) # diagonal of Hessian wrt lower slack at intermediate shooting nodes (1 to N-1) + (ns,) + ) # diagonal of Hessian wrt lower slack at intermediate shooting nodes (1 to N-1) ocp.cost.zu = 1000 * np.ones((ns,)) ocp.cost.Zu = 1 * np.ones((ns,)) @@ -175,7 +175,7 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: ocp.constraints.x0 = X0 # Set initialize parameters - init_contact_status = np.array([1., 1., 1., 1.]) + init_contact_status = np.array([1.0, 1.0, 1.0, 1.0]) init_mu = np.array([0.5]) init_stance_proximity = np.array([0, 0, 0, 0]) init_base_position = np.array([0, 0, 0]) @@ -184,59 +184,68 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: init_inertia = config.inertia.reshape((9,)) init_mass = np.array([config.mass]) - ocp.parameter_values = np.concatenate((init_contact_status, init_mu, init_stance_proximity, - init_base_position, init_base_yaw, init_external_wrench, - init_inertia, init_mass)) + ocp.parameter_values = np.concatenate( + ( + init_contact_status, + init_mu, + init_stance_proximity, + init_base_position, + init_base_yaw, + init_external_wrench, + init_inertia, + init_mass, + ) + ) # Set options ocp.solver_options.qp_solver = "PARTIAL_CONDENSING_HPIPM" # FULL_CONDENSING_QPOASES PARTIAL_CONDENSING_OSQP # PARTIAL_CONDENSING_HPIPM ocp.solver_options.hessian_approx = "GAUSS_NEWTON" # 'GAUSS_NEWTON', 'EXACT' ocp.solver_options.integrator_type = "ERK" # ERK IRK GNSF DISCRETE - if (self.use_DDP): - ocp.solver_options.nlp_solver_type = 'DDP' - ocp.solver_options.nlp_solver_max_iter = config.mpc_params['num_qp_iterations'] + if self.use_DDP: + ocp.solver_options.nlp_solver_type = "DDP" + ocp.solver_options.nlp_solver_max_iter = config.mpc_params["num_qp_iterations"] # ocp.solver_options.globalization = 'MERIT_BACKTRACKING' ocp.solver_options.with_adaptive_levenberg_marquardt = True - ocp.cost.cost_type = 'NONLINEAR_LS' - ocp.cost.cost_type_e = 'NONLINEAR_LS' + ocp.cost.cost_type = "NONLINEAR_LS" + ocp.cost.cost_type_e = "NONLINEAR_LS" ocp.model.cost_y_expr = cs.vertcat(ocp.model.x, ocp.model.u) ocp.model.cost_y_expr_e = ocp.model.x ocp.translate_to_feasibility_problem(keep_x0=True, keep_cost=True) - elif (self.use_RTI): + elif self.use_RTI: ocp.solver_options.nlp_solver_type = "SQP_RTI" ocp.solver_options.nlp_solver_max_iter = 1 - # Set the RTI type for the advanced RTI method + # Set the RTI type for the advanced RTI method # (see https://arxiv.org/pdf/2403.07101.pdf) - if (config.mpc_params['as_rti_type'] == "AS-RTI-A"): + if config.mpc_params["as_rti_type"] == "AS-RTI-A": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 0 - elif (config.mpc_params['as_rti_type'] == "AS-RTI-B"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-B": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 1 - elif (config.mpc_params['as_rti_type'] == "AS-RTI-C"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-C": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 2 - elif (config.mpc_params['as_rti_type'] == "AS-RTI-D"): + elif config.mpc_params["as_rti_type"] == "AS-RTI-D": ocp.solver_options.as_rti_iter = 1 ocp.solver_options.as_rti_level = 3 else: ocp.solver_options.nlp_solver_type = "SQP" - ocp.solver_options.nlp_solver_max_iter = config.mpc_params['num_qp_iterations'] + ocp.solver_options.nlp_solver_max_iter = config.mpc_params["num_qp_iterations"] # ocp.solver_options.globalization = "MERIT_BACKTRACKING" # FIXED_STEP, MERIT_BACKTRACKING - if (config.mpc_params['solver_mode'] == "balance"): + if config.mpc_params["solver_mode"] == "balance": ocp.solver_options.hpipm_mode = "BALANCE" - elif (config.mpc_params['solver_mode'] == "robust"): + elif config.mpc_params["solver_mode"] == "robust": ocp.solver_options.hpipm_mode = "ROBUST" - elif (config.mpc_params['solver_mode'] == "fast"): + elif config.mpc_params["solver_mode"] == "fast": ocp.solver_options.qp_solver_iter_max = 10 ocp.solver_options.hpipm_mode = "SPEED" - elif (config.mpc_params['solver_mode'] == "crazy_speed"): + elif config.mpc_params["solver_mode"] == "crazy_speed": ocp.solver_options.qp_solver_iter_max = 5 ocp.solver_options.hpipm_mode = "SPEED_ABS" @@ -249,11 +258,13 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: ocp.solver_options.tf = self.T_horizon # Nonuniform discretization - if (config.mpc_params['use_nonuniform_discretization']): - time_steps_fine_grained = np.tile(config.mpc_params['dt_fine_grained'], - config.mpc_params['horizon_fine_grained']) + if config.mpc_params["use_nonuniform_discretization"]: + time_steps_fine_grained = np.tile( + config.mpc_params["dt_fine_grained"], config.mpc_params["horizon_fine_grained"] + ) time_steps = np.concatenate( - (time_steps_fine_grained, np.tile(self.dt, self.horizon - config.mpc_params['horizon_fine_grained']))) + (time_steps_fine_grained, np.tile(self.dt, self.horizon - config.mpc_params["horizon_fine_grained"])) + ) shooting_nodes = np.zeros((self.horizon + 1,)) for i in range(len(time_steps)): shooting_nodes[i + 1] = shooting_nodes[i] + time_steps[i] @@ -262,8 +273,9 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: return ocp # Create a constraint for stability (COM, ZMP or CP inside support polygon) - def create_stability_constraints(self, ) -> None: - + def create_stability_constraints( + self, + ) -> None: base_w = self.centroidal_model.states[0:3] base_vel_w = self.centroidal_model.states[3:6] @@ -285,7 +297,7 @@ def create_stability_constraints(self, ) -> None: RL[0:2] = h_R_w @ (RL[0:2] - base_w[0:2]) RR[0:2] = h_R_w @ (RR[0:2] - base_w[0:2]) - if (self.use_static_stability): + if self.use_static_stability: x = 0.0 y = 0.0 else: @@ -363,14 +375,16 @@ def create_stability_constraints(self, ) -> None: ub[5] = 0 lb[5] = -ACADOS_INFTY - Jb = cs.vertcat(constraint_FL_FR, constraint_FR_RR, - constraint_RR_RL, constraint_RL_FL, - constraint_FL_RR, constraint_FR_RL) + Jb = cs.vertcat( + constraint_FL_FR, constraint_FR_RR, constraint_RR_RL, constraint_RL_FL, constraint_FL_RR, constraint_FR_RL + ) return Jb, ub, lb # Create a standard foothold box constraint - def create_foothold_constraints(self, ): + def create_foothold_constraints( + self, + ): """ This function calculates the symbolic foothold constraints for the centroidal NMPC problem. @@ -412,12 +426,13 @@ def create_foothold_constraints(self, ): foot_position_rr[0:2] = h_R_w @ cs.vertcat(self.centroidal_model.states[21:23] - base[0:2]) foot_position_rr[2] = self.centroidal_model.states[23] - Jbu = cs.vertcat(foot_position_fl, foot_position_fr, - foot_position_rl, foot_position_rr) + Jbu = cs.vertcat(foot_position_fl, foot_position_fr, foot_position_rl, foot_position_rr) return Jbu, ubu, lbu # Create the friction cone constraint - def create_friction_cone_constraints(self, ) -> None: + def create_friction_cone_constraints( + self, + ) -> None: """ This function calculates the symbolic friction cone constraints for the centroidal NMPC problem. @@ -430,8 +445,8 @@ def create_friction_cone_constraints(self, ) -> None: t = np.array([1, 0, 0]) b = np.array([0, 1, 0]) mu = self.centroidal_model.mu_friction - f_max = config.mpc_params['grf_max'] - f_min = config.mpc_params['grf_min'] + f_max = config.mpc_params["grf_max"] + f_min = config.mpc_params["grf_min"] # Derivation can be found in the paper # "High-slope terrain locomotion for torque-controlled quadruped robots", @@ -504,35 +519,55 @@ def set_weight(self, nx, nu): Q_pitch_integral_integral = np.array([10]) # integral of pitch R_foot_vel = np.array([0.0001, 0.0001, 0.00001]) # v_x, v_y, v_z (should be 4 times this, once per foot) - if (config.robot == "hyqreal"): + if config.robot == "hyqreal": R_foot_force = np.array( - [0.00001, 0.00001, 0.00001]) # f_x, f_y, f_z (should be 4 times this, once per foot) + [0.00001, 0.00001, 0.00001] + ) # f_x, f_y, f_z (should be 4 times this, once per foot) else: R_foot_force = np.array([0.001, 0.001, 0.001]) - Q_mat = np.diag(np.concatenate((Q_position, Q_velocity, - Q_base_angle, Q_base_angle_rates, - Q_foot_pos, Q_foot_pos, Q_foot_pos, Q_foot_pos, - Q_com_position_z_integral, Q_com_velocity_x_integral, - Q_com_velocity_y_integral, Q_com_velocity_z_integral, - Q_roll_integral_integral, Q_pitch_integral_integral))) - - R_mat = np.diag(np.concatenate((R_foot_vel, R_foot_vel, R_foot_vel, R_foot_vel, - R_foot_force, R_foot_force, R_foot_force, R_foot_force))) + Q_mat = np.diag( + np.concatenate( + ( + Q_position, + Q_velocity, + Q_base_angle, + Q_base_angle_rates, + Q_foot_pos, + Q_foot_pos, + Q_foot_pos, + Q_foot_pos, + Q_com_position_z_integral, + Q_com_velocity_x_integral, + Q_com_velocity_y_integral, + Q_com_velocity_z_integral, + Q_roll_integral_integral, + Q_pitch_integral_integral, + ) + ) + ) + + R_mat = np.diag( + np.concatenate( + (R_foot_vel, R_foot_vel, R_foot_vel, R_foot_vel, R_foot_force, R_foot_force, R_foot_force, R_foot_force) + ) + ) return Q_mat, R_mat def reset(self): self.acados_ocp_solver.reset() - self.acados_ocp_solver = AcadosOcpSolver(self.ocp, - json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + - ".json", - build=False, generate=False) + self.acados_ocp_solver = AcadosOcpSolver( + self.ocp, + json_file=self.ocp.code_export_directory + "/centroidal_nmpc" + ".json", + build=False, + generate=False, + ) def set_stage_constraint(self, constraint, state, reference, contact_sequence, h_R_w, stance_proximity): """ Set the stage constraint for the centroidal NMPC problem. We only consider the stance constraint, and the swing - constraint up to 2 maximum references. + constraint up to 2 maximum references. Args: constraint (numpy.ndarray or None): Constraint passed from outside (e.g. vision). If None, @@ -547,7 +582,6 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h None """ try: - # Take the array of the contact sequence and split # it in 4 arrays for clarity FL_contact_sequence = contact_sequence[0] @@ -573,7 +607,7 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h h_R_w = h_R_w.reshape((2, 2)) # Divide the constraint in upper and lower bound. The constraint are - # represented by 4 vertex, but we only use box constraint hence + # represented by 4 vertex, but we only use box constraint hence # we need only 2 vertex for each constraint (first an last) # For the leg in stance now, we simply enlarge the actual position as a box constraint @@ -625,8 +659,8 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h # Constraint for the first footholds at the next touchdown. If constraint == True # we have some constraint passed from outside (e.g. vision), otherwise we use the - # nominal foothold enlarged as we do previously - if (constraint is not None): + # nominal foothold enlarged as we do previously + if constraint is not None: # From the VFA first_up_constraint_FL = np.array([constraint[0][0], constraint[1][0], constraint[2][0] + 0.002]) first_up_constraint_FR = np.array([constraint[0][1], constraint[1][1], constraint[2][1] + 0.002]) @@ -675,38 +709,46 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h # FL first touchdown constraint first_up_constraint_FL = np.array( - [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] + 0.002]) + [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] + 0.002] + ) first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - base[0:2]) + 0.15 first_low_constraint_FL = np.array( - [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] - 0.002]) + [FL_reference_foot[0][0], FL_reference_foot[0][1], FL_reference_foot[0][2] - 0.002] + ) first_low_constraint_FL[0:2] = h_R_w @ (first_low_constraint_FL[0:2] - base[0:2]) - 0.15 # FR first touchdown constraint first_up_constraint_FR = np.array( - [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] + 0.002]) + [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] + 0.002] + ) first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - base[0:2]) + 0.15 first_low_constraint_FR = np.array( - [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] - 0.002]) + [FR_reference_foot[0][0], FR_reference_foot[0][1], FR_reference_foot[0][2] - 0.002] + ) first_low_constraint_FR[0:2] = h_R_w @ (first_low_constraint_FR[0:2] - base[0:2]) - 0.15 # RL first touchdown constraint first_up_constraint_RL = np.array( - [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] + 0.002]) + [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] + 0.002] + ) first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - base[0:2]) + 0.15 first_low_constraint_RL = np.array( - [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] - 0.002]) + [RL_reference_foot[0][0], RL_reference_foot[0][1], RL_reference_foot[0][2] - 0.002] + ) first_low_constraint_RL[0:2] = h_R_w @ (first_low_constraint_RL[0:2] - base[0:2]) - 0.15 # RR first touchdown constraint first_up_constraint_RR = np.array( - [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] + 0.002]) + [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] + 0.002] + ) first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - base[0:2]) + 0.15 first_low_constraint_RR = np.array( - [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] - 0.002]) + [RR_reference_foot[0][0], RR_reference_foot[0][1], RR_reference_foot[0][2] - 0.002] + ) first_low_constraint_RR[0:2] = h_R_w @ (first_low_constraint_RR[0:2] - base[0:2]) - 0.15 # we stack all the constraint we have for now @@ -724,52 +766,60 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h # We do not expect more than two reference footholds... # FL second touchdown constraint - if (FL_reference_foot.shape[0] == 2): + if FL_reference_foot.shape[0] == 2: second_up_constraint_FL = np.array( - [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] + 0.002]) + [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] + 0.002] + ) second_up_constraint_FL[0:2] = h_R_w @ (second_up_constraint_FL[0:2] - base[0:2]) + 0.15 second_low_constraint_FL = np.array( - [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] - 0.002]) + [FL_reference_foot[1][0], FL_reference_foot[1][1], FL_reference_foot[1][2] - 0.002] + ) second_low_constraint_FL[0:2] = h_R_w @ (second_low_constraint_FL[0:2] - base[0:2]) - 0.15 up_constraint_FL = np.vstack((up_constraint_FL, second_up_constraint_FL)) low_constraint_FL = np.vstack((low_constraint_FL, second_low_constraint_FL)) # FR second touchdown constraint - if (FR_reference_foot.shape[0] == 2): + if FR_reference_foot.shape[0] == 2: second_up_constraint_FR = np.array( - [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] + 0.002]) + [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] + 0.002] + ) second_up_constraint_FR[0:2] = h_R_w @ (second_up_constraint_FR[0:2] - base[0:2]) + 0.15 second_low_constraint_FR = np.array( - [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] - 0.002]) + [FR_reference_foot[1][0], FR_reference_foot[1][1], FR_reference_foot[1][2] - 0.002] + ) second_low_constraint_FR[0:2] = h_R_w @ (second_low_constraint_FR[0:2] - base[0:2]) - 0.15 up_constraint_FR = np.vstack((up_constraint_FR, second_up_constraint_FR)) low_constraint_FR = np.vstack((low_constraint_FR, second_low_constraint_FR)) # RL second touchdown constraint - if (RL_reference_foot.shape[0] == 2): + if RL_reference_foot.shape[0] == 2: second_up_constraint_RL = np.array( - [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] + 0.002]) + [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] + 0.002] + ) second_up_constraint_RL[0:2] = h_R_w @ (second_up_constraint_RL[0:2] - base[0:2]) + 0.15 second_low_constraint_RL = np.array( - [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] - 0.002]) + [RL_reference_foot[1][0], RL_reference_foot[1][1], RL_reference_foot[1][2] - 0.002] + ) second_low_constraint_RL[0:2] = h_R_w @ (second_low_constraint_RL[0:2] - base[0:2]) - 0.15 up_constraint_RL = np.vstack((up_constraint_RL, second_up_constraint_RL)) low_constraint_RL = np.vstack((low_constraint_RL, second_low_constraint_RL)) # RR second touchdown constraint - if (RR_reference_foot.shape[0] == 2): + if RR_reference_foot.shape[0] == 2: second_up_constraint_RR = np.array( - [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] + 0.002]) + [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] + 0.002] + ) second_up_constraint_RR[0:2] = h_R_w @ (second_up_constraint_RR[0:2] - base[0:2]) + 0.15 second_low_constraint_RR = np.array( - [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] - 0.002]) + [RR_reference_foot[1][0], RR_reference_foot[1][1], RR_reference_foot[1][2] - 0.002] + ) second_low_constraint_RR[0:2] = h_R_w @ (second_low_constraint_RR[0:2] - base[0:2]) - 0.15 up_constraint_RR = np.vstack((up_constraint_RR, second_up_constraint_RR)) @@ -782,17 +832,17 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h # If the foothold is in swing, the idx of the constraint start from 1 idx_constraint = np.array([0, 0, 0, 0]) - if (FL_contact_sequence[0] == 0): + if FL_contact_sequence[0] == 0: idx_constraint[0] = 1 - if (FR_contact_sequence[0] == 0): + if FR_contact_sequence[0] == 0: idx_constraint[1] = 1 - if (RL_contact_sequence[0] == 0): + if RL_contact_sequence[0] == 0: idx_constraint[2] = 1 - if (RR_contact_sequence[0] == 0): + if RR_contact_sequence[0] == 0: idx_constraint[3] = 1 for j in range(0, self.horizon): - # take the constraint for the current timestep + # take the constraint for the current timestep ub_foot_FL = up_constraint_FL[idx_constraint[0]] lb_foot_FL = low_constraint_FL[idx_constraint[0]] @@ -806,11 +856,9 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h lb_foot_RR = low_constraint_RR[idx_constraint[3]] # Concatenate the friction and foothold constraint - ub_foot = copy.deepcopy(np.concatenate((ub_foot_FL, ub_foot_FR, - ub_foot_RL, ub_foot_RR))) - lb_foot = copy.deepcopy(np.concatenate((lb_foot_FL, lb_foot_FR, - lb_foot_RL, lb_foot_RR))) - if (self.use_foothold_constraints): + ub_foot = copy.deepcopy(np.concatenate((ub_foot_FL, ub_foot_FR, ub_foot_RL, ub_foot_RR))) + lb_foot = copy.deepcopy(np.concatenate((lb_foot_FL, lb_foot_FR, lb_foot_RL, lb_foot_RR))) + if self.use_foothold_constraints: ub_total = np.concatenate((ub_friction, ub_foot)) lb_total = np.concatenate((lb_friction, lb_foot)) else: @@ -819,7 +867,7 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h # Constraints for the support polygon depending on the leg in stance # all disabled at the beginning!! - if (self.use_stability_constraints): + if self.use_stability_constraints: ub_support_FL_FR = ACADOS_INFTY lb_support_FL_FR = -ACADOS_INFTY @@ -839,10 +887,12 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h lb_support_FR_RL = -ACADOS_INFTY # We have 4 cases for the stability constraint: trot, pace, crawl, full stance - if (FL_contact_sequence[j] == 1 and - FR_contact_sequence[j] == 1 and - RL_contact_sequence[j] == 1 and - RR_contact_sequence[j] == 1): + if ( + FL_contact_sequence[j] == 1 + and FR_contact_sequence[j] == 1 + and RL_contact_sequence[j] == 1 + and RR_contact_sequence[j] == 1 + ): # FULL STANCE TODO ub_support_FL_FR = ACADOS_INFTY lb_support_FL_FR = -ACADOS_INFTY @@ -862,95 +912,113 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h ub_support_FR_RL = ACADOS_INFTY lb_support_FR_RL = -ACADOS_INFTY - elif (np.array_equal(FL_contact_sequence, RR_contact_sequence) - and np.array_equal(FR_contact_sequence, RL_contact_sequence)): + elif np.array_equal(FL_contact_sequence, RR_contact_sequence) and np.array_equal( + FR_contact_sequence, RL_contact_sequence + ): # TROT - stability_margin = config.mpc_params['trot_stability_margin'] - if (FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0): + stability_margin = config.mpc_params["trot_stability_margin"] + if FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0: ub_support_FL_RR = 0 + stability_margin lb_support_FL_RR = 0 - stability_margin - if (FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0): + if FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0: ub_support_FR_RL = 0 + stability_margin lb_support_FR_RL = 0 - stability_margin - elif (np.array_equal(FL_contact_sequence, RL_contact_sequence) - and np.array_equal(FR_contact_sequence, RR_contact_sequence)): + elif np.array_equal(FL_contact_sequence, RL_contact_sequence) and np.array_equal( + FR_contact_sequence, RR_contact_sequence + ): # PACE - stability_margin = config.mpc_params['pace_stability_margin'] - if (FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0): + stability_margin = config.mpc_params["pace_stability_margin"] + if FL_contact_sequence[j] == 1 and FR_contact_sequence[j] == 0: ub_support_RL_FL = 0 + stability_margin lb_support_RL_FL = 0 - stability_margin - if (FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0): + if FR_contact_sequence[j] == 1 and FL_contact_sequence[j] == 0: ub_support_FR_RR = 0 + stability_margin lb_support_FR_RR = 0 - stability_margin else: # CRAWL BACKDIAGONALCRAWL ONLY - stability_margin = config.mpc_params['crawl_stability_margin'] + stability_margin = config.mpc_params["crawl_stability_margin"] - if (FL_contact_sequence[j] == 1): - if (FR_contact_sequence[j] == 1): + if FL_contact_sequence[j] == 1: + if FR_contact_sequence[j] == 1: ub_support_FL_FR = -0.0 - stability_margin lb_support_FL_FR = -ACADOS_INFTY else: ub_support_FL_RR = ACADOS_INFTY lb_support_FL_RR = 0.0 + stability_margin - if (FR_contact_sequence[j] == 1): - if (RR_contact_sequence[j] == 1): + if FR_contact_sequence[j] == 1: + if RR_contact_sequence[j] == 1: ub_support_FR_RR = ACADOS_INFTY lb_support_FR_RR = 0.0 + stability_margin else: ub_support_FR_RL = ACADOS_INFTY lb_support_FR_RL = 0.0 + stability_margin - if (RR_contact_sequence[j] == 1): - if (RL_contact_sequence[j] == 1): + if RR_contact_sequence[j] == 1: + if RL_contact_sequence[j] == 1: ub_support_RR_RL = ACADOS_INFTY lb_support_RR_RL = 0.0 + stability_margin else: ub_support_FL_RR = -0.0 - stability_margin lb_support_FL_RR = -ACADOS_INFTY - if (RL_contact_sequence[j] == 1): - if (FL_contact_sequence[j] == 1): + if RL_contact_sequence[j] == 1: + if FL_contact_sequence[j] == 1: ub_support_RL_FL = -0.0 - stability_margin lb_support_RL_FL = -ACADOS_INFTY else: ub_support_FR_RL = -0.0 - stability_margin lb_support_FR_RL = -ACADOS_INFTY - ub_support = np.array([ub_support_FL_FR, ub_support_FR_RR, ub_support_RR_RL, ub_support_RL_FL, - ub_support_FL_RR, ub_support_FR_RL]) - lb_support = np.array([lb_support_FL_FR, lb_support_FR_RR, lb_support_RR_RL, lb_support_RL_FL, - lb_support_FL_RR, lb_support_FR_RL]) + ub_support = np.array( + [ + ub_support_FL_FR, + ub_support_FR_RR, + ub_support_RR_RL, + ub_support_RL_FL, + ub_support_FL_RR, + ub_support_FR_RL, + ] + ) + lb_support = np.array( + [ + lb_support_FL_FR, + lb_support_FR_RR, + lb_support_RR_RL, + lb_support_RL_FL, + lb_support_FL_RR, + lb_support_FR_RL, + ] + ) ub_total = np.concatenate((ub_total, ub_support)) lb_total = np.concatenate((lb_total, lb_support)) # No friction constraint at the end! we don't have u_N - if (j == self.horizon): - if (self.use_foothold_constraints): - if (self.use_stability_constraints): + if j == self.horizon: + if self.use_foothold_constraints: + if self.use_stability_constraints: ub_total = np.concatenate((ub_foot, ub_support)) lb_total = np.concatenate((lb_foot, lb_support)) else: ub_total = ub_foot lb_total = lb_foot else: - if (self.use_stability_constraints): + if self.use_stability_constraints: ub_total = ub_support lb_total = lb_support else: continue # Only friction costraints at the beginning - if (j == 0): + if j == 0: self.acados_ocp_solver.constraints_set(j, "uh", ub_friction) self.acados_ocp_solver.constraints_set(j, "lh", lb_friction) - if (j > 0): + if j > 0: self.acados_ocp_solver.constraints_set(j, "uh", ub_total) self.acados_ocp_solver.constraints_set(j, "lh", lb_total) @@ -959,33 +1027,40 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h self.lower_bound[j] = lb_total.tolist() # ugly procedure to update the idx of the constraint - if (j >= 1): - if (FL_contact_sequence[j] == 0 and FL_contact_sequence[j - 1] == 1): - if (idx_constraint[0] < up_constraint_FL.shape[0] - 1): + if j >= 1: + if FL_contact_sequence[j] == 0 and FL_contact_sequence[j - 1] == 1: + if idx_constraint[0] < up_constraint_FL.shape[0] - 1: idx_constraint[0] += 1 - if (FR_contact_sequence[j] == 0 and FR_contact_sequence[j - 1] == 1): - if (idx_constraint[1] < up_constraint_FR.shape[0] - 1): + if FR_contact_sequence[j] == 0 and FR_contact_sequence[j - 1] == 1: + if idx_constraint[1] < up_constraint_FR.shape[0] - 1: idx_constraint[1] += 1 - if (RL_contact_sequence[j] == 0 and RL_contact_sequence[j - 1] == 1): - if (idx_constraint[2] < up_constraint_RL.shape[0] - 1): + if RL_contact_sequence[j] == 0 and RL_contact_sequence[j - 1] == 1: + if idx_constraint[2] < up_constraint_RL.shape[0] - 1: idx_constraint[2] += 1 - if (RR_contact_sequence[j] == 0 and RR_contact_sequence[j - 1] == 1): - if (idx_constraint[3] < up_constraint_RR.shape[0] - 1): + if RR_contact_sequence[j] == 0 and RR_contact_sequence[j - 1] == 1: + if idx_constraint[3] < up_constraint_RR.shape[0] - 1: idx_constraint[3] += 1 except: - if(self.verbose): + if self.verbose: print("###WARNING: error in setting the constraints") return - def set_warm_start(self, state_acados, reference, FL_contact_sequence, FR_contact_sequence, RL_contact_sequence, - RR_contact_sequence): + def set_warm_start( + self, + state_acados, + reference, + FL_contact_sequence, + FR_contact_sequence, + RL_contact_sequence, + RR_contact_sequence, + ): """ - Sets the warm start for the optimization problem. This could be useful in the case + Sets the warm start for the optimization problem. This could be useful in the case some old guess is outside some new constraints.. Maybe we could have some infeasibility problem. Still to investigate.. @@ -1001,41 +1076,41 @@ def set_warm_start(self, state_acados, reference, FL_contact_sequence, FR_contac for j in range(self.horizon): # We only want to modify the warm start of the footholds... - # hence i take the current warm start from acados and + # hence i take the current warm start from acados and # modify onlya subset of it warm_start = copy.deepcopy(self.acados_ocp_solver.get(j, "x")) # Update the idx_ref_foot_to_assign. Every time there is a change in the contact phase # between 1 and 0, it means that the leg go into swing and a new reference is needed!!! - if (j > 1 and j < self.horizon - 1): - if (FL_contact_sequence[j] == 0 and FL_contact_sequence[j - 1] == 1): - if (reference['ref_foot_FL'].shape[0] > idx_ref_foot_to_assign[0] + 1): + if j > 1 and j < self.horizon - 1: + if FL_contact_sequence[j] == 0 and FL_contact_sequence[j - 1] == 1: + if reference["ref_foot_FL"].shape[0] > idx_ref_foot_to_assign[0] + 1: idx_ref_foot_to_assign[0] += 1 - if (FR_contact_sequence[j] == 0 and FR_contact_sequence[j - 1] == 1): - if (reference['ref_foot_FR'].shape[0] > idx_ref_foot_to_assign[1] + 1): + if FR_contact_sequence[j] == 0 and FR_contact_sequence[j - 1] == 1: + if reference["ref_foot_FR"].shape[0] > idx_ref_foot_to_assign[1] + 1: idx_ref_foot_to_assign[1] += 1 - if (RL_contact_sequence[j] == 0 and RL_contact_sequence[j - 1] == 1): - if (reference['ref_foot_RL'].shape[0] > idx_ref_foot_to_assign[2] + 1): + if RL_contact_sequence[j] == 0 and RL_contact_sequence[j - 1] == 1: + if reference["ref_foot_RL"].shape[0] > idx_ref_foot_to_assign[2] + 1: idx_ref_foot_to_assign[2] += 1 - if (RR_contact_sequence[j] == 0 and RR_contact_sequence[j - 1] == 1): - if (reference['ref_foot_RR'].shape[0] > idx_ref_foot_to_assign[3] + 1): + if RR_contact_sequence[j] == 0 and RR_contact_sequence[j - 1] == 1: + if reference["ref_foot_RR"].shape[0] > idx_ref_foot_to_assign[3] + 1: idx_ref_foot_to_assign[3] += 1 warm_start[8] = state_acados[8] - if (idx_ref_foot_to_assign[0] == 0): + if idx_ref_foot_to_assign[0] == 0: warm_start[12:15] = state_acados[12:15].reshape((3,)) else: warm_start[12:15] = reference["ref_foot_FL"][0] - if (idx_ref_foot_to_assign[1] == 0): + if idx_ref_foot_to_assign[1] == 0: warm_start[15:18] = state_acados[15:18].reshape((3,)) else: warm_start[15:18] = reference["ref_foot_FR"][0] - if (idx_ref_foot_to_assign[2] == 0): + if idx_ref_foot_to_assign[2] == 0: warm_start[18:21] = state_acados[18:21].reshape((3,)) else: warm_start[18:21] = reference["ref_foot_RL"][0] - if (idx_ref_foot_to_assign[3] == 0): + if idx_ref_foot_to_assign[3] == 0: warm_start[21:24] = state_acados[21:24].reshape((3,)) else: warm_start[21:24] = reference["ref_foot_RR"][0] @@ -1065,9 +1140,16 @@ def perform_scaling(self, state, reference, constraint=None): return state, reference, constraint # Main loop for computing the control - def compute_control(self, state, reference, contact_sequence, constraint=None, external_wrenches=np.zeros((6,)), - inertia=config.inertia.reshape((9,)), mass=config.mass): - + def compute_control( + self, + state, + reference, + contact_sequence, + constraint=None, + external_wrenches=np.zeros((6,)), + inertia=config.inertia.reshape((9,)), + mass=config.mass, + ): # Take the array of the contact sequence and split it in 4 arrays, # one for each leg FL_contact_sequence = contact_sequence[0] @@ -1081,46 +1163,44 @@ def compute_control(self, state, reference, contact_sequence, constraint=None, e RR_previous_contact_sequence = self.previous_contact_sequence[3] # Perform the scaling of the states and the reference - state, \ - reference, \ - constraint = self.perform_scaling(state, reference, constraint) + state, reference, constraint = self.perform_scaling(state, reference, constraint) # Fill reference (self.states_dim+self.inputs_dim) idx_ref_foot_to_assign = np.array([0, 0, 0, 0]) for j in range(self.horizon): - yref = np.zeros(shape=(self.states_dim + self.inputs_dim,)) - yref[0:3] = reference['ref_position'] - yref[3:6] = reference['ref_linear_velocity'] - yref[6:9] = reference['ref_orientation'] - yref[9:12] = reference['ref_angular_velocity'] - yref[12:15] = reference['ref_foot_FL'][idx_ref_foot_to_assign[0]] - yref[15:18] = reference['ref_foot_FR'][idx_ref_foot_to_assign[1]] - yref[18:21] = reference['ref_foot_RL'][idx_ref_foot_to_assign[2]] - yref[21:24] = reference['ref_foot_RR'][idx_ref_foot_to_assign[3]] + yref[0:3] = reference["ref_position"] + yref[3:6] = reference["ref_linear_velocity"] + yref[6:9] = reference["ref_orientation"] + yref[9:12] = reference["ref_angular_velocity"] + yref[12:15] = reference["ref_foot_FL"][idx_ref_foot_to_assign[0]] + yref[15:18] = reference["ref_foot_FR"][idx_ref_foot_to_assign[1]] + yref[18:21] = reference["ref_foot_RL"][idx_ref_foot_to_assign[2]] + yref[21:24] = reference["ref_foot_RR"][idx_ref_foot_to_assign[3]] # Update the idx_ref_foot_to_assign. Every time there is a change in the contact phase # between 1 and 0, it means that the leg go into swing and a new reference is needed!!! - if (j > 1 and j < self.horizon - 1): - if (FL_contact_sequence[j + 1] == 0 and FL_contact_sequence[j] == 1): - if (reference['ref_foot_FL'].shape[0] > idx_ref_foot_to_assign[0] + 1): + if j > 1 and j < self.horizon - 1: + if FL_contact_sequence[j + 1] == 0 and FL_contact_sequence[j] == 1: + if reference["ref_foot_FL"].shape[0] > idx_ref_foot_to_assign[0] + 1: idx_ref_foot_to_assign[0] += 1 - if (FR_contact_sequence[j + 1] == 0 and FR_contact_sequence[j] == 1): - if (reference['ref_foot_FR'].shape[0] > idx_ref_foot_to_assign[1] + 1): + if FR_contact_sequence[j + 1] == 0 and FR_contact_sequence[j] == 1: + if reference["ref_foot_FR"].shape[0] > idx_ref_foot_to_assign[1] + 1: idx_ref_foot_to_assign[1] += 1 - if (RL_contact_sequence[j + 1] == 0 and RL_contact_sequence[j] == 1): - if (reference['ref_foot_RL'].shape[0] > idx_ref_foot_to_assign[2] + 1): + if RL_contact_sequence[j + 1] == 0 and RL_contact_sequence[j] == 1: + if reference["ref_foot_RL"].shape[0] > idx_ref_foot_to_assign[2] + 1: idx_ref_foot_to_assign[2] += 1 - if (RR_contact_sequence[j + 1] == 0 and RR_contact_sequence[j] == 1): - if (reference['ref_foot_RR'].shape[0] > idx_ref_foot_to_assign[3] + 1): + if RR_contact_sequence[j + 1] == 0 and RR_contact_sequence[j] == 1: + if reference["ref_foot_RR"].shape[0] > idx_ref_foot_to_assign[3] + 1: idx_ref_foot_to_assign[3] += 1 # Calculate the reference force z for the leg in stance # It's simply mass*acc/number_of_legs_in_stance!! # Force x and y are always 0 - number_of_legs_in_stance = np.array([FL_contact_sequence[j], FR_contact_sequence[j], - RL_contact_sequence[j], RR_contact_sequence[j]]).sum() - if (number_of_legs_in_stance == 0): + number_of_legs_in_stance = np.array( + [FL_contact_sequence[j], FR_contact_sequence[j], RL_contact_sequence[j], RR_contact_sequence[j]] + ).sum() + if number_of_legs_in_stance == 0: reference_force_stance_legs = 0 else: reference_force_stance_legs = (mass * 9.81) / number_of_legs_in_stance @@ -1135,33 +1215,40 @@ def compute_control(self, state, reference, contact_sequence, constraint=None, e yref[53] = reference_force_rr_z # Setting the reference to acados - if (self.use_DDP): - if (j == 0): + if self.use_DDP: + if j == 0: num_l2_penalties = self.ocp.model.cost_y_expr_0.shape[0] - (self.states_dim + self.inputs_dim) else: num_l2_penalties = self.ocp.model.cost_y_expr.shape[0] - (self.states_dim + self.inputs_dim) - yref_tot = np.concatenate((yref, np.zeros(num_l2_penalties, ))) + yref_tot = np.concatenate( + ( + yref, + np.zeros( + num_l2_penalties, + ), + ) + ) self.acados_ocp_solver.set(j, "yref", yref_tot) else: self.acados_ocp_solver.set(j, "yref", yref) # Fill last step horizon reference (self.states_dim - no control action!!) yref_N = np.zeros(shape=(self.states_dim,)) - yref_N[0:3] = reference['ref_position'] - yref_N[3:6] = reference['ref_linear_velocity'] - yref_N[6:9] = reference['ref_orientation'] - yref_N[9:12] = reference['ref_angular_velocity'] - yref_N[12:15] = reference['ref_foot_FL'][idx_ref_foot_to_assign[0]] - yref_N[15:18] = reference['ref_foot_FR'][idx_ref_foot_to_assign[1]] - yref_N[18:21] = reference['ref_foot_RL'][idx_ref_foot_to_assign[2]] - yref_N[21:24] = reference['ref_foot_RR'][idx_ref_foot_to_assign[3]] + yref_N[0:3] = reference["ref_position"] + yref_N[3:6] = reference["ref_linear_velocity"] + yref_N[6:9] = reference["ref_orientation"] + yref_N[9:12] = reference["ref_angular_velocity"] + yref_N[12:15] = reference["ref_foot_FL"][idx_ref_foot_to_assign[0]] + yref_N[15:18] = reference["ref_foot_FR"][idx_ref_foot_to_assign[1]] + yref_N[18:21] = reference["ref_foot_RL"][idx_ref_foot_to_assign[2]] + yref_N[21:24] = reference["ref_foot_RR"][idx_ref_foot_to_assign[3]] # Setting the reference to acados self.acados_ocp_solver.set(self.horizon, "yref", yref_N) # Fill stance param, friction and stance proximity # (stance proximity will disable foothold optimization near a stance!!) - mu = config.mpc_params['mu'] + mu = config.mpc_params["mu"] yaw = state["orientation"][2] # Stance Proximity ugly routine. Basically we disable foothold optimization @@ -1173,91 +1260,116 @@ def compute_control(self, state, reference, contact_sequence, constraint=None, e stance_proximity_RR = np.zeros((self.horizon,)) for j in range(self.horizon): - if (FL_contact_sequence[j] == 0): + if FL_contact_sequence[j] == 0: if (j + 1) < self.horizon: - if (FL_contact_sequence[j + 1] == 1): + if FL_contact_sequence[j + 1] == 1: stance_proximity_FL[j] = 1 * 0 if (j + 2) < self.horizon: - if (FL_contact_sequence[j + 1] == 0 and FL_contact_sequence[j + 2] == 1): + if FL_contact_sequence[j + 1] == 0 and FL_contact_sequence[j + 2] == 1: stance_proximity_FL[j] = 1 * 0 - if (FR_contact_sequence[j] == 0): + if FR_contact_sequence[j] == 0: if (j + 1) < self.horizon: - if (FR_contact_sequence[j + 1] == 1): + if FR_contact_sequence[j + 1] == 1: stance_proximity_FR[j] = 1 * 0 if (j + 2) < self.horizon: - if (FR_contact_sequence[j + 1] == 0 and FR_contact_sequence[j + 2] == 1): + if FR_contact_sequence[j + 1] == 0 and FR_contact_sequence[j + 2] == 1: stance_proximity_FR[j] = 1 * 0 - if (RL_contact_sequence[j] == 0): + if RL_contact_sequence[j] == 0: if (j + 1) < self.horizon: - if (RL_contact_sequence[j + 1] == 1): + if RL_contact_sequence[j + 1] == 1: stance_proximity_RL[j] = 1 * 0 if (j + 2) < self.horizon: - if (RL_contact_sequence[j + 1] == 0 and RL_contact_sequence[j + 2] == 1): + if RL_contact_sequence[j + 1] == 0 and RL_contact_sequence[j + 2] == 1: stance_proximity_RL[j] = 1 * 0 - if (RR_contact_sequence[j] == 0): + if RR_contact_sequence[j] == 0: if (j + 1) < self.horizon: - if (RR_contact_sequence[j + 1] == 1): + if RR_contact_sequence[j + 1] == 1: stance_proximity_RR[j] = 1 * 0 if (j + 2) < self.horizon: - if (RR_contact_sequence[j + 1] == 0 and RR_contact_sequence[j + 2] == 1): + if RR_contact_sequence[j + 1] == 0 and RR_contact_sequence[j + 2] == 1: stance_proximity_RR[j] = 1 * 0 # Set the parameters to acados for j in range(self.horizon): # If we have estimated an external wrench, we can compensate it for all steps # or less (maybe the disturbance is not costant along the horizon!) - if (config.mpc_params['external_wrenches_compensation'] and - config.mpc_params['external_wrenches_compensation_num_step'] and - j < config.mpc_params['external_wrenches_compensation_num_step']): + if ( + config.mpc_params["external_wrenches_compensation"] + and config.mpc_params["external_wrenches_compensation_num_step"] + and j < config.mpc_params["external_wrenches_compensation_num_step"] + ): external_wrenches_estimated_param = copy.deepcopy(external_wrenches) external_wrenches_estimated_param = external_wrenches_estimated_param.reshape((6,)) else: external_wrenches_estimated_param = np.zeros((6,)) - param = np.array([FL_contact_sequence[j], FR_contact_sequence[j], - RL_contact_sequence[j], RR_contact_sequence[j], mu, - stance_proximity_FL[j], - stance_proximity_FR[j], - stance_proximity_RL[j], - stance_proximity_RR[j], - state["position"][0], state["position"][1], - state["position"][2], state["orientation"][2], - external_wrenches_estimated_param[0], external_wrenches_estimated_param[1], - external_wrenches_estimated_param[2], external_wrenches_estimated_param[3], - external_wrenches_estimated_param[4], external_wrenches_estimated_param[5], - inertia[0], inertia[1], inertia[2], inertia[3], inertia[4], inertia[5], - inertia[6], inertia[7], inertia[8], mass]) + param = np.array( + [ + FL_contact_sequence[j], + FR_contact_sequence[j], + RL_contact_sequence[j], + RR_contact_sequence[j], + mu, + stance_proximity_FL[j], + stance_proximity_FR[j], + stance_proximity_RL[j], + stance_proximity_RR[j], + state["position"][0], + state["position"][1], + state["position"][2], + state["orientation"][2], + external_wrenches_estimated_param[0], + external_wrenches_estimated_param[1], + external_wrenches_estimated_param[2], + external_wrenches_estimated_param[3], + external_wrenches_estimated_param[4], + external_wrenches_estimated_param[5], + inertia[0], + inertia[1], + inertia[2], + inertia[3], + inertia[4], + inertia[5], + inertia[6], + inertia[7], + inertia[8], + mass, + ] + ) self.acados_ocp_solver.set(j, "p", copy.deepcopy(param)) # Set initial state constraint. We teleported the robot foothold # to the previous optimal foothold. This is done to avoid the optimization # of a foothold that is not considered at all at touchdown! In any case, - # the height cames always from the VFA - if (FL_contact_sequence[0] == 0): + # the height cames always from the VFA + if FL_contact_sequence[0] == 0: state["foot_FL"] = reference["ref_foot_FL"][0] - if (FR_contact_sequence[0] == 0): + if FR_contact_sequence[0] == 0: state["foot_FR"] = reference["ref_foot_FR"][0] - if (RL_contact_sequence[0] == 0): + if RL_contact_sequence[0] == 0: state["foot_RL"] = reference["ref_foot_RL"][0] - if (RR_contact_sequence[0] == 0): + if RR_contact_sequence[0] == 0: state["foot_RR"] = reference["ref_foot_RR"][0] - if (self.use_integrators): + if self.use_integrators: # Compute error for integral action alpha_integrator = config.mpc_params["alpha_integrator"] self.integral_errors[0] += (state["position"][2] - reference["ref_position"][2]) * alpha_integrator - self.integral_errors[1] += (state["linear_velocity"][0] - reference["ref_linear_velocity"][ - 0]) * alpha_integrator - self.integral_errors[2] += (state["linear_velocity"][1] - reference["ref_linear_velocity"][ - 1]) * alpha_integrator - self.integral_errors[3] += (state["linear_velocity"][2] - reference["ref_linear_velocity"][ - 2]) * alpha_integrator + self.integral_errors[1] += ( + state["linear_velocity"][0] - reference["ref_linear_velocity"][0] + ) * alpha_integrator + self.integral_errors[2] += ( + state["linear_velocity"][1] - reference["ref_linear_velocity"][1] + ) * alpha_integrator + self.integral_errors[3] += ( + state["linear_velocity"][2] - reference["ref_linear_velocity"][2] + ) * alpha_integrator self.integral_errors[4] += (state["orientation"][0] - reference["ref_orientation"][0]) * (alpha_integrator) self.integral_errors[5] += (state["orientation"][1] - reference["ref_orientation"][1]) * alpha_integrator @@ -1268,72 +1380,98 @@ def compute_control(self, state, reference, contact_sequence, constraint=None, e cap_integrator_roll = config.mpc_params["integrator_cap"][4] cap_integrator_pitch = config.mpc_params["integrator_cap"][5] - self.integral_errors[0] = np.where(np.abs(self.integral_errors[0]) > cap_integrator_z, - cap_integrator_z * np.sign(self.integral_errors[0]), - self.integral_errors[0]) - self.integral_errors[1] = np.where(np.abs(self.integral_errors[1]) > cap_integrator_x_dot, - cap_integrator_x_dot * np.sign(self.integral_errors[1]), - self.integral_errors[1]) - self.integral_errors[2] = np.where(np.abs(self.integral_errors[2]) > cap_integrator_y_dot, - cap_integrator_y_dot * np.sign(self.integral_errors[2]), - self.integral_errors[2]) - self.integral_errors[3] = np.where(np.abs(self.integral_errors[3]) > cap_integrator_z_dot, - cap_integrator_z_dot * np.sign(self.integral_errors[3]), - self.integral_errors[3]) - self.integral_errors[4] = np.where(np.abs(self.integral_errors[4]) > cap_integrator_roll, - cap_integrator_roll * np.sign(self.integral_errors[4]), - self.integral_errors[4]) - self.integral_errors[5] = np.where(np.abs(self.integral_errors[5]) > cap_integrator_pitch, - cap_integrator_pitch * np.sign(self.integral_errors[5]), - self.integral_errors[5]) - - if(self.verbose): + self.integral_errors[0] = np.where( + np.abs(self.integral_errors[0]) > cap_integrator_z, + cap_integrator_z * np.sign(self.integral_errors[0]), + self.integral_errors[0], + ) + self.integral_errors[1] = np.where( + np.abs(self.integral_errors[1]) > cap_integrator_x_dot, + cap_integrator_x_dot * np.sign(self.integral_errors[1]), + self.integral_errors[1], + ) + self.integral_errors[2] = np.where( + np.abs(self.integral_errors[2]) > cap_integrator_y_dot, + cap_integrator_y_dot * np.sign(self.integral_errors[2]), + self.integral_errors[2], + ) + self.integral_errors[3] = np.where( + np.abs(self.integral_errors[3]) > cap_integrator_z_dot, + cap_integrator_z_dot * np.sign(self.integral_errors[3]), + self.integral_errors[3], + ) + self.integral_errors[4] = np.where( + np.abs(self.integral_errors[4]) > cap_integrator_roll, + cap_integrator_roll * np.sign(self.integral_errors[4]), + self.integral_errors[4], + ) + self.integral_errors[5] = np.where( + np.abs(self.integral_errors[5]) > cap_integrator_pitch, + cap_integrator_pitch * np.sign(self.integral_errors[5]), + self.integral_errors[5], + ) + + if self.verbose: print("self.integral_errors: ", self.integral_errors) # Set initial state constraint acados, converting first the dictionary to np array - state_acados = np.concatenate((state["position"], state["linear_velocity"], - state["orientation"], state["angular_velocity"], - state["foot_FL"], state["foot_FR"], - state["foot_RL"], state["foot_RR"], - self.integral_errors)).reshape((self.states_dim, 1)) + state_acados = np.concatenate( + ( + state["position"], + state["linear_velocity"], + state["orientation"], + state["angular_velocity"], + state["foot_FL"], + state["foot_FR"], + state["foot_RL"], + state["foot_RR"], + self.integral_errors, + ) + ).reshape((self.states_dim, 1)) self.acados_ocp_solver.set(0, "lbx", state_acados) self.acados_ocp_solver.set(0, "ubx", state_acados) # Set Warm start in case... - if (self.use_warm_start): - self.set_warm_start(state_acados, reference, FL_contact_sequence, FR_contact_sequence, RL_contact_sequence, - RR_contact_sequence) + if self.use_warm_start: + self.set_warm_start( + state_acados, + reference, + FL_contact_sequence, + FR_contact_sequence, + RL_contact_sequence, + RR_contact_sequence, + ) # Set stage constraint - h_R_w = np.array([np.cos(yaw), np.sin(yaw), - -np.sin(yaw), np.cos(yaw)]) - if (self.use_foothold_constraints or self.use_stability_constraints): - stance_proximity = np.vstack((stance_proximity_FL, stance_proximity_FR, - stance_proximity_RL, stance_proximity_RR)) + h_R_w = np.array([np.cos(yaw), np.sin(yaw), -np.sin(yaw), np.cos(yaw)]) + if self.use_foothold_constraints or self.use_stability_constraints: + stance_proximity = np.vstack( + (stance_proximity_FL, stance_proximity_FR, stance_proximity_RL, stance_proximity_RR) + ) self.set_stage_constraint(constraint, state, reference, contact_sequence, h_R_w, stance_proximity) # Solve ocp via RTI or normal ocp if self.use_RTI: # feedback phase - self.acados_ocp_solver.options_set('rti_phase', 2) + self.acados_ocp_solver.options_set("rti_phase", 2) status = self.acados_ocp_solver.solve() - if(self.verbose): - print("feedback phase time: ", self.acados_ocp_solver.get_stats('time_tot')) + if self.verbose: + print("feedback phase time: ", self.acados_ocp_solver.get_stats("time_tot")) else: status = self.acados_ocp_solver.solve() - if(self.verbose): - print("ocp time: ", self.acados_ocp_solver.get_stats('time_tot')) + if self.verbose: + print("ocp time: ", self.acados_ocp_solver.get_stats("time_tot")) # Take the solution control = self.acados_ocp_solver.get(0, "u") optimal_GRF = control[12:] optimal_foothold = np.zeros((4, 3)) - optimal_footholds_assigned = np.zeros((4,), dtype='bool') + optimal_footholds_assigned = np.zeros((4,), dtype="bool") # We need to provide the next touchdown foothold position. # We first take the foothold in stance now (they are not optimized!) - # and flag them as True (aka "assigned") + # and flag them as True (aka "assigned") if FL_contact_sequence[0] == 1: optimal_foothold[0] = state["foot_FL"] optimal_footholds_assigned[0] = True @@ -1348,23 +1486,34 @@ def compute_control(self, state, reference, contact_sequence, constraint=None, e optimal_footholds_assigned[3] = True # Then we take the foothold at the next touchdown from the one - # that are not flagged as True from before, and saturate them!! + # that are not flagged as True from before, and saturate them!! # P.S. The saturation is in the horizontal frame h_R_w = h_R_w.reshape((2, 2)) for j in range(1, self.horizon): - if (FL_contact_sequence[j] != FL_contact_sequence[j - 1] and not optimal_footholds_assigned[0]): + if FL_contact_sequence[j] != FL_contact_sequence[j - 1] and not optimal_footholds_assigned[0]: optimal_foothold[0] = self.acados_ocp_solver.get(j, "x")[12:15] optimal_footholds_assigned[0] = True - if (constraint is None): - first_up_constraint_FL = np.array([reference["ref_foot_FL"][0][0], reference["ref_foot_FL"][0][1], - reference["ref_foot_FL"][0][2] + 0.002]) - first_low_constraint_FL = np.array([reference["ref_foot_FL"][0][0], reference["ref_foot_FL"][0][1], - reference["ref_foot_FL"][0][2] - 0.002]) + if constraint is None: + first_up_constraint_FL = np.array( + [ + reference["ref_foot_FL"][0][0], + reference["ref_foot_FL"][0][1], + reference["ref_foot_FL"][0][2] + 0.002, + ] + ) + first_low_constraint_FL = np.array( + [ + reference["ref_foot_FL"][0][0], + reference["ref_foot_FL"][0][1], + reference["ref_foot_FL"][0][2] - 0.002, + ] + ) first_up_constraint_FL[0:2] = h_R_w @ (first_up_constraint_FL[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_FL[0:2] = h_R_w @ ( - first_low_constraint_FL[0:2] - state["position"][0:2]) - 0.15 + first_low_constraint_FL[0:2] = ( + h_R_w @ (first_low_constraint_FL[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_FL = np.array([constraint[0][0], constraint[1][0], constraint[2][0] + 0.002]) first_low_constraint_FL = np.array([constraint[9][0], constraint[10][0], constraint[11][0] - 0.002]) @@ -1373,23 +1522,35 @@ def compute_control(self, state, reference, contact_sequence, constraint=None, e first_low_constraint_FL[0:2] = h_R_w @ (first_low_constraint_FL[0:2] - state["position"][0:2]) optimal_foothold[0][0:2] = h_R_w @ (optimal_foothold[0][0:2] - state["position"][0:2]) - optimal_foothold[0][0:2] = np.clip(optimal_foothold[0][0:2], first_low_constraint_FL[0:2], - first_up_constraint_FL[0:2]) + optimal_foothold[0][0:2] = np.clip( + optimal_foothold[0][0:2], first_low_constraint_FL[0:2], first_up_constraint_FL[0:2] + ) optimal_foothold[0][0:2] = h_R_w.T @ optimal_foothold[0][0:2] + state["position"][0:2] - if (FR_contact_sequence[j] != FR_contact_sequence[j - 1] and not optimal_footholds_assigned[1]): + if FR_contact_sequence[j] != FR_contact_sequence[j - 1] and not optimal_footholds_assigned[1]: optimal_foothold[1] = self.acados_ocp_solver.get(j, "x")[15:18] optimal_footholds_assigned[1] = True - if (constraint is None): - first_up_constraint_FR = np.array([reference["ref_foot_FR"][0][0], reference["ref_foot_FR"][0][1], - reference["ref_foot_FR"][0][2] + 0.002]) - first_low_constraint_FR = np.array([reference["ref_foot_FR"][0][0], reference["ref_foot_FR"][0][1], - reference["ref_foot_FR"][0][2] - 0.002]) + if constraint is None: + first_up_constraint_FR = np.array( + [ + reference["ref_foot_FR"][0][0], + reference["ref_foot_FR"][0][1], + reference["ref_foot_FR"][0][2] + 0.002, + ] + ) + first_low_constraint_FR = np.array( + [ + reference["ref_foot_FR"][0][0], + reference["ref_foot_FR"][0][1], + reference["ref_foot_FR"][0][2] - 0.002, + ] + ) first_up_constraint_FR[0:2] = h_R_w @ (first_up_constraint_FR[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_FR[0:2] = h_R_w @ ( - first_low_constraint_FR[0:2] - state["position"][0:2]) - 0.15 + first_low_constraint_FR[0:2] = ( + h_R_w @ (first_low_constraint_FR[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_FR = np.array([constraint[0][1], constraint[1][1], constraint[2][1] + 0.002]) first_low_constraint_FR = np.array([constraint[9][1], constraint[10][1], constraint[11][1] - 0.002]) @@ -1398,23 +1559,35 @@ def compute_control(self, state, reference, contact_sequence, constraint=None, e first_low_constraint_FR[0:2] = h_R_w @ (first_low_constraint_FR[0:2] - state["position"][0:2]) optimal_foothold[1][0:2] = h_R_w @ (optimal_foothold[1][0:2] - state["position"][0:2]) - optimal_foothold[1][0:2] = np.clip(optimal_foothold[1][0:2], first_low_constraint_FR[0:2], - first_up_constraint_FR[0:2]) + optimal_foothold[1][0:2] = np.clip( + optimal_foothold[1][0:2], first_low_constraint_FR[0:2], first_up_constraint_FR[0:2] + ) optimal_foothold[1][0:2] = h_R_w.T @ optimal_foothold[1][0:2] + state["position"][0:2] - if (RL_contact_sequence[j] != RL_contact_sequence[j - 1] and not optimal_footholds_assigned[2]): + if RL_contact_sequence[j] != RL_contact_sequence[j - 1] and not optimal_footholds_assigned[2]: optimal_foothold[2] = self.acados_ocp_solver.get(j, "x")[18:21] optimal_footholds_assigned[2] = True - if (constraint is None): - first_up_constraint_RL = np.array([reference["ref_foot_RL"][0][0], reference["ref_foot_RL"][0][1], - reference["ref_foot_RL"][0][2] + 0.002]) - first_low_constraint_RL = np.array([reference["ref_foot_RL"][0][0], reference["ref_foot_RL"][0][1], - reference["ref_foot_RL"][0][2] - 0.002]) + if constraint is None: + first_up_constraint_RL = np.array( + [ + reference["ref_foot_RL"][0][0], + reference["ref_foot_RL"][0][1], + reference["ref_foot_RL"][0][2] + 0.002, + ] + ) + first_low_constraint_RL = np.array( + [ + reference["ref_foot_RL"][0][0], + reference["ref_foot_RL"][0][1], + reference["ref_foot_RL"][0][2] - 0.002, + ] + ) first_up_constraint_RL[0:2] = h_R_w @ (first_up_constraint_RL[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_RL[0:2] = h_R_w @ ( - first_low_constraint_RL[0:2] - state["position"][0:2]) - 0.15 + first_low_constraint_RL[0:2] = ( + h_R_w @ (first_low_constraint_RL[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_RL = np.array([constraint[0][2], constraint[1][2], constraint[2][2] + 0.002]) first_low_constraint_RL = np.array([constraint[9][2], constraint[10][2], constraint[11][2] - 0.002]) @@ -1423,23 +1596,35 @@ def compute_control(self, state, reference, contact_sequence, constraint=None, e first_low_constraint_RL[0:2] = h_R_w @ (first_low_constraint_RL[0:2] - state["position"][0:2]) optimal_foothold[2][0:2] = h_R_w @ (optimal_foothold[2][0:2] - state["position"][0:2]) - optimal_foothold[2][0:2] = np.clip(optimal_foothold[2][0:2], first_low_constraint_RL[0:2], - first_up_constraint_RL[0:2]) + optimal_foothold[2][0:2] = np.clip( + optimal_foothold[2][0:2], first_low_constraint_RL[0:2], first_up_constraint_RL[0:2] + ) optimal_foothold[2][0:2] = h_R_w.T @ optimal_foothold[2][0:2] + state["position"][0:2] - if (RR_contact_sequence[j] != RR_contact_sequence[j - 1] and not optimal_footholds_assigned[3]): + if RR_contact_sequence[j] != RR_contact_sequence[j - 1] and not optimal_footholds_assigned[3]: optimal_foothold[3] = self.acados_ocp_solver.get(j, "x")[21:24] optimal_footholds_assigned[3] = True - if (constraint is None): - first_up_constraint_RR = np.array([reference["ref_foot_RR"][0][0], reference["ref_foot_RR"][0][1], - reference["ref_foot_RR"][0][2] + 0.002]) - first_low_constraint_RR = np.array([reference["ref_foot_RR"][0][0], reference["ref_foot_RR"][0][1], - reference["ref_foot_RR"][0][2] - 0.002]) + if constraint is None: + first_up_constraint_RR = np.array( + [ + reference["ref_foot_RR"][0][0], + reference["ref_foot_RR"][0][1], + reference["ref_foot_RR"][0][2] + 0.002, + ] + ) + first_low_constraint_RR = np.array( + [ + reference["ref_foot_RR"][0][0], + reference["ref_foot_RR"][0][1], + reference["ref_foot_RR"][0][2] - 0.002, + ] + ) first_up_constraint_RR[0:2] = h_R_w @ (first_up_constraint_RR[0:2] - state["position"][0:2]) + 0.15 - first_low_constraint_RR[0:2] = h_R_w @ ( - first_low_constraint_RR[0:2] - state["position"][0:2]) - 0.15 + first_low_constraint_RR[0:2] = ( + h_R_w @ (first_low_constraint_RR[0:2] - state["position"][0:2]) - 0.15 + ) else: first_up_constraint_RR = np.array([constraint[0][3], constraint[1][3], constraint[2][3] + 0.002]) first_low_constraint_RR = np.array([constraint[9][3], constraint[10][3], constraint[11][3] - 0.002]) @@ -1448,36 +1633,38 @@ def compute_control(self, state, reference, contact_sequence, constraint=None, e first_low_constraint_RR[0:2] = h_R_w @ (first_low_constraint_RR[0:2] - state["position"][0:2]) optimal_foothold[3][0:2] = h_R_w @ (optimal_foothold[3][0:2] - state["position"][0:2]) - optimal_foothold[3][0:2] = np.clip(optimal_foothold[3][0:2], first_low_constraint_RR[0:2], - first_up_constraint_RR[0:2]) + optimal_foothold[3][0:2] = np.clip( + optimal_foothold[3][0:2], first_low_constraint_RR[0:2], first_up_constraint_RR[0:2] + ) optimal_foothold[3][0:2] = h_R_w.T @ optimal_foothold[3][0:2] + state["position"][0:2] # If in the prediction horizon, the foot is never in stance, we replicate the reference # to not confuse the swing controller - if (optimal_footholds_assigned[0] == False): + if optimal_footholds_assigned[0] == False: optimal_foothold[0] = reference["ref_foot_FL"][0] - if (optimal_footholds_assigned[1] == False): + if optimal_footholds_assigned[1] == False: optimal_foothold[1] = reference["ref_foot_FR"][0] - if (optimal_footholds_assigned[2] == False): + if optimal_footholds_assigned[2] == False: optimal_foothold[2] = reference["ref_foot_RL"][0] - if (optimal_footholds_assigned[3] == False): + if optimal_footholds_assigned[3] == False: optimal_foothold[3] = reference["ref_foot_RR"][0] - if (config.mpc_params['dt'] <= 0.02 or ( - config.mpc_params['use_nonuniform_discretization'] and config.mpc_params['dt_fine_grained'] <= 0.02)): + if config.mpc_params["dt"] <= 0.02 or ( + config.mpc_params["use_nonuniform_discretization"] and config.mpc_params["dt_fine_grained"] <= 0.02 + ): optimal_next_state_index = 2 else: optimal_next_state_index = 1 optimal_next_state = self.acados_ocp_solver.get(optimal_next_state_index, "x")[0:24] self.optimal_next_state = optimal_next_state - if(self.verbose): + if self.verbose: self.acados_ocp_solver.print_statistics() # Check if QPs converged, if not just use the reference footholds # and a GRF over Z distribuited between the leg in stance - if (status == 1 or status == 4): - if(self.verbose): + if status == 1 or status == 4: + if self.verbose: print("status", status) if FL_contact_sequence[0] == 0: optimal_foothold[0] = reference["ref_foot_FL"][0] @@ -1488,9 +1675,10 @@ def compute_control(self, state, reference, contact_sequence, constraint=None, e if RR_contact_sequence[0] == 0: optimal_foothold[3] = reference["ref_foot_RR"][0] - number_of_legs_in_stance = np.array([FL_contact_sequence[0], FR_contact_sequence[0], - RL_contact_sequence[0], RR_contact_sequence[0]]).sum() - if (number_of_legs_in_stance == 0): + number_of_legs_in_stance = np.array( + [FL_contact_sequence[0], FR_contact_sequence[0], RL_contact_sequence[0], RR_contact_sequence[0]] + ).sum() + if number_of_legs_in_stance == 0: reference_force_stance_legs = 0 else: reference_force_stance_legs = (mass * 9.81) / number_of_legs_in_stance diff --git a/quadruped_pympc/controllers/sampling/centroidal_model_jax.py b/quadruped_pympc/controllers/sampling/centroidal_model_jax.py index 03298d0..a43101c 100644 --- a/quadruped_pympc/controllers/sampling/centroidal_model_jax.py +++ b/quadruped_pympc/controllers/sampling/centroidal_model_jax.py @@ -3,31 +3,28 @@ # Authors: Giulio Turrisi - +import os import time -import os dir_path = os.path.dirname(os.path.realpath(__file__)) import sys -sys.path.append(dir_path) -sys.path.append(dir_path + '/../') -from quadruped_pympc import config +sys.path.append(dir_path) +sys.path.append(dir_path + "/../") -import jax.numpy as jnp -from jax import jit import jax -from jax import random - - +import jax.numpy as jnp +from jax import jit, random -dtype_general='float32' +from quadruped_pympc import config +dtype_general = "float32" # Class that defines the prediction model of the NMPC class Centroidal_Model_JAX: - def __init__(self,dt,device) -> None: + def __init__(self, dt, device) -> None: """ This method initializes the foothold generator Centroidal_Model, which creates the prediction model of the NMPC. @@ -35,47 +32,44 @@ def __init__(self,dt,device) -> None: self.dt = dt - if(device=="gpu"): + if device == "gpu": try: - self.device = jax.devices('gpu')[0] + self.device = jax.devices("gpu")[0] except: - self.device = jax.devices('cpu')[0] + self.device = jax.devices("cpu")[0] print("GPU not available, using CPU") else: - self.device = jax.devices('cpu')[0] - + self.device = jax.devices("cpu")[0] # Mass and Inertia robot dependant self.mass = config.mass self.inertia = jnp.array(config.inertia) - # Nonuniform discretization - if(config.mpc_params['use_nonuniform_discretization']): - time_steps_fine_grained = jnp.tile(config.mpc_params['dt_fine_grained'], config.mpc_params['horizon_fine_grained']) - self.dts = jnp.concatenate((time_steps_fine_grained, jnp.tile(self.dt, config.mpc_params['horizon'] - config.mpc_params['horizon_fine_grained']))) + if config.mpc_params["use_nonuniform_discretization"]: + time_steps_fine_grained = jnp.tile( + config.mpc_params["dt_fine_grained"], config.mpc_params["horizon_fine_grained"] + ) + self.dts = jnp.concatenate( + ( + time_steps_fine_grained, + jnp.tile(self.dt, config.mpc_params["horizon"] - config.mpc_params["horizon_fine_grained"]), + ) + ) else: - self.dts = jnp.tile(self.dt, config.mpc_params['horizon']) - - - + self.dts = jnp.tile(self.dt, config.mpc_params["horizon"]) # We precompute the inverse of the inertia self.inertia_inv = self.calculate_inverse(self.inertia) - # State and input dimensions self.state_dim = 12 self.input_dim = 24 - - - # this is useful to manage the fact that df depends on self.mass and self.inertia which jit does not really like - # by compiling it here we ensure that the parameters will always stay the same which is what jit likes + # this is useful to manage the fact that df depends on self.mass and self.inertia which jit does not really like + # by compiling it here we ensure that the parameters will always stay the same which is what jit likes vectorized_integrate_jax = jax.vmap(self.integrate_jax, in_axes=(None, 0, None), out_axes=0) - self.compiled_integrate_jax = jit(vectorized_integrate_jax, device = self.device) - - + self.compiled_integrate_jax = jit(vectorized_integrate_jax, device=self.device) def calculate_inverse(self, A): a11 = A[0, 0] @@ -92,23 +86,25 @@ def calculate_inverse(self, A): DET = a11 * (a33 * a22 - a32 * a23) - a21 * (a33 * a12 - a32 * a13) + a31 * (a23 * a12 - a22 * a13) # Calculate the inverse of A - return jnp.array([ - [(a33 * a22 - a32 * a23), -(a33 * a12 - a32 * a13), (a23 * a12 - a22 * a13)], - [-(a33 * a21 - a31 * a23), (a33 * a11 - a31 * a13), -(a23 * a11 - a21 * a13)], - [(a32 * a21 - a31 * a22), -(a32 * a11 - a31 * a12), (a22 * a11 - a21 * a12)] - ]) / DET + return ( + jnp.array( + [ + [(a33 * a22 - a32 * a23), -(a33 * a12 - a32 * a13), (a23 * a12 - a22 * a13)], + [-(a33 * a21 - a31 * a23), (a33 * a11 - a31 * a13), -(a23 * a11 - a21 * a13)], + [(a32 * a21 - a31 * a22), -(a32 * a11 - a31 * a12), (a22 * a11 - a21 * a12)], + ] + ) + / DET + ) - def fd(self, states: jnp.ndarray, inputs: jnp.ndarray, contact_status: jnp.ndarray): """ This method computes the state derivative of the system. """ - + def skew(v): - return jnp.array([[0, -v[2], v[1]], - [v[2], 0, -v[0]], - [-v[1], v[0], 0]]) - + return jnp.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]]) + # Extracting variables for clarity foot_position_fl, foot_position_fr, foot_position_rl, foot_position_rr = jnp.split(states[12:], 4) foot_force_fl, foot_force_fr, foot_force_rl, foot_force_rr = jnp.split(inputs[12:], 4) @@ -117,140 +113,167 @@ def skew(v): # Compute linear_com_vel linear_com_vel = states[3:6] - # Compute linear_com_acc - temp = jnp.dot(foot_force_fl, stanceFL) + jnp.dot(foot_force_fr, stanceFR) + jnp.dot(foot_force_rl, stanceRL) + jnp.dot(foot_force_rr, stanceRR) + temp = ( + jnp.dot(foot_force_fl, stanceFL) + + jnp.dot(foot_force_fr, stanceFR) + + jnp.dot(foot_force_rl, stanceRL) + + jnp.dot(foot_force_rr, stanceRR) + ) gravity = jnp.array([jnp.float32(0), jnp.float32(0), jnp.float32(-9.81)]) linear_com_acc = jnp.dot(jnp.float32(1) / self.mass, temp) + gravity - # Compute euler_rates_base and angular_acc_base w = states[9:12] roll, pitch, yaw = states[6:9] - - conj_euler_rates = jnp.array([ - [jnp.float32(1), jnp.float32(0), -jnp.sin(pitch)], - [jnp.float32(0), jnp.cos(roll), jnp.cos(pitch) * jnp.sin(roll)], - [jnp.float32(0), -jnp.sin(roll), jnp.cos(pitch) * jnp.cos(roll)] - ]) + conj_euler_rates = jnp.array( + [ + [jnp.float32(1), jnp.float32(0), -jnp.sin(pitch)], + [jnp.float32(0), jnp.cos(roll), jnp.cos(pitch) * jnp.sin(roll)], + [jnp.float32(0), -jnp.sin(roll), jnp.cos(pitch) * jnp.cos(roll)], + ] + ) - - temp2 = jnp.dot(skew(foot_position_fl - com_position), foot_force_fl) * stanceFL temp2 += jnp.dot(skew(foot_position_fr - com_position), foot_force_fr) * stanceFR temp2 += jnp.dot(skew(foot_position_rl - com_position), foot_force_rl) * stanceRL temp2 += jnp.dot(skew(foot_position_rr - com_position), foot_force_rr) * stanceRR - euler_rates_base = jnp.dot(self.calculate_inverse(conj_euler_rates), w) - # FINAL angular_acc_base STATE (4) - #Z Y X rotations! - b_R_w = jnp.array([ - [jnp.cos(pitch)*jnp.cos(yaw), jnp.cos(pitch)*jnp.sin(yaw), -jnp.sin(pitch)], - [jnp.sin(roll)*jnp.sin(pitch)*jnp.cos(yaw)-jnp.cos(roll)*jnp.sin(yaw), jnp.sin(roll)*jnp.sin(pitch)*jnp.sin(yaw)+jnp.cos(roll)*jnp.cos(yaw), jnp.sin(roll)*jnp.cos(pitch)], - [jnp.cos(roll)*jnp.sin(pitch)*jnp.cos(yaw)+jnp.sin(roll)*jnp.sin(yaw), jnp.cos(roll)*jnp.sin(pitch)*jnp.sin(yaw)-jnp.sin(roll)*jnp.cos(yaw), jnp.cos(roll)*jnp.cos(pitch)] - ]) - - - - - angular_acc_base = -jnp.dot(self.inertia_inv, jnp.dot(skew(w), jnp.dot(self.inertia, w))) + jnp.dot(self.inertia_inv, jnp.dot(b_R_w, temp2)) - - + # Z Y X rotations! + b_R_w = jnp.array( + [ + [jnp.cos(pitch) * jnp.cos(yaw), jnp.cos(pitch) * jnp.sin(yaw), -jnp.sin(pitch)], + [ + jnp.sin(roll) * jnp.sin(pitch) * jnp.cos(yaw) - jnp.cos(roll) * jnp.sin(yaw), + jnp.sin(roll) * jnp.sin(pitch) * jnp.sin(yaw) + jnp.cos(roll) * jnp.cos(yaw), + jnp.sin(roll) * jnp.cos(pitch), + ], + [ + jnp.cos(roll) * jnp.sin(pitch) * jnp.cos(yaw) + jnp.sin(roll) * jnp.sin(yaw), + jnp.cos(roll) * jnp.sin(pitch) * jnp.sin(yaw) - jnp.sin(roll) * jnp.cos(yaw), + jnp.cos(roll) * jnp.cos(pitch), + ], + ] + ) + + angular_acc_base = -jnp.dot(self.inertia_inv, jnp.dot(skew(w), jnp.dot(self.inertia, w))) + jnp.dot( + self.inertia_inv, jnp.dot(b_R_w, temp2) + ) + # Returning the results return jnp.concatenate([linear_com_vel, linear_com_acc, euler_rates_base, angular_acc_base]) - - - def integrate_jax(self, state, inputs, contact_status, n): """ This method computes the forward evolution of the system. """ fd = self.fd(state, inputs, contact_status) - # Simple euler! dt = self.dts[n] - new_state = state[0:12] + fd*dt + new_state = state[0:12] + fd * dt return jnp.concatenate([new_state, state[12:]]) - - - - -if __name__=="__main__": - model_jax = Centroidal_Model_JAX(dt = jnp.float32(0.04), device="gpu") +if __name__ == "__main__": + model_jax = Centroidal_Model_JAX(dt=jnp.float32(0.04), device="gpu") + + # all stance, friction, stance proximity + param = jnp.array([1.0, 1.0, 1.0, 1.0, 0.6, 0.0, 0.0, 0.0, 0.0], dtype=dtype_general) + + state = jnp.array( + [ + 0.0, + 0.0, + 0.0, # com position + 0.0, + 0.0, + 0.0, # com velocity + 0.0, + 0.0, + 0.0, # euler angles + 0.0, + 0.0, + 0.0, + ], + dtype=dtype_general, + ) # euler rates + + input = jnp.array( + [ + 0.0, + 0.0, + 0.0, # foot position fl + 0.0, + 0.0, + 0.0, # foot position fr + 0.0, + 0.0, + 0.0, # foot position rl + 0.0, + 0.0, + 0.0, # foot position rr + 0.0, + 0.0, + 0.0, # foot force fl + 0.0, + 0.0, + 0.0, # foot force fr + 0.0, + 0.0, + 0.0, # foot force rl + 0.0, + 0.0, + 0.0, + ], + dtype=dtype_general, + ) # foot force rr - # all stance, friction, stance proximity - param = jnp.array([1.0, 1.0 ,1.0 , 1.0 , 0.6, 0.0, 0.0, 0.0, 0.0], dtype=dtype_general) - - - state = jnp.array([0.0, 0.0, 0.0, # com position - 0.0, 0.0, 0.0, # com velocity - 0.0, 0.0, 0.0, # euler angles - 0.0, 0.0, 0.0], dtype=dtype_general) # euler rates - - input = jnp.array([0.0, 0.0, 0.0, # foot position fl - 0.0, 0.0, 0.0, # foot position fr - 0.0, 0.0, 0.0, # foot position rl - 0.0, 0.0, 0.0, # foot position rr - 0.0, 0.0, 0.0, # foot force fl - 0.0, 0.0, 0.0, # foot force fr - 0.0, 0.0, 0.0, # foot force rl - 0.0, 0.0, 0.0], dtype=dtype_general) # foot force rr - # test fd - acc = model_jax.fd(state,input,param) - + acc = model_jax.fd(state, input, param) # test integrated print("testing SINGLE integration PYTHON") time_start = time.time() state_integrated = model_jax.integrate_jax(state, input, param) - print("computation time: ", time.time()-time_start) - + print("computation time: ", time.time() - time_start) # test compiled integrated print("\ntesting SINGLE integration COMPILED-FIRST TIME") - compiled_integrated_jax_single = jit(model_jax.integrate_jax, device = model_jax.device) + compiled_integrated_jax_single = jit(model_jax.integrate_jax, device=model_jax.device) time_start = time.time() state_integrated = compiled_integrated_jax_single(state, input, param).block_until_ready() - print("computation time: ", time.time()-time_start) - + print("computation time: ", time.time() - time_start) print("\ntesting SINGLE integration COMPILED-SECOND TIME") time_start = time.time() state_integrated = compiled_integrated_jax_single(state, input, param).block_until_ready() - print("computation time: ", time.time()-time_start) - - + print("computation time: ", time.time() - time_start) print("\ntesting MULTIPLE integration COMPILED-FIRST TIME") threads = 10000 - + key = random.PRNGKey(42) - input_vec = random.randint(key,(model_jax.input_dim*threads,), minval=-2, maxval=2).reshape(threads,model_jax.input_dim) - + input_vec = random.randint(key, (model_jax.input_dim * threads,), minval=-2, maxval=2).reshape( + threads, model_jax.input_dim + ) + time_start = time.time() model_jax.compiled_integrate_jax(state, input_vec, param).block_until_ready() - print("computation time: ", time.time()-time_start) - + print("computation time: ", time.time() - time_start) print("\ntesting MULTIPLE integration SECOND TIME") time_start = time.time() key = random.PRNGKey(50) - input_vec = random.randint(key,(model_jax.input_dim*threads,), minval=-2, maxval=2).reshape(threads,model_jax.input_dim) + input_vec = random.randint(key, (model_jax.input_dim * threads,), minval=-2, maxval=2).reshape( + threads, model_jax.input_dim + ) time_start = time.time() model_jax.compiled_integrate_jax(state, input_vec, param).block_until_ready() - print("computation time: ", time.time()-time_start) - - - - \ No newline at end of file + print("computation time: ", time.time() - time_start) diff --git a/quadruped_pympc/controllers/sampling/centroidal_nmpc_jax.py b/quadruped_pympc/controllers/sampling/centroidal_nmpc_jax.py index dd4125d..10eed9e 100644 --- a/quadruped_pympc/controllers/sampling/centroidal_nmpc_jax.py +++ b/quadruped_pympc/controllers/sampling/centroidal_nmpc_jax.py @@ -1,159 +1,162 @@ import numpy as np -np.set_printoptions(precision=3, suppress = True) -import os +np.set_printoptions(precision=3, suppress=True) + +import os + dir_path = os.path.dirname(os.path.realpath(__file__)) -os.environ['XLA_FLAGS'] = ('--xla_gpu_triton_gemm_any=True') +os.environ["XLA_FLAGS"] = "--xla_gpu_triton_gemm_any=True" +import sys + import jax import jax.numpy as jnp -from jax import jit, random +from jax import random -import sys sys.path.append(dir_path) -sys.path.append(dir_path + '/../') +sys.path.append(dir_path + "/../") -from quadruped_pympc import config -from centroidal_model_jax import Centroidal_Model_JAX - -import time import copy +from centroidal_model_jax import Centroidal_Model_JAX +from quadruped_pympc import config - -dtype_general='float32' +dtype_general = "float32" class Sampling_MPC: """This is a small class that implements a sampling based control law""" - - def __init__(self, horizon = 200, dt = 0.01, num_parallel_computations = 10000, sampling_method = 'random_sampling', control_parametrization = "linear_spline_1", device="gpu"): + def __init__( + self, + horizon=200, + dt=0.01, + num_parallel_computations=10000, + sampling_method="random_sampling", + control_parametrization="linear_spline_1", + device="gpu", + ): """ Args: - horizon (int): how much to look into the future for optimizing the gains + horizon (int): how much to look into the future for optimizing the gains dt (int): desidered sampling time """ - self.num_parallel_computations = config.mpc_params['num_parallel_computations'] - self.sampling_method = config.mpc_params['sampling_method'] - self.control_parametrization = config.mpc_params['control_parametrization'] - self.num_sampling_iterations = config.mpc_params['num_sampling_iterations'] - self.dt = config.mpc_params['dt'] - self.horizon = config.mpc_params['horizon'] - + self.num_parallel_computations = config.mpc_params["num_parallel_computations"] + self.sampling_method = config.mpc_params["sampling_method"] + self.control_parametrization = config.mpc_params["control_parametrization"] + self.num_sampling_iterations = config.mpc_params["num_sampling_iterations"] + self.dt = config.mpc_params["dt"] + self.horizon = config.mpc_params["horizon"] self.state_dim = 24 self.control_dim = 24 self.reference_dim = self.state_dim - + self.max_sampling_forces = 30 - - if(device=="gpu"): + if device == "gpu": try: - self.device = jax.devices('gpu')[0] + self.device = jax.devices("gpu")[0] except: - self.device = jax.devices('cpu')[0] + self.device = jax.devices("cpu")[0] print("GPU not available, using CPU") else: - self.device = jax.devices('cpu')[0] - + self.device = jax.devices("cpu")[0] - - if(self.control_parametrization == "linear_spline_1"): + if self.control_parametrization == "linear_spline_1": # Along the horizon, we have only 1 spline per control input (3 forces) # Each spline has 2 parameters - self.num_control_parameters_single_leg = 2*3 - + self.num_control_parameters_single_leg = 2 * 3 + # In totale we have 4 legs - self.num_control_parameters = self.num_control_parameters_single_leg*4 - + self.num_control_parameters = self.num_control_parameters_single_leg * 4 + # We have 4 different spline functions, one for each leg self.spline_fun_FL = self.compute_linear_spline_1 self.spline_fun_FR = self.compute_linear_spline_1 self.spline_fun_RL = self.compute_linear_spline_1 self.spline_fun_RR = self.compute_linear_spline_1 - - elif(self.control_parametrization == "linear_spline_2"): + + elif self.control_parametrization == "linear_spline_2": # Along the horizon, we have 2 splines per control input (3 forces) # Each spline has 2 parameters, but one is shared between the two splines - self.num_control_parameters_single_leg = 3*3 - + self.num_control_parameters_single_leg = 3 * 3 + # In totale we have 4 legs - self.num_control_parameters = self.num_control_parameters_single_leg*4 - + self.num_control_parameters = self.num_control_parameters_single_leg * 4 + # We have 4 different spline functions, one for each leg self.spline_fun_FL = self.compute_linear_spline_2 self.spline_fun_FR = self.compute_linear_spline_2 self.spline_fun_RL = self.compute_linear_spline_2 self.spline_fun_RR = self.compute_linear_spline_2 - elif(self.control_parametrization == "linear_spline_N"): + elif self.control_parametrization == "linear_spline_N": # Along the horizon, we have 2 splines per control input (3 forces) # Each spline has 2 parameters, but one is shared between the two splines - self.num_spline = config.mpc_params['num_splines'] - self.num_control_parameters_single_leg = (self.num_spline + 1)*3 - + self.num_spline = config.mpc_params["num_splines"] + self.num_control_parameters_single_leg = (self.num_spline + 1) * 3 + # In totale we have 4 legs - self.num_control_parameters = self.num_control_parameters_single_leg*4 - + self.num_control_parameters = self.num_control_parameters_single_leg * 4 + # We have 4 different spline functions, one for each leg self.spline_fun_FL = self.compute_linear_spline_N self.spline_fun_FR = self.compute_linear_spline_N self.spline_fun_RL = self.compute_linear_spline_N self.spline_fun_RR = self.compute_linear_spline_N - elif(self.control_parametrization == "cubic_spline_1"): + elif self.control_parametrization == "cubic_spline_1": # Along the horizon, we have 1 splines per control input (3 forces) # Each spline has 3 parameters - self.num_control_parameters_single_leg = 4*3 + self.num_control_parameters_single_leg = 4 * 3 # In totale we have 4 legs - self.num_control_parameters = self.num_control_parameters_single_leg*4 + self.num_control_parameters = self.num_control_parameters_single_leg * 4 # We have 4 different spline functions, one for each leg self.spline_fun_FL = self.compute_cubic_spline self.spline_fun_FR = self.compute_cubic_spline self.spline_fun_RL = self.compute_cubic_spline self.spline_fun_RR = self.compute_cubic_spline - - elif(self.control_parametrization == "cubic_spline_2"): + + elif self.control_parametrization == "cubic_spline_2": # Along the horizon, we have 1 splines per control input (3 forces) # Each spline has 3 parameters - self.num_control_parameters_single_leg = 4*3 + 4*3 + self.num_control_parameters_single_leg = 4 * 3 + 4 * 3 # In totale we have 4 legs - self.num_control_parameters = self.num_control_parameters_single_leg*4 + self.num_control_parameters = self.num_control_parameters_single_leg * 4 # We have 4 different spline functions, one for each leg self.spline_fun_FL = self.compute_cubic_spline_2 self.spline_fun_FR = self.compute_cubic_spline_2 self.spline_fun_RL = self.compute_cubic_spline_2 - self.spline_fun_RR = self.compute_cubic_spline_2 + self.spline_fun_RR = self.compute_cubic_spline_2 - elif(self.control_parametrization == "cubic_spline_N"): + elif self.control_parametrization == "cubic_spline_N": # Along the horizon, we have 1 splines per control input (3 forces) # Each spline has 3 parameters - self.num_spline = config.mpc_params['num_splines'] - self.num_control_parameters_single_leg = 4*3*self.num_spline + self.num_spline = config.mpc_params["num_splines"] + self.num_control_parameters_single_leg = 4 * 3 * self.num_spline # In totale we have 4 legs - self.num_control_parameters = self.num_control_parameters_single_leg*4 + self.num_control_parameters = self.num_control_parameters_single_leg * 4 # We have 4 different spline functions, one for each leg self.spline_fun_FL = self.compute_cubic_spline_N self.spline_fun_FR = self.compute_cubic_spline_N self.spline_fun_RL = self.compute_cubic_spline_N - self.spline_fun_RR = self.compute_cubic_spline_N + self.spline_fun_RR = self.compute_cubic_spline_N else: # We have 1 parameters for every 3 force direction (x,y,z)...for each time horizon!! - self.num_control_parameters_single_leg = self.horizon*3 - + self.num_control_parameters_single_leg = self.horizon * 3 + # In totale we have 4 legs - self.num_control_parameters = self.num_control_parameters_single_leg*4 + self.num_control_parameters = self.num_control_parameters_single_leg * 4 # We have 4 different spline functions, one for each leg self.spline_fun_FL = self.compute_zero_order_spline @@ -161,17 +164,17 @@ def __init__(self, horizon = 200, dt = 0.01, num_parallel_computations = 10000, self.spline_fun_RL = self.compute_zero_order_spline self.spline_fun_RR = self.compute_zero_order_spline - - - if(self.sampling_method == 'random_sampling'): + if self.sampling_method == "random_sampling": self.compute_control = self.compute_control_random_sampling - self.sigma_random_sampling = config.mpc_params['sigma_random_sampling'] - elif(self.sampling_method == 'mppi'): - self.compute_control = self.compute_control_mppi - self.sigma_mppi = config.mpc_params['sigma_mppi'] - elif(self.sampling_method == 'cem_mppi'): + self.sigma_random_sampling = config.mpc_params["sigma_random_sampling"] + elif self.sampling_method == "mppi": + self.compute_control = self.compute_control_mppi + self.sigma_mppi = config.mpc_params["sigma_mppi"] + elif self.sampling_method == "cem_mppi": self.compute_control = self.compute_control_cem_mppi - self.sigma_cem_mppi = jnp.ones(self.num_control_parameters, dtype=dtype_general) * config.mpc_params['sigma_cem_mppi'] + self.sigma_cem_mppi = ( + jnp.ones(self.num_control_parameters, dtype=dtype_general) * config.mpc_params["sigma_cem_mppi"] + ) else: # return error and stop execution print("Error: sampling method not recognized") @@ -179,284 +182,297 @@ def __init__(self, horizon = 200, dt = 0.01, num_parallel_computations = 10000, self.jitted_compute_control = jax.jit(self.compute_control, device=self.device) - # Initialize the robot model - self.robot = Centroidal_Model_JAX(self.dt,self.device) - + self.robot = Centroidal_Model_JAX(self.dt, self.device) # Initialize the cost function matrices - self.Q = jnp.identity(self.state_dim, dtype=dtype_general)*0 - self.Q = self.Q.at[0,0].set(0.0) - self.Q = self.Q.at[1,1].set(0.0) - self.Q = self.Q.at[2,2].set(111500) #com_z - self.Q = self.Q.at[3,3].set(5000) #com_vel_x - self.Q = self.Q.at[4,4].set(5000) #com_vel_y - self.Q = self.Q.at[5,5].set(200) #com_vel_z - self.Q = self.Q.at[6,6].set(11200) #base_angle_roll - self.Q = self.Q.at[7,7].set(11200) #base_angle_pitch - self.Q = self.Q.at[8,8].set(0.0) #base_angle_yaw - self.Q = self.Q.at[9,9].set(20) #base_angle_rates_x - self.Q = self.Q.at[10,10].set(20) #base_angle_rates_y - self.Q = self.Q.at[11,11].set(600) #base_angle_rates_z - - + self.Q = jnp.identity(self.state_dim, dtype=dtype_general) * 0 + self.Q = self.Q.at[0, 0].set(0.0) + self.Q = self.Q.at[1, 1].set(0.0) + self.Q = self.Q.at[2, 2].set(111500) # com_z + self.Q = self.Q.at[3, 3].set(5000) # com_vel_x + self.Q = self.Q.at[4, 4].set(5000) # com_vel_y + self.Q = self.Q.at[5, 5].set(200) # com_vel_z + self.Q = self.Q.at[6, 6].set(11200) # base_angle_roll + self.Q = self.Q.at[7, 7].set(11200) # base_angle_pitch + self.Q = self.Q.at[8, 8].set(0.0) # base_angle_yaw + self.Q = self.Q.at[9, 9].set(20) # base_angle_rates_x + self.Q = self.Q.at[10, 10].set(20) # base_angle_rates_y + self.Q = self.Q.at[11, 11].set(600) # base_angle_rates_z self.R = jnp.identity(self.control_dim, dtype=dtype_general) - self.R = self.R.at[0,0].set(0.0) #foot_pos_x_FL - self.R = self.R.at[1,1].set(0.0) #foot_pos_y_FL - self.R = self.R.at[2,2].set(0.0) #foot_pos_z_FL - self.R = self.R.at[3,3].set(0.0) #foot_pos_x_FR - self.R = self.R.at[4,4].set(0.0) #foot_pos_y_FR - self.R = self.R.at[5,5].set(0.0) #foot_pos_z_FR - self.R = self.R.at[6,6].set(0.0) #foot_pos_x_RL - self.R = self.R.at[7,7].set(0.0) #foot_pos_y_RL - self.R = self.R.at[8,8].set(0.0) #foot_pos_z_RL - self.R = self.R.at[9,9].set(0.0) #foot_pos_x_RR - self.R = self.R.at[10,10].set(0.0) #foot_pos_y_RR - self.R = self.R.at[11,11].set(0.0) #foot_pos_z_RR - - self.R = self.R.at[12,12].set(0.1) #foot_force_x_FL - self.R = self.R.at[13,13].set(0.1) #foot_force_y_FL - self.R = self.R.at[14,14].set(0.001) #foot_force_z_FL - self.R = self.R.at[15,15].set(0.1) #foot_force_x_FR - self.R = self.R.at[16,16].set(0.1) #foot_force_y_FR - self.R = self.R.at[17,17].set(0.001) #foot_force_z_FR - self.R = self.R.at[18,18].set(0.1) #foot_force_x_RL - self.R = self.R.at[19,19].set(0.1) #foot_force_y_RL - self.R = self.R.at[20,20].set(0.001) #foot_force_z_RL - self.R = self.R.at[21,21].set(0.1) #foot_force_x_RR - self.R = self.R.at[22,22].set(0.1) #foot_force_y_RR - self.R = self.R.at[23,23].set(0.001) #foot_force_z_RR - - # mu is the friction coefficient - self.mu = config.mpc_params['mu'] + self.R = self.R.at[0, 0].set(0.0) # foot_pos_x_FL + self.R = self.R.at[1, 1].set(0.0) # foot_pos_y_FL + self.R = self.R.at[2, 2].set(0.0) # foot_pos_z_FL + self.R = self.R.at[3, 3].set(0.0) # foot_pos_x_FR + self.R = self.R.at[4, 4].set(0.0) # foot_pos_y_FR + self.R = self.R.at[5, 5].set(0.0) # foot_pos_z_FR + self.R = self.R.at[6, 6].set(0.0) # foot_pos_x_RL + self.R = self.R.at[7, 7].set(0.0) # foot_pos_y_RL + self.R = self.R.at[8, 8].set(0.0) # foot_pos_z_RL + self.R = self.R.at[9, 9].set(0.0) # foot_pos_x_RR + self.R = self.R.at[10, 10].set(0.0) # foot_pos_y_RR + self.R = self.R.at[11, 11].set(0.0) # foot_pos_z_RR + + self.R = self.R.at[12, 12].set(0.1) # foot_force_x_FL + self.R = self.R.at[13, 13].set(0.1) # foot_force_y_FL + self.R = self.R.at[14, 14].set(0.001) # foot_force_z_FL + self.R = self.R.at[15, 15].set(0.1) # foot_force_x_FR + self.R = self.R.at[16, 16].set(0.1) # foot_force_y_FR + self.R = self.R.at[17, 17].set(0.001) # foot_force_z_FR + self.R = self.R.at[18, 18].set(0.1) # foot_force_x_RL + self.R = self.R.at[19, 19].set(0.1) # foot_force_y_RL + self.R = self.R.at[20, 20].set(0.001) # foot_force_z_RL + self.R = self.R.at[21, 21].set(0.1) # foot_force_x_RR + self.R = self.R.at[22, 22].set(0.1) # foot_force_y_RR + self.R = self.R.at[23, 23].set(0.001) # foot_force_z_RR - # maximum allowed z contact forces - self.f_z_max = config.mpc_params['grf_max'] - self.f_z_min = config.mpc_params['grf_min'] + # mu is the friction coefficient + self.mu = config.mpc_params["mu"] + # maximum allowed z contact forces + self.f_z_max = config.mpc_params["grf_max"] + self.f_z_min = config.mpc_params["grf_min"] self.best_control_parameters = jnp.zeros((self.num_control_parameters,), dtype=dtype_general) self.master_key = jax.random.PRNGKey(42) - self.initial_random_parameters = jax.random.uniform(key=self.master_key, minval=-self.max_sampling_forces, maxval=self.max_sampling_forces, shape=(self.num_parallel_computations, self.num_control_parameters )) - + self.initial_random_parameters = jax.random.uniform( + key=self.master_key, + minval=-self.max_sampling_forces, + maxval=self.max_sampling_forces, + shape=(self.num_parallel_computations, self.num_control_parameters), + ) - # jitting the vmap function! self.vectorized_rollout = jax.vmap(self.compute_rollout, in_axes=(None, None, 0, None), out_axes=0) self.jit_vectorized_rollout = jax.jit(self.vectorized_rollout, device=self.device) - # the first call of jax is very slow, hence we should do this since the beginning + # the first call of jax is very slow, hence we should do this since the beginning # creating a fake initial state, reference and contact sequence initial_state = jnp.zeros((self.state_dim,), dtype=dtype_general) initial_reference = jnp.zeros((self.reference_dim,), dtype=dtype_general) contact_sequence = jnp.ones((4, self.horizon), dtype=dtype_general) - - self.control_parameters_vec = random.uniform(self.master_key, (self.num_control_parameters*self.num_parallel_computations, ), minval=-100., maxval=100.) - self.jit_vectorized_rollout(initial_state, initial_reference, - self.control_parameters_vec.reshape(self.num_parallel_computations, self.num_control_parameters), - contact_sequence) + self.control_parameters_vec = random.uniform( + self.master_key, + (self.num_control_parameters * self.num_parallel_computations,), + minval=-100.0, + maxval=100.0, + ) + self.jit_vectorized_rollout( + initial_state, + initial_reference, + self.control_parameters_vec.reshape(self.num_parallel_computations, self.num_control_parameters), + contact_sequence, + ) - - - def compute_linear_spline_1(self, parameters, step, horizon_leg): """ Compute a linear spline parametrization of the GRF - """ + """ - tau = step/(horizon_leg) + tau = step / (horizon_leg) - q = (tau - 0.0)/(1.0-0.0) - + q = (tau - 0.0) / (1.0 - 0.0) index = 0 - f_x = (1-q)*parameters[index+0] + q*parameters[index+1] - f_y = (1-q)*parameters[index+2] + q*parameters[index+3] - f_z = (1-q)*parameters[index+4] + q*parameters[index+5] - - return f_x, f_y, f_z - + f_x = (1 - q) * parameters[index + 0] + q * parameters[index + 1] + f_y = (1 - q) * parameters[index + 2] + q * parameters[index + 3] + f_z = (1 - q) * parameters[index + 4] + q * parameters[index + 5] + return f_x, f_y, f_z def compute_linear_spline_2(self, parameters, step, horizon_leg): """ Compute the linear spline parametrization of the GRF (2 splines) - """ + """ + + tau = step / (horizon_leg) - tau = step/(horizon_leg) - index = 0 - index = jax.numpy.where(step > self.horizon/2, 1, index) + index = jax.numpy.where(step > self.horizon / 2, 1, index) - tau = jax.numpy.where(step > self.horizon/2, tau-1, tau) - q = (tau - 0.0)/(1.0-0.0) + tau = jax.numpy.where(step > self.horizon / 2, tau - 1, tau) + q = (tau - 0.0) / (1.0 - 0.0) - - f_x = (1-q)*parameters[index+0] + q*parameters[index+1] - f_y = (1-q)*parameters[index+3] + q*parameters[index+4] - f_z = (1-q)*parameters[index+6] + q*parameters[index+7] + f_x = (1 - q) * parameters[index + 0] + q * parameters[index + 1] + f_y = (1 - q) * parameters[index + 3] + q * parameters[index + 4] + f_z = (1 - q) * parameters[index + 6] + q * parameters[index + 7] - return f_x, f_y, f_z - + return f_x, f_y, f_z def compute_linear_spline_N(self, parameters, step, horizon_leg): """ Compute the linear spline parametrization of the GRF (N splines) - """ + """ + + tau = step / (horizon_leg) - tau = step/(horizon_leg) - # Adding the last boundary for the case when step is exactly self.horizon chunk_boundaries = jnp.linspace(0, self.horizon, self.num_spline + 1) # Find the chunk index by checking in which interval the step falls index = jnp.max(jnp.where(step >= chunk_boundaries, jnp.arange(self.num_spline + 1), 0)) - tau = jax.numpy.where(jnp.isin(step, chunk_boundaries[1:]), tau - 1, tau) - q = (tau - 0.0)/(1.0-0.0) + q = (tau - 0.0) / (1.0 - 0.0) - - shift = (self.num_spline + 1) - f_x = (1-q)*parameters[index+0] + q*parameters[index+1] - f_y = (1-q)*parameters[index+shift] + q*parameters[index+shift+1] - f_z = (1-q)*parameters[index+shift*2] + q*parameters[index+shift*2+1] + shift = self.num_spline + 1 + f_x = (1 - q) * parameters[index + 0] + q * parameters[index + 1] + f_y = (1 - q) * parameters[index + shift] + q * parameters[index + shift + 1] + f_z = (1 - q) * parameters[index + shift * 2] + q * parameters[index + shift * 2 + 1] - return f_x, f_y, f_z - + return f_x, f_y, f_z - - - def compute_cubic_spline(self, parameters, step, horizon_leg): """ Compute the cubic spline parametrization of the GRF - """ + """ - tau = step/(horizon_leg) - q = (tau - 0.0)/(1.0-0.0) - a = 2*q*q*q - 3*q*q + 1 - b = (q*q*q - 2*q*q + q)*1.0 - c = -2*q*q*q + 3*q*q - d = (q*q*q - q*q)*1.0 + tau = step / (horizon_leg) + q = (tau - 0.0) / (1.0 - 0.0) + a = 2 * q * q * q - 3 * q * q + 1 + b = (q * q * q - 2 * q * q + q) * 1.0 + c = -2 * q * q * q + 3 * q * q + d = (q * q * q - q * q) * 1.0 - phi = (1./2.)*(((parameters[2] - parameters[1])/1.0) + ((parameters[1] - parameters[0])/1.0)) - phi_next = (1./2.)*(((parameters[3] - parameters[2])/1.0) + ((parameters[2] - parameters[1])/1.0)) - f_x = a*parameters[1] + b*phi + c*parameters[2] + d*phi_next + phi = (1.0 / 2.0) * (((parameters[2] - parameters[1]) / 1.0) + ((parameters[1] - parameters[0]) / 1.0)) + phi_next = (1.0 / 2.0) * (((parameters[3] - parameters[2]) / 1.0) + ((parameters[2] - parameters[1]) / 1.0)) + f_x = a * parameters[1] + b * phi + c * parameters[2] + d * phi_next - phi = (1./2.)*(((parameters[6] - parameters[5])/1.0) + ((parameters[5] - parameters[4])/1.0)) - phi_next = (1./2.)*(((parameters[7] - parameters[6])/1.0) + ((parameters[6] - parameters[5])/1.0)) - f_y = a*parameters[5] + b*phi + c*parameters[6] + d*phi_next + phi = (1.0 / 2.0) * (((parameters[6] - parameters[5]) / 1.0) + ((parameters[5] - parameters[4]) / 1.0)) + phi_next = (1.0 / 2.0) * (((parameters[7] - parameters[6]) / 1.0) + ((parameters[6] - parameters[5]) / 1.0)) + f_y = a * parameters[5] + b * phi + c * parameters[6] + d * phi_next + phi = (1.0 / 2.0) * (((parameters[10] - parameters[9]) / 1.0) + ((parameters[9] - parameters[8]) / 1.0)) + phi_next = (1.0 / 2.0) * (((parameters[11] - parameters[10]) / 1.0) + ((parameters[10] - parameters[9]) / 1.0)) + f_z = a * parameters[9] + b * phi + c * parameters[10] + d * phi_next - phi = (1./2.)*(((parameters[10] - parameters[9])/1.0) + ((parameters[9] - parameters[8])/1.0)) - phi_next = (1./2.)*(((parameters[11] - parameters[10])/1.0) + ((parameters[10] - parameters[9])/1.0)) - f_z = a*parameters[9] + b*phi + c*parameters[10] + d*phi_next - return f_x, f_y, f_z - def compute_cubic_spline_2(self, parameters, step, horizon_leg): """ Compute the cubic spline parametrization of the GRF - """ + """ - tau = step/(horizon_leg) + tau = step / (horizon_leg) index = 0 - index = jax.numpy.where(step > self.horizon/2, 1, index) - tau = jax.numpy.where(step > self.horizon/2, tau-1, tau) - - start_index = 10*index - + index = jax.numpy.where(step > self.horizon / 2, 1, index) + tau = jax.numpy.where(step > self.horizon / 2, tau - 1, tau) + + start_index = 10 * index + + q = (tau - 0.0) / (1.0 - 0.0) + a = 2 * q * q * q - 3 * q * q + 1 + b = (q * q * q - 2 * q * q + q) * 1.0 + c = -2 * q * q * q + 3 * q * q + d = (q * q * q - q * q) * 1.0 + + phi = (1.0 / 2.0) * ( + ((parameters[start_index + 2] - parameters[start_index + 1]) / 1.0) + + ((parameters[start_index + 1] - parameters[start_index + 0]) / 1.0) + ) + phi_next = (1.0 / 2.0) * ( + ((parameters[start_index + 3] - parameters[start_index + 2]) / 1.0) + + ((parameters[start_index + 2] - parameters[start_index + 1]) / 1.0) + ) + f_x = a * parameters[start_index + 1] + b * phi + c * parameters[start_index + 2] + d * phi_next + + phi = (1.0 / 2.0) * ( + ((parameters[start_index + 6] - parameters[start_index + 5]) / 1.0) + + ((parameters[start_index + 5] - parameters[start_index + 4]) / 1.0) + ) + phi_next = (1.0 / 2.0) * ( + ((parameters[start_index + 7] - parameters[start_index + 6]) / 1.0) + + ((parameters[start_index + 6] - parameters[start_index + 5]) / 1.0) + ) + f_y = a * parameters[start_index + 5] + b * phi + c * parameters[start_index + 6] + d * phi_next + + phi = (1.0 / 2.0) * ( + ((parameters[start_index + 10] - parameters[start_index + 9]) / 1.0) + + ((parameters[start_index + 9] - parameters[start_index + 8]) / 1.0) + ) + phi_next = (1.0 / 2.0) * ( + ((parameters[start_index + 11] - parameters[start_index + 10]) / 1.0) + + ((parameters[start_index + 10] - parameters[start_index + 9]) / 1.0) + ) + f_z = a * parameters[start_index + 9] + b * phi + c * parameters[start_index + 10] + d * phi_next - q = (tau - 0.0)/(1.0-0.0) - a = 2*q*q*q - 3*q*q + 1 - b = (q*q*q - 2*q*q + q)*1.0 - c = -2*q*q*q + 3*q*q - d = (q*q*q - q*q)*1.0 - - phi = (1./2.)*(((parameters[start_index+2] - parameters[start_index+1])/1.0) + ((parameters[start_index+1] - parameters[start_index+0])/1.0)) - phi_next = (1./2.)*(((parameters[start_index+3] - parameters[start_index+2])/1.0) + ((parameters[start_index+2] - parameters[start_index+1])/1.0)) - f_x = a*parameters[start_index+1] + b*phi + c*parameters[start_index+2] + d*phi_next - - phi = (1./2.)*(((parameters[start_index+6] - parameters[start_index+5])/1.0) + ((parameters[start_index+5] - parameters[start_index+4])/1.0)) - phi_next = (1./2.)*(((parameters[start_index+7] - parameters[start_index+6])/1.0) + ((parameters[start_index+6] - parameters[start_index+5])/1.0)) - f_y = a*parameters[start_index+5] + b*phi + c*parameters[start_index+6] + d*phi_next - - - phi = (1./2.)*(((parameters[start_index+10] - parameters[start_index+9])/1.0) + ((parameters[start_index+9] - parameters[start_index+8])/1.0)) - phi_next = (1./2.)*(((parameters[start_index+11] - parameters[start_index+10])/1.0) + ((parameters[start_index+10] - parameters[start_index+9])/1.0)) - f_z = a*parameters[start_index+9] + b*phi + c*parameters[start_index+10] + d*phi_next - return f_x, f_y, f_z - - def compute_cubic_spline_N(self, parameters, step, horizon_leg): """ Compute the cubic spline parametrization of the GRF (N splines) - """ + """ + + tau = step / (horizon_leg) - tau = step/(horizon_leg) - # Adding the last boundary for the case when step is exactly self.horizon chunk_boundaries = jnp.linspace(0, self.horizon, self.num_spline + 1) # Find the chunk index by checking in which interval the step falls index = jnp.max(jnp.where(step >= chunk_boundaries, jnp.arange(self.num_spline + 1), 0)) - tau = jax.numpy.where(jnp.isin(step, chunk_boundaries[1:]), tau - 1, tau) - q = (tau - 0.0)/(1.0-0.0) - - start_index = 10*index - + q = (tau - 0.0) / (1.0 - 0.0) + + start_index = 10 * index + + q = (tau - 0.0) / (1.0 - 0.0) + a = 2 * q * q * q - 3 * q * q + 1 + b = (q * q * q - 2 * q * q + q) * 1.0 + c = -2 * q * q * q + 3 * q * q + d = (q * q * q - q * q) * 1.0 + + phi = (1.0 / 2.0) * ( + ((parameters[start_index + 2] - parameters[start_index + 1]) / 1.0) + + ((parameters[start_index + 1] - parameters[start_index + 0]) / 1.0) + ) + phi_next = (1.0 / 2.0) * ( + ((parameters[start_index + 3] - parameters[start_index + 2]) / 1.0) + + ((parameters[start_index + 2] - parameters[start_index + 1]) / 1.0) + ) + f_x = a * parameters[start_index + 1] + b * phi + c * parameters[start_index + 2] + d * phi_next + + phi = (1.0 / 2.0) * ( + ((parameters[start_index + 6] - parameters[start_index + 5]) / 1.0) + + ((parameters[start_index + 5] - parameters[start_index + 4]) / 1.0) + ) + phi_next = (1.0 / 2.0) * ( + ((parameters[start_index + 7] - parameters[start_index + 6]) / 1.0) + + ((parameters[start_index + 6] - parameters[start_index + 5]) / 1.0) + ) + f_y = a * parameters[start_index + 5] + b * phi + c * parameters[start_index + 6] + d * phi_next + + phi = (1.0 / 2.0) * ( + ((parameters[start_index + 10] - parameters[start_index + 9]) / 1.0) + + ((parameters[start_index + 9] - parameters[start_index + 8]) / 1.0) + ) + phi_next = (1.0 / 2.0) * ( + ((parameters[start_index + 11] - parameters[start_index + 10]) / 1.0) + + ((parameters[start_index + 10] - parameters[start_index + 9]) / 1.0) + ) + f_z = a * parameters[start_index + 9] + b * phi + c * parameters[start_index + 10] + d * phi_next - q = (tau - 0.0)/(1.0-0.0) - a = 2*q*q*q - 3*q*q + 1 - b = (q*q*q - 2*q*q + q)*1.0 - c = -2*q*q*q + 3*q*q - d = (q*q*q - q*q)*1.0 - - phi = (1./2.)*(((parameters[start_index+2] - parameters[start_index+1])/1.0) + ((parameters[start_index+1] - parameters[start_index+0])/1.0)) - phi_next = (1./2.)*(((parameters[start_index+3] - parameters[start_index+2])/1.0) + ((parameters[start_index+2] - parameters[start_index+1])/1.0)) - f_x = a*parameters[start_index+1] + b*phi + c*parameters[start_index+2] + d*phi_next - - phi = (1./2.)*(((parameters[start_index+6] - parameters[start_index+5])/1.0) + ((parameters[start_index+5] - parameters[start_index+4])/1.0)) - phi_next = (1./2.)*(((parameters[start_index+7] - parameters[start_index+6])/1.0) + ((parameters[start_index+6] - parameters[start_index+5])/1.0)) - f_y = a*parameters[start_index+5] + b*phi + c*parameters[start_index+6] + d*phi_next - - - phi = (1./2.)*(((parameters[start_index+10] - parameters[start_index+9])/1.0) + ((parameters[start_index+9] - parameters[start_index+8])/1.0)) - phi_next = (1./2.)*(((parameters[start_index+11] - parameters[start_index+10])/1.0) + ((parameters[start_index+10] - parameters[start_index+9])/1.0)) - f_z = a*parameters[start_index+9] + b*phi + c*parameters[start_index+10] + d*phi_next - return f_x, f_y, f_z - - def compute_zero_order_spline(self, parameters, step, horizon_leg): """ Compute the zero-order control input - """ - + """ + index = jnp.int16(step) f_x = parameters[index] - f_y = parameters[index+self.horizon] - f_z = parameters[index+self.horizon*2] - return f_x, f_y, f_z - - + f_y = parameters[index + self.horizon] + f_z = parameters[index + self.horizon * 2] + return f_x, f_y, f_z - def enforce_force_constraints(self, f_x_FL, f_y_FL, f_z_FL, - f_x_FR, f_y_FR, f_z_FR, - f_x_RL, f_y_RL, f_z_RL, - f_x_RR, f_y_RR, f_z_RR): + def enforce_force_constraints( + self, f_x_FL, f_y_FL, f_z_FL, f_x_FR, f_y_FR, f_z_FR, f_x_RL, f_y_RL, f_z_RL, f_x_RR, f_y_RR, f_z_RR + ): """ Enforce the friction cone and the force limits constraints - """ - + """ + # Enforce push-only of the ground! f_z_FL = jax.numpy.where(f_z_FL > self.f_z_min, f_z_FL, self.f_z_min) f_z_FR = jax.numpy.where(f_z_FR > self.f_z_min, f_z_FR, self.f_z_min) @@ -464,45 +480,37 @@ def enforce_force_constraints(self, f_x_FL, f_y_FL, f_z_FL, f_z_RR = jax.numpy.where(f_z_RR > self.f_z_min, f_z_RR, self.f_z_min) # Enforce maximum force per leg! - f_z_FL = jax.numpy.where(f_z_FL-self.mu*f_z_FL, f_x_FL,-self.mu*f_z_FL) - f_x_FL = jax.numpy.where(f_x_FL-self.mu*f_z_FL, f_y_FL,-self.mu*f_z_FL) - f_y_FL = jax.numpy.where(f_y_FL-self.mu*f_z_FR, f_x_FR,-self.mu*f_z_FR) - f_x_FR = jax.numpy.where(f_x_FR-self.mu*f_z_FR, f_y_FR,-self.mu*f_z_FR) - f_y_FR = jax.numpy.where(f_y_FR-self.mu*f_z_RL, f_x_RL,-self.mu*f_z_RL) - f_x_RL = jax.numpy.where(f_x_RL-self.mu*f_z_RL, f_y_RL,-self.mu*f_z_RL) - f_y_RL = jax.numpy.where(f_y_RL-self.mu*f_z_RR, f_x_RR,-self.mu*f_z_RR) - f_x_RR = jax.numpy.where(f_x_RR-self.mu*f_z_RR, f_y_RR,-self.mu*f_z_RR) - f_y_RR = jax.numpy.where(f_y_RR -self.mu * f_z_FL, f_x_FL, -self.mu * f_z_FL) + f_x_FL = jax.numpy.where(f_x_FL < self.mu * f_z_FL, f_x_FL, self.mu * f_z_FL) + f_y_FL = jax.numpy.where(f_y_FL > -self.mu * f_z_FL, f_y_FL, -self.mu * f_z_FL) + f_y_FL = jax.numpy.where(f_y_FL < self.mu * f_z_FL, f_y_FL, self.mu * f_z_FL) + + f_x_FR = jax.numpy.where(f_x_FR > -self.mu * f_z_FR, f_x_FR, -self.mu * f_z_FR) + f_x_FR = jax.numpy.where(f_x_FR < self.mu * f_z_FR, f_x_FR, self.mu * f_z_FR) + f_y_FR = jax.numpy.where(f_y_FR > -self.mu * f_z_FR, f_y_FR, -self.mu * f_z_FR) + f_y_FR = jax.numpy.where(f_y_FR < self.mu * f_z_FR, f_y_FR, self.mu * f_z_FR) + + f_x_RL = jax.numpy.where(f_x_RL > -self.mu * f_z_RL, f_x_RL, -self.mu * f_z_RL) + f_x_RL = jax.numpy.where(f_x_RL < self.mu * f_z_RL, f_x_RL, self.mu * f_z_RL) + f_y_RL = jax.numpy.where(f_y_RL > -self.mu * f_z_RL, f_y_RL, -self.mu * f_z_RL) + f_y_RL = jax.numpy.where(f_y_RL < self.mu * f_z_RL, f_y_RL, self.mu * f_z_RL) + + f_x_RR = jax.numpy.where(f_x_RR > -self.mu * f_z_RR, f_x_RR, -self.mu * f_z_RR) + f_x_RR = jax.numpy.where(f_x_RR < self.mu * f_z_RR, f_x_RR, self.mu * f_z_RR) + f_y_RR = jax.numpy.where(f_y_RR > -self.mu * f_z_RR, f_y_RR, -self.mu * f_z_RR) + f_y_RR = jax.numpy.where(f_y_RR < self.mu * f_z_RR, f_y_RR, self.mu * f_z_RR) + + return f_x_FL, f_y_FL, f_z_FL, f_x_FR, f_y_FR, f_z_FR, f_x_RL, f_y_RL, f_z_RL, f_x_RR, f_y_RR, f_z_RR def compute_rollout(self, initial_state, reference, control_parameters, contact_sequence): """Calculate cost of a rollout of the dynamics given random parameters @@ -513,17 +521,16 @@ def compute_rollout(self, initial_state, reference, control_parameters, contact_ parameters (np.array): parameters for the simplified dynamics Returns: (float): cost of the rollout - """ + """ state = initial_state cost = jnp.float32(0.0) - n_ = jnp.array([-1,-1,-1,-1]) + n_ = jnp.array([-1, -1, -1, -1]) - - #FL_num_of_contact = jnp.sum(contact_sequence[0])+1 - #FR_num_of_contact = jnp.sum(contact_sequence[1])+1 - #RL_num_of_contact = jnp.sum(contact_sequence[2])+1 - #RR_num_of_contact = jnp.sum(contact_sequence[3])+1 + # FL_num_of_contact = jnp.sum(contact_sequence[0])+1 + # FR_num_of_contact = jnp.sum(contact_sequence[1])+1 + # RL_num_of_contact = jnp.sum(contact_sequence[2])+1 + # RR_num_of_contact = jnp.sum(contact_sequence[3])+1 FL_num_of_contact = self.horizon FR_num_of_contact = self.horizon @@ -533,291 +540,368 @@ def compute_rollout(self, initial_state, reference, control_parameters, contact_ def iterate_fun(n, carry): cost, state, reference, n_ = carry - - - #n_ = n_.at[0].set(n_[0]+1*contact_sequence[0][n]) - #n_ = n_.at[1].set(n_[1]+1*contact_sequence[1][n]) - #n_ = n_.at[2].set(n_[2]+1*contact_sequence[2][n]) - #n_ = n_.at[3].set(n_[3]+1*contact_sequence[3][n]) + # n_ = n_.at[0].set(n_[0]+1*contact_sequence[0][n]) + # n_ = n_.at[1].set(n_[1]+1*contact_sequence[1][n]) + # n_ = n_.at[2].set(n_[2]+1*contact_sequence[2][n]) + # n_ = n_.at[3].set(n_[3]+1*contact_sequence[3][n]) n_ = n_.at[0].set(n) n_ = n_.at[1].set(n) n_ = n_.at[2].set(n) n_ = n_.at[3].set(n) - - f_x_FL, f_y_FL, f_z_FL = self.spline_fun_FL(control_parameters[0:self.num_control_parameters_single_leg], n_[0], FL_num_of_contact) - f_x_FR, f_y_FR, f_z_FR = self.spline_fun_FR(control_parameters[self.num_control_parameters_single_leg:self.num_control_parameters_single_leg*2], n_[1], FR_num_of_contact) - f_x_RL, f_y_RL, f_z_RL = self.spline_fun_RL(control_parameters[self.num_control_parameters_single_leg*2:self.num_control_parameters_single_leg*3], n_[2], RL_num_of_contact) - f_x_RR, f_y_RR, f_z_RR = self.spline_fun_RR(control_parameters[self.num_control_parameters_single_leg*3:self.num_control_parameters_single_leg*4], n_[3], RR_num_of_contact) - - - - + f_x_FL, f_y_FL, f_z_FL = self.spline_fun_FL( + control_parameters[0 : self.num_control_parameters_single_leg], n_[0], FL_num_of_contact + ) + f_x_FR, f_y_FR, f_z_FR = self.spline_fun_FR( + control_parameters[self.num_control_parameters_single_leg : self.num_control_parameters_single_leg * 2], + n_[1], + FR_num_of_contact, + ) + f_x_RL, f_y_RL, f_z_RL = self.spline_fun_RL( + control_parameters[ + self.num_control_parameters_single_leg * 2 : self.num_control_parameters_single_leg * 3 + ], + n_[2], + RL_num_of_contact, + ) + f_x_RR, f_y_RR, f_z_RR = self.spline_fun_RR( + control_parameters[ + self.num_control_parameters_single_leg * 3 : self.num_control_parameters_single_leg * 4 + ], + n_[3], + RR_num_of_contact, + ) # The sampling over f_z is a delta over gravity compensation (only for the leg in stance!) - number_of_legs_in_stance = contact_sequence[0][n] + contact_sequence[1][n] + contact_sequence[2][n] + contact_sequence[3][n] + number_of_legs_in_stance = ( + contact_sequence[0][n] + contact_sequence[1][n] + contact_sequence[2][n] + contact_sequence[3][n] + ) reference_force_stance_legs = (self.robot.mass * 9.81) / number_of_legs_in_stance - f_z_FL = reference_force_stance_legs + f_z_FL f_z_FR = reference_force_stance_legs + f_z_FR f_z_RL = reference_force_stance_legs + f_z_RL f_z_RR = reference_force_stance_legs + f_z_RR - # Foot in swing (contact sequence = 0) have zero force - f_x_FL = f_x_FL*contact_sequence[0][n] - f_y_FL = f_y_FL*contact_sequence[0][n] - f_z_FL = f_z_FL*contact_sequence[0][n] - - f_x_FR = f_x_FR*contact_sequence[1][n] - f_y_FR = f_y_FR*contact_sequence[1][n] - f_z_FR = f_z_FR*contact_sequence[1][n] + f_x_FL = f_x_FL * contact_sequence[0][n] + f_y_FL = f_y_FL * contact_sequence[0][n] + f_z_FL = f_z_FL * contact_sequence[0][n] - f_x_RL = f_x_RL*contact_sequence[2][n] - f_y_RL = f_y_RL*contact_sequence[2][n] - f_z_RL = f_z_RL*contact_sequence[2][n] + f_x_FR = f_x_FR * contact_sequence[1][n] + f_y_FR = f_y_FR * contact_sequence[1][n] + f_z_FR = f_z_FR * contact_sequence[1][n] - f_x_RR = f_x_RR*contact_sequence[3][n] - f_y_RR = f_y_RR*contact_sequence[3][n] - f_z_RR = f_z_RR*contact_sequence[3][n] + f_x_RL = f_x_RL * contact_sequence[2][n] + f_y_RL = f_y_RL * contact_sequence[2][n] + f_z_RL = f_z_RL * contact_sequence[2][n] + f_x_RR = f_x_RR * contact_sequence[3][n] + f_y_RR = f_y_RR * contact_sequence[3][n] + f_z_RR = f_z_RR * contact_sequence[3][n] # Enforce force constraints - f_x_FL, f_y_FL, f_z_FL, \ - f_x_FR, f_y_FR, f_z_FR, \ - f_x_RL, f_y_RL, f_z_RL, \ - f_x_RR, f_y_RR, f_z_RR = self.enforce_force_constraints(f_x_FL, f_y_FL, f_z_FL, - f_x_FR, f_y_FR, f_z_FR, - f_x_RL, f_y_RL, f_z_RL, - f_x_RR, f_y_RR, f_z_RR) - - - - - input = jnp.array([jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - f_x_FL, f_y_FL, f_z_FL, # foot position fl - f_x_FR, f_y_FR, f_z_FR, # foot position fr - f_x_RL, f_y_RL, f_z_RL, # foot position rl - f_x_RR, f_y_RR, f_z_RR, # foot position rr - ], dtype=dtype_general) - - + f_x_FL, f_y_FL, f_z_FL, f_x_FR, f_y_FR, f_z_FR, f_x_RL, f_y_RL, f_z_RL, f_x_RR, f_y_RR, f_z_RR = ( + self.enforce_force_constraints( + f_x_FL, f_y_FL, f_z_FL, f_x_FR, f_y_FR, f_z_FR, f_x_RL, f_y_RL, f_z_RL, f_x_RR, f_y_RR, f_z_RR + ) + ) + + input = jnp.array( + [ + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + f_x_FL, + f_y_FL, + f_z_FL, # foot position fl + f_x_FR, + f_y_FR, + f_z_FR, # foot position fr + f_x_RL, + f_y_RL, + f_z_RL, # foot position rl + f_x_RR, + f_y_RR, + f_z_RR, # foot position rr + ], + dtype=dtype_general, + ) + # Integrate the dynamics - current_contact = jnp.array([contact_sequence[0][n], contact_sequence[1][n], - contact_sequence[2][n], contact_sequence[3][n]], dtype=dtype_general) + current_contact = jnp.array( + [contact_sequence[0][n], contact_sequence[1][n], contact_sequence[2][n], contact_sequence[3][n]], + dtype=dtype_general, + ) state_next = self.robot.integrate_jax(state, input, current_contact, n) - - - # Compute the cost - - # Calculate cost regulation state - state_error = state_next - reference[0:self.state_dim] - error_cost = state_error.T@self.Q@state_error - input_for_cost = jnp.array([jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - f_x_FL, f_y_FL, f_z_FL - reference_force_stance_legs, - f_x_FR, f_y_FR, f_z_FR - reference_force_stance_legs, - f_x_RL, f_y_RL, f_z_RL - reference_force_stance_legs, - f_x_RR, f_y_RR, f_z_RR - reference_force_stance_legs, - ], dtype=dtype_general) + # Compute the cost + # Calculate cost regulation state + state_error = state_next - reference[0 : self.state_dim] + error_cost = state_error.T @ self.Q @ state_error + + input_for_cost = jnp.array( + [ + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + f_x_FL, + f_y_FL, + f_z_FL - reference_force_stance_legs, + f_x_FR, + f_y_FR, + f_z_FR - reference_force_stance_legs, + f_x_RL, + f_y_RL, + f_z_RL - reference_force_stance_legs, + f_x_RR, + f_y_RR, + f_z_RR - reference_force_stance_legs, + ], + dtype=dtype_general, + ) # Calculate cost regulation input - #error_cost += input_for_cost.T@self.R@input_for_cost - - - + # error_cost += input_for_cost.T@self.R@input_for_cost # saturate in the case of NaN - #state_next = jnp.where(jnp.isnan(state_next), 100, state_next) - #error_cost = jnp.where(jnp.isnan(error_cost), 1000, error_cost) - #cost = jnp.where(jnp.isnan(cost), 10000, cost) - - + # state_next = jnp.where(jnp.isnan(state_next), 100, state_next) + # error_cost = jnp.where(jnp.isnan(error_cost), 1000, error_cost) + # cost = jnp.where(jnp.isnan(cost), 10000, cost) return (cost + error_cost, state_next, reference, n_) carry = (cost, state, reference, n_) cost, state, reference, n_ = jax.lax.fori_loop(0, self.horizon, iterate_fun, carry) - - return cost - - + return cost def with_newkey(self): newkey, subkey = jax.random.split(self.master_key) self.master_key = newkey return self - + def get_key(self): return self.master_key - + def with_newsigma(self, sigma): self.sigma_cem_mppi = sigma return self - + def get_sigma(self): return self.sigma_cem_mppi - - - def shift_solution(self, best_control_parameters, step): """ This function shift the control parameter ahed - """ - + """ best_control_parameters = np.array(best_control_parameters) - FL_control = copy.deepcopy(best_control_parameters[0:self.num_control_parameters_single_leg]) - FR_control = copy.deepcopy(best_control_parameters[self.num_control_parameters_single_leg:self.num_control_parameters_single_leg*2]) - RL_control = copy.deepcopy(best_control_parameters[self.num_control_parameters_single_leg*2:self.num_control_parameters_single_leg*3]) - RR_control = copy.deepcopy(best_control_parameters[self.num_control_parameters_single_leg*3:self.num_control_parameters_single_leg*4]) - + FL_control = copy.deepcopy(best_control_parameters[0 : self.num_control_parameters_single_leg]) + FR_control = copy.deepcopy( + best_control_parameters[self.num_control_parameters_single_leg : self.num_control_parameters_single_leg * 2] + ) + RL_control = copy.deepcopy( + best_control_parameters[ + self.num_control_parameters_single_leg * 2 : self.num_control_parameters_single_leg * 3 + ] + ) + RR_control = copy.deepcopy( + best_control_parameters[ + self.num_control_parameters_single_leg * 3 : self.num_control_parameters_single_leg * 4 + ] + ) FL_control_temp = copy.deepcopy(FL_control) FL_control[0], FL_control[2], FL_control[4] = self.spline_fun_FL(FL_control_temp, step) - #FL_control[1], FL_control[3], FL_control[5] = controller.spline_fun_FL(FL_control_temp, controller.horizon+step) + # FL_control[1], FL_control[3], FL_control[5] = controller.spline_fun_FL(FL_control_temp, controller.horizon+step) FR_control_temp = copy.deepcopy(FR_control) FR_control[0], FR_control[2], FR_control[4] = self.spline_fun_FR(FR_control_temp, step) - #FR_control[1], FR_control[3], FR_control[5] = controller.spline_fun_FR(FR_control_temp, controller.horizon+step) - - + # FR_control[1], FR_control[3], FR_control[5] = controller.spline_fun_FR(FR_control_temp, controller.horizon+step) + RL_control_temp = copy.deepcopy(RL_control) RL_control[0], RL_control[2], RL_control[4] = self.spline_fun_RL(RL_control_temp, step) - #RL_control[1], RL_control[3], RL_control[5] = controller.spline_fun_RL(RL_control_temp, controller.horizon+step) - - + # RL_control[1], RL_control[3], RL_control[5] = controller.spline_fun_RL(RL_control_temp, controller.horizon+step) + RR_control_temp = copy.deepcopy(RR_control) RR_control[0], RR_control[2], RR_control[4] = self.spline_fun_RR(RR_control_temp, step) - #RR_control[1], RR_control[3], RR_control[5] = controller.spline_fun_RR(RR_control_temp, controller.horizon+step) - - - best_control_parameters[0:self.num_control_parameters_single_leg] = FL_control - best_control_parameters[self.num_control_parameters_single_leg:self.num_control_parameters_single_leg*2] = FR_control - best_control_parameters[self.num_control_parameters_single_leg*2:self.num_control_parameters_single_leg*3] = RL_control - best_control_parameters[self.num_control_parameters_single_leg*3:self.num_control_parameters_single_leg*4] = RR_control - - return best_control_parameters - + # RR_control[1], RR_control[3], RR_control[5] = controller.spline_fun_RR(RR_control_temp, controller.horizon+step) + + best_control_parameters[0 : self.num_control_parameters_single_leg] = FL_control + best_control_parameters[self.num_control_parameters_single_leg : self.num_control_parameters_single_leg * 2] = ( + FR_control + ) + best_control_parameters[ + self.num_control_parameters_single_leg * 2 : self.num_control_parameters_single_leg * 3 + ] = RL_control + best_control_parameters[ + self.num_control_parameters_single_leg * 3 : self.num_control_parameters_single_leg * 4 + ] = RR_control + return best_control_parameters - def prepare_state_and_reference(self, state_current, reference_state, - current_contact, previous_contact, mpc_frequency=100): + def prepare_state_and_reference( + self, state_current, reference_state, current_contact, previous_contact, mpc_frequency=100 + ): """ This function jaxify the current state and reference for further processing. - """ + """ # Shift the previous solution ahead - if (config.mpc_params['shift_solution']): - index_shift = 1./mpc_frequency + if config.mpc_params["shift_solution"]: + index_shift = 1.0 / mpc_frequency self.best_control_parameters = self.shift_solution(self.best_control_parameters, index_shift) - - - state_current_jax = np.concatenate((state_current["position"], state_current["linear_velocity"], - state_current["orientation"], state_current["angular_velocity"], - state_current["foot_FL"], state_current["foot_FR"], - state_current["foot_RL"], state_current["foot_RR"])).reshape((24, )) - - if(current_contact[0] == 0.): + state_current_jax = np.concatenate( + ( + state_current["position"], + state_current["linear_velocity"], + state_current["orientation"], + state_current["angular_velocity"], + state_current["foot_FL"], + state_current["foot_FR"], + state_current["foot_RL"], + state_current["foot_RR"], + ) + ).reshape((24,)) + + if current_contact[0] == 0.0: state_current_jax[12:15] = reference_state["ref_foot_FL"].reshape((3,)) - if(current_contact[1] == 0.): + if current_contact[1] == 0.0: state_current_jax[15:18] = reference_state["ref_foot_FR"].reshape((3,)) - if(current_contact[2] == 0.): - state_current_jax[18:21] = reference_state["ref_foot_RL"].reshape((3,)) - if(current_contact[3] == 0.): + if current_contact[2] == 0.0: + state_current_jax[18:21] = reference_state["ref_foot_RL"].reshape((3,)) + if current_contact[3] == 0.0: state_current_jax[21:24] = reference_state["ref_foot_RR"].reshape((3,)) - - reference_state_jax = np.concatenate((reference_state["ref_position"], reference_state["ref_linear_velocity"], - reference_state["ref_orientation"], reference_state["ref_angular_velocity"], - reference_state["ref_foot_FL"].reshape((3,)), reference_state["ref_foot_FR"].reshape((3,)), - reference_state["ref_foot_RL"].reshape((3,)), reference_state["ref_foot_RR"].reshape((3,)))).reshape((24, )) - + reference_state_jax = np.concatenate( + ( + reference_state["ref_position"], + reference_state["ref_linear_velocity"], + reference_state["ref_orientation"], + reference_state["ref_angular_velocity"], + reference_state["ref_foot_FL"].reshape((3,)), + reference_state["ref_foot_FR"].reshape((3,)), + reference_state["ref_foot_RL"].reshape((3,)), + reference_state["ref_foot_RR"].reshape((3,)), + ) + ).reshape((24,)) self.best_control_parameters = np.array(self.best_control_parameters) - - if(previous_contact[0] == 1 and current_contact[0] == 0): - self.best_control_parameters[0:self.num_control_parameters_single_leg] = 0.0 - if(previous_contact[1] == 1 and current_contact[1] == 0): - self.best_control_parameters[self.num_control_parameters_single_leg:self.num_control_parameters_single_leg*2] = 0.0 - if(previous_contact[2] == 1 and current_contact[2] == 0): - self.best_control_parameters[self.num_control_parameters_single_leg*2:self.num_control_parameters_single_leg*3] = 0.0 - if(previous_contact[3] == 1 and current_contact[3] == 0): - self.best_control_parameters[self.num_control_parameters_single_leg*3:self.num_control_parameters_single_leg*4] = 0.0 - - return state_current_jax, reference_state_jax - + if previous_contact[0] == 1 and current_contact[0] == 0: + self.best_control_parameters[0 : self.num_control_parameters_single_leg] = 0.0 + if previous_contact[1] == 1 and current_contact[1] == 0: + self.best_control_parameters[ + self.num_control_parameters_single_leg : self.num_control_parameters_single_leg * 2 + ] = 0.0 + if previous_contact[2] == 1 and current_contact[2] == 0: + self.best_control_parameters[ + self.num_control_parameters_single_leg * 2 : self.num_control_parameters_single_leg * 3 + ] = 0.0 + if previous_contact[3] == 1 and current_contact[3] == 0: + self.best_control_parameters[ + self.num_control_parameters_single_leg * 3 : self.num_control_parameters_single_leg * 4 + ] = 0.0 + return state_current_jax, reference_state_jax - def compute_control_random_sampling(self, state, reference, contact_sequence, best_control_parameters, key, timing, nominal_step_frequency, optimize_swing): + def compute_control_random_sampling( + self, + state, + reference, + contact_sequence, + best_control_parameters, + key, + timing, + nominal_step_frequency, + optimize_swing, + ): """ This function computes the control parameters by sampling from a Gaussian and a uniform distribution. - """ - - - + """ + # Generate random parameters - - # The first control parameters is the old best one, so we add zero noise there - additional_random_parameters = self.initial_random_parameters*0.0 + # The first control parameters is the old best one, so we add zero noise there + additional_random_parameters = self.initial_random_parameters * 0.0 # FIRST GAUSSIAN - num_sample_gaussian_1 = (1 + int(self.num_parallel_computations/3)) - 1 + num_sample_gaussian_1 = (1 + int(self.num_parallel_computations / 3)) - 1 sigma_gaussian_1 = self.sigma_random_sampling[0] - additional_random_parameters = additional_random_parameters.at[1:1 + int(self.num_parallel_computations/3)].set(sigma_gaussian_1*jax.random.normal(key=key, shape=(num_sample_gaussian_1, self.num_control_parameters))) - + additional_random_parameters = additional_random_parameters.at[ + 1 : 1 + int(self.num_parallel_computations / 3) + ].set(sigma_gaussian_1 * jax.random.normal(key=key, shape=(num_sample_gaussian_1, self.num_control_parameters))) # SECOND GAUSSIAN - num_sample_gaussian_2 = (1 + int(self.num_parallel_computations/3)*2) - (1 + int(self.num_parallel_computations/3)) + num_sample_gaussian_2 = (1 + int(self.num_parallel_computations / 3) * 2) - ( + 1 + int(self.num_parallel_computations / 3) + ) sigma_gaussian_2 = self.sigma_random_sampling[1] - additional_random_parameters = additional_random_parameters.at[1 + int(self.num_parallel_computations/3):1 + int(self.num_parallel_computations/3)*2].set(sigma_gaussian_2*jax.random.normal(key=key, shape=(num_sample_gaussian_2, self.num_control_parameters))) - + additional_random_parameters = additional_random_parameters.at[ + 1 + int(self.num_parallel_computations / 3) : 1 + int(self.num_parallel_computations / 3) * 2 + ].set(sigma_gaussian_2 * jax.random.normal(key=key, shape=(num_sample_gaussian_2, self.num_control_parameters))) # UNIFORM max_sampling_forces = self.sigma_random_sampling[2] - num_samples_uniform = int(self.num_parallel_computations) - (1 + int(self.num_parallel_computations/3)*2) - additional_random_parameters = additional_random_parameters.at[1 + int(self.num_parallel_computations/3)*2:int(self.num_parallel_computations)].set(jax.random.uniform(key=key, minval=-max_sampling_forces, maxval=max_sampling_forces, shape=(num_samples_uniform, self.num_control_parameters ))) - + num_samples_uniform = int(self.num_parallel_computations) - (1 + int(self.num_parallel_computations / 3) * 2) + additional_random_parameters = additional_random_parameters.at[ + 1 + int(self.num_parallel_computations / 3) * 2 : int(self.num_parallel_computations) + ].set( + jax.random.uniform( + key=key, + minval=-max_sampling_forces, + maxval=max_sampling_forces, + shape=(num_samples_uniform, self.num_control_parameters), + ) + ) # Add sampling to the best old control parameters control_parameters_vec = best_control_parameters + additional_random_parameters - - # Do rollout costs = self.jit_vectorized_rollout(state, reference, control_parameters_vec, contact_sequence) - # Saturate the cost in case of NaN or inf costs = jnp.where(jnp.isnan(costs), 1000000, costs) costs = jnp.where(jnp.isinf(costs), 1000000, costs) - # Take the best found control parameters best_index = jnp.nanargmin(costs) best_cost = costs.take(best_index) best_control_parameters = control_parameters_vec[best_index] - # and redistribute it to each leg - best_control_parameters_FL = best_control_parameters[0:self.num_control_parameters_single_leg] - best_control_parameters_FR = best_control_parameters[self.num_control_parameters_single_leg:self.num_control_parameters_single_leg*2] - best_control_parameters_RL = best_control_parameters[self.num_control_parameters_single_leg*2:self.num_control_parameters_single_leg*3] - best_control_parameters_RR = best_control_parameters[self.num_control_parameters_single_leg*3:self.num_control_parameters_single_leg*4] - + best_control_parameters_FL = best_control_parameters[0 : self.num_control_parameters_single_leg] + best_control_parameters_FR = best_control_parameters[ + self.num_control_parameters_single_leg : self.num_control_parameters_single_leg * 2 + ] + best_control_parameters_RL = best_control_parameters[ + self.num_control_parameters_single_leg * 2 : self.num_control_parameters_single_leg * 3 + ] + best_control_parameters_RR = best_control_parameters[ + self.num_control_parameters_single_leg * 3 : self.num_control_parameters_single_leg * 4 + ] # Compute the GRF associated to the best parameter fx_FL, fy_FL, fz_FL = self.spline_fun_FL(best_control_parameters_FL, 0.0, 1) @@ -825,126 +909,143 @@ def compute_control_random_sampling(self, state, reference, contact_sequence, be fx_RL, fy_RL, fz_RL = self.spline_fun_RL(best_control_parameters_RL, 0.0, 1) fx_RR, fy_RR, fz_RR = self.spline_fun_RR(best_control_parameters_RR, 0.0, 1) - - # Add the gravity compensation to the stance legs and put to zero + # Add the gravity compensation to the stance legs and put to zero # the GRF of the swing legs - number_of_legs_in_stance = contact_sequence[0][0] + contact_sequence[1][0] + contact_sequence[2][0] + contact_sequence[3][0] + number_of_legs_in_stance = ( + contact_sequence[0][0] + contact_sequence[1][0] + contact_sequence[2][0] + contact_sequence[3][0] + ) reference_force_stance_legs = (self.robot.mass * 9.81) / number_of_legs_in_stance fz_FL = reference_force_stance_legs + fz_FL fz_FR = reference_force_stance_legs + fz_FR fz_RL = reference_force_stance_legs + fz_RL fz_RR = reference_force_stance_legs + fz_RR - - fx_FL = fx_FL*contact_sequence[0][0] - fy_FL = fy_FL*contact_sequence[0][0] - fz_FL = fz_FL*contact_sequence[0][0] - - fx_FR = fx_FR*contact_sequence[1][0] - fy_FR = fy_FR*contact_sequence[1][0] - fz_FR = fz_FR*contact_sequence[1][0] - fx_RL = fx_RL*contact_sequence[2][0] - fy_RL = fy_RL*contact_sequence[2][0] - fz_RL = fz_RL*contact_sequence[2][0] + fx_FL = fx_FL * contact_sequence[0][0] + fy_FL = fy_FL * contact_sequence[0][0] + fz_FL = fz_FL * contact_sequence[0][0] - fx_RR = fx_RR*contact_sequence[3][0] - fy_RR = fy_RR*contact_sequence[3][0] - fz_RR = fz_RR*contact_sequence[3][0] + fx_FR = fx_FR * contact_sequence[1][0] + fy_FR = fy_FR * contact_sequence[1][0] + fz_FR = fz_FR * contact_sequence[1][0] + fx_RL = fx_RL * contact_sequence[2][0] + fy_RL = fy_RL * contact_sequence[2][0] + fz_RL = fz_RL * contact_sequence[2][0] + + fx_RR = fx_RR * contact_sequence[3][0] + fy_RR = fy_RR * contact_sequence[3][0] + fz_RR = fz_RR * contact_sequence[3][0] # Enforce force constraints - fx_FL, fy_FL, fz_FL, \ - fx_FR, fy_FR, fz_FR, \ - fx_RL, fy_RL, fz_RL, \ - fx_RR, fy_RR, fz_RR =self.enforce_force_constraints(fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR) - - - - nmpc_GRFs = jnp.array([fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR]) - nmpc_footholds = jnp.array([0, 0, 0, - 0, 0, 0, - 0, 0, 0, - 0, 0, 0]) - + fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR = ( + self.enforce_force_constraints( + fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR + ) + ) + + nmpc_GRFs = jnp.array([fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR]) + nmpc_footholds = jnp.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) # Compute predicted state for IK - input = jnp.array([jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR, - ], dtype=dtype_general) - current_contact = jnp.array([contact_sequence[0][0], contact_sequence[1][0], - contact_sequence[2][0], contact_sequence[3][0]], dtype=dtype_general) + input = jnp.array( + [ + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + fx_FL, + fy_FL, + fz_FL, + fx_FR, + fy_FR, + fz_FR, + fx_RL, + fy_RL, + fz_RL, + fx_RR, + fy_RR, + fz_RR, + ], + dtype=dtype_general, + ) + current_contact = jnp.array( + [contact_sequence[0][0], contact_sequence[1][0], contact_sequence[2][0], contact_sequence[3][0]], + dtype=dtype_general, + ) nmpc_predicted_state = self.robot.integrate_jax(state, input, current_contact, 0) - best_freq = 1.4 return nmpc_GRFs, nmpc_footholds, nmpc_predicted_state, best_control_parameters, best_cost, best_freq, costs - - - - def compute_control_mppi(self, state, reference, contact_sequence, best_control_parameters, key, timing, nominal_step_frequency, optimize_swing): + def compute_control_mppi( + self, + state, + reference, + contact_sequence, + best_control_parameters, + key, + timing, + nominal_step_frequency, + optimize_swing, + ): """ This function computes the control parameters by applying MPPI. - """ - + """ + # Generate random parameters # The first control parameters is the old best one, so we add zero noise there - additional_random_parameters = self.initial_random_parameters*0.0 - + additional_random_parameters = self.initial_random_parameters * 0.0 # GAUSSIAN - num_sample_gaussian_1 = self.num_parallel_computations-1 - additional_random_parameters = additional_random_parameters.at[1:self.num_parallel_computations].set(self.sigma_mppi*jax.random.normal(key=key, shape=(num_sample_gaussian_1, self.num_control_parameters))) - - - control_parameters_vec = best_control_parameters + additional_random_parameters + num_sample_gaussian_1 = self.num_parallel_computations - 1 + additional_random_parameters = additional_random_parameters.at[1 : self.num_parallel_computations].set( + self.sigma_mppi * jax.random.normal(key=key, shape=(num_sample_gaussian_1, self.num_control_parameters)) + ) - - + control_parameters_vec = best_control_parameters + additional_random_parameters # Do rollout costs = self.jit_vectorized_rollout(state, reference, control_parameters_vec, contact_sequence) - # Saturate the cost in case of NaN or inf costs = jnp.where(jnp.isnan(costs), 1000000, costs) costs = jnp.where(jnp.isinf(costs), 1000000, costs) - # Take the best found control parameters best_index = jnp.nanargmin(costs) best_cost = costs.take(best_index) - # Compute MPPI update beta = best_cost - temperature = 1. - exp_costs = jnp.exp((-1./temperature) * (costs - beta)) + temperature = 1.0 + exp_costs = jnp.exp((-1.0 / temperature) * (costs - beta)) denom = np.sum(exp_costs) - weights = exp_costs/denom - weighted_inputs = weights[:, jnp.newaxis, jnp.newaxis] * additional_random_parameters.reshape((self.num_parallel_computations,self.num_control_parameters,1)) - best_control_parameters += jnp.sum(weighted_inputs, axis=0).reshape((self.num_control_parameters, )) - + weights = exp_costs / denom + weighted_inputs = weights[:, jnp.newaxis, jnp.newaxis] * additional_random_parameters.reshape( + (self.num_parallel_computations, self.num_control_parameters, 1) + ) + best_control_parameters += jnp.sum(weighted_inputs, axis=0).reshape((self.num_control_parameters,)) # And redistribute it to each leg - best_control_parameters_FL = best_control_parameters[0:self.num_control_parameters_single_leg] - best_control_parameters_FR = best_control_parameters[self.num_control_parameters_single_leg:self.num_control_parameters_single_leg*2] - best_control_parameters_RL = best_control_parameters[self.num_control_parameters_single_leg*2:self.num_control_parameters_single_leg*3] - best_control_parameters_RR = best_control_parameters[self.num_control_parameters_single_leg*3:self.num_control_parameters_single_leg*4] - + best_control_parameters_FL = best_control_parameters[0 : self.num_control_parameters_single_leg] + best_control_parameters_FR = best_control_parameters[ + self.num_control_parameters_single_leg : self.num_control_parameters_single_leg * 2 + ] + best_control_parameters_RL = best_control_parameters[ + self.num_control_parameters_single_leg * 2 : self.num_control_parameters_single_leg * 3 + ] + best_control_parameters_RR = best_control_parameters[ + self.num_control_parameters_single_leg * 3 : self.num_control_parameters_single_leg * 4 + ] # Compute the GRF associated to the best parameter fx_FL, fy_FL, fz_FL = self.spline_fun_FL(best_control_parameters_FL, 0.0, 1) @@ -952,127 +1053,145 @@ def compute_control_mppi(self, state, reference, contact_sequence, best_control_ fx_RL, fy_RL, fz_RL = self.spline_fun_RL(best_control_parameters_RL, 0.0, 1) fx_RR, fy_RR, fz_RR = self.spline_fun_RR(best_control_parameters_RR, 0.0, 1) - - # Add the gravity compensation to the stance legs and put to zero + # Add the gravity compensation to the stance legs and put to zero # the GRF of the swing legs - number_of_legs_in_stance = contact_sequence[0][0] + contact_sequence[1][0] + contact_sequence[2][0] + contact_sequence[3][0] + number_of_legs_in_stance = ( + contact_sequence[0][0] + contact_sequence[1][0] + contact_sequence[2][0] + contact_sequence[3][0] + ) reference_force_stance_legs = (self.robot.mass * 9.81) / number_of_legs_in_stance - fz_FL = reference_force_stance_legs + fz_FL fz_FR = reference_force_stance_legs + fz_FR fz_RL = reference_force_stance_legs + fz_RL fz_RR = reference_force_stance_legs + fz_RR - - fx_FL = fx_FL*contact_sequence[0][0] - fy_FL = fy_FL*contact_sequence[0][0] - fz_FL = fz_FL*contact_sequence[0][0] - - fx_FR = fx_FR*contact_sequence[1][0] - fy_FR = fy_FR*contact_sequence[1][0] - fz_FR = fz_FR*contact_sequence[1][0] - fx_RL = fx_RL*contact_sequence[2][0] - fy_RL = fy_RL*contact_sequence[2][0] - fz_RL = fz_RL*contact_sequence[2][0] + fx_FL = fx_FL * contact_sequence[0][0] + fy_FL = fy_FL * contact_sequence[0][0] + fz_FL = fz_FL * contact_sequence[0][0] + + fx_FR = fx_FR * contact_sequence[1][0] + fy_FR = fy_FR * contact_sequence[1][0] + fz_FR = fz_FR * contact_sequence[1][0] - fx_RR = fx_RR*contact_sequence[3][0] - fy_RR = fy_RR*contact_sequence[3][0] - fz_RR = fz_RR*contact_sequence[3][0] + fx_RL = fx_RL * contact_sequence[2][0] + fy_RL = fy_RL * contact_sequence[2][0] + fz_RL = fz_RL * contact_sequence[2][0] + fx_RR = fx_RR * contact_sequence[3][0] + fy_RR = fy_RR * contact_sequence[3][0] + fz_RR = fz_RR * contact_sequence[3][0] # Enforce force constraints - fx_FL, fy_FL, fz_FL, \ - fx_FR, fy_FR, fz_FR, \ - fx_RL, fy_RL, fz_RL, \ - fx_RR, fy_RR, fz_RR =self.enforce_force_constraints(fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR) - - - nmpc_GRFs = jnp.array([fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR]) - nmpc_footholds = jnp.array([0, 0, 0, - 0, 0, 0, - 0, 0, 0, - 0, 0, 0]) - + fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR = ( + self.enforce_force_constraints( + fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR + ) + ) + + nmpc_GRFs = jnp.array([fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR]) + nmpc_footholds = jnp.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) # Compute predicted state for IK - input = jnp.array([jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR, - ], dtype=dtype_general) - current_contact = jnp.array([contact_sequence[0][0], contact_sequence[1][0], - contact_sequence[2][0], contact_sequence[3][0]], dtype=dtype_general) + input = jnp.array( + [ + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + fx_FL, + fy_FL, + fz_FL, + fx_FR, + fy_FR, + fz_FR, + fx_RL, + fy_RL, + fz_RL, + fx_RR, + fy_RR, + fz_RR, + ], + dtype=dtype_general, + ) + current_contact = jnp.array( + [contact_sequence[0][0], contact_sequence[1][0], contact_sequence[2][0], contact_sequence[3][0]], + dtype=dtype_general, + ) nmpc_predicted_state = self.robot.integrate_jax(state, input, current_contact, 0) - - - best_freq = 1.4 - - return nmpc_GRFs, nmpc_footholds, nmpc_predicted_state, best_control_parameters, best_cost, best_freq, costs - + best_freq = 1.4 + return nmpc_GRFs, nmpc_footholds, nmpc_predicted_state, best_control_parameters, best_cost, best_freq, costs - def compute_control_cem_mppi(self, state, reference, contact_sequence, best_control_parameters, key, sigma, timing = None, nominal_step_frequency = None): + def compute_control_cem_mppi( + self, + state, + reference, + contact_sequence, + best_control_parameters, + key, + sigma, + timing=None, + nominal_step_frequency=None, + ): """ This function computes the control parameters by applying CEM-MPPI. - """ - + """ + # Generate random parameters # The first control parameters is the old best one, so we add zero noise there - additional_random_parameters = self.initial_random_parameters*0.0 - + additional_random_parameters = self.initial_random_parameters * 0.0 # GAUSSIAN - num_sample_gaussian_1 = self.num_parallel_computations-1 + num_sample_gaussian_1 = self.num_parallel_computations - 1 - additional_random_parameters = additional_random_parameters.at[1:self.num_parallel_computations].set(jax.random.normal(key=key, shape=(num_sample_gaussian_1, self.num_control_parameters))*sigma) - - - control_parameters_vec = best_control_parameters + additional_random_parameters + additional_random_parameters = additional_random_parameters.at[1 : self.num_parallel_computations].set( + jax.random.normal(key=key, shape=(num_sample_gaussian_1, self.num_control_parameters)) * sigma + ) + control_parameters_vec = best_control_parameters + additional_random_parameters - # Do rollout costs = self.jit_vectorized_rollout(state, reference, control_parameters_vec, contact_sequence) - # Saturate the cost in case of NaN or inf costs = jnp.where(jnp.isnan(costs), 1000000, costs) costs = jnp.where(jnp.isinf(costs), 1000000, costs) - # Take the best found control parameters best_index = jnp.nanargmin(costs) best_cost = costs.take(best_index) - # Compute MPPI update beta = best_cost - temperature = 1. - exp_costs = jnp.exp((-1./temperature) * (costs - beta)) + temperature = 1.0 + exp_costs = jnp.exp((-1.0 / temperature) * (costs - beta)) denom = np.sum(exp_costs) - weights = exp_costs/denom - weighted_inputs = weights[:, jnp.newaxis, jnp.newaxis] * additional_random_parameters.reshape((self.num_parallel_computations,self.num_control_parameters,1)) - best_control_parameters += jnp.sum(weighted_inputs, axis=0).reshape((self.num_control_parameters, )) - + weights = exp_costs / denom + weighted_inputs = weights[:, jnp.newaxis, jnp.newaxis] * additional_random_parameters.reshape( + (self.num_parallel_computations, self.num_control_parameters, 1) + ) + best_control_parameters += jnp.sum(weighted_inputs, axis=0).reshape((self.num_control_parameters,)) # And redistribute it to each leg - best_control_parameters_FL = best_control_parameters[0:self.num_control_parameters_single_leg] - best_control_parameters_FR = best_control_parameters[self.num_control_parameters_single_leg:self.num_control_parameters_single_leg*2] - best_control_parameters_RL = best_control_parameters[self.num_control_parameters_single_leg*2:self.num_control_parameters_single_leg*3] - best_control_parameters_RR = best_control_parameters[self.num_control_parameters_single_leg*3:self.num_control_parameters_single_leg*4] - + best_control_parameters_FL = best_control_parameters[0 : self.num_control_parameters_single_leg] + best_control_parameters_FR = best_control_parameters[ + self.num_control_parameters_single_leg : self.num_control_parameters_single_leg * 2 + ] + best_control_parameters_RL = best_control_parameters[ + self.num_control_parameters_single_leg * 2 : self.num_control_parameters_single_leg * 3 + ] + best_control_parameters_RR = best_control_parameters[ + self.num_control_parameters_single_leg * 3 : self.num_control_parameters_single_leg * 4 + ] # Compute the GRF associated to the best parameter fx_FL, fy_FL, fz_FL = self.spline_fun_FL(best_control_parameters_FL, 0.0, 1) @@ -1080,81 +1199,99 @@ def compute_control_cem_mppi(self, state, reference, contact_sequence, best_cont fx_RL, fy_RL, fz_RL = self.spline_fun_RL(best_control_parameters_RL, 0.0, 1) fx_RR, fy_RR, fz_RR = self.spline_fun_RR(best_control_parameters_RR, 0.0, 1) - - # Add the gravity compensation to the stance legs and put to zero + # Add the gravity compensation to the stance legs and put to zero # the GRF of the swing legs - number_of_legs_in_stance = contact_sequence[0][0] + contact_sequence[1][0] + contact_sequence[2][0] + contact_sequence[3][0] + number_of_legs_in_stance = ( + contact_sequence[0][0] + contact_sequence[1][0] + contact_sequence[2][0] + contact_sequence[3][0] + ) reference_force_stance_legs = (self.robot.mass * 9.81) / number_of_legs_in_stance fz_FL = reference_force_stance_legs + fz_FL fz_FR = reference_force_stance_legs + fz_FR fz_RL = reference_force_stance_legs + fz_RL fz_RR = reference_force_stance_legs + fz_RR - - fx_FL = fx_FL*contact_sequence[0][0] - fy_FL = fy_FL*contact_sequence[0][0] - fz_FL = fz_FL*contact_sequence[0][0] - - fx_FR = fx_FR*contact_sequence[1][0] - fy_FR = fy_FR*contact_sequence[1][0] - fz_FR = fz_FR*contact_sequence[1][0] - fx_RL = fx_RL*contact_sequence[2][0] - fy_RL = fy_RL*contact_sequence[2][0] - fz_RL = fz_RL*contact_sequence[2][0] + fx_FL = fx_FL * contact_sequence[0][0] + fy_FL = fy_FL * contact_sequence[0][0] + fz_FL = fz_FL * contact_sequence[0][0] - fx_RR = fx_RR*contact_sequence[3][0] - fy_RR = fy_RR*contact_sequence[3][0] - fz_RR = fz_RR*contact_sequence[3][0] + fx_FR = fx_FR * contact_sequence[1][0] + fy_FR = fy_FR * contact_sequence[1][0] + fz_FR = fz_FR * contact_sequence[1][0] + fx_RL = fx_RL * contact_sequence[2][0] + fy_RL = fy_RL * contact_sequence[2][0] + fz_RL = fz_RL * contact_sequence[2][0] + fx_RR = fx_RR * contact_sequence[3][0] + fy_RR = fy_RR * contact_sequence[3][0] + fz_RR = fz_RR * contact_sequence[3][0] # Enforce force constraints - fx_FL, fy_FL, fz_FL, \ - fx_FR, fy_FR, fz_FR, \ - fx_RL, fy_RL, fz_RL, \ - fx_RR, fy_RR, fz_RR =self.enforce_force_constraints(fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR) - - - nmpc_GRFs = jnp.array([fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR]) - nmpc_footholds = jnp.array([0, 0, 0, - 0, 0, 0, - 0, 0, 0, - 0, 0, 0]) + fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR = ( + self.enforce_force_constraints( + fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR + ) + ) + + nmpc_GRFs = jnp.array([fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR]) + nmpc_footholds = jnp.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) # Compute predicted state for IK - input = jnp.array([jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR, - ], dtype=dtype_general) - current_contact = jnp.array([contact_sequence[0][0], contact_sequence[1][0], - contact_sequence[2][0], contact_sequence[3][0]], dtype=dtype_general) + input = jnp.array( + [ + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + fx_FL, + fy_FL, + fz_FL, + fx_FR, + fy_FR, + fz_FR, + fx_RL, + fy_RL, + fz_RL, + fx_RR, + fy_RR, + fz_RR, + ], + dtype=dtype_general, + ) + current_contact = jnp.array( + [contact_sequence[0][0], contact_sequence[1][0], contact_sequence[2][0], contact_sequence[3][0]], + dtype=dtype_general, + ) nmpc_predicted_state = self.robot.integrate_jax(state, input, current_contact, 0) - indexes = jnp.argsort(costs)[:10] elite = additional_random_parameters[indexes] - new_sigma_cem_mppi = jnp.cov(elite, rowvar=False) + np.eye(self.num_control_parameters)*1e-8 + new_sigma_cem_mppi = jnp.cov(elite, rowvar=False) + np.eye(self.num_control_parameters) * 1e-8 new_sigma_cem_mppi = jnp.diag(new_sigma_cem_mppi) new_sigma_cem_mppi = jnp.sqrt(new_sigma_cem_mppi) new_sigma_cem_mppi = jnp.where(new_sigma_cem_mppi > 5, 5, new_sigma_cem_mppi) new_sigma_cem_mppi = jnp.where(new_sigma_cem_mppi < 0.2, 0.2, new_sigma_cem_mppi) - best_freq = 1.65 - - return nmpc_GRFs, nmpc_footholds, nmpc_predicted_state, best_control_parameters, best_cost, best_freq, costs, new_sigma_cem_mppi - + return ( + nmpc_GRFs, + nmpc_footholds, + nmpc_predicted_state, + best_control_parameters, + best_cost, + best_freq, + costs, + new_sigma_cem_mppi, + ) + def reset(self): - print("Resetting the controller") \ No newline at end of file + print("Resetting the controller") diff --git a/quadruped_pympc/controllers/sampling/centroidal_nmpc_jax_gait_adaptive.py b/quadruped_pympc/controllers/sampling/centroidal_nmpc_jax_gait_adaptive.py index 2b63cf2..081072e 100644 --- a/quadruped_pympc/controllers/sampling/centroidal_nmpc_jax_gait_adaptive.py +++ b/quadruped_pympc/controllers/sampling/centroidal_nmpc_jax_gait_adaptive.py @@ -1,163 +1,156 @@ import numpy as np -np.set_printoptions(precision=3, suppress = True) -import os +np.set_printoptions(precision=3, suppress=True) + +import os + dir_path = os.path.dirname(os.path.realpath(__file__)) -os.environ['XLA_FLAGS'] = ('--xla_gpu_triton_gemm_any=True') +os.environ["XLA_FLAGS"] = "--xla_gpu_triton_gemm_any=True" +import sys + import jax import jax.numpy as jnp -from jax import jit, random +from jax import random -import sys sys.path.append(dir_path) -sys.path.append(dir_path + '/../') -sys.path.append(dir_path + '/../helpers/') +sys.path.append(dir_path + "/../") +sys.path.append(dir_path + "/../helpers/") -from quadruped_pympc import config -from centroidal_model_jax import Centroidal_Model_JAX -from quadruped_pympc.helpers.periodic_gait_generator_jax import PeriodicGaitGeneratorJax - -import time import copy +from centroidal_model_jax import Centroidal_Model_JAX +from quadruped_pympc import config +from quadruped_pympc.helpers.periodic_gait_generator_jax import PeriodicGaitGeneratorJax -dtype_general='float32' +dtype_general = "float32" class Sampling_MPC: """This is a small class that implements a sampling based control law""" - def __init__(self, device="gpu"): """ Args: - horizon (int): how much to look into the future for optimizing the gains + horizon (int): how much to look into the future for optimizing the gains dt (int): desidered sampling time """ - self.num_parallel_computations = config.mpc_params['num_parallel_computations'] - self.sampling_method = config.mpc_params['sampling_method'] - self.control_parametrization = config.mpc_params['control_parametrization'] - self.num_sampling_iterations = config.mpc_params['num_sampling_iterations'] - self.dt = config.mpc_params['dt'] - self.horizon = config.mpc_params['horizon'] - + self.num_parallel_computations = config.mpc_params["num_parallel_computations"] + self.sampling_method = config.mpc_params["sampling_method"] + self.control_parametrization = config.mpc_params["control_parametrization"] + self.num_sampling_iterations = config.mpc_params["num_sampling_iterations"] + self.dt = config.mpc_params["dt"] + self.horizon = config.mpc_params["horizon"] self.state_dim = 24 self.control_dim = 24 self.reference_dim = self.state_dim - - self.max_sampling_forces = 30 - - + self.max_sampling_forces = 30 - - if(device=="gpu"): + if device == "gpu": try: - self.device = jax.devices('gpu')[0] + self.device = jax.devices("gpu")[0] except: - self.device = jax.devices('cpu')[0] + self.device = jax.devices("cpu")[0] print("GPU not available, using CPU") else: - self.device = jax.devices('cpu')[0] - + self.device = jax.devices("cpu")[0] - - if(self.control_parametrization == "linear_spline_1"): + if self.control_parametrization == "linear_spline_1": # Along the horizon, we have only 1 spline per control input (3 forces) # Each spline has 2 parameters - self.num_control_parameters_single_leg = 2*3 - + self.num_control_parameters_single_leg = 2 * 3 + # In totale we have 4 legs - self.num_control_parameters = self.num_control_parameters_single_leg*4 - + self.num_control_parameters = self.num_control_parameters_single_leg * 4 + # We have 4 different spline functions, one for each leg self.spline_fun_FL = self.compute_linear_spline_1 self.spline_fun_FR = self.compute_linear_spline_1 self.spline_fun_RL = self.compute_linear_spline_1 self.spline_fun_RR = self.compute_linear_spline_1 - - elif(self.control_parametrization == "linear_spline_2"): + + elif self.control_parametrization == "linear_spline_2": # Along the horizon, we have 2 splines per control input (3 forces) # Each spline has 2 parameters, but one is shared between the two splines - self.num_control_parameters_single_leg = 3*3 - + self.num_control_parameters_single_leg = 3 * 3 + # In totale we have 4 legs - self.num_control_parameters = self.num_control_parameters_single_leg*4 - + self.num_control_parameters = self.num_control_parameters_single_leg * 4 + # We have 4 different spline functions, one for each leg self.spline_fun_FL = self.compute_linear_spline_2 self.spline_fun_FR = self.compute_linear_spline_2 self.spline_fun_RL = self.compute_linear_spline_2 self.spline_fun_RR = self.compute_linear_spline_2 - elif(self.control_parametrization == "linear_spline_N"): + elif self.control_parametrization == "linear_spline_N": # Along the horizon, we have 2 splines per control input (3 forces) # Each spline has 2 parameters, but one is shared between the two splines - self.num_spline = config.mpc_params['num_splines'] - self.num_control_parameters_single_leg = (self.num_spline + 1)*3 - + self.num_spline = config.mpc_params["num_splines"] + self.num_control_parameters_single_leg = (self.num_spline + 1) * 3 + # In totale we have 4 legs - self.num_control_parameters = self.num_control_parameters_single_leg*4 - + self.num_control_parameters = self.num_control_parameters_single_leg * 4 + # We have 4 different spline functions, one for each leg self.spline_fun_FL = self.compute_linear_spline_N self.spline_fun_FR = self.compute_linear_spline_N self.spline_fun_RL = self.compute_linear_spline_N self.spline_fun_RR = self.compute_linear_spline_N - elif(self.control_parametrization == "cubic_spline_1"): + elif self.control_parametrization == "cubic_spline_1": # Along the horizon, we have 1 splines per control input (3 forces) # Each spline has 3 parameters - self.num_control_parameters_single_leg = 4*3 + self.num_control_parameters_single_leg = 4 * 3 # In totale we have 4 legs - self.num_control_parameters = self.num_control_parameters_single_leg*4 + self.num_control_parameters = self.num_control_parameters_single_leg * 4 # We have 4 different spline functions, one for each leg self.spline_fun_FL = self.compute_cubic_spline self.spline_fun_FR = self.compute_cubic_spline self.spline_fun_RL = self.compute_cubic_spline self.spline_fun_RR = self.compute_cubic_spline - - elif(self.control_parametrization == "cubic_spline_2"): + + elif self.control_parametrization == "cubic_spline_2": # Along the horizon, we have 1 splines per control input (3 forces) # Each spline has 3 parameters - self.num_control_parameters_single_leg = 4*3 + 4*3 + self.num_control_parameters_single_leg = 4 * 3 + 4 * 3 # In totale we have 4 legs - self.num_control_parameters = self.num_control_parameters_single_leg*4 + self.num_control_parameters = self.num_control_parameters_single_leg * 4 # We have 4 different spline functions, one for each leg self.spline_fun_FL = self.compute_cubic_spline_2 self.spline_fun_FR = self.compute_cubic_spline_2 self.spline_fun_RL = self.compute_cubic_spline_2 - self.spline_fun_RR = self.compute_cubic_spline_2 + self.spline_fun_RR = self.compute_cubic_spline_2 - elif(self.control_parametrization == "cubic_spline_N"): + elif self.control_parametrization == "cubic_spline_N": # Along the horizon, we have 1 splines per control input (3 forces) # Each spline has 3 parameters - self.num_spline = config.mpc_params['num_splines'] - self.num_control_parameters_single_leg = 4*3*self.num_spline + self.num_spline = config.mpc_params["num_splines"] + self.num_control_parameters_single_leg = 4 * 3 * self.num_spline # In totale we have 4 legs - self.num_control_parameters = self.num_control_parameters_single_leg*4 + self.num_control_parameters = self.num_control_parameters_single_leg * 4 # We have 4 different spline functions, one for each leg self.spline_fun_FL = self.compute_cubic_spline_N self.spline_fun_FR = self.compute_cubic_spline_N self.spline_fun_RL = self.compute_cubic_spline_N - self.spline_fun_RR = self.compute_cubic_spline_N + self.spline_fun_RR = self.compute_cubic_spline_N else: # We have 1 parameters for every 3 force direction (x,y,z)...for each time horizon!! - self.num_control_parameters_single_leg = self.horizon*3 - + self.num_control_parameters_single_leg = self.horizon * 3 + # In totale we have 4 legs - self.num_control_parameters = self.num_control_parameters_single_leg*4 + self.num_control_parameters = self.num_control_parameters_single_leg * 4 # We have 4 different spline functions, one for each leg self.spline_fun_FL = self.compute_zero_order_spline @@ -165,18 +158,17 @@ def __init__(self, device="gpu"): self.spline_fun_RL = self.compute_zero_order_spline self.spline_fun_RR = self.compute_zero_order_spline - - - - if(self.sampling_method == 'random_sampling'): + if self.sampling_method == "random_sampling": self.compute_control = self.compute_control_random_sampling - self.sigma_random_sampling = config.mpc_params['sigma_random_sampling'] - elif(self.sampling_method == 'mppi'): - self.compute_control = self.compute_control_mppi - self.sigma_mppi = config.mpc_params['sigma_mppi'] - elif(self.sampling_method == 'cem_mppi'): + self.sigma_random_sampling = config.mpc_params["sigma_random_sampling"] + elif self.sampling_method == "mppi": + self.compute_control = self.compute_control_mppi + self.sigma_mppi = config.mpc_params["sigma_mppi"] + elif self.sampling_method == "cem_mppi": self.compute_control = self.compute_control_cem_mppi - self.sigma_cem_mppi = jnp.ones(self.num_control_parameters, dtype=dtype_general) * config.mpc_params['sigma_cem_mppi'] + self.sigma_cem_mppi = ( + jnp.ones(self.num_control_parameters, dtype=dtype_general) * config.mpc_params["sigma_cem_mppi"] + ) else: # return error and stop execution print("Error: sampling method not recognized") @@ -184,345 +176,343 @@ def __init__(self, device="gpu"): self.jitted_compute_control = jax.jit(self.compute_control, device=self.device) - # Initialize the robot model - self.robot = Centroidal_Model_JAX(self.dt,self.device) - + self.robot = Centroidal_Model_JAX(self.dt, self.device) # Initialize the cost function matrices - self.Q = jnp.identity(self.state_dim, dtype=dtype_general)*0 - self.Q = self.Q.at[0,0].set(0.0) - self.Q = self.Q.at[1,1].set(0.0) - self.Q = self.Q.at[2,2].set(111500) #com_z - self.Q = self.Q.at[3,3].set(5000) #com_vel_x - self.Q = self.Q.at[4,4].set(5000) #com_vel_y - self.Q = self.Q.at[5,5].set(200) #com_vel_z - self.Q = self.Q.at[6,6].set(11200) #base_angle_roll - self.Q = self.Q.at[7,7].set(11200) #base_angle_pitch - self.Q = self.Q.at[8,8].set(0.0) #base_angle_yaw - self.Q = self.Q.at[9,9].set(20) #base_angle_rates_x - self.Q = self.Q.at[10,10].set(20) #base_angle_rates_y - self.Q = self.Q.at[11,11].set(600) #base_angle_rates_z - - + self.Q = jnp.identity(self.state_dim, dtype=dtype_general) * 0 + self.Q = self.Q.at[0, 0].set(0.0) + self.Q = self.Q.at[1, 1].set(0.0) + self.Q = self.Q.at[2, 2].set(111500) # com_z + self.Q = self.Q.at[3, 3].set(5000) # com_vel_x + self.Q = self.Q.at[4, 4].set(5000) # com_vel_y + self.Q = self.Q.at[5, 5].set(200) # com_vel_z + self.Q = self.Q.at[6, 6].set(11200) # base_angle_roll + self.Q = self.Q.at[7, 7].set(11200) # base_angle_pitch + self.Q = self.Q.at[8, 8].set(0.0) # base_angle_yaw + self.Q = self.Q.at[9, 9].set(20) # base_angle_rates_x + self.Q = self.Q.at[10, 10].set(20) # base_angle_rates_y + self.Q = self.Q.at[11, 11].set(600) # base_angle_rates_z self.R = jnp.identity(self.control_dim, dtype=dtype_general) - self.R = self.R.at[0,0].set(0.0) #foot_pos_x_FL - self.R = self.R.at[1,1].set(0.0) #foot_pos_y_FL - self.R = self.R.at[2,2].set(0.0) #foot_pos_z_FL - self.R = self.R.at[3,3].set(0.0) #foot_pos_x_FR - self.R = self.R.at[4,4].set(0.0) #foot_pos_y_FR - self.R = self.R.at[5,5].set(0.0) #foot_pos_z_FR - self.R = self.R.at[6,6].set(0.0) #foot_pos_x_RL - self.R = self.R.at[7,7].set(0.0) #foot_pos_y_RL - self.R = self.R.at[8,8].set(0.0) #foot_pos_z_RL - self.R = self.R.at[9,9].set(0.0) #foot_pos_x_RR - self.R = self.R.at[10,10].set(0.0) #foot_pos_y_RR - self.R = self.R.at[11,11].set(0.0) #foot_pos_z_RR - - self.R = self.R.at[12,12].set(0.1) #foot_force_x_FL - self.R = self.R.at[13,13].set(0.1) #foot_force_y_FL - self.R = self.R.at[14,14].set(0.001) #foot_force_z_FL - self.R = self.R.at[15,15].set(0.1) #foot_force_x_FR - self.R = self.R.at[16,16].set(0.1) #foot_force_y_FR - self.R = self.R.at[17,17].set(0.001) #foot_force_z_FR - self.R = self.R.at[18,18].set(0.1) #foot_force_x_RL - self.R = self.R.at[19,19].set(0.1) #foot_force_y_RL - self.R = self.R.at[20,20].set(0.001) #foot_force_z_RL - self.R = self.R.at[21,21].set(0.1) #foot_force_x_RR - self.R = self.R.at[22,22].set(0.1) #foot_force_y_RR - self.R = self.R.at[23,23].set(0.001) #foot_force_z_RR - - # mu is the friction coefficient - self.mu = config.mpc_params['mu'] - - # maximum allowed z contact forces - self.f_z_max = config.mpc_params['grf_max'] - self.f_z_min = config.mpc_params['grf_min'] + self.R = self.R.at[0, 0].set(0.0) # foot_pos_x_FL + self.R = self.R.at[1, 1].set(0.0) # foot_pos_y_FL + self.R = self.R.at[2, 2].set(0.0) # foot_pos_z_FL + self.R = self.R.at[3, 3].set(0.0) # foot_pos_x_FR + self.R = self.R.at[4, 4].set(0.0) # foot_pos_y_FR + self.R = self.R.at[5, 5].set(0.0) # foot_pos_z_FR + self.R = self.R.at[6, 6].set(0.0) # foot_pos_x_RL + self.R = self.R.at[7, 7].set(0.0) # foot_pos_y_RL + self.R = self.R.at[8, 8].set(0.0) # foot_pos_z_RL + self.R = self.R.at[9, 9].set(0.0) # foot_pos_x_RR + self.R = self.R.at[10, 10].set(0.0) # foot_pos_y_RR + self.R = self.R.at[11, 11].set(0.0) # foot_pos_z_RR + + self.R = self.R.at[12, 12].set(0.1) # foot_force_x_FL + self.R = self.R.at[13, 13].set(0.1) # foot_force_y_FL + self.R = self.R.at[14, 14].set(0.001) # foot_force_z_FL + self.R = self.R.at[15, 15].set(0.1) # foot_force_x_FR + self.R = self.R.at[16, 16].set(0.1) # foot_force_y_FR + self.R = self.R.at[17, 17].set(0.001) # foot_force_z_FR + self.R = self.R.at[18, 18].set(0.1) # foot_force_x_RL + self.R = self.R.at[19, 19].set(0.1) # foot_force_y_RL + self.R = self.R.at[20, 20].set(0.001) # foot_force_z_RL + self.R = self.R.at[21, 21].set(0.1) # foot_force_x_RR + self.R = self.R.at[22, 22].set(0.1) # foot_force_y_RR + self.R = self.R.at[23, 23].set(0.001) # foot_force_z_RR + # mu is the friction coefficient + self.mu = config.mpc_params["mu"] + # maximum allowed z contact forces + self.f_z_max = config.mpc_params["grf_max"] + self.f_z_min = config.mpc_params["grf_min"] self.best_control_parameters = jnp.zeros((self.num_control_parameters,), dtype=dtype_general) self.master_key = jax.random.PRNGKey(42) - self.initial_random_parameters = jax.random.uniform(key=self.master_key, minval=-self.max_sampling_forces, maxval=self.max_sampling_forces, shape=(self.num_parallel_computations, self.num_control_parameters )) - + self.initial_random_parameters = jax.random.uniform( + key=self.master_key, + minval=-self.max_sampling_forces, + maxval=self.max_sampling_forces, + shape=(self.num_parallel_computations, self.num_control_parameters), + ) # Initialize the periodic gait generator and jit it! - self.pgg = PeriodicGaitGeneratorJax(duty_factor=0.65, step_freq=1.65, horizon=self.horizon, mpc_dt = self.dt) + self.pgg = PeriodicGaitGeneratorJax(duty_factor=0.65, step_freq=1.65, horizon=self.horizon, mpc_dt=self.dt) _, _ = self.pgg.compute_contact_sequence(simulation_dt=0.002, t=self.pgg.get_t(), step_freq=1.65) self.jitted_compute_contact_sequence = jax.jit(self.pgg.compute_contact_sequence, device=self.device) - # jitting the vmap function! self.vectorized_rollout = jax.vmap(self.compute_rollout, in_axes=(None, None, None, 0, 0), out_axes=0) self.jit_vectorized_rollout = jax.jit(self.vectorized_rollout, device=self.device) - # the first call of jax is very slow, hence we should do this since the beginning + # the first call of jax is very slow, hence we should do this since the beginning # creating a fake initial state, reference and contact sequence initial_state = jnp.zeros((self.state_dim,), dtype=dtype_general) initial_reference = jnp.zeros((self.reference_dim,), dtype=dtype_general) - - self.control_parameters_vec = random.uniform(self.master_key, (self.num_control_parameters*self.num_parallel_computations, ), minval=-100., maxval=100.) + self.control_parameters_vec = random.uniform( + self.master_key, + (self.num_control_parameters * self.num_parallel_computations,), + minval=-100.0, + maxval=100.0, + ) step_frequencies_vec = jax.random.uniform(self.master_key, shape=(self.num_parallel_computations, 1)) - - + temp = self.pgg.get_t() temp2 = self.control_parameters_vec.reshape(self.num_parallel_computations, self.num_control_parameters) - temp3 = step_frequencies_vec.reshape(self.num_parallel_computations, ) - self.jit_vectorized_rollout(initial_state, initial_reference, - temp, - temp2, - temp3) - - - self.step_freq_delta = jnp.array(config.mpc_params['step_freq_available']) - - - - + temp3 = step_frequencies_vec.reshape( + self.num_parallel_computations, + ) + self.jit_vectorized_rollout(initial_state, initial_reference, temp, temp2, temp3) + self.step_freq_delta = jnp.array(config.mpc_params["step_freq_available"]) def compute_linear_spline_1(self, parameters, step, horizon_leg): """ Compute a linear spline parametrization of the GRF - """ + """ - tau = step/(horizon_leg) + tau = step / (horizon_leg) - q = (tau - 0.0)/(1.0-0.0) - + q = (tau - 0.0) / (1.0 - 0.0) index = 0 - f_x = (1-q)*parameters[index+0] + q*parameters[index+1] - f_y = (1-q)*parameters[index+2] + q*parameters[index+3] - f_z = (1-q)*parameters[index+4] + q*parameters[index+5] - - return f_x, f_y, f_z - + f_x = (1 - q) * parameters[index + 0] + q * parameters[index + 1] + f_y = (1 - q) * parameters[index + 2] + q * parameters[index + 3] + f_z = (1 - q) * parameters[index + 4] + q * parameters[index + 5] + return f_x, f_y, f_z def compute_linear_spline_2(self, parameters, step, horizon_leg): """ Compute the linear spline parametrization of the GRF (2 splines) - """ + """ + + tau = step / (horizon_leg) - tau = step/(horizon_leg) - index = 0 - index = jax.numpy.where(step > self.horizon/2, 1, index) + index = jax.numpy.where(step > self.horizon / 2, 1, index) - tau = jax.numpy.where(step > self.horizon/2, tau-1, tau) - q = (tau - 0.0)/(1.0-0.0) + tau = jax.numpy.where(step > self.horizon / 2, tau - 1, tau) + q = (tau - 0.0) / (1.0 - 0.0) - - f_x = (1-q)*parameters[index+0] + q*parameters[index+1] - f_y = (1-q)*parameters[index+3] + q*parameters[index+4] - f_z = (1-q)*parameters[index+6] + q*parameters[index+7] + f_x = (1 - q) * parameters[index + 0] + q * parameters[index + 1] + f_y = (1 - q) * parameters[index + 3] + q * parameters[index + 4] + f_z = (1 - q) * parameters[index + 6] + q * parameters[index + 7] - return f_x, f_y, f_z - + return f_x, f_y, f_z def compute_linear_spline_N(self, parameters, step, horizon_leg): """ Compute the linear spline parametrization of the GRF (N splines) - """ + """ + + tau = step / (horizon_leg) - tau = step/(horizon_leg) - # Adding the last boundary for the case when step is exactly self.horizon chunk_boundaries = jnp.linspace(0, self.horizon, self.num_spline + 1) # Find the chunk index by checking in which interval the step falls index = jnp.max(jnp.where(step >= chunk_boundaries, jnp.arange(self.num_spline + 1), 0)) - tau = jax.numpy.where(jnp.isin(step, chunk_boundaries[1:]), tau - 1, tau) - q = (tau - 0.0)/(1.0-0.0) - - - shift = (self.num_spline + 1) - f_x = (1-q)*parameters[index+0] + q*parameters[index+1] - f_y = (1-q)*parameters[index+shift] + q*parameters[index+shift+1] - f_z = (1-q)*parameters[index+shift*2] + q*parameters[index+shift*2+1] - - return f_x, f_y, f_z - - - + q = (tau - 0.0) / (1.0 - 0.0) + + shift = self.num_spline + 1 + f_x = (1 - q) * parameters[index + 0] + q * parameters[index + 1] + f_y = (1 - q) * parameters[index + shift] + q * parameters[index + shift + 1] + f_z = (1 - q) * parameters[index + shift * 2] + q * parameters[index + shift * 2 + 1] + + return f_x, f_y, f_z + def compute_cubic_spline(self, parameters, step, horizon_leg): """ Compute the cubic spline parametrization of the GRF - """ + """ - tau = step/(horizon_leg) - q = (tau - 0.0)/(1.0-0.0) - a = 2*q*q*q - 3*q*q + 1 - b = (q*q*q - 2*q*q + q)*1.0 - c = -2*q*q*q + 3*q*q - d = (q*q*q - q*q)*1.0 + tau = step / (horizon_leg) + q = (tau - 0.0) / (1.0 - 0.0) + a = 2 * q * q * q - 3 * q * q + 1 + b = (q * q * q - 2 * q * q + q) * 1.0 + c = -2 * q * q * q + 3 * q * q + d = (q * q * q - q * q) * 1.0 - phi = (1./2.)*(((parameters[2] - parameters[1])/1.0) + ((parameters[1] - parameters[0])/1.0)) - phi_next = (1./2.)*(((parameters[3] - parameters[2])/1.0) + ((parameters[2] - parameters[1])/1.0)) - f_x = a*parameters[1] + b*phi + c*parameters[2] + d*phi_next + phi = (1.0 / 2.0) * (((parameters[2] - parameters[1]) / 1.0) + ((parameters[1] - parameters[0]) / 1.0)) + phi_next = (1.0 / 2.0) * (((parameters[3] - parameters[2]) / 1.0) + ((parameters[2] - parameters[1]) / 1.0)) + f_x = a * parameters[1] + b * phi + c * parameters[2] + d * phi_next - phi = (1./2.)*(((parameters[6] - parameters[5])/1.0) + ((parameters[5] - parameters[4])/1.0)) - phi_next = (1./2.)*(((parameters[7] - parameters[6])/1.0) + ((parameters[6] - parameters[5])/1.0)) - f_y = a*parameters[5] + b*phi + c*parameters[6] + d*phi_next + phi = (1.0 / 2.0) * (((parameters[6] - parameters[5]) / 1.0) + ((parameters[5] - parameters[4]) / 1.0)) + phi_next = (1.0 / 2.0) * (((parameters[7] - parameters[6]) / 1.0) + ((parameters[6] - parameters[5]) / 1.0)) + f_y = a * parameters[5] + b * phi + c * parameters[6] + d * phi_next + phi = (1.0 / 2.0) * (((parameters[10] - parameters[9]) / 1.0) + ((parameters[9] - parameters[8]) / 1.0)) + phi_next = (1.0 / 2.0) * (((parameters[11] - parameters[10]) / 1.0) + ((parameters[10] - parameters[9]) / 1.0)) + f_z = a * parameters[9] + b * phi + c * parameters[10] + d * phi_next - phi = (1./2.)*(((parameters[10] - parameters[9])/1.0) + ((parameters[9] - parameters[8])/1.0)) - phi_next = (1./2.)*(((parameters[11] - parameters[10])/1.0) + ((parameters[10] - parameters[9])/1.0)) - f_z = a*parameters[9] + b*phi + c*parameters[10] + d*phi_next - return f_x, f_y, f_z - def compute_cubic_spline_2(self, parameters, step, horizon_leg): """ Compute the cubic spline parametrization of the GRF - """ + """ - tau = step/(horizon_leg) + tau = step / (horizon_leg) index = 0 - index = jax.numpy.where(step > self.horizon/2, 1, index) - tau = jax.numpy.where(step > self.horizon/2, tau-1, tau) - - start_index = 10*index - - - q = (tau - 0.0)/(1.0-0.0) - a = 2*q*q*q - 3*q*q + 1 - b = (q*q*q - 2*q*q + q)*1.0 - c = -2*q*q*q + 3*q*q - d = (q*q*q - q*q)*1.0 - - phi = (1./2.)*(((parameters[start_index+2] - parameters[start_index+1])/1.0) + ((parameters[start_index+1] - parameters[start_index+0])/1.0)) - phi_next = (1./2.)*(((parameters[start_index+3] - parameters[start_index+2])/1.0) + ((parameters[start_index+2] - parameters[start_index+1])/1.0)) - f_x = a*parameters[start_index+1] + b*phi + c*parameters[start_index+2] + d*phi_next - - phi = (1./2.)*(((parameters[start_index+6] - parameters[start_index+5])/1.0) + ((parameters[start_index+5] - parameters[start_index+4])/1.0)) - phi_next = (1./2.)*(((parameters[start_index+7] - parameters[start_index+6])/1.0) + ((parameters[start_index+6] - parameters[start_index+5])/1.0)) - f_y = a*parameters[start_index+5] + b*phi + c*parameters[start_index+6] + d*phi_next - + index = jax.numpy.where(step > self.horizon / 2, 1, index) + tau = jax.numpy.where(step > self.horizon / 2, tau - 1, tau) + + start_index = 10 * index + + q = (tau - 0.0) / (1.0 - 0.0) + a = 2 * q * q * q - 3 * q * q + 1 + b = (q * q * q - 2 * q * q + q) * 1.0 + c = -2 * q * q * q + 3 * q * q + d = (q * q * q - q * q) * 1.0 + + phi = (1.0 / 2.0) * ( + ((parameters[start_index + 2] - parameters[start_index + 1]) / 1.0) + + ((parameters[start_index + 1] - parameters[start_index + 0]) / 1.0) + ) + phi_next = (1.0 / 2.0) * ( + ((parameters[start_index + 3] - parameters[start_index + 2]) / 1.0) + + ((parameters[start_index + 2] - parameters[start_index + 1]) / 1.0) + ) + f_x = a * parameters[start_index + 1] + b * phi + c * parameters[start_index + 2] + d * phi_next + + phi = (1.0 / 2.0) * ( + ((parameters[start_index + 6] - parameters[start_index + 5]) / 1.0) + + ((parameters[start_index + 5] - parameters[start_index + 4]) / 1.0) + ) + phi_next = (1.0 / 2.0) * ( + ((parameters[start_index + 7] - parameters[start_index + 6]) / 1.0) + + ((parameters[start_index + 6] - parameters[start_index + 5]) / 1.0) + ) + f_y = a * parameters[start_index + 5] + b * phi + c * parameters[start_index + 6] + d * phi_next + + phi = (1.0 / 2.0) * ( + ((parameters[start_index + 10] - parameters[start_index + 9]) / 1.0) + + ((parameters[start_index + 9] - parameters[start_index + 8]) / 1.0) + ) + phi_next = (1.0 / 2.0) * ( + ((parameters[start_index + 11] - parameters[start_index + 10]) / 1.0) + + ((parameters[start_index + 10] - parameters[start_index + 9]) / 1.0) + ) + f_z = a * parameters[start_index + 9] + b * phi + c * parameters[start_index + 10] + d * phi_next - phi = (1./2.)*(((parameters[start_index+10] - parameters[start_index+9])/1.0) + ((parameters[start_index+9] - parameters[start_index+8])/1.0)) - phi_next = (1./2.)*(((parameters[start_index+11] - parameters[start_index+10])/1.0) + ((parameters[start_index+10] - parameters[start_index+9])/1.0)) - f_z = a*parameters[start_index+9] + b*phi + c*parameters[start_index+10] + d*phi_next - return f_x, f_y, f_z - def compute_cubic_spline_N(self, parameters, step, horizon_leg): """ Compute the cubic spline parametrization of the GRF (N splines) - """ + """ + + tau = step / (horizon_leg) - tau = step/(horizon_leg) - # Adding the last boundary for the case when step is exactly self.horizon chunk_boundaries = jnp.linspace(0, self.horizon, self.num_spline + 1) # Find the chunk index by checking in which interval the step falls index = jnp.max(jnp.where(step >= chunk_boundaries, jnp.arange(self.num_spline + 1), 0)) - tau = jax.numpy.where(jnp.isin(step, chunk_boundaries[1:]), tau - 1, tau) - q = (tau - 0.0)/(1.0-0.0) - - start_index = 10*index - - - q = (tau - 0.0)/(1.0-0.0) - a = 2*q*q*q - 3*q*q + 1 - b = (q*q*q - 2*q*q + q)*1.0 - c = -2*q*q*q + 3*q*q - d = (q*q*q - q*q)*1.0 + q = (tau - 0.0) / (1.0 - 0.0) + + start_index = 10 * index + + q = (tau - 0.0) / (1.0 - 0.0) + a = 2 * q * q * q - 3 * q * q + 1 + b = (q * q * q - 2 * q * q + q) * 1.0 + c = -2 * q * q * q + 3 * q * q + d = (q * q * q - q * q) * 1.0 + + phi = (1.0 / 2.0) * ( + ((parameters[start_index + 2] - parameters[start_index + 1]) / 1.0) + + ((parameters[start_index + 1] - parameters[start_index + 0]) / 1.0) + ) + phi_next = (1.0 / 2.0) * ( + ((parameters[start_index + 3] - parameters[start_index + 2]) / 1.0) + + ((parameters[start_index + 2] - parameters[start_index + 1]) / 1.0) + ) + f_x = a * parameters[start_index + 1] + b * phi + c * parameters[start_index + 2] + d * phi_next + + phi = (1.0 / 2.0) * ( + ((parameters[start_index + 6] - parameters[start_index + 5]) / 1.0) + + ((parameters[start_index + 5] - parameters[start_index + 4]) / 1.0) + ) + phi_next = (1.0 / 2.0) * ( + ((parameters[start_index + 7] - parameters[start_index + 6]) / 1.0) + + ((parameters[start_index + 6] - parameters[start_index + 5]) / 1.0) + ) + f_y = a * parameters[start_index + 5] + b * phi + c * parameters[start_index + 6] + d * phi_next + + phi = (1.0 / 2.0) * ( + ((parameters[start_index + 10] - parameters[start_index + 9]) / 1.0) + + ((parameters[start_index + 9] - parameters[start_index + 8]) / 1.0) + ) + phi_next = (1.0 / 2.0) * ( + ((parameters[start_index + 11] - parameters[start_index + 10]) / 1.0) + + ((parameters[start_index + 10] - parameters[start_index + 9]) / 1.0) + ) + f_z = a * parameters[start_index + 9] + b * phi + c * parameters[start_index + 10] + d * phi_next - phi = (1./2.)*(((parameters[start_index+2] - parameters[start_index+1])/1.0) + ((parameters[start_index+1] - parameters[start_index+0])/1.0)) - phi_next = (1./2.)*(((parameters[start_index+3] - parameters[start_index+2])/1.0) + ((parameters[start_index+2] - parameters[start_index+1])/1.0)) - f_x = a*parameters[start_index+1] + b*phi + c*parameters[start_index+2] + d*phi_next - - phi = (1./2.)*(((parameters[start_index+6] - parameters[start_index+5])/1.0) + ((parameters[start_index+5] - parameters[start_index+4])/1.0)) - phi_next = (1./2.)*(((parameters[start_index+7] - parameters[start_index+6])/1.0) + ((parameters[start_index+6] - parameters[start_index+5])/1.0)) - f_y = a*parameters[start_index+5] + b*phi + c*parameters[start_index+6] + d*phi_next - - - phi = (1./2.)*(((parameters[start_index+10] - parameters[start_index+9])/1.0) + ((parameters[start_index+9] - parameters[start_index+8])/1.0)) - phi_next = (1./2.)*(((parameters[start_index+11] - parameters[start_index+10])/1.0) + ((parameters[start_index+10] - parameters[start_index+9])/1.0)) - f_z = a*parameters[start_index+9] + b*phi + c*parameters[start_index+10] + d*phi_next - return f_x, f_y, f_z - def compute_zero_order_spline(self, parameters, step, horizon_leg): """ Compute the zero-order control input - """ + """ index = jnp.int16(step) f_x = parameters[index] - f_y = parameters[index+self.horizon] - f_z = parameters[index+self.horizon*2] - return f_x, f_y, f_z - - - + f_y = parameters[index + self.horizon] + f_z = parameters[index + self.horizon * 2] + return f_x, f_y, f_z - def enforce_force_constraints(self, f_x_FL, f_y_FL, f_z_FL, - f_x_FR, f_y_FR, f_z_FR, - f_x_RL, f_y_RL, f_z_RL, - f_x_RR, f_y_RR, f_z_RR): + def enforce_force_constraints( + self, f_x_FL, f_y_FL, f_z_FL, f_x_FR, f_y_FR, f_z_FR, f_x_RL, f_y_RL, f_z_RL, f_x_RR, f_y_RR, f_z_RR + ): """ Enforce the friction cone and the force limits constraints - """ - + """ + # Enforce push-only of the ground! f_z_FL = jax.numpy.where(f_z_FL > self.f_z_min, f_z_FL, self.f_z_min) f_z_FR = jax.numpy.where(f_z_FR > self.f_z_min, f_z_FR, self.f_z_min) f_z_RL = jax.numpy.where(f_z_RL > self.f_z_min, f_z_RL, self.f_z_min) f_z_RR = jax.numpy.where(f_z_RR > self.f_z_min, f_z_RR, self.f_z_min) - # Enforce maximum force per leg! - f_z_FL = jax.numpy.where(f_z_FL-self.mu*f_z_FL, f_x_FL,-self.mu*f_z_FL) - f_x_FL = jax.numpy.where(f_x_FL-self.mu*f_z_FL, f_y_FL,-self.mu*f_z_FL) - f_y_FL = jax.numpy.where(f_y_FL-self.mu*f_z_FR, f_x_FR,-self.mu*f_z_FR) - f_x_FR = jax.numpy.where(f_x_FR-self.mu*f_z_FR, f_y_FR,-self.mu*f_z_FR) - f_y_FR = jax.numpy.where(f_y_FR-self.mu*f_z_RL, f_x_RL,-self.mu*f_z_RL) - f_x_RL = jax.numpy.where(f_x_RL-self.mu*f_z_RL, f_y_RL,-self.mu*f_z_RL) - f_y_RL = jax.numpy.where(f_y_RL-self.mu*f_z_RR, f_x_RR,-self.mu*f_z_RR) - f_x_RR = jax.numpy.where(f_x_RR-self.mu*f_z_RR, f_y_RR,-self.mu*f_z_RR) - f_y_RR = jax.numpy.where(f_y_RR -self.mu * f_z_FL, f_x_FL, -self.mu * f_z_FL) + f_x_FL = jax.numpy.where(f_x_FL < self.mu * f_z_FL, f_x_FL, self.mu * f_z_FL) + f_y_FL = jax.numpy.where(f_y_FL > -self.mu * f_z_FL, f_y_FL, -self.mu * f_z_FL) + f_y_FL = jax.numpy.where(f_y_FL < self.mu * f_z_FL, f_y_FL, self.mu * f_z_FL) + + f_x_FR = jax.numpy.where(f_x_FR > -self.mu * f_z_FR, f_x_FR, -self.mu * f_z_FR) + f_x_FR = jax.numpy.where(f_x_FR < self.mu * f_z_FR, f_x_FR, self.mu * f_z_FR) + f_y_FR = jax.numpy.where(f_y_FR > -self.mu * f_z_FR, f_y_FR, -self.mu * f_z_FR) + f_y_FR = jax.numpy.where(f_y_FR < self.mu * f_z_FR, f_y_FR, self.mu * f_z_FR) + + f_x_RL = jax.numpy.where(f_x_RL > -self.mu * f_z_RL, f_x_RL, -self.mu * f_z_RL) + f_x_RL = jax.numpy.where(f_x_RL < self.mu * f_z_RL, f_x_RL, self.mu * f_z_RL) + f_y_RL = jax.numpy.where(f_y_RL > -self.mu * f_z_RL, f_y_RL, -self.mu * f_z_RL) + f_y_RL = jax.numpy.where(f_y_RL < self.mu * f_z_RL, f_y_RL, self.mu * f_z_RL) + + f_x_RR = jax.numpy.where(f_x_RR > -self.mu * f_z_RR, f_x_RR, -self.mu * f_z_RR) + f_x_RR = jax.numpy.where(f_x_RR < self.mu * f_z_RR, f_x_RR, self.mu * f_z_RR) + f_y_RR = jax.numpy.where(f_y_RR > -self.mu * f_z_RR, f_y_RR, -self.mu * f_z_RR) + f_y_RR = jax.numpy.where(f_y_RR < self.mu * f_z_RR, f_y_RR, self.mu * f_z_RR) + + return f_x_FL, f_y_FL, f_z_FL, f_x_FR, f_y_FR, f_z_FR, f_x_RL, f_y_RL, f_z_RL, f_x_RR, f_y_RR, f_z_RR def compute_rollout(self, initial_state, reference, timing, control_parameters, step_frequency): """Calculate cost of a rollout of the dynamics given random parameters @@ -533,292 +523,371 @@ def compute_rollout(self, initial_state, reference, timing, control_parameters, parameters (np.array): parameters for the simplified dynamics Returns: (float): cost of the rollout - """ + """ state = initial_state cost = jnp.float32(0.0) - n_ = jnp.array([-1,-1,-1,-1]) + n_ = jnp.array([-1, -1, -1, -1]) + + contact_sequence, _ = self.jitted_compute_contact_sequence( + simulation_dt=0.002, t=timing, step_freq=step_frequency + ) - contact_sequence, _ = self.jitted_compute_contact_sequence(simulation_dt=0.002, t=timing, step_freq=step_frequency) + FL_num_of_contact = jnp.sum(contact_sequence[0]) + 1 + FR_num_of_contact = jnp.sum(contact_sequence[1]) + 1 + RL_num_of_contact = jnp.sum(contact_sequence[2]) + 1 + RR_num_of_contact = jnp.sum(contact_sequence[3]) + 1 - FL_num_of_contact = jnp.sum(contact_sequence[0])+1 - FR_num_of_contact = jnp.sum(contact_sequence[1])+1 - RL_num_of_contact = jnp.sum(contact_sequence[2])+1 - RR_num_of_contact = jnp.sum(contact_sequence[3])+1 - def iterate_fun(n, carry): cost, state, reference, n_ = carry - - n_ = n_.at[0].set(n_[0]+1*contact_sequence[0][n]) - n_ = n_.at[1].set(n_[1]+1*contact_sequence[1][n]) - n_ = n_.at[2].set(n_[2]+1*contact_sequence[2][n]) - n_ = n_.at[3].set(n_[3]+1*contact_sequence[3][n]) - - f_x_FL, f_y_FL, f_z_FL = self.spline_fun_FL(control_parameters[0:self.num_control_parameters_single_leg], n_[0], FL_num_of_contact) - f_x_FR, f_y_FR, f_z_FR = self.spline_fun_FR(control_parameters[self.num_control_parameters_single_leg:self.num_control_parameters_single_leg*2], n_[1], FR_num_of_contact) - f_x_RL, f_y_RL, f_z_RL = self.spline_fun_RL(control_parameters[self.num_control_parameters_single_leg*2:self.num_control_parameters_single_leg*3], n_[2], RL_num_of_contact) - f_x_RR, f_y_RR, f_z_RR = self.spline_fun_RR(control_parameters[self.num_control_parameters_single_leg*3:self.num_control_parameters_single_leg*4], n_[3], RR_num_of_contact) - - - + n_ = n_.at[0].set(n_[0] + 1 * contact_sequence[0][n]) + n_ = n_.at[1].set(n_[1] + 1 * contact_sequence[1][n]) + n_ = n_.at[2].set(n_[2] + 1 * contact_sequence[2][n]) + n_ = n_.at[3].set(n_[3] + 1 * contact_sequence[3][n]) + + f_x_FL, f_y_FL, f_z_FL = self.spline_fun_FL( + control_parameters[0 : self.num_control_parameters_single_leg], n_[0], FL_num_of_contact + ) + f_x_FR, f_y_FR, f_z_FR = self.spline_fun_FR( + control_parameters[self.num_control_parameters_single_leg : self.num_control_parameters_single_leg * 2], + n_[1], + FR_num_of_contact, + ) + f_x_RL, f_y_RL, f_z_RL = self.spline_fun_RL( + control_parameters[ + self.num_control_parameters_single_leg * 2 : self.num_control_parameters_single_leg * 3 + ], + n_[2], + RL_num_of_contact, + ) + f_x_RR, f_y_RR, f_z_RR = self.spline_fun_RR( + control_parameters[ + self.num_control_parameters_single_leg * 3 : self.num_control_parameters_single_leg * 4 + ], + n_[3], + RR_num_of_contact, + ) # The sampling over f_z is a delta over gravity compensation (only for the leg in stance!) - number_of_legs_in_stance = contact_sequence[0][n] + contact_sequence[1][n] + contact_sequence[2][n] + contact_sequence[3][n] + number_of_legs_in_stance = ( + contact_sequence[0][n] + contact_sequence[1][n] + contact_sequence[2][n] + contact_sequence[3][n] + ) reference_force_stance_legs = (self.robot.mass * 9.81) / number_of_legs_in_stance - - f_z_FL = reference_force_stance_legs + f_z_FL f_z_FR = reference_force_stance_legs + f_z_FR f_z_RL = reference_force_stance_legs + f_z_RL f_z_RR = reference_force_stance_legs + f_z_RR - # Foot in swing (contact sequence = 0) have zero force - f_x_FL = f_x_FL*contact_sequence[0][n] - f_y_FL = f_y_FL*contact_sequence[0][n] - f_z_FL = f_z_FL*contact_sequence[0][n] - - f_x_FR = f_x_FR*contact_sequence[1][n] - f_y_FR = f_y_FR*contact_sequence[1][n] - f_z_FR = f_z_FR*contact_sequence[1][n] + f_x_FL = f_x_FL * contact_sequence[0][n] + f_y_FL = f_y_FL * contact_sequence[0][n] + f_z_FL = f_z_FL * contact_sequence[0][n] - f_x_RL = f_x_RL*contact_sequence[2][n] - f_y_RL = f_y_RL*contact_sequence[2][n] - f_z_RL = f_z_RL*contact_sequence[2][n] + f_x_FR = f_x_FR * contact_sequence[1][n] + f_y_FR = f_y_FR * contact_sequence[1][n] + f_z_FR = f_z_FR * contact_sequence[1][n] - f_x_RR = f_x_RR*contact_sequence[3][n] - f_y_RR = f_y_RR*contact_sequence[3][n] - f_z_RR = f_z_RR*contact_sequence[3][n] + f_x_RL = f_x_RL * contact_sequence[2][n] + f_y_RL = f_y_RL * contact_sequence[2][n] + f_z_RL = f_z_RL * contact_sequence[2][n] - + f_x_RR = f_x_RR * contact_sequence[3][n] + f_y_RR = f_y_RR * contact_sequence[3][n] + f_z_RR = f_z_RR * contact_sequence[3][n] # Enforce force constraints - f_x_FL, f_y_FL, f_z_FL, \ - f_x_FR, f_y_FR, f_z_FR, \ - f_x_RL, f_y_RL, f_z_RL, \ - f_x_RR, f_y_RR, f_z_RR = self.enforce_force_constraints(f_x_FL, f_y_FL, f_z_FL, - f_x_FR, f_y_FR, f_z_FR, - f_x_RL, f_y_RL, f_z_RL, - f_x_RR, f_y_RR, f_z_RR) - - - - input = jnp.array([jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - f_x_FL, f_y_FL, f_z_FL, # foot position fl - f_x_FR, f_y_FR, f_z_FR, # foot position fr - f_x_RL, f_y_RL, f_z_RL, # foot position rl - f_x_RR, f_y_RR, f_z_RR, # foot position rr - ], dtype=dtype_general) - - + f_x_FL, f_y_FL, f_z_FL, f_x_FR, f_y_FR, f_z_FR, f_x_RL, f_y_RL, f_z_RL, f_x_RR, f_y_RR, f_z_RR = ( + self.enforce_force_constraints( + f_x_FL, f_y_FL, f_z_FL, f_x_FR, f_y_FR, f_z_FR, f_x_RL, f_y_RL, f_z_RL, f_x_RR, f_y_RR, f_z_RR + ) + ) + + input = jnp.array( + [ + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + f_x_FL, + f_y_FL, + f_z_FL, # foot position fl + f_x_FR, + f_y_FR, + f_z_FR, # foot position fr + f_x_RL, + f_y_RL, + f_z_RL, # foot position rl + f_x_RR, + f_y_RR, + f_z_RR, # foot position rr + ], + dtype=dtype_general, + ) + # Integrate the dynamics - current_contact = jnp.array([contact_sequence[0][n], contact_sequence[1][n], - contact_sequence[2][n], contact_sequence[3][n]], dtype=dtype_general) + current_contact = jnp.array( + [contact_sequence[0][n], contact_sequence[1][n], contact_sequence[2][n], contact_sequence[3][n]], + dtype=dtype_general, + ) state_next = self.robot.integrate_jax(state, input, current_contact, n) - - - + # Calculate cost regulation state - state_error = state_next - reference[0:self.state_dim] - error_cost = state_error.T@self.Q@state_error - - input_for_cost = jnp.array([jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - f_x_FL, f_y_FL, f_z_FL - reference_force_stance_legs, - f_x_FR, f_y_FR, f_z_FR - reference_force_stance_legs, - f_x_RL, f_y_RL, f_z_RL - reference_force_stance_legs, - f_x_RR, f_y_RR, f_z_RR - reference_force_stance_legs, - ], dtype=dtype_general) - - # Calculate cost regulation input - #error_cost += input_for_cost.T@self.R@input_for_cost + state_error = state_next - reference[0 : self.state_dim] + error_cost = state_error.T @ self.Q @ state_error + + input_for_cost = jnp.array( + [ + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + f_x_FL, + f_y_FL, + f_z_FL - reference_force_stance_legs, + f_x_FR, + f_y_FR, + f_z_FR - reference_force_stance_legs, + f_x_RL, + f_y_RL, + f_z_RL - reference_force_stance_legs, + f_x_RR, + f_y_RR, + f_z_RR - reference_force_stance_legs, + ], + dtype=dtype_general, + ) - + # Calculate cost regulation input + # error_cost += input_for_cost.T@self.R@input_for_cost # saturate in the case of NaN - #state_next = jnp.where(jnp.isnan(state_next), 100, state_next) - #error_cost = jnp.where(jnp.isnan(error_cost), 1000, error_cost) - #cost = jnp.where(jnp.isnan(cost), 10000, cost) - - + # state_next = jnp.where(jnp.isnan(state_next), 100, state_next) + # error_cost = jnp.where(jnp.isnan(error_cost), 1000, error_cost) + # cost = jnp.where(jnp.isnan(cost), 10000, cost) return (cost + error_cost, state_next, reference, n_) carry = (cost, state, reference, n_) cost, state, reference, n_ = jax.lax.fori_loop(0, self.horizon, iterate_fun, carry) - - cost += (step_frequency - 1.3)*100*(step_frequency - 1.3) - return cost + cost += (step_frequency - 1.3) * 100 * (step_frequency - 1.3) + return cost - def with_newkey(self): newkey, subkey = jax.random.split(self.master_key) self.master_key = newkey return self - + def get_key(self): return self.master_key - + def with_newsigma(self, sigma): self.sigma_cem_mppi = sigma return self - + def get_sigma(self): return self.sigma_cem_mppi - - - def shift_solution(self, best_control_parameters, step): """ This function shift the control parameter ahed - """ - + """ + best_control_parameters = np.array(best_control_parameters) - FL_control = copy.deepcopy(best_control_parameters[0:self.num_control_parameters_single_leg]) - FR_control = copy.deepcopy(best_control_parameters[self.num_control_parameters_single_leg:self.num_control_parameters_single_leg*2]) - RL_control = copy.deepcopy(best_control_parameters[self.num_control_parameters_single_leg*2:self.num_control_parameters_single_leg*3]) - RR_control = copy.deepcopy(best_control_parameters[self.num_control_parameters_single_leg*3:self.num_control_parameters_single_leg*4]) - + FL_control = copy.deepcopy(best_control_parameters[0 : self.num_control_parameters_single_leg]) + FR_control = copy.deepcopy( + best_control_parameters[self.num_control_parameters_single_leg : self.num_control_parameters_single_leg * 2] + ) + RL_control = copy.deepcopy( + best_control_parameters[ + self.num_control_parameters_single_leg * 2 : self.num_control_parameters_single_leg * 3 + ] + ) + RR_control = copy.deepcopy( + best_control_parameters[ + self.num_control_parameters_single_leg * 3 : self.num_control_parameters_single_leg * 4 + ] + ) FL_control_temp = copy.deepcopy(FL_control) FL_control[0], FL_control[2], FL_control[4] = self.spline_fun_FL(FL_control_temp, step) - #FL_control[1], FL_control[3], FL_control[5] = controller.spline_fun_FL(FL_control_temp, controller.horizon+step) + # FL_control[1], FL_control[3], FL_control[5] = controller.spline_fun_FL(FL_control_temp, controller.horizon+step) FR_control_temp = copy.deepcopy(FR_control) FR_control[0], FR_control[2], FR_control[4] = self.spline_fun_FR(FR_control_temp, step) - #FR_control[1], FR_control[3], FR_control[5] = controller.spline_fun_FR(FR_control_temp, controller.horizon+step) - - + # FR_control[1], FR_control[3], FR_control[5] = controller.spline_fun_FR(FR_control_temp, controller.horizon+step) + RL_control_temp = copy.deepcopy(RL_control) RL_control[0], RL_control[2], RL_control[4] = self.spline_fun_RL(RL_control_temp, step) - #RL_control[1], RL_control[3], RL_control[5] = controller.spline_fun_RL(RL_control_temp, controller.horizon+step) - - + # RL_control[1], RL_control[3], RL_control[5] = controller.spline_fun_RL(RL_control_temp, controller.horizon+step) + RR_control_temp = copy.deepcopy(RR_control) RR_control[0], RR_control[2], RR_control[4] = self.spline_fun_RR(RR_control_temp, step) - #RR_control[1], RR_control[3], RR_control[5] = controller.spline_fun_RR(RR_control_temp, controller.horizon+step) - - - best_control_parameters[0:self.num_control_parameters_single_leg] = FL_control - best_control_parameters[self.num_control_parameters_single_leg:self.num_control_parameters_single_leg*2] = FR_control - best_control_parameters[self.num_control_parameters_single_leg*2:self.num_control_parameters_single_leg*3] = RL_control - best_control_parameters[self.num_control_parameters_single_leg*3:self.num_control_parameters_single_leg*4] = RR_control - - return best_control_parameters - - + # RR_control[1], RR_control[3], RR_control[5] = controller.spline_fun_RR(RR_control_temp, controller.horizon+step) + + best_control_parameters[0 : self.num_control_parameters_single_leg] = FL_control + best_control_parameters[self.num_control_parameters_single_leg : self.num_control_parameters_single_leg * 2] = ( + FR_control + ) + best_control_parameters[ + self.num_control_parameters_single_leg * 2 : self.num_control_parameters_single_leg * 3 + ] = RL_control + best_control_parameters[ + self.num_control_parameters_single_leg * 3 : self.num_control_parameters_single_leg * 4 + ] = RR_control + return best_control_parameters - def prepare_state_and_reference(self, state_current, reference_state, - current_contact, previous_contact, mpc_frequency=100): + def prepare_state_and_reference( + self, state_current, reference_state, current_contact, previous_contact, mpc_frequency=100 + ): """ This function jaxify the current state and reference for further processing. - """ + """ # Shift the previous solution ahead - if (config.mpc_params['shift_solution']): - index_shift = 1./mpc_frequency + if config.mpc_params["shift_solution"]: + index_shift = 1.0 / mpc_frequency self.best_control_parameters = self.shift_solution(self.best_control_parameters, index_shift) - - - state_current_jax = np.concatenate((state_current["position"], state_current["linear_velocity"], - state_current["orientation"], state_current["angular_velocity"], - state_current["foot_FL"], state_current["foot_FR"], - state_current["foot_RL"], state_current["foot_RR"])).reshape((24, )) - - if(current_contact[0] == 0.): + state_current_jax = np.concatenate( + ( + state_current["position"], + state_current["linear_velocity"], + state_current["orientation"], + state_current["angular_velocity"], + state_current["foot_FL"], + state_current["foot_FR"], + state_current["foot_RL"], + state_current["foot_RR"], + ) + ).reshape((24,)) + + if current_contact[0] == 0.0: state_current_jax[12:15] = reference_state["ref_foot_FL"].reshape((3,)) - if(current_contact[1] == 0.): + if current_contact[1] == 0.0: state_current_jax[15:18] = reference_state["ref_foot_FR"].reshape((3,)) - if(current_contact[2] == 0.): - state_current_jax[18:21] = reference_state["ref_foot_RL"].reshape((3,)) - if(current_contact[3] == 0.): + if current_contact[2] == 0.0: + state_current_jax[18:21] = reference_state["ref_foot_RL"].reshape((3,)) + if current_contact[3] == 0.0: state_current_jax[21:24] = reference_state["ref_foot_RR"].reshape((3,)) - - reference_state_jax = np.concatenate((reference_state["ref_position"], reference_state["ref_linear_velocity"], - reference_state["ref_orientation"], reference_state["ref_angular_velocity"], - reference_state["ref_foot_FL"].reshape((3,)), reference_state["ref_foot_FR"].reshape((3,)), - reference_state["ref_foot_RL"].reshape((3,)), reference_state["ref_foot_RR"].reshape((3,)))).reshape((24, )) - + reference_state_jax = np.concatenate( + ( + reference_state["ref_position"], + reference_state["ref_linear_velocity"], + reference_state["ref_orientation"], + reference_state["ref_angular_velocity"], + reference_state["ref_foot_FL"].reshape((3,)), + reference_state["ref_foot_FR"].reshape((3,)), + reference_state["ref_foot_RL"].reshape((3,)), + reference_state["ref_foot_RR"].reshape((3,)), + ) + ).reshape((24,)) self.best_control_parameters = np.array(self.best_control_parameters) - - if(previous_contact[0] == 1 and current_contact[0] == 0): - self.best_control_parameters[0:self.num_control_parameters_single_leg] = 0.0 - if(previous_contact[1] == 1 and current_contact[1] == 0): - self.best_control_parameters[self.num_control_parameters_single_leg:self.num_control_parameters_single_leg*2] = 0.0 - if(previous_contact[2] == 1 and current_contact[2] == 0): - self.best_control_parameters[self.num_control_parameters_single_leg*2:self.num_control_parameters_single_leg*3] = 0.0 - if(previous_contact[3] == 1 and current_contact[3] == 0): - self.best_control_parameters[self.num_control_parameters_single_leg*3:self.num_control_parameters_single_leg*4] = 0.0 - - return state_current_jax, reference_state_jax - - + if previous_contact[0] == 1 and current_contact[0] == 0: + self.best_control_parameters[0 : self.num_control_parameters_single_leg] = 0.0 + if previous_contact[1] == 1 and current_contact[1] == 0: + self.best_control_parameters[ + self.num_control_parameters_single_leg : self.num_control_parameters_single_leg * 2 + ] = 0.0 + if previous_contact[2] == 1 and current_contact[2] == 0: + self.best_control_parameters[ + self.num_control_parameters_single_leg * 2 : self.num_control_parameters_single_leg * 3 + ] = 0.0 + if previous_contact[3] == 1 and current_contact[3] == 0: + self.best_control_parameters[ + self.num_control_parameters_single_leg * 3 : self.num_control_parameters_single_leg * 4 + ] = 0.0 + return state_current_jax, reference_state_jax - def compute_control_random_sampling(self, state, reference, contact_sequence, best_control_parameters, key, timing, nominal_step_frequency, optimize_swing): + def compute_control_random_sampling( + self, + state, + reference, + contact_sequence, + best_control_parameters, + key, + timing, + nominal_step_frequency, + optimize_swing, + ): """ This function computes the control parameters by sampling from a Gaussian and a uniform distribution. - """ - + """ + # Generate random parameters - - # The first control parameters is the old best one, so we add zero noise there - additional_random_parameters = self.initial_random_parameters*0.0 + # The first control parameters is the old best one, so we add zero noise there + additional_random_parameters = self.initial_random_parameters * 0.0 # FIRST GAUSSIAN - num_sample_gaussian_1 = (1 + int(self.num_parallel_computations/3)) - 1 + num_sample_gaussian_1 = (1 + int(self.num_parallel_computations / 3)) - 1 sigma_gaussian_1 = self.sigma_random_sampling[0] - additional_random_parameters = additional_random_parameters.at[1:1 + int(self.num_parallel_computations/3)].set(sigma_gaussian_1*jax.random.normal(key=key, shape=(num_sample_gaussian_1, self.num_control_parameters))) - + additional_random_parameters = additional_random_parameters.at[ + 1 : 1 + int(self.num_parallel_computations / 3) + ].set(sigma_gaussian_1 * jax.random.normal(key=key, shape=(num_sample_gaussian_1, self.num_control_parameters))) # SECOND GAUSSIAN - num_sample_gaussian_2 = (1 + int(self.num_parallel_computations/3)*2) - (1 + int(self.num_parallel_computations/3)) + num_sample_gaussian_2 = (1 + int(self.num_parallel_computations / 3) * 2) - ( + 1 + int(self.num_parallel_computations / 3) + ) sigma_gaussian_2 = self.sigma_random_sampling[1] - additional_random_parameters = additional_random_parameters.at[1 + int(self.num_parallel_computations/3):1 + int(self.num_parallel_computations/3)*2].set(sigma_gaussian_2*jax.random.normal(key=key, shape=(num_sample_gaussian_2, self.num_control_parameters))) - + additional_random_parameters = additional_random_parameters.at[ + 1 + int(self.num_parallel_computations / 3) : 1 + int(self.num_parallel_computations / 3) * 2 + ].set(sigma_gaussian_2 * jax.random.normal(key=key, shape=(num_sample_gaussian_2, self.num_control_parameters))) # UNIFORM max_sampling_forces = self.sigma_random_sampling[2] - num_samples_uniform = int(self.num_parallel_computations) - (1 + int(self.num_parallel_computations/3)*2) - additional_random_parameters = additional_random_parameters.at[1 + int(self.num_parallel_computations/3)*2:int(self.num_parallel_computations)].set(jax.random.uniform(key=key, minval=-max_sampling_forces, maxval=max_sampling_forces, shape=(num_samples_uniform, self.num_control_parameters ))) - + num_samples_uniform = int(self.num_parallel_computations) - (1 + int(self.num_parallel_computations / 3) * 2) + additional_random_parameters = additional_random_parameters.at[ + 1 + int(self.num_parallel_computations / 3) * 2 : int(self.num_parallel_computations) + ].set( + jax.random.uniform( + key=key, + minval=-max_sampling_forces, + maxval=max_sampling_forces, + shape=(num_samples_uniform, self.num_control_parameters), + ) + ) # Add sampling to the best old control parameters control_parameters_vec = best_control_parameters + additional_random_parameters - # Sampling step frequency available_freq_increment = jnp.where(optimize_swing, self.step_freq_delta, nominal_step_frequency) - #available_freq_increment = self.step_freq_delta - - - #step_frequencies_vec = jax.random.choice(key, available_freq_increment, shape=(self.num_parallel_computations, ))*optimize_swing + nominal_step_frequency - step_frequencies_vec = jax.random.choice(key, available_freq_increment, shape=(self.num_parallel_computations, )) + # available_freq_increment = self.step_freq_delta + # step_frequencies_vec = jax.random.choice(key, available_freq_increment, shape=(self.num_parallel_computations, ))*optimize_swing + nominal_step_frequency + step_frequencies_vec = jax.random.choice(key, available_freq_increment, shape=(self.num_parallel_computations,)) # Do rollout costs = self.jit_vectorized_rollout(state, reference, timing, control_parameters_vec, step_frequencies_vec) - # Saturate the cost in case of NaN or inf costs = jnp.where(jnp.isnan(costs), 1000000, costs) costs = jnp.where(jnp.isinf(costs), 1000000, costs) - # Take the best found control parameters best_index = jnp.nanargmin(costs) @@ -826,13 +895,17 @@ def compute_control_random_sampling(self, state, reference, contact_sequence, be best_control_parameters = control_parameters_vec[best_index] best_step_frequency = step_frequencies_vec[best_index] - # and redistribute it to each leg - best_control_parameters_FL = best_control_parameters[0:self.num_control_parameters_single_leg] - best_control_parameters_FR = best_control_parameters[self.num_control_parameters_single_leg:self.num_control_parameters_single_leg*2] - best_control_parameters_RL = best_control_parameters[self.num_control_parameters_single_leg*2:self.num_control_parameters_single_leg*3] - best_control_parameters_RR = best_control_parameters[self.num_control_parameters_single_leg*3:self.num_control_parameters_single_leg*4] - + best_control_parameters_FL = best_control_parameters[0 : self.num_control_parameters_single_leg] + best_control_parameters_FR = best_control_parameters[ + self.num_control_parameters_single_leg : self.num_control_parameters_single_leg * 2 + ] + best_control_parameters_RL = best_control_parameters[ + self.num_control_parameters_single_leg * 2 : self.num_control_parameters_single_leg * 3 + ] + best_control_parameters_RR = best_control_parameters[ + self.num_control_parameters_single_leg * 3 : self.num_control_parameters_single_leg * 4 + ] # Compute the GRF associated to the best parameter fx_FL, fy_FL, fz_FL = self.spline_fun_FL(best_control_parameters_FL, 0.0, 1) @@ -840,10 +913,11 @@ def compute_control_random_sampling(self, state, reference, contact_sequence, be fx_RL, fy_RL, fz_RL = self.spline_fun_RL(best_control_parameters_RL, 0.0, 1) fx_RR, fy_RR, fz_RR = self.spline_fun_RR(best_control_parameters_RR, 0.0, 1) - - # Add the gravity compensation to the stance legs and put to zero + # Add the gravity compensation to the stance legs and put to zero # the GRF of the swing legs - number_of_legs_in_stance = contact_sequence[0][0] + contact_sequence[1][0] + contact_sequence[2][0] + contact_sequence[3][0] + number_of_legs_in_stance = ( + contact_sequence[0][0] + contact_sequence[1][0] + contact_sequence[2][0] + contact_sequence[3][0] + ) reference_force_stance_legs = (self.robot.mass * 9.81) / number_of_legs_in_stance fz_FL = reference_force_stance_legs + fz_FL @@ -851,119 +925,143 @@ def compute_control_random_sampling(self, state, reference, contact_sequence, be fz_RL = reference_force_stance_legs + fz_RL fz_RR = reference_force_stance_legs + fz_RR + fx_FL = fx_FL * contact_sequence[0][0] + fy_FL = fy_FL * contact_sequence[0][0] + fz_FL = fz_FL * contact_sequence[0][0] + fx_FR = fx_FR * contact_sequence[1][0] + fy_FR = fy_FR * contact_sequence[1][0] + fz_FR = fz_FR * contact_sequence[1][0] + fx_RL = fx_RL * contact_sequence[2][0] + fy_RL = fy_RL * contact_sequence[2][0] + fz_RL = fz_RL * contact_sequence[2][0] - fx_FL = fx_FL*contact_sequence[0][0] - fy_FL = fy_FL*contact_sequence[0][0] - fz_FL = fz_FL*contact_sequence[0][0] - - fx_FR = fx_FR*contact_sequence[1][0] - fy_FR = fy_FR*contact_sequence[1][0] - fz_FR = fz_FR*contact_sequence[1][0] - - fx_RL = fx_RL*contact_sequence[2][0] - fy_RL = fy_RL*contact_sequence[2][0] - fz_RL = fz_RL*contact_sequence[2][0] - - fx_RR = fx_RR*contact_sequence[3][0] - fy_RR = fy_RR*contact_sequence[3][0] - fz_RR = fz_RR*contact_sequence[3][0] - + fx_RR = fx_RR * contact_sequence[3][0] + fy_RR = fy_RR * contact_sequence[3][0] + fz_RR = fz_RR * contact_sequence[3][0] # Enforce force constraints - fx_FL, fy_FL, fz_FL, \ - fx_FR, fy_FR, fz_FR, \ - fx_RL, fy_RL, fz_RL, \ - fx_RR, fy_RR, fz_RR =self.enforce_force_constraints(fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR) - - - - nmpc_GRFs = jnp.array([fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR]) - nmpc_footholds = jnp.array([0, 0, 0, - 0, 0, 0, - 0, 0, 0, - 0, 0, 0]) + fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR = ( + self.enforce_force_constraints( + fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR + ) + ) + + nmpc_GRFs = jnp.array([fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR]) + nmpc_footholds = jnp.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) # Compute predicted state for IK - input = jnp.array([jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR, - ], dtype=dtype_general) - current_contact = jnp.array([contact_sequence[0][0], contact_sequence[1][0], - contact_sequence[2][0], contact_sequence[3][0]], dtype=dtype_general) + input = jnp.array( + [ + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + fx_FL, + fy_FL, + fz_FL, + fx_FR, + fy_FR, + fz_FR, + fx_RL, + fy_RL, + fz_RL, + fx_RR, + fy_RR, + fz_RR, + ], + dtype=dtype_general, + ) + current_contact = jnp.array( + [contact_sequence[0][0], contact_sequence[1][0], contact_sequence[2][0], contact_sequence[3][0]], + dtype=dtype_general, + ) nmpc_predicted_state = self.robot.integrate_jax(state, input, current_contact, 0) - - return nmpc_GRFs, nmpc_footholds, nmpc_predicted_state, best_control_parameters, best_cost, best_step_frequency, costs - - - - - def compute_control_mppi(self, state, reference, contact_sequence, best_control_parameters, key, timing, nominal_step_frequency, optimize_swing): + return ( + nmpc_GRFs, + nmpc_footholds, + nmpc_predicted_state, + best_control_parameters, + best_cost, + best_step_frequency, + costs, + ) + + def compute_control_mppi( + self, + state, + reference, + contact_sequence, + best_control_parameters, + key, + timing, + nominal_step_frequency, + optimize_swing, + ): """ This function computes the control parameters by applying MPPI. - """ - - - # The first control parameters is the old best one, so we add zero noise there - additional_random_parameters = self.initial_random_parameters*0.0 - + """ + + # The first control parameters is the old best one, so we add zero noise there + additional_random_parameters = self.initial_random_parameters * 0.0 # GAUSSIAN - num_sample_gaussian_1 = self.num_parallel_computations-1 - additional_random_parameters = additional_random_parameters.at[1:self.num_parallel_computations].set(self.sigma_mppi*jax.random.normal(key=key, shape=(num_sample_gaussian_1, self.num_control_parameters))) - - control_parameters_vec = best_control_parameters + additional_random_parameters + num_sample_gaussian_1 = self.num_parallel_computations - 1 + additional_random_parameters = additional_random_parameters.at[1 : self.num_parallel_computations].set( + self.sigma_mppi * jax.random.normal(key=key, shape=(num_sample_gaussian_1, self.num_control_parameters)) + ) + control_parameters_vec = best_control_parameters + additional_random_parameters # Sampling step frequency available_freq_increment = self.step_freq_delta - step_frequencies_vec = jax.random.choice(key, available_freq_increment, shape=(self.num_parallel_computations, ))#*optimize_swing + nominal_step_frequency - + step_frequencies_vec = jax.random.choice( + key, available_freq_increment, shape=(self.num_parallel_computations,) + ) # *optimize_swing + nominal_step_frequency # Do rollout costs = self.jit_vectorized_rollout(state, reference, timing, control_parameters_vec, step_frequencies_vec) - # Saturate the cost in case of NaN or inf costs = jnp.where(jnp.isnan(costs), 1000000, costs) costs = jnp.where(jnp.isinf(costs), 1000000, costs) - # Take the best found control parameters best_index = jnp.nanargmin(costs) best_cost = costs.take(best_index) - # Compute MPPI update beta = best_cost - temperature = 1. - exp_costs = jnp.exp((-1./temperature) * (costs - beta)) + temperature = 1.0 + exp_costs = jnp.exp((-1.0 / temperature) * (costs - beta)) denom = np.sum(exp_costs) - weights = exp_costs/denom - weighted_inputs = weights[:, jnp.newaxis, jnp.newaxis] * additional_random_parameters.reshape((self.num_parallel_computations,self.num_control_parameters,1)) - best_control_parameters += jnp.sum(weighted_inputs, axis=0).reshape((self.num_control_parameters, )) + weights = exp_costs / denom + weighted_inputs = weights[:, jnp.newaxis, jnp.newaxis] * additional_random_parameters.reshape( + (self.num_parallel_computations, self.num_control_parameters, 1) + ) + best_control_parameters += jnp.sum(weighted_inputs, axis=0).reshape((self.num_control_parameters,)) best_step_frequency = step_frequencies_vec[best_index] - - # And redistribute it to each leg - best_control_parameters_FL = best_control_parameters[0:self.num_control_parameters_single_leg] - best_control_parameters_FR = best_control_parameters[self.num_control_parameters_single_leg:self.num_control_parameters_single_leg*2] - best_control_parameters_RL = best_control_parameters[self.num_control_parameters_single_leg*2:self.num_control_parameters_single_leg*3] - best_control_parameters_RR = best_control_parameters[self.num_control_parameters_single_leg*3:self.num_control_parameters_single_leg*4] - + best_control_parameters_FL = best_control_parameters[0 : self.num_control_parameters_single_leg] + best_control_parameters_FR = best_control_parameters[ + self.num_control_parameters_single_leg : self.num_control_parameters_single_leg * 2 + ] + best_control_parameters_RL = best_control_parameters[ + self.num_control_parameters_single_leg * 2 : self.num_control_parameters_single_leg * 3 + ] + best_control_parameters_RR = best_control_parameters[ + self.num_control_parameters_single_leg * 3 : self.num_control_parameters_single_leg * 4 + ] # Compute the GRF associated to the best parameter fx_FL, fy_FL, fz_FL = self.spline_fun_FL(best_control_parameters_FL, 0.0, 1) @@ -971,133 +1069,160 @@ def compute_control_mppi(self, state, reference, contact_sequence, best_control_ fx_RL, fy_RL, fz_RL = self.spline_fun_RL(best_control_parameters_RL, 0.0, 1) fx_RR, fy_RR, fz_RR = self.spline_fun_RR(best_control_parameters_RR, 0.0, 1) - - # Add the gravity compensation to the stance legs and put to zero + # Add the gravity compensation to the stance legs and put to zero # the GRF of the swing legs - number_of_legs_in_stance = contact_sequence[0][0] + contact_sequence[1][0] + contact_sequence[2][0] + contact_sequence[3][0] + number_of_legs_in_stance = ( + contact_sequence[0][0] + contact_sequence[1][0] + contact_sequence[2][0] + contact_sequence[3][0] + ) reference_force_stance_legs = (self.robot.mass * 9.81) / number_of_legs_in_stance - fz_FL = reference_force_stance_legs + fz_FL fz_FR = reference_force_stance_legs + fz_FR fz_RL = reference_force_stance_legs + fz_RL fz_RR = reference_force_stance_legs + fz_RR - - fx_FL = fx_FL*contact_sequence[0][0] - fy_FL = fy_FL*contact_sequence[0][0] - fz_FL = fz_FL*contact_sequence[0][0] - - fx_FR = fx_FR*contact_sequence[1][0] - fy_FR = fy_FR*contact_sequence[1][0] - fz_FR = fz_FR*contact_sequence[1][0] - fx_RL = fx_RL*contact_sequence[2][0] - fy_RL = fy_RL*contact_sequence[2][0] - fz_RL = fz_RL*contact_sequence[2][0] + fx_FL = fx_FL * contact_sequence[0][0] + fy_FL = fy_FL * contact_sequence[0][0] + fz_FL = fz_FL * contact_sequence[0][0] + + fx_FR = fx_FR * contact_sequence[1][0] + fy_FR = fy_FR * contact_sequence[1][0] + fz_FR = fz_FR * contact_sequence[1][0] - fx_RR = fx_RR*contact_sequence[3][0] - fy_RR = fy_RR*contact_sequence[3][0] - fz_RR = fz_RR*contact_sequence[3][0] + fx_RL = fx_RL * contact_sequence[2][0] + fy_RL = fy_RL * contact_sequence[2][0] + fz_RL = fz_RL * contact_sequence[2][0] + fx_RR = fx_RR * contact_sequence[3][0] + fy_RR = fy_RR * contact_sequence[3][0] + fz_RR = fz_RR * contact_sequence[3][0] # Enforce force constraints - fx_FL, fy_FL, fz_FL, \ - fx_FR, fy_FR, fz_FR, \ - fx_RL, fy_RL, fz_RL, \ - fx_RR, fy_RR, fz_RR =self.enforce_force_constraints(fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR) - - - - nmpc_GRFs = jnp.array([fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR]) - nmpc_footholds = jnp.array([0, 0, 0, - 0, 0, 0, - 0, 0, 0, - 0, 0, 0]) + fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR = ( + self.enforce_force_constraints( + fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR + ) + ) + + nmpc_GRFs = jnp.array([fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR]) + nmpc_footholds = jnp.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) # Compute predicted state for IK - input = jnp.array([jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR, - ], dtype=dtype_general) - current_contact = jnp.array([contact_sequence[0][0], contact_sequence[1][0], - contact_sequence[2][0], contact_sequence[3][0]], dtype=dtype_general) + input = jnp.array( + [ + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + fx_FL, + fy_FL, + fz_FL, + fx_FR, + fy_FR, + fz_FR, + fx_RL, + fy_RL, + fz_RL, + fx_RR, + fy_RR, + fz_RR, + ], + dtype=dtype_general, + ) + current_contact = jnp.array( + [contact_sequence[0][0], contact_sequence[1][0], contact_sequence[2][0], contact_sequence[3][0]], + dtype=dtype_general, + ) nmpc_predicted_state = self.robot.integrate_jax(state, input, current_contact, 0) - - return nmpc_GRFs, nmpc_footholds, nmpc_predicted_state, best_control_parameters, best_cost, best_step_frequency, costs - - - - def compute_control_cem_mppi(self, state, reference, contact_sequence, best_control_parameters, key, sigma, timing, nominal_step_frequency, optimize_swing): + return ( + nmpc_GRFs, + nmpc_footholds, + nmpc_predicted_state, + best_control_parameters, + best_cost, + best_step_frequency, + costs, + ) + + def compute_control_cem_mppi( + self, + state, + reference, + contact_sequence, + best_control_parameters, + key, + sigma, + timing, + nominal_step_frequency, + optimize_swing, + ): """ This function computes the control parameters by applying CEM-MPPI. - """ - - + """ + # Generate random parameters # The first control parameters is the old best one, so we add zero noise there - additional_random_parameters = self.initial_random_parameters*0.0 - + additional_random_parameters = self.initial_random_parameters * 0.0 # GAUSSIAN - num_sample_gaussian_1 = self.num_parallel_computations-1 + num_sample_gaussian_1 = self.num_parallel_computations - 1 - #jax.debug.breakpoint() - additional_random_parameters = additional_random_parameters.at[1:self.num_parallel_computations].set(jax.random.normal(key=key, shape=(num_sample_gaussian_1, self.num_control_parameters))*sigma) - - - control_parameters_vec = best_control_parameters + additional_random_parameters + # jax.debug.breakpoint() + additional_random_parameters = additional_random_parameters.at[1 : self.num_parallel_computations].set( + jax.random.normal(key=key, shape=(num_sample_gaussian_1, self.num_control_parameters)) * sigma + ) - + control_parameters_vec = best_control_parameters + additional_random_parameters # Sampling step frequency available_freq_increment = jnp.array([0.0, 0.2, 0.4]) - step_frequencies_vec = jax.random.choice(key, available_freq_increment, shape=(self.num_parallel_computations, ))*optimize_swing + nominal_step_frequency - - + step_frequencies_vec = ( + jax.random.choice(key, available_freq_increment, shape=(self.num_parallel_computations,)) * optimize_swing + + nominal_step_frequency + ) # Do rollout costs = self.jit_vectorized_rollout(state, reference, timing, control_parameters_vec, step_frequencies_vec) - # Saturate the cost in case of NaN or inf costs = jnp.where(jnp.isnan(costs), 1000000, costs) costs = jnp.where(jnp.isinf(costs), 1000000, costs) - # Take the best found control parameters best_index = jnp.nanargmin(costs) best_cost = costs.take(best_index) - # Compute MPPI update beta = best_cost - temperature = 1. - exp_costs = jnp.exp((-1./temperature) * (costs - beta)) + temperature = 1.0 + exp_costs = jnp.exp((-1.0 / temperature) * (costs - beta)) denom = np.sum(exp_costs) - weights = exp_costs/denom - weighted_inputs = weights[:, jnp.newaxis, jnp.newaxis] * additional_random_parameters.reshape((self.num_parallel_computations,self.num_control_parameters,1)) - best_control_parameters += jnp.sum(weighted_inputs, axis=0).reshape((self.num_control_parameters, )) + weights = exp_costs / denom + weighted_inputs = weights[:, jnp.newaxis, jnp.newaxis] * additional_random_parameters.reshape( + (self.num_parallel_computations, self.num_control_parameters, 1) + ) + best_control_parameters += jnp.sum(weighted_inputs, axis=0).reshape((self.num_control_parameters,)) best_step_frequency = step_frequencies_vec[best_index] - - # And redistribute it to each leg - best_control_parameters_FL = best_control_parameters[0:self.num_control_parameters_single_leg] - best_control_parameters_FR = best_control_parameters[self.num_control_parameters_single_leg:self.num_control_parameters_single_leg*2] - best_control_parameters_RL = best_control_parameters[self.num_control_parameters_single_leg*2:self.num_control_parameters_single_leg*3] - best_control_parameters_RR = best_control_parameters[self.num_control_parameters_single_leg*3:self.num_control_parameters_single_leg*4] - + best_control_parameters_FL = best_control_parameters[0 : self.num_control_parameters_single_leg] + best_control_parameters_FR = best_control_parameters[ + self.num_control_parameters_single_leg : self.num_control_parameters_single_leg * 2 + ] + best_control_parameters_RL = best_control_parameters[ + self.num_control_parameters_single_leg * 2 : self.num_control_parameters_single_leg * 3 + ] + best_control_parameters_RR = best_control_parameters[ + self.num_control_parameters_single_leg * 3 : self.num_control_parameters_single_leg * 4 + ] # Compute the GRF associated to the best parameter fx_FL, fy_FL, fz_FL = self.spline_fun_FL(best_control_parameters_FL, 0.0, 1) @@ -1105,79 +1230,96 @@ def compute_control_cem_mppi(self, state, reference, contact_sequence, best_cont fx_RL, fy_RL, fz_RL = self.spline_fun_RL(best_control_parameters_RL, 0.0, 1) fx_RR, fy_RR, fz_RR = self.spline_fun_RR(best_control_parameters_RR, 0.0, 1) - - # Add the gravity compensation to the stance legs and put to zero + # Add the gravity compensation to the stance legs and put to zero # the GRF of the swing legs - number_of_legs_in_stance = contact_sequence[0][0] + contact_sequence[1][0] + contact_sequence[2][0] + contact_sequence[3][0] + number_of_legs_in_stance = ( + contact_sequence[0][0] + contact_sequence[1][0] + contact_sequence[2][0] + contact_sequence[3][0] + ) reference_force_stance_legs = (self.robot.mass * 9.81) / number_of_legs_in_stance - fz_FL = reference_force_stance_legs + fz_FL fz_FR = reference_force_stance_legs + fz_FR fz_RL = reference_force_stance_legs + fz_RL fz_RR = reference_force_stance_legs + fz_RR - - fx_FL = fx_FL*contact_sequence[0][0] - fy_FL = fy_FL*contact_sequence[0][0] - fz_FL = fz_FL*contact_sequence[0][0] - - fx_FR = fx_FR*contact_sequence[1][0] - fy_FR = fy_FR*contact_sequence[1][0] - fz_FR = fz_FR*contact_sequence[1][0] - fx_RL = fx_RL*contact_sequence[2][0] - fy_RL = fy_RL*contact_sequence[2][0] - fz_RL = fz_RL*contact_sequence[2][0] + fx_FL = fx_FL * contact_sequence[0][0] + fy_FL = fy_FL * contact_sequence[0][0] + fz_FL = fz_FL * contact_sequence[0][0] + + fx_FR = fx_FR * contact_sequence[1][0] + fy_FR = fy_FR * contact_sequence[1][0] + fz_FR = fz_FR * contact_sequence[1][0] - fx_RR = fx_RR*contact_sequence[3][0] - fy_RR = fy_RR*contact_sequence[3][0] - fz_RR = fz_RR*contact_sequence[3][0] + fx_RL = fx_RL * contact_sequence[2][0] + fy_RL = fy_RL * contact_sequence[2][0] + fz_RL = fz_RL * contact_sequence[2][0] + fx_RR = fx_RR * contact_sequence[3][0] + fy_RR = fy_RR * contact_sequence[3][0] + fz_RR = fz_RR * contact_sequence[3][0] # Enforce force constraints - fx_FL, fy_FL, fz_FL, \ - fx_FR, fy_FR, fz_FR, \ - fx_RL, fy_RL, fz_RL, \ - fx_RR, fy_RR, fz_RR =self.enforce_force_constraints(fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR) - - - nmpc_GRFs = jnp.array([fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR]) - nmpc_footholds = jnp.array([0, 0, 0, - 0, 0, 0, - 0, 0, 0, - 0, 0, 0]) + fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR = ( + self.enforce_force_constraints( + fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR + ) + ) + + nmpc_GRFs = jnp.array([fx_FL, fy_FL, fz_FL, fx_FR, fy_FR, fz_FR, fx_RL, fy_RL, fz_RL, fx_RR, fy_RR, fz_RR]) + nmpc_footholds = jnp.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) # Compute predicted state for IK - input = jnp.array([jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - jnp.float32(0), jnp.float32(0), jnp.float32(0), - fx_FL, fy_FL, fz_FL, - fx_FR, fy_FR, fz_FR, - fx_RL, fy_RL, fz_RL, - fx_RR, fy_RR, fz_RR, - ], dtype=dtype_general) - current_contact = jnp.array([contact_sequence[0][0], contact_sequence[1][0], - contact_sequence[2][0], contact_sequence[3][0]], dtype=dtype_general) + input = jnp.array( + [ + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + jnp.float32(0), + fx_FL, + fy_FL, + fz_FL, + fx_FR, + fy_FR, + fz_FR, + fx_RL, + fy_RL, + fz_RL, + fx_RR, + fy_RR, + fz_RR, + ], + dtype=dtype_general, + ) + current_contact = jnp.array( + [contact_sequence[0][0], contact_sequence[1][0], contact_sequence[2][0], contact_sequence[3][0]], + dtype=dtype_general, + ) nmpc_predicted_state = self.robot.integrate_jax(state, input, current_contact, 0) - - + indexes = jnp.argsort(costs)[:200] elite = additional_random_parameters[indexes] - new_sigma_cem_mppi = jnp.cov(elite, rowvar=False) + np.eye(self.num_control_parameters)*1e-8 + new_sigma_cem_mppi = jnp.cov(elite, rowvar=False) + np.eye(self.num_control_parameters) * 1e-8 new_sigma_cem_mppi = jnp.diag(new_sigma_cem_mppi) new_sigma_cem_mppi = jnp.where(new_sigma_cem_mppi > 3, 3, new_sigma_cem_mppi) new_sigma_cem_mppi = jnp.where(new_sigma_cem_mppi < 0.5, 0.5, new_sigma_cem_mppi) - - return nmpc_GRFs, nmpc_footholds, nmpc_predicted_state, best_control_parameters, best_cost, costs, new_sigma_cem_mppi - + return ( + nmpc_GRFs, + nmpc_footholds, + nmpc_predicted_state, + best_control_parameters, + best_cost, + costs, + new_sigma_cem_mppi, + ) def reset(self): - print("Resetting the controller") \ No newline at end of file + print("Resetting the controller") diff --git a/quadruped_pympc/helpers/foothold_reference_generator.py b/quadruped_pympc/helpers/foothold_reference_generator.py index 39da609..b44dff7 100644 --- a/quadruped_pympc/helpers/foothold_reference_generator.py +++ b/quadruped_pympc/helpers/foothold_reference_generator.py @@ -1,21 +1,20 @@ import collections +import copy import mujoco -import copy import numpy as np +from gym_quadruped.utils.quadruped_utils import LegsAttr from scipy.spatial.transform import Rotation -from gym_quadruped.utils.quadruped_utils import LegsAttr from quadruped_pympc.helpers.quadruped_utils import GaitType # Class for the generation of the reference footholds # TODO: @Giulio Should we convert this to a single function instead of a class? Stance time, can be passed as argument class FootholdReferenceGenerator: - - def __init__(self, stance_time: float, lift_off_positions: LegsAttr, - vel_moving_average_length=20, - hip_height: float = None) -> None: + def __init__( + self, stance_time: float, lift_off_positions: LegsAttr, vel_moving_average_length=20, hip_height: float = None + ) -> None: """This method initializes the foothold generator class, which computes the reference foothold for the nonlinear MPC. @@ -30,27 +29,29 @@ def __init__(self, stance_time: float, lift_off_positions: LegsAttr, self.touch_down_positions = copy.deepcopy(lift_off_positions) # The offset of the COM position wrt the estimated one. HACK compensation - self.com_pos_offset_b = np.zeros((3, )) - self.com_pos_offset_w = np.zeros((3, )) + self.com_pos_offset_b = np.zeros((3,)) + self.com_pos_offset_w = np.zeros((3,)) """R_W2H = np.array([np.cos(yaw), np.sin(yaw), -np.sin(yaw), np.cos(yaw)]) R_W2H = R_W2H.reshape((2, 2)) self.lift_off_positions_h = R_W2H @ (self.lift_off_positions - base_position[0:2])""" - self.lift_off_positions_h = copy.deepcopy(lift_off_positions) #TODO wrong - self.touch_down_positions_h = copy.deepcopy(lift_off_positions) #TODO wrong + self.lift_off_positions_h = copy.deepcopy(lift_off_positions) # TODO wrong + self.touch_down_positions_h = copy.deepcopy(lift_off_positions) # TODO wrong # The footholds are wrt the hip position, so if we want to change # the default foothold, we need to use a variable to add an offset self.hip_offset = 0.1 - def compute_footholds_reference(self, - base_position: np.ndarray, - base_ori_euler_xyz: np.ndarray, - base_xy_lin_vel: np.ndarray, - ref_base_xy_lin_vel: np.ndarray, - hips_position: LegsAttr, - com_height_nominal: np.float32) -> LegsAttr: + def compute_footholds_reference( + self, + base_position: np.ndarray, + base_ori_euler_xyz: np.ndarray, + base_xy_lin_vel: np.ndarray, + ref_base_xy_lin_vel: np.ndarray, + hips_position: LegsAttr, + com_height_nominal: np.float32, + ) -> LegsAttr: """Compute the reference footholds for a quadruped robot, using simple geometric heuristics. TODO: This function should be adapted to: @@ -72,13 +73,13 @@ def compute_footholds_reference(self, ------- ref_feet: (LegsAttr) The reference footholds for the robot in world frame. """ - assert base_xy_lin_vel.shape == (2,) and ref_base_xy_lin_vel.shape == (2,), \ + assert base_xy_lin_vel.shape == (2,) and ref_base_xy_lin_vel.shape == (2,), ( f"Expected shape (2,):=[x_dot, y_dot], got {base_xy_lin_vel.shape} and {ref_base_xy_lin_vel.shape}." + ) # Get the rotation matrix to transform from world to horizontal frame (hip-centric) yaw = base_ori_euler_xyz[2] - R_W2H = np.array([np.cos(yaw), np.sin(yaw), - -np.sin(yaw), np.cos(yaw)]) + R_W2H = np.array([np.cos(yaw), np.sin(yaw), -np.sin(yaw), np.cos(yaw)]) R_W2H = R_W2H.reshape((2, 2)) # Compute desired and error velocity compensation values for all legs @@ -89,22 +90,19 @@ def compute_footholds_reference(self, self.base_vel_hist.append(base_lin_vel_H) base_vel_mvg = np.mean(list(self.base_vel_hist), axis=0) # Compensation due to average velocity - #delta_ref_H = (self.stance_time / 2.) * base_vel_mvg - + # delta_ref_H = (self.stance_time / 2.) * base_vel_mvg + # Compensation due to desired velocity - delta_ref_H = (self.stance_time / 2.) * ref_base_lin_vel_H + delta_ref_H = (self.stance_time / 2.0) * ref_base_lin_vel_H delta_ref_H = np.clip(delta_ref_H, -self.hip_height * 1.5, self.hip_height * 1.5) vel_offset = np.concatenate((delta_ref_H, np.zeros(1))) - # Compensation for the error in velocity tracking - error_compensation = np.sqrt(com_height_nominal/9.81)*(base_vel_mvg - ref_base_lin_vel_H) + error_compensation = np.sqrt(com_height_nominal / 9.81) * (base_vel_mvg - ref_base_lin_vel_H) error_compensation = np.where(error_compensation > 0.05, 0.05, error_compensation) error_compensation = np.where(error_compensation < -0.05, -0.05, error_compensation) error_compensation = np.concatenate((error_compensation, np.zeros(1))) - - # Reference footholds in the horizontal frame ref_feet = LegsAttr(*[np.zeros(3) for _ in range(4)]) @@ -123,19 +121,17 @@ def compute_footholds_reference(self, ref_feet.RL[1] += self.hip_offset ref_feet.RR[1] -= self.hip_offset - - # Add the velocity compensation and desired velocity to the feet positions ref_feet += vel_offset + error_compensation # Add offset to all feet # Reference footholds in world frame - ref_feet.FL[0:2] = R_W2H.T @ ref_feet.FL[:2] + base_position[0:2] + ref_feet.FL[0:2] = R_W2H.T @ ref_feet.FL[:2] + base_position[0:2] ref_feet.FR[0:2] = R_W2H.T @ ref_feet.FR[:2] + base_position[0:2] ref_feet.RL[0:2] = R_W2H.T @ ref_feet.RL[:2] + base_position[0:2] ref_feet.RR[0:2] = R_W2H.T @ ref_feet.RR[:2] + base_position[0:2] # Add offset to the feet positions for manual com offset - R_B2W = Rotation.from_euler('xyz', base_ori_euler_xyz).as_matrix() + R_B2W = Rotation.from_euler("xyz", base_ori_euler_xyz).as_matrix() self.com_pos_offset_w = R_B2W @ self.com_pos_offset_b ref_feet.FL[0:2] += self.com_pos_offset_w[0:2] ref_feet.FR[0:2] += self.com_pos_offset_w[0:2] @@ -144,78 +140,67 @@ def compute_footholds_reference(self, # TODO: we should rotate them considering the terrain estimator maybe # or we can just do exteroceptive height adjustement... - for leg_id in ['FL', 'FR', 'RL', 'RR']: - ref_feet[leg_id][2] = self.lift_off_positions[leg_id][2]# - 0.02 + for leg_id in ["FL", "FR", "RL", "RR"]: + ref_feet[leg_id][2] = self.lift_off_positions[leg_id][2] # - 0.02 return ref_feet - - def update_lift_off_positions(self, previous_contact, current_contact, feet_pos, legs_order, gait_type, - base_position, base_ori_euler_xyz): - + def update_lift_off_positions( + self, previous_contact, current_contact, feet_pos, legs_order, gait_type, base_position, base_ori_euler_xyz + ): yaw = base_ori_euler_xyz[2] - R_W2H = np.array([np.cos(yaw), np.sin(yaw), 0, - -np.sin(yaw), np.cos(yaw), 0, - 0, 0, 1]) + R_W2H = np.array([np.cos(yaw), np.sin(yaw), 0, -np.sin(yaw), np.cos(yaw), 0, 0, 0, 1]) R_W2H = R_W2H.reshape((3, 3)) for leg_id, leg_name in enumerate(legs_order): - - if(gait_type == GaitType.FULL_STANCE.value): + if gait_type == GaitType.FULL_STANCE.value: self.lift_off_positions[leg_name] = feet_pos[leg_name] continue # Set lif-offs in world frame and base frame if previous_contact[leg_id] == 1 and current_contact[leg_id] == 0: self.lift_off_positions[leg_name] = feet_pos[leg_name] - self.lift_off_positions_h[leg_name] = R_W2H @ (self.lift_off_positions[leg_name] - base_position) - - elif(previous_contact[leg_id] == 0 and current_contact[leg_id] == 0): + self.lift_off_positions_h[leg_name] = R_W2H @ (self.lift_off_positions[leg_name] - base_position) + + elif previous_contact[leg_id] == 0 and current_contact[leg_id] == 0: # Update lift-offs in world frame self.lift_off_positions[leg_name] = R_W2H.T @ self.lift_off_positions_h[leg_name] + base_position - - - def update_touch_down_positions(self, previous_contact, current_contact, feet_pos, legs_order, gait_type, - base_position, base_ori_euler_xyz): - + def update_touch_down_positions( + self, previous_contact, current_contact, feet_pos, legs_order, gait_type, base_position, base_ori_euler_xyz + ): yaw = base_ori_euler_xyz[2] - R_W2H = np.array([np.cos(yaw), np.sin(yaw), 0, - -np.sin(yaw), np.cos(yaw), 0, - 0, 0, 1]) + R_W2H = np.array([np.cos(yaw), np.sin(yaw), 0, -np.sin(yaw), np.cos(yaw), 0, 0, 0, 1]) R_W2H = R_W2H.reshape((3, 3)) for leg_id, leg_name in enumerate(legs_order): - if(gait_type == GaitType.FULL_STANCE.value): + if gait_type == GaitType.FULL_STANCE.value: self.touch_down_positions[leg_name] = feet_pos[leg_name] continue # Set touch-downs if previous_contact[leg_id] == 0 and current_contact[leg_id] == 1: self.touch_down_positions[leg_name] = feet_pos[leg_name] - self.touch_down_positions_h[leg_name] = R_W2H @ (self.touch_down_positions[leg_name] - base_position) - - elif(previous_contact[leg_id] == 1 and current_contact[leg_id] == 1): + self.touch_down_positions_h[leg_name] = R_W2H @ (self.touch_down_positions[leg_name] - base_position) + + elif previous_contact[leg_id] == 1 and current_contact[leg_id] == 1: # Update touch-downs in world frame self.touch_down_positions[leg_name] = R_W2H.T @ self.touch_down_positions_h[leg_name] + base_position - - if __name__ == "__main__": - m = mujoco.MjModel.from_xml_path('./../simulation/unitree_go1/scene.xml') + m = mujoco.MjModel.from_xml_path("./../simulation/unitree_go1/scene.xml") d = mujoco.MjData(m) mujoco.mj_step(m, d) - FL_hip_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, 'FL_hip') - FR_hip_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, 'FR_hip') - RL_hip_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, 'RL_hip') - RR_hip_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, 'RR_hip') + FL_hip_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, "FL_hip") + FR_hip_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, "FR_hip") + RL_hip_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, "RL_hip") + RR_hip_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, "RR_hip") - hip_pos = np.array(([d.body(FL_hip_id).xpos], - [d.body(FR_hip_id).xpos], - [d.body(RL_hip_id).xpos], - [d.body(RR_hip_id).xpos])) + hip_pos = np.array( + ([d.body(FL_hip_id).xpos], [d.body(FR_hip_id).xpos], [d.body(RL_hip_id).xpos], [d.body(RR_hip_id).xpos]) + ) print("hip_pos", hip_pos) @@ -225,8 +210,8 @@ def update_touch_down_positions(self, previous_contact, current_contact, feet_po com_height = d.qpos[2] foothold_generator = FootholdReferenceGenerator(stance_time) - footholds_reference = foothold_generator.compute_footholds_reference(linear_com_velocity[0:2], - desired_linear_com_velocity[0:2], - hip_pos, com_height) + footholds_reference = foothold_generator.compute_footholds_reference( + linear_com_velocity[0:2], desired_linear_com_velocity[0:2], hip_pos, com_height + ) print("iniztial hip_pos: ", hip_pos) print("footholds_reference: ", footholds_reference) diff --git a/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_numeric.py b/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_numeric.py index f494660..ae88b69 100644 --- a/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_numeric.py +++ b/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_numeric.py @@ -1,34 +1,35 @@ import numpy as np -np.set_printoptions(precision=3, suppress = True) -from numpy.linalg import norm, solve + +np.set_printoptions(precision=3, suppress=True) +# import example_robot_data as robex +import copy +import os import time + import casadi as cs -#import example_robot_data as robex -import copy +import gym_quadruped # Mujoco magic import mujoco import mujoco.viewer +from adam import Representations # Adam and Liecasadi magic from adam.casadi import KinDynComputations -from adam import Representations from liecasadi import SO3 -import gym_quadruped -import os dir_path = os.path.dirname(os.path.realpath(__file__)) gym_quadruped_path = os.path.dirname(gym_quadruped.__file__) - -from quadruped_pympc import config - +from quadruped_pympc import config # Class for solving a generic inverse kinematics problem class InverseKinematicsNumeric: - def __init__(self,) -> None: + def __init__( + self, + ) -> None: """ This method initializes the inverse kinematics solver class. @@ -36,39 +37,43 @@ def __init__(self,) -> None: """ + if config.robot == "go2": + urdf_filename = gym_quadruped_path + "/robot_model/go2/go2.urdf" + xml_filename = gym_quadruped_path + "/robot_model/go2/go2.xml" + if config.robot == "go1": + urdf_filename = gym_quadruped_path + "/robot_model/go1/go1.urdf" + xml_filename = gym_quadruped_path + "/robot_model/go1/go1.xml" + elif config.robot == "aliengo": + urdf_filename = gym_quadruped_path + "/robot_model/aliengo/aliengo.urdf" + xml_filename = gym_quadruped_path + "/robot_model/aliengo/aliengo.xml" + elif config.robot == "b2": + urdf_filename = gym_quadruped_path + "/robot_model/b2/b2.urdf" + xml_filename = gym_quadruped_path + "/robot_model/b2/b2.xml" + elif config.robot == "hyqreal": + urdf_filename = gym_quadruped_path + "/robot_model/hyqreal/hyqreal.urdf" + xml_filename = gym_quadruped_path + "/robot_model/hyqreal/hyqreal.xml" + elif config.robot == "mini_cheetah": + urdf_filename = gym_quadruped_path + "/robot_model/mini_cheetah/mini_cheetah.urdf" + xml_filename = gym_quadruped_path + "/robot_model/mini_cheetah/mini_cheetah.xml" + + joint_list = [ + "FL_hip_joint", + "FL_thigh_joint", + "FL_calf_joint", + "FR_hip_joint", + "FR_thigh_joint", + "FR_calf_joint", + "RL_hip_joint", + "RL_thigh_joint", + "RL_calf_joint", + "RR_hip_joint", + "RR_thigh_joint", + "RR_calf_joint", + ] - if(config.robot == 'go2'): - urdf_filename = gym_quadruped_path + '/robot_model/go2/go2.urdf' - xml_filename = gym_quadruped_path + '/robot_model/go2/go2.xml' - if(config.robot == 'go1'): - urdf_filename = gym_quadruped_path + '/robot_model/go1/go1.urdf' - xml_filename = gym_quadruped_path + '/robot_model/go1/go1.xml' - elif(config.robot == 'aliengo'): - urdf_filename = gym_quadruped_path + '/robot_model/aliengo/aliengo.urdf' - xml_filename = gym_quadruped_path + '/robot_model/aliengo/aliengo.xml' - elif(config.robot == 'b2'): - urdf_filename = gym_quadruped_path + '/robot_model/b2/b2.urdf' - xml_filename = gym_quadruped_path + '/robot_model/b2/b2.xml' - elif(config.robot == 'hyqreal'): - urdf_filename = gym_quadruped_path + '/robot_model/hyqreal/hyqreal.urdf' - xml_filename = gym_quadruped_path + '/robot_model/hyqreal/hyqreal.xml' - elif(config.robot == 'mini_cheetah'): - urdf_filename = gym_quadruped_path + '/robot_model/mini_cheetah/mini_cheetah.urdf' - xml_filename = gym_quadruped_path + '/robot_model/mini_cheetah/mini_cheetah.xml' - - - - joint_list = ['FL_hip_joint', 'FL_thigh_joint', 'FL_calf_joint', - 'FR_hip_joint', 'FR_thigh_joint', 'FR_calf_joint', - 'RL_hip_joint', 'RL_thigh_joint', 'RL_calf_joint', - 'RR_hip_joint', 'RR_thigh_joint', 'RR_calf_joint'] - - self.kindyn = KinDynComputations(urdfstring=urdf_filename, joints_name_list=joint_list) self.kindyn.set_frame_velocity_representation(representation=Representations.MIXED_REPRESENTATION) - - self.forward_kinematics_FL_fun = self.kindyn.forward_kinematics_fun("FL_foot") self.forward_kinematics_FR_fun = self.kindyn.forward_kinematics_fun("FR_foot") self.forward_kinematics_RL_fun = self.kindyn.forward_kinematics_fun("RL_foot") @@ -78,23 +83,29 @@ def __init__(self,) -> None: self.jacobian_FR_fun = self.kindyn.jacobian_fun("FR_foot") self.jacobian_RL_fun = self.kindyn.jacobian_fun("RL_foot") self.jacobian_RR_fun = self.kindyn.jacobian_fun("RR_foot") - - - q = cs.SX.sym('q', 12+7) - FL_foot_target_position = cs.SX.sym('FL_foot_target_position', 3) - FR_foot_target_position = cs.SX.sym('FR_foot_target_position', 3) - RL_foot_target_position = cs.SX.sym('RL_foot_target_position', 3) - RR_foot_target_position = cs.SX.sym('RR_foot_target_position', 3) - ik = self.compute_solution(q, FL_foot_target_position, FR_foot_target_position, - RL_foot_target_position, RR_foot_target_position) - self.fun_compute_solution = cs.Function('fun_ik', [q, FL_foot_target_position, FR_foot_target_position, RL_foot_target_position, RR_foot_target_position], [ik]) - - - - - def compute_solution(self, q: np.ndarray, FL_foot_target_position: np.ndarray, FR_foot_target_position: np.ndarray, - RL_foot_target_position: np.ndarray, RR_foot_target_position: np.ndarray) -> np.ndarray: + q = cs.SX.sym("q", 12 + 7) + FL_foot_target_position = cs.SX.sym("FL_foot_target_position", 3) + FR_foot_target_position = cs.SX.sym("FR_foot_target_position", 3) + RL_foot_target_position = cs.SX.sym("RL_foot_target_position", 3) + RR_foot_target_position = cs.SX.sym("RR_foot_target_position", 3) + ik = self.compute_solution( + q, FL_foot_target_position, FR_foot_target_position, RL_foot_target_position, RR_foot_target_position + ) + self.fun_compute_solution = cs.Function( + "fun_ik", + [q, FL_foot_target_position, FR_foot_target_position, RL_foot_target_position, RR_foot_target_position], + [ik], + ) + + def compute_solution( + self, + q: np.ndarray, + FL_foot_target_position: np.ndarray, + FR_foot_target_position: np.ndarray, + RL_foot_target_position: np.ndarray, + RR_foot_target_position: np.ndarray, + ) -> np.ndarray: """ This method computes the forward kinematics from initial joint angles and desired foot target positions. @@ -108,116 +119,105 @@ def compute_solution(self, q: np.ndarray, FL_foot_target_position: np.ndarray, F Returns: np.ndarray: The joint angles that achieve the desired foot positions. """ - + eps = 1e-2 IT_MAX = 5 DT = 1e-2 damp = 1e-2 damp_matrix = damp * np.eye(12) - - i=0 - err_FL = cs.SX.zeros(3,1) - err_FR = cs.SX.zeros(3,1) - err_RL = cs.SX.zeros(3,1) - err_RR = cs.SX.zeros(3,1) + i = 0 + + err_FL = cs.SX.zeros(3, 1) + err_FR = cs.SX.zeros(3, 1) + err_RL = cs.SX.zeros(3, 1) + err_RR = cs.SX.zeros(3, 1) err = cs.SX.zeros(3) q_joint = q[7:] quaternion = q[3:7] quaternion = np.array([quaternion[1], quaternion[2], quaternion[3], quaternion[0]]) R = SO3.from_quat(quaternion).as_matrix() - #H = cs.DM.eye(4) + # H = cs.DM.eye(4) H = cs.SX.eye(4) - + H[0:3, 0:3] = R H[0:3, 3] = q[0:3] - while i <= IT_MAX: - FL_foot_actual_pos = self.forward_kinematics_FL_fun(H, q_joint)[0:3, 3] FR_foot_actual_pos = self.forward_kinematics_FR_fun(H, q_joint)[0:3, 3] RL_foot_actual_pos = self.forward_kinematics_RL_fun(H, q_joint)[0:3, 3] RR_foot_actual_pos = self.forward_kinematics_RR_fun(H, q_joint)[0:3, 3] - - err_FL = (FL_foot_target_position - FL_foot_actual_pos) - err_FR = (FR_foot_target_position - FR_foot_actual_pos) - err_RL = (RL_foot_target_position - RL_foot_actual_pos) - err_RR = (RR_foot_target_position - RR_foot_actual_pos) - + err_FL = FL_foot_target_position - FL_foot_actual_pos + err_FR = FR_foot_target_position - FR_foot_actual_pos + err_RL = RL_foot_target_position - RL_foot_actual_pos + err_RR = RR_foot_target_position - RR_foot_actual_pos err = err_FL + err_FR + err_RL + err_RR norm_err = cs.norm_2(err) - #if(norm_err < eps): + # if(norm_err < eps): # success = True # break - J_FL = self.jacobian_FL_fun(H, q_joint)[0:3, 6:] J_FR = self.jacobian_FR_fun(H, q_joint)[0:3, 6:] J_RL = self.jacobian_RL_fun(H, q_joint)[0:3, 6:] J_RR = self.jacobian_RR_fun(H, q_joint)[0:3, 6:] - total_jac = cs.vertcat(J_FL, J_FR, J_RL, J_RR) total_err = 100.0 * cs.vertcat(err_FL, err_FR, err_RL, err_RR) - damped_pinv = cs.inv(total_jac.T@total_jac + damp_matrix)@total_jac.T - v = damped_pinv @ total_err + damped_pinv = cs.inv(total_jac.T @ total_jac + damp_matrix) @ total_jac.T + v = damped_pinv @ total_err q_joint = q_joint + DT * v - i += 1 - return q_joint - - - if __name__ == "__main__": - - if(config.robot == 'go2'): - urdf_filename = dir_path + '/../../gym-quadruped/gym_quadruped/robot_model/go2/go2.urdf' - xml_filename = dir_path + '/../../gym-quadruped/gym_quadruped/robot_model/go2/go2.xml' - if(config.robot == 'go1'): - urdf_filename = dir_path + '/../../gym-quadruped/gym_quadruped/robot_model/go1/go1.urdf' - xml_filename = dir_path + '/../../gym-quadruped/gym_quadruped/robot_model/go1/go1.xml' - elif(config.robot == 'aliengo'): - urdf_filename = dir_path + '/../../gym-quadruped/gym_quadruped/robot_model/aliengo/aliengo.urdf' - xml_filename = dir_path + '/../../gym-quadruped/gym_quadruped/robot_model/aliengo/aliengo.xml' - elif(config.robot == 'hyqreal'): - urdf_filename = dir_path + '/../../../gym-quadruped/gym_quadruped/robot_model/hyqreal/hyqreal.urdf' - xml_filename = dir_path + '/../../../gym-quadruped/gym_quadruped/robot_model/hyqreal/hyqreal.xml' - elif(config.robot == 'mini_cheetah'): - urdf_filename = dir_path + '/../../../gym-quadruped/gym_quadruped/robot_model/mini_cheetah/mini_cheetah.urdf' - xml_filename = dir_path + '/../../../gym-quadruped/gym_quadruped/robot_model/mini_cheetah/mini_cheetah.xml' - - - ik = InverseKinematicsNumeric() - - + if config.robot == "go2": + urdf_filename = dir_path + "/../../gym-quadruped/gym_quadruped/robot_model/go2/go2.urdf" + xml_filename = dir_path + "/../../gym-quadruped/gym_quadruped/robot_model/go2/go2.xml" + if config.robot == "go1": + urdf_filename = dir_path + "/../../gym-quadruped/gym_quadruped/robot_model/go1/go1.urdf" + xml_filename = dir_path + "/../../gym-quadruped/gym_quadruped/robot_model/go1/go1.xml" + elif config.robot == "aliengo": + urdf_filename = dir_path + "/../../gym-quadruped/gym_quadruped/robot_model/aliengo/aliengo.urdf" + xml_filename = dir_path + "/../../gym-quadruped/gym_quadruped/robot_model/aliengo/aliengo.xml" + elif config.robot == "hyqreal": + urdf_filename = dir_path + "/../../../gym-quadruped/gym_quadruped/robot_model/hyqreal/hyqreal.urdf" + xml_filename = dir_path + "/../../../gym-quadruped/gym_quadruped/robot_model/hyqreal/hyqreal.xml" + elif config.robot == "mini_cheetah": + urdf_filename = dir_path + "/../../../gym-quadruped/gym_quadruped/robot_model/mini_cheetah/mini_cheetah.urdf" + xml_filename = dir_path + "/../../../gym-quadruped/gym_quadruped/robot_model/mini_cheetah/mini_cheetah.xml" + + ik = InverseKinematicsNumeric() # Check consistency in mujoco m = mujoco.MjModel.from_xml_path(xml_filename) d = mujoco.MjData(m) - random_q_joint = np.random.rand(12,) + random_q_joint = np.random.rand( + 12, + ) d.qpos[7:] = random_q_joint # random quaternion - rand_quat = np.random.rand(4,) + rand_quat = np.random.rand( + 4, + ) rand_quat = rand_quat / np.linalg.norm(rand_quat) d.qpos[3:7] = rand_quat mujoco.mj_step(m, d) - FL_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, 'FL') - FR_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, 'FR') - RL_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, 'RL') - RR_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, 'RR') + FL_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "FL") + FR_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "FR") + RL_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "RL") + RR_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "RR") FL_foot_target_position = d.geom_xpos[FL_id] FR_foot_target_position = d.geom_xpos[FR_id] RL_foot_target_position = d.geom_xpos[RL_id] @@ -228,10 +228,10 @@ def compute_solution(self, q: np.ndarray, FL_foot_target_position: np.ndarray, F print("RL foot target position: ", RL_foot_target_position) print("RR foot target position: ", RR_foot_target_position) - - initial_q = copy.deepcopy(d.qpos) - initial_q[7:] = np.random.rand(12,) + initial_q[7:] = np.random.rand( + 12, + ) quaternion = d.qpos[3:7] quaternion = np.array([quaternion[1], quaternion[2], quaternion[3], quaternion[0]]) R = SO3.from_quat(quaternion).as_matrix() @@ -242,15 +242,13 @@ def compute_solution(self, q: np.ndarray, FL_foot_target_position: np.ndarray, F print("FR foot start position", ik.forward_kinematics_FR_fun(H, initial_q[7:])[0:3, 3]) print("RL foot start position", ik.forward_kinematics_RL_fun(H, initial_q[7:])[0:3, 3]) print("RR foot start position", ik.forward_kinematics_RR_fun(H, initial_q[7:])[0:3, 3]) - - + initial_time = time.time() - solution = ik.fun_compute_solution(initial_q, FL_foot_target_position, FR_foot_target_position, - RL_foot_target_position, RR_foot_target_position) + solution = ik.fun_compute_solution( + initial_q, FL_foot_target_position, FR_foot_target_position, RL_foot_target_position, RR_foot_target_position + ) print("time: ", time.time() - initial_time) - - print("\n") print("MUJOCO SOLUTION") foot_position_FL = d.geom_xpos[FL_id] @@ -262,8 +260,7 @@ def compute_solution(self, q: np.ndarray, FL_foot_target_position: np.ndarray, F print("FR foot position: ", foot_position_FR) print("RL foot position: ", foot_position_RL) print("RR foot position: ", foot_position_RR) - - + print("\n") print("ADAM SOLUTION") print("joints: ", solution) @@ -278,7 +275,6 @@ def compute_solution(self, q: np.ndarray, FL_foot_target_position: np.ndarray, F print("RL foot position", ik.forward_kinematics_RL_fun(H, solution)[0:3, 3]) print("RR foot position", ik.forward_kinematics_RR_fun(H, solution)[0:3, 3]) - with mujoco.viewer.launch_passive(m, d) as viewer: while True: - viewer.sync() \ No newline at end of file + viewer.sync() diff --git a/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_numeric_mujoco.py b/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_numeric_mujoco.py index d773a5a..69ee5ba 100644 --- a/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_numeric_mujoco.py +++ b/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_numeric_mujoco.py @@ -1,31 +1,26 @@ import numpy as np -np.set_printoptions(precision=3, suppress = True) -from numpy.linalg import norm, solve -import time -import casadi as cs -#import example_robot_data as robex + +np.set_printoptions(precision=3, suppress=True) +# import example_robot_data as robex import copy +import os +import time + +import gym_quadruped # Mujoco magic import mujoco import mujoco.viewer # Adam and Liecasadi magic -from adam.casadi import KinDynComputations -from adam import Representations -from liecasadi import SO3 -import gym_quadruped -import os dir_path = os.path.dirname(os.path.realpath(__file__)) gym_quadruped_path = os.path.dirname(gym_quadruped.__file__) -from quadruped_pympc import config as cfg - from gym_quadruped.quadruped_env import QuadrupedEnv - +from quadruped_pympc import config as cfg IT_MAX = 5 DT = 1e-2 @@ -35,7 +30,9 @@ # Class for solving a generic inverse kinematics problem class InverseKinematicsNumeric: - def __init__(self,) -> None: + def __init__( + self, + ) -> None: """ This method initializes the inverse kinematics solver class. @@ -43,28 +40,29 @@ def __init__(self,) -> None: """ - robot_name = cfg.robot hip_height = cfg.hip_height robot_leg_joints = cfg.robot_leg_joints robot_feet_geom_names = cfg.robot_feet_geom_names - #scene_name = cfg.simulation_params['scene'] - #simulation_dt = cfg.simulation_params['dt'] - - - + # scene_name = cfg.simulation_params['scene'] + # simulation_dt = cfg.simulation_params['dt'] # Create the quadruped robot environment ----------------------------------------------------------- - self.env = QuadrupedEnv(robot=robot_name, - hip_height=hip_height, - legs_joint_names=robot_leg_joints, # Joint names of the legs DoF - feet_geom_name=robot_feet_geom_names, # Geom/Frame id of feet - ) - - - - def compute_solution(self, q: np.ndarray, FL_foot_target_position: np.ndarray, FR_foot_target_position: np.ndarray, - RL_foot_target_position: np.ndarray, RR_foot_target_position: np.ndarray) -> np.ndarray: + self.env = QuadrupedEnv( + robot=robot_name, + hip_height=hip_height, + legs_joint_names=robot_leg_joints, # Joint names of the legs DoF + feet_geom_name=robot_feet_geom_names, # Geom/Frame id of feet + ) + + def compute_solution( + self, + q: np.ndarray, + FL_foot_target_position: np.ndarray, + FR_foot_target_position: np.ndarray, + RL_foot_target_position: np.ndarray, + RR_foot_target_position: np.ndarray, + ) -> np.ndarray: """ This method computes the forward kinematics from initial joint angles and desired foot target positions. @@ -78,42 +76,36 @@ def compute_solution(self, q: np.ndarray, FL_foot_target_position: np.ndarray, F Returns: np.ndarray: The joint angles that achieve the desired foot positions. """ - - # Set the initial states self.env.mjData.qpos = q mujoco.mj_fwdPosition(self.env.mjModel, self.env.mjData) - for j in range(IT_MAX): + feet_pos = self.env.feet_pos(frame="world") - feet_pos = self.env.feet_pos(frame='world') - FL_foot_actual_pos = feet_pos.FL FR_foot_actual_pos = feet_pos.FR RL_foot_actual_pos = feet_pos.RL RR_foot_actual_pos = feet_pos.RR - - err_FL = (FL_foot_target_position - FL_foot_actual_pos) - err_FR = (FR_foot_target_position - FR_foot_actual_pos) - err_RL = (RL_foot_target_position - RL_foot_actual_pos) - err_RR = (RR_foot_target_position - RR_foot_actual_pos) - + err_FL = FL_foot_target_position - FL_foot_actual_pos + err_FR = FR_foot_target_position - FR_foot_actual_pos + err_RL = RL_foot_target_position - RL_foot_actual_pos + err_RR = RR_foot_target_position - RR_foot_actual_pos # Compute feet jacobian - feet_jac = self.env.feet_jacobians(frame='world', return_rot_jac=False) - - J_FL = feet_jac.FL[:,6:] - J_FR = feet_jac.FR[:,6:] - J_RL = feet_jac.RL[:,6:] - J_RR = feet_jac.RR[:,6:] - + feet_jac = self.env.feet_jacobians(frame="world", return_rot_jac=False) + + J_FL = feet_jac.FL[:, 6:] + J_FR = feet_jac.FR[:, 6:] + J_RL = feet_jac.RL[:, 6:] + J_RR = feet_jac.RR[:, 6:] + total_jac = np.vstack((J_FL, J_FR, J_RL, J_RR)) total_err = np.hstack((err_FL, err_FR, err_RL, err_RR)) - #breakpoint() - + # breakpoint() + # Solve the IK problem dq = total_jac.T @ np.linalg.solve(total_jac @ total_jac.T + damp_matrix, total_err) @@ -123,53 +115,48 @@ def compute_solution(self, q: np.ndarray, FL_foot_target_position: np.ndarray, F print("joint step", self.env.mjData.qpos) self.env.mjData.qpos[7:] = q_joint - #mujoco.mj_fwdPosition(self.env.mjModel, self.env.mjData) - #mujoco.mj_kinematics(self.env.mjModel, self.env.mjData) - - + # mujoco.mj_fwdPosition(self.env.mjModel, self.env.mjData) + # mujoco.mj_kinematics(self.env.mjModel, self.env.mjData) return q_joint - - - if __name__ == "__main__": - - if(cfg.robot == 'go2'): - xml_filename = gym_quadruped_path + '/robot_model/go2/go2.xml' - if(cfg.robot == 'go1'): - xml_filename = gym_quadruped_path + '/robot_model/go1/go1.xml' - elif(cfg.robot == 'aliengo'): - xml_filename = gym_quadruped_path + '/robot_model/aliengo/aliengo.xml' - elif(cfg.robot == 'hyqreal'): - xml_filename = gym_quadruped_path + '/robot_model/hyqreal/hyqreal.xml' - elif(cfg.robot == 'mini_cheetah'): - xml_filename = gym_quadruped_path + '/robot_model/mini_cheetah/mini_cheetah.xml' - - - ik = InverseKinematicsNumeric() - - + if cfg.robot == "go2": + xml_filename = gym_quadruped_path + "/robot_model/go2/go2.xml" + if cfg.robot == "go1": + xml_filename = gym_quadruped_path + "/robot_model/go1/go1.xml" + elif cfg.robot == "aliengo": + xml_filename = gym_quadruped_path + "/robot_model/aliengo/aliengo.xml" + elif cfg.robot == "hyqreal": + xml_filename = gym_quadruped_path + "/robot_model/hyqreal/hyqreal.xml" + elif cfg.robot == "mini_cheetah": + xml_filename = gym_quadruped_path + "/robot_model/mini_cheetah/mini_cheetah.xml" + + ik = InverseKinematicsNumeric() # Check consistency in mujoco m = mujoco.MjModel.from_xml_path(xml_filename) d = mujoco.MjData(m) - random_q_joint = np.random.rand(12,) + random_q_joint = np.random.rand( + 12, + ) d.qpos[7:] = random_q_joint # random quaternion - rand_quat = np.random.rand(4,) + rand_quat = np.random.rand( + 4, + ) rand_quat = rand_quat / np.linalg.norm(rand_quat) d.qpos[3:7] = rand_quat mujoco.mj_fwdPosition(m, d) - FL_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, 'FL') - FR_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, 'FR') - RL_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, 'RL') - RR_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, 'RR') + FL_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "FL") + FR_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "FR") + RL_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "RL") + RR_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "RR") FL_foot_target_position = d.geom_xpos[FL_id] FR_foot_target_position = d.geom_xpos[FR_id] RL_foot_target_position = d.geom_xpos[RL_id] @@ -180,30 +167,27 @@ def compute_solution(self, q: np.ndarray, FL_foot_target_position: np.ndarray, F print("RL foot target position: ", RL_foot_target_position) print("RR foot target position: ", RR_foot_target_position) - - initial_q = copy.deepcopy(d.qpos) - initial_q[7:] = np.random.rand(12,) - + initial_q[7:] = np.random.rand( + 12, + ) + ik.env.mjData.qpos = initial_q mujoco.mj_fwdPosition(ik.env.mjModel, ik.env.mjData) - feet = ik.env.feet_pos(frame='world') + feet = ik.env.feet_pos(frame="world") print("joints start position: ", initial_q) print("FL foot start position", feet.FL) print("FR foot start position", feet.FR) print("RL foot start position", feet.RL) print("RR foot start position", feet.RR) - - + initial_time = time.time() - solution = ik.compute_solution(initial_q, - FL_foot_target_position, FR_foot_target_position, - RL_foot_target_position, RR_foot_target_position) + solution = ik.compute_solution( + initial_q, FL_foot_target_position, FR_foot_target_position, RL_foot_target_position, RR_foot_target_position + ) print("time: ", time.time() - initial_time) - - print("\n") print("DESIRED SOLUTION") foot_position_FL = d.geom_xpos[FL_id] @@ -215,13 +199,12 @@ def compute_solution(self, q: np.ndarray, FL_foot_target_position: np.ndarray, F print("FR foot desired position: ", foot_position_FR) print("RL foot desired position: ", foot_position_RL) print("RR foot desired position: ", foot_position_RR) - - + print("\n") print("MUJOCO IK SOLUTION") ik.env.mjData.qpos[7:] = solution mujoco.mj_fwdPosition(ik.env.mjModel, ik.env.mjData) - feet = ik.env.feet_pos(frame='world') + feet = ik.env.feet_pos(frame="world") print("joints solution: ", ik.env.mjData.qpos) print("FL foot solution position", feet.FL) @@ -229,9 +212,6 @@ def compute_solution(self, q: np.ndarray, FL_foot_target_position: np.ndarray, F print("RL foot solution position", feet.RL) print("RR foot solution position", feet.RR) - - - with mujoco.viewer.launch_passive(m, d) as viewer: while True: - viewer.sync() \ No newline at end of file + viewer.sync() diff --git a/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_qp.py b/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_qp.py index ffd6697..f180050 100644 --- a/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_qp.py +++ b/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_qp.py @@ -1,21 +1,23 @@ import numpy as np -np.set_printoptions(precision=3, suppress = True) -from numpy.linalg import norm + +np.set_printoptions(precision=3, suppress=True) +# import example_robot_data as robex +import os import time -import unittest + import casadi as cs -#import example_robot_data as robex -import os dir_path = os.path.dirname(os.path.realpath(__file__)) import sys + sys.path.append(dir_path) -sys.path.append(dir_path + '/../../') +sys.path.append(dir_path + "/../../") -#from jnrh2023.utils.meshcat_viewer_wrapper import MeshcatVisualizer +# from jnrh2023.utils.meshcat_viewer_wrapper import MeshcatVisualizer # Mujoco magic + import mujoco import mujoco.viewer @@ -23,9 +25,6 @@ import pinocchio as pin from pinocchio import casadi as cpin -import copy - - # Class for solving a generic inverse kinematics problem class InverseKinematicsQP: @@ -40,24 +39,23 @@ def __init__(self, robot: pin.robot_wrapper.RobotWrapper, use_viewer: bool = Non self.robot = robot self.model = self.robot.model self.data = self.robot.data - + # generate the casadi graph cmodel = cpin.Model(self.model) cdata = cmodel.createData() cq = cs.SX.sym("q", self.model.nq, 1) - + # precompute the forward kinematics graph cpin.framesForwardKinematics(cmodel, cdata, cq) # initialize the viewer if requested - self.use_viewer = use_viewer + self.use_viewer = use_viewer """if(self.use_viewer): # Open the viewer self.viz = MeshcatVisualizer(self.robot) self.viz.display(self.robot.q0) self.viewer = self.viz.viewer time.sleep(10)""" - # takes the ID of the feet, and generate a casadi function for a generic forward kinematics self.FL_foot_id = self.model.getFrameId("FL_foot_fixed") @@ -72,11 +70,8 @@ def __init__(self, robot: pin.robot_wrapper.RobotWrapper, use_viewer: bool = Non self.RR_foot_id = self.model.getFrameId("RR_foot_fixed") self.RR_foot_position = cs.Function("RR_foot_pos", [cq], [cdata.oMf[self.RR_foot_id].translation]) - # create the NLP for computing the forward kinematics self.create_nlp_ik() - - def create_nlp_ik(self) -> None: """ @@ -84,48 +79,47 @@ def create_nlp_ik(self) -> None: """ # create NLP self.opti = cs.Opti() - # casadi param to be updated at each request - self.base_pose = self.opti.parameter(7) #7 is the number of DoF of the base, position + quaternion - self.FL_foot_target_position = self.opti.parameter(3) #3 is the number of DoF of the foot, position - self.FR_foot_target_position = self.opti.parameter(3) #3 is the number of DoF of the foot, position - self.RL_foot_target_position = self.opti.parameter(3) #3 is the number of DoF of the foot, position - self.RR_foot_target_position = self.opti.parameter(3) #3 is the number of DoF of the foot, position - - - # define the configuration variables (base pose + joints)) - self.var_q = self.opti.variable(self.model.nq) + self.base_pose = self.opti.parameter(7) # 7 is the number of DoF of the base, position + quaternion + self.FL_foot_target_position = self.opti.parameter(3) # 3 is the number of DoF of the foot, position + self.FR_foot_target_position = self.opti.parameter(3) # 3 is the number of DoF of the foot, position + self.RL_foot_target_position = self.opti.parameter(3) # 3 is the number of DoF of the foot, position + self.RR_foot_target_position = self.opti.parameter(3) # 3 is the number of DoF of the foot, position + # define the configuration variables (base pose + joints)) + self.var_q = self.opti.variable(self.model.nq) # define the cost function (it's parametric!!) - totalcost = cs.sumsqr(self.FL_foot_position(self.var_q) - self.FL_foot_target_position) +\ - cs.sumsqr(self.FR_foot_position(self.var_q) - self.FR_foot_target_position) +\ - cs.sumsqr(self.RL_foot_position(self.var_q) - self.RL_foot_target_position) +\ - cs.sumsqr(self.RR_foot_position(self.var_q) - self.RR_foot_target_position) + totalcost = ( + cs.sumsqr(self.FL_foot_position(self.var_q) - self.FL_foot_target_position) + + cs.sumsqr(self.FR_foot_position(self.var_q) - self.FR_foot_target_position) + + cs.sumsqr(self.RL_foot_position(self.var_q) - self.RL_foot_target_position) + + cs.sumsqr(self.RR_foot_position(self.var_q) - self.RR_foot_target_position) + ) self.opti.minimize(totalcost) - # define the solver - p_opts = dict(print_time=False, verbose=False) + p_opts = dict(print_time=False, verbose=False) s_opts = dict(print_level=0) self.opti.solver("ipopt", p_opts, s_opts) - - # define the parametric constraints for the base, it's fixed, only the leg can move! self.opti.subject_to(self.var_q[0:7] == self.base_pose) - # if use_viewer is yes, you can see the different solution iteration by iteration # in the browser - if(self.use_viewer): + if self.use_viewer: self.opti.callback(lambda i: self.callback(self.opti.debug.value(self.var_q))) - - - def compute_solution(self, q: np.ndarray, FL_foot_target_position: np.ndarray, FR_foot_target_position: np.ndarray, - RL_foot_target_position: np.ndarray, RR_foot_target_position: np.ndarray) -> np.ndarray: + def compute_solution( + self, + q: np.ndarray, + FL_foot_target_position: np.ndarray, + FR_foot_target_position: np.ndarray, + RL_foot_target_position: np.ndarray, + RR_foot_target_position: np.ndarray, + ) -> np.ndarray: """ This method computes the inverse kinematics from initial joint angles and desired foot target positions. @@ -139,8 +133,8 @@ def compute_solution(self, q: np.ndarray, FL_foot_target_position: np.ndarray, F Returns: np.ndarray: The joint angles that achieve the desired foot positions. """ - - #print("initial state", q) + + # print("initial state", q) # set the value for the constraints self.opti.set_value(self.base_pose, q[0:7]) @@ -158,23 +152,21 @@ def compute_solution(self, q: np.ndarray, FL_foot_target_position: np.ndarray, F try: sol = self.opti.solve_limited() sol_q = self.opti.value(self.var_q) - #print("final q: \n", sol_q) + # print("final q: \n", sol_q) return sol_q except: print("ERROR in convergence, plotting debug info.") sol_q = self.opti.debug.value(self.var_q) - - def callback(self, q: np.ndarray) -> None: """ - This method is called by the solver at each iteration (if use_viewer is TRUE) + This method is called by the solver at each iteration (if use_viewer is TRUE) and displays the current joint angles and foot positions. """ pin.framesForwardKinematics(self.model, self.data, q) - #transform_frame_to_world = self.data.oMf[self.FL_foot_id] - #self.viewer["target"].set_transform(self.transform_target_to_world.np) - #self.viewer["current"].set_transform(transform_frame_to_world.np) + # transform_frame_to_world = self.data.oMf[self.FL_foot_id] + # self.viewer["target"].set_transform(self.transform_target_to_world.np) + # self.viewer["current"].set_transform(transform_frame_to_world.np) self.viz.display(q) print("q: \n", q) time.sleep(0.5) @@ -189,32 +181,28 @@ def callback(self, q: np.ndarray) -> None: RL_foot_target_position = np.array([-0.12, 0, 0.06]) RR_foot_target_position = np.array([0, 0.2, 0]) - - initial_time = time.time() - solution = ik.compute_solution(robot.q0, FL_foot_target_position, FR_foot_target_position, - RL_foot_target_position, RR_foot_target_position) + solution = ik.compute_solution( + robot.q0, FL_foot_target_position, FR_foot_target_position, RL_foot_target_position, RR_foot_target_position + ) print("time: ", time.time() - initial_time) - # Check consistency in mujoco - m = mujoco.MjModel.from_xml_path('./../simulation/robot_model/unitree_go1/scene.xml') + m = mujoco.MjModel.from_xml_path("./../simulation/robot_model/unitree_go1/scene.xml") d = mujoco.MjData(m) d.qpos[2] = robot.q0[2] d.qpos[7:] = robot.q0[7:] + FL_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "FL") + FR_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "FR") + RL_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "RL") + RR_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "RR") - FL_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, 'FL') - FR_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, 'FR') - RL_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, 'RL') - RR_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, 'RR') - - joint_FL = solution[7:10] joint_FR = solution[10:13] joint_RL = solution[13:16] joint_RR = solution[16:19] - + d.qpos[7:] = np.concatenate((joint_FR, joint_FL, joint_RR, joint_RL)) mujoco.mj_step(m, d) print("\n") @@ -228,7 +216,7 @@ def callback(self, q: np.ndarray) -> None: print("FR foot position: ", foot_position_FR) print("RL foot position: ", foot_position_RL) print("RR foot position: ", foot_position_RR) - + print("\n") print("PINOCCHIO SOLUTION") print("joints: ", solution[7:]) @@ -239,4 +227,4 @@ def callback(self, q: np.ndarray) -> None: with mujoco.viewer.launch_passive(m, d) as viewer: while True: - viewer.sync() \ No newline at end of file + viewer.sync() diff --git a/quadruped_pympc/helpers/math_utils.py b/quadruped_pympc/helpers/math_utils.py index b150430..947041e 100644 --- a/quadruped_pympc/helpers/math_utils.py +++ b/quadruped_pympc/helpers/math_utils.py @@ -3,9 +3,7 @@ def skew(x): """Skew symmetric matrix from a 3D vector.""" - return np.array([[0, -x[2], x[1]], - [x[2], 0, -x[0]], - [-x[1], x[0], 0]]) + return np.array([[0, -x[2], x[1]], [x[2], 0, -x[0]], [-x[1], x[0], 0]]) def homogenous_transform(vec: np.ndarray, X: np.ndarray) -> np.ndarray: diff --git a/quadruped_pympc/helpers/periodic_gait_generator.py b/quadruped_pympc/helpers/periodic_gait_generator.py index a8c4fe9..d883cec 100644 --- a/quadruped_pympc/helpers/periodic_gait_generator.py +++ b/quadruped_pympc/helpers/periodic_gait_generator.py @@ -1,19 +1,20 @@ -import numpy as np import copy -from quadruped_pympc.helpers.quadruped_utils import GaitType +import numpy as np from gym_quadruped.utils.quadruped_utils import LegsAttr -class PeriodicGaitGenerator: +from quadruped_pympc.helpers.quadruped_utils import GaitType + +class PeriodicGaitGenerator: def __init__(self, duty_factor, step_freq, gait_type: GaitType, horizon): self.duty_factor = duty_factor self.step_freq = step_freq self.horizon = horizon self.gait_type = gait_type self.previous_gait_type = copy.deepcopy(gait_type) - - self.start_and_stop_activated = False # If True, the robot will start and stop walking to save energy + + self.start_and_stop_activated = False # If True, the robot will start and stop walking to save energy # Private variables self._phase_signal, self._init = None, None @@ -39,8 +40,8 @@ def reset(self): self.phase_offset = [0.0, 0.5, 0.5, 0.0] # Set random gait_signal respecting the gait phase offset between legs - #self._phase_signal = (np.asarray(self.phase_offset) + np.random.rand()) % 1.0 - self._phase_signal = (np.asarray(self.phase_offset)) + # self._phase_signal = (np.asarray(self.phase_offset) + np.random.rand()) % 1.0 + self._phase_signal = np.asarray(self.phase_offset) self._init = [False] * len(self.phase_offset) self.n_contact = len(self.phase_offset) self.time_before_switch_freq = 0 @@ -48,7 +49,6 @@ def reset(self): def run(self, dt, new_step_freq): contact = np.zeros(self.n_contact) for leg in range(self.n_contact): - # Increase time by dt # self.t[leg] += dt*self.step_freq self._phase_signal[leg] += dt * new_step_freq @@ -93,7 +93,7 @@ def phase_signal(self): def compute_contact_sequence(self, contact_sequence_dts, contact_sequence_lenghts): # TODO: This function can be vectorized and computed with numpy vectorized operations - if (self.gait_type == GaitType.FULL_STANCE.value): + if self.gait_type == GaitType.FULL_STANCE.value: contact_sequence = np.ones((4, self.horizon * 2)) self.reset() return contact_sequence @@ -106,40 +106,42 @@ def compute_contact_sequence(self, contact_sequence_dts, contact_sequence_lenght # the first value is simply the current predicted contact by the timer contact_sequence[:, 0] = self.run(0.0, self.step_freq) - + # contact_sequence_dts contains a list of dt (usefull for nonuniform sampling) # contact_sequence_lenghts contains the number of steps for each dt j = 0 for i in range(1, self.horizon): - if(i >= contact_sequence_lenghts[j]): + if i >= contact_sequence_lenghts[j]: j += 1 dt = contact_sequence_dts[j] contact_sequence[:, i] = self.run(dt, self.step_freq) self.set_phase_signal(t_init, init_init) return contact_sequence - def set_full_stance(self): self.gait_type = GaitType.FULL_STANCE.value self.reset() - def restore_previous_gait(self): self.gait_type = copy.deepcopy(self.previous_gait_type) self.reset() - - def update_start_and_stop(self, feet_pos, hip_pos, hip_offset, - base_pos, base_ori_euler_xyz, - base_lin_vel, base_ang_vel, - ref_base_lin_vel, ref_base_ang_vel, - current_contact): - - + def update_start_and_stop( + self, + feet_pos, + hip_pos, + hip_offset, + base_pos, + base_ori_euler_xyz, + base_lin_vel, + base_ang_vel, + ref_base_lin_vel, + ref_base_ang_vel, + current_contact, + ): # Get the rotation matrix to transform from world to horizontal frame (hip-centric) yaw = base_ori_euler_xyz[2] - R_W2H = np.array([np.cos(yaw), np.sin(yaw), - -np.sin(yaw), np.cos(yaw)]) + R_W2H = np.array([np.cos(yaw), np.sin(yaw), -np.sin(yaw), np.cos(yaw)]) R_W2H = R_W2H.reshape((2, 2)) # Check if the feet are close to the hip enough to stop the motion @@ -159,23 +161,37 @@ def update_start_and_stop(self, feet_pos, hip_pos, hip_offset, hip_pos_h.FR[:2] = R_W2H @ (hip_pos.FR[:2] - base_pos[0:2]) hip_pos_h.RL[:2] = R_W2H @ (hip_pos.RL[:2] - base_pos[0:2]) hip_pos_h.RR[:2] = R_W2H @ (hip_pos.RR[:2] - base_pos[0:2]) - - feet_to_hip_distance_h_FL = np.sqrt(np.square(feet_pos_h.FL[0] - hip_pos_h.FL[0]) + np.square(feet_pos_h.FL[1] - hip_pos_h.FL[1])) - feet_to_hip_distance_h_FR = np.sqrt(np.square(feet_pos_h.FR[0] - hip_pos_h.FR[0]) + np.square(feet_pos_h.FR[1] - hip_pos_h.FR[1])) - feet_to_hip_distance_h_RL = np.sqrt(np.square(feet_pos_h.RL[0] - hip_pos_h.RL[0]) + np.square(feet_pos_h.RL[1] - hip_pos_h.RL[1])) - feet_to_hip_distance_h_RR = np.sqrt(np.square(feet_pos_h.RR[0] - hip_pos_h.RR[0]) + np.square(feet_pos_h.RR[1] - hip_pos_h.RR[1])) - feet_to_hip_distance_h = np.mean([feet_to_hip_distance_h_FL, feet_to_hip_distance_h_FR, feet_to_hip_distance_h_RL, feet_to_hip_distance_h_RR]) - + + feet_to_hip_distance_h_FL = np.sqrt( + np.square(feet_pos_h.FL[0] - hip_pos_h.FL[0]) + np.square(feet_pos_h.FL[1] - hip_pos_h.FL[1]) + ) + feet_to_hip_distance_h_FR = np.sqrt( + np.square(feet_pos_h.FR[0] - hip_pos_h.FR[0]) + np.square(feet_pos_h.FR[1] - hip_pos_h.FR[1]) + ) + feet_to_hip_distance_h_RL = np.sqrt( + np.square(feet_pos_h.RL[0] - hip_pos_h.RL[0]) + np.square(feet_pos_h.RL[1] - hip_pos_h.RL[1]) + ) + feet_to_hip_distance_h_RR = np.sqrt( + np.square(feet_pos_h.RR[0] - hip_pos_h.RR[0]) + np.square(feet_pos_h.RR[1] - hip_pos_h.RR[1]) + ) + feet_to_hip_distance_h = np.mean( + [feet_to_hip_distance_h_FL, feet_to_hip_distance_h_FR, feet_to_hip_distance_h_RL, feet_to_hip_distance_h_RR] + ) + # Some safety checks to safely stop motion - if(np.linalg.norm(ref_base_lin_vel) == 0.0 and np.linalg.norm(ref_base_ang_vel) == 0.0 and - np.linalg.norm(base_lin_vel) < 0.1 and np.linalg.norm(base_ang_vel) < 0.1 and - np.abs(base_ori_euler_xyz[0]) < 0.05 and np.abs(base_ori_euler_xyz[1]) < 0.05 and - np.sum(current_contact) == 4 and - feet_to_hip_distance_h_FL < 0.06 and - feet_to_hip_distance_h_FR < 0.06 and - feet_to_hip_distance_h_RL < 0.06 and - feet_to_hip_distance_h_RR < 0.06): - self.set_full_stance() - elif(self.gait_type == GaitType.FULL_STANCE.value): + if ( + np.linalg.norm(ref_base_lin_vel) == 0.0 + and np.linalg.norm(ref_base_ang_vel) == 0.0 + and np.linalg.norm(base_lin_vel) < 0.1 + and np.linalg.norm(base_ang_vel) < 0.1 + and np.abs(base_ori_euler_xyz[0]) < 0.05 + and np.abs(base_ori_euler_xyz[1]) < 0.05 + and np.sum(current_contact) == 4 + and feet_to_hip_distance_h_FL < 0.06 + and feet_to_hip_distance_h_FR < 0.06 + and feet_to_hip_distance_h_RL < 0.06 + and feet_to_hip_distance_h_RR < 0.06 + ): + self.set_full_stance() + elif self.gait_type == GaitType.FULL_STANCE.value: self.restore_previous_gait() - diff --git a/quadruped_pympc/helpers/periodic_gait_generator_jax.py b/quadruped_pympc/helpers/periodic_gait_generator_jax.py index 520893f..fc0615e 100644 --- a/quadruped_pympc/helpers/periodic_gait_generator_jax.py +++ b/quadruped_pympc/helpers/periodic_gait_generator_jax.py @@ -1,13 +1,13 @@ -import copy -import numpy as np +import os + import jax import jax.numpy as jnp -import os dir_path = os.path.dirname(os.path.realpath(__file__)) import sys -sys.path.append(dir_path + '/../') + +sys.path.append(dir_path + "/../") # Parameters for both MPC and simulation from quadruped_pympc import config @@ -19,17 +19,14 @@ def __init__(self, duty_factor, step_freq, horizon, mpc_dt): self.step_freq = step_freq self.horizon = horizon self.mpc_dt = mpc_dt - # Only Trot for now self.delta = jnp.array([0.5, 1.0, 1.0, 0.5]) self.t = jnp.zeros(len(self.delta)) - #self.t = jnp.array([0.0, 0.5, 0.5, 0.0]) + # self.t = jnp.array([0.0, 0.5, 0.5, 0.0]) self.n_contact = len(self.delta) - - """def run(self, t, step_freq): contact = jnp.zeros(self.n_contact) #for leg in range(self.n_contact): @@ -63,36 +60,30 @@ def __init__(self, duty_factor, step_freq, horizon, mpc_dt): contact = contact.at[3].set(jnp.where(t[3] > (1-self.duty_factor), 1.0, 0.0)) return contact, t""" - def run(self, t, step_freq): contact = jnp.zeros(self.n_contact) - #for leg in range(self.n_contact): - - #restart condition + # for leg in range(self.n_contact): + + # restart condition t = t.at[0].set(jnp.where(t[0] >= 1.0, 0, t[0])) t = t.at[1].set(jnp.where(t[1] >= 1.0, 0, t[1])) t = t.at[2].set(jnp.where(t[2] >= 1.0, 0, t[2])) t = t.at[3].set(jnp.where(t[3] >= 1.0, 0, t[3])) - - - #increase time by dt - t = t.at[0].set(t[0] + self.mpc_dt*step_freq) - t = t.at[1].set(t[1] + self.mpc_dt*step_freq) - t = t.at[2].set(t[2] + self.mpc_dt*step_freq) - t = t.at[3].set(t[3] + self.mpc_dt*step_freq) + # increase time by dt + t = t.at[0].set(t[0] + self.mpc_dt * step_freq) + t = t.at[1].set(t[1] + self.mpc_dt * step_freq) + t = t.at[2].set(t[2] + self.mpc_dt * step_freq) + t = t.at[3].set(t[3] + self.mpc_dt * step_freq) contact = contact.at[0].set(jnp.where(t[0] < self.duty_factor, 1.0, 0.0)) contact = contact.at[1].set(jnp.where(t[1] < self.duty_factor, 1.0, 0.0)) contact = contact.at[2].set(jnp.where(t[2] < self.duty_factor, 1.0, 0.0)) contact = contact.at[3].set(jnp.where(t[3] < self.duty_factor, 1.0, 0.0)) - - return contact, t - - + def set(self, t): self.t = t @@ -103,7 +94,6 @@ def with_newt(self, t): def get_t(self): return self.t - """def compute_contact_sequence(self, simulation_dt, t, step_freq): t_init = jnp.array(t) @@ -138,31 +128,24 @@ def body_fn(n, carry): contact_sequence = contact_sequence.at[:, 0].set(contact) return contact_sequence, new_t""" - + def compute_contact_sequence(self, simulation_dt, t, step_freq): t_init = jnp.array(t) - + contact_sequence = jnp.zeros((self.n_contact, self.horizon)) new_t = t_init - - def body_fn(n, carry): new_t, contact_sequence = carry new_contact_sequence, new_t = self.run(new_t, step_freq) contact_sequence = contact_sequence.at[:, n].set(new_contact_sequence) - return (new_t, contact_sequence)#, None + return (new_t, contact_sequence) # , None - init_carry = (new_t, contact_sequence) new_t, contact_sequence = jax.lax.fori_loop(0, self.horizon, body_fn, init_carry) - - return contact_sequence, new_t - - class Gait: TROT = 0 @@ -176,16 +159,16 @@ class Gait: if __name__ == "__main__": # Periodic gait generator - - mpc_dt = config.mpc_params['dt'] - horizon = config.mpc_params['horizon'] - gait = config.simulation_params['gait'] - simulation_dt = config.simulation_params['dt'] - if(gait == "trot"): + + mpc_dt = config.mpc_params["dt"] + horizon = config.mpc_params["horizon"] + gait = config.simulation_params["gait"] + simulation_dt = config.simulation_params["dt"] + if gait == "trot": step_frequency = 2.65 duty_factor = 0.65 p_gait = Gait.TROT - elif(gait == "crawl"): + elif gait == "crawl": step_frequency = 0.4 duty_factor = 0.9 p_gait = Gait.BACKDIAGONALCRAWL @@ -193,13 +176,17 @@ class Gait: step_frequency = 2 duty_factor = 0.7 p_gait = Gait.PACE - + pgg = PeriodicGaitGenerator(duty_factor=duty_factor, step_freq=step_frequency, gait_type=p_gait, horizon=horizon) - pgg_jax = PeriodicGaitGeneratorJax(duty_factor=duty_factor, step_freq=step_frequency, horizon=horizon, mpc_dt=mpc_dt) + pgg_jax = PeriodicGaitGeneratorJax( + duty_factor=duty_factor, step_freq=step_frequency, horizon=horizon, mpc_dt=mpc_dt + ) jitted_pgg = jax.jit(pgg_jax.compute_contact_sequence) for j in range(100): contact_sequence = pgg.compute_contact_sequence(mpc_dt=mpc_dt, simulation_dt=simulation_dt) - contact_sequence_jax = pgg_jax.compute_contact_sequence(simulation_dt=simulation_dt, t = pgg.get_t(), step_freq=step_frequency) + contact_sequence_jax = pgg_jax.compute_contact_sequence( + simulation_dt=simulation_dt, t=pgg.get_t(), step_freq=step_frequency + ) contact_sequence_jax_jitted = jitted_pgg(simulation_dt, pgg.get_t(), step_frequency) print("contact_sequence:\n", contact_sequence) print("contact_sequence_jax:\n", contact_sequence_jax[0]) @@ -207,7 +194,4 @@ class Gait: pgg.run(simulation_dt) print("############") - - - - breakpoint() \ No newline at end of file + breakpoint() diff --git a/quadruped_pympc/helpers/quadruped_utils.py b/quadruped_pympc/helpers/quadruped_utils.py index 1b88c3b..91ee2b5 100644 --- a/quadruped_pympc/helpers/quadruped_utils.py +++ b/quadruped_pympc/helpers/quadruped_utils.py @@ -23,15 +23,16 @@ class GaitType(Enum): FULL_STANCE = 7 -def plot_swing_mujoco(viewer: Handle, - swing_traj_controller: SwingTrajectoryController, - swing_period, - swing_time: namedtuple, - lift_off_positions: namedtuple, - nmpc_footholds: namedtuple, - ref_feet_pos: namedtuple, - geom_ids: namedtuple = None, - ): +def plot_swing_mujoco( + viewer: Handle, + swing_traj_controller: SwingTrajectoryController, + swing_period, + swing_time: namedtuple, + lift_off_positions: namedtuple, + nmpc_footholds: namedtuple, + ref_feet_pos: namedtuple, + geom_ids: namedtuple = None, +): """Function to plot the desired foot swing trajectory in Mujoco. Args: @@ -58,40 +59,41 @@ def plot_swing_mujoco(viewer: Handle, if geom_ids is None: geom_ids = LegsAttr(FL=[], FR=[], RL=[], RR=[]) # Instantiate a new geometry - for leg_id, leg_name in enumerate(['FL', 'FR', 'RL', 'RR']): + for leg_id, leg_name in enumerate(["FL", "FR", "RL", "RR"]): viewer.user_scn.ngeom += NUM_TRAJ_POINTS geom_ids[leg_name] = list(range(viewer.user_scn.ngeom - NUM_TRAJ_POINTS - 1, viewer.user_scn.ngeom - 1)) # viewer.user_scn.ngeom = 1 # We first draw the trajectory of the feet des_foot_traj = LegsAttr(FL=[], FR=[], RL=[], RR=[]) - for leg_id, leg_name in enumerate(['FL', 'FR', 'RL', 'RR']): + for leg_id, leg_name in enumerate(["FL", "FR", "RL", "RR"]): # TODO: This function should be vectorized rather than queried sequentially if swing_time[leg_name] == 0.0: continue for point_idx, foot_swing_time in enumerate(np.linspace(swing_time[leg_name], swing_period, NUM_TRAJ_POINTS)): ref_foot_pos, _, _ = swing_traj_controller.swing_generator.compute_trajectory_references( - foot_swing_time, - lift_off_positions[leg_name], - nmpc_footholds[leg_name]) + foot_swing_time, lift_off_positions[leg_name], nmpc_footholds[leg_name] + ) des_foot_traj[leg_name].append(ref_foot_pos.squeeze()) for point_idx in range(NUM_TRAJ_POINTS - 1): - render_line(viewer=viewer, - initial_point=des_foot_traj[leg_name][point_idx], - target_point=des_foot_traj[leg_name][point_idx + 1], - width=.005, - color=np.array([1, 0, 0, 1]), - geom_id=geom_ids[leg_name][point_idx] - ) + render_line( + viewer=viewer, + initial_point=des_foot_traj[leg_name][point_idx], + target_point=des_foot_traj[leg_name][point_idx + 1], + width=0.005, + color=np.array([1, 0, 0, 1]), + geom_id=geom_ids[leg_name][point_idx], + ) # Add a sphere at the the ref_feet_pos - render_sphere(viewer=viewer, - position=ref_feet_pos[leg_name], - diameter=0.04, - color=np.array([0, 1, 0, .5]), - geom_id=geom_ids[leg_name][-1] - ) + render_sphere( + viewer=viewer, + position=ref_feet_pos[leg_name], + diameter=0.04, + color=np.array([0, 1, 0, 0.5]), + geom_id=geom_ids[leg_name][-1], + ) return geom_ids @@ -99,15 +101,15 @@ def check_zmp_constraint_satisfaction(state, contact_status, forces): # TODO: This import should go from quadruped_pympc import config - base_w = copy.deepcopy(state['position']) - base_vel_w = copy.deepcopy(state['linear_velocity']) + base_w = copy.deepcopy(state["position"]) + base_vel_w = copy.deepcopy(state["linear_velocity"]) - FL = copy.deepcopy(state['foot_FL']) - FR = copy.deepcopy(state['foot_FR']) - RL = copy.deepcopy(state['foot_RL']) - RR = copy.deepcopy(state['foot_RR']) + FL = copy.deepcopy(state["foot_FL"]) + FR = copy.deepcopy(state["foot_FR"]) + RL = copy.deepcopy(state["foot_RL"]) + RR = copy.deepcopy(state["foot_RR"]) - yaw = copy.deepcopy(state['orientation'][2]) + yaw = copy.deepcopy(state["orientation"][2]) h_R_w = np.zeros((2, 2)) h_R_w[0, 0] = np.cos(yaw) h_R_w[0, 1] = np.sin(yaw) @@ -127,7 +129,7 @@ def check_zmp_constraint_satisfaction(state, contact_status, forces): gravity = np.array([0, 0, -9.81]) linear_com_acc = (1 / config.mass) * temp + gravity - if (config.mpc_params['use_zmp_stability']): + if config.mpc_params["use_zmp_stability"]: gravity_z = 9.81 robotHeight = base_w[2] zmp = base_w[0:2] - linear_com_acc[0:2] * (robotHeight / gravity_z) @@ -176,24 +178,24 @@ def check_zmp_constraint_satisfaction(state, contact_status, forces): violation = 0 - if (FL_contact == 1): - if (FR_contact == 1): + if FL_contact == 1: + if FR_contact == 1: # ub_support_FL_FR = -0.0 # lb_support_FL_FR = -1000 - if (constraint_FL_FR > 0): + if constraint_FL_FR > 0: violation = violation + 1 else: # ub_support_FL_RR = 1000 # lb_support_FL_RR = 0.0 - if (constraint_FL_RR < 0): + if constraint_FL_RR < 0: violation = violation + 1 - if (FR_contact == 1): - if (RR_contact == 1): + if FR_contact == 1: + if RR_contact == 1: # ub_support_FR_RR = 1000 # lb_support_FR_RR = 0.0 - if (constraint_FR_RR < 0): + if constraint_FR_RR < 0: violation = violation + 1 else: @@ -203,35 +205,35 @@ def check_zmp_constraint_satisfaction(state, contact_status, forces): # violation = violation + 1 # TOCHECK - if (RR_contact == 1): - if (RL_contact == 1): + if RR_contact == 1: + if RL_contact == 1: # ub_support_RR_RL = 1000 # lb_support_RR_RL = 0.0 - if (constraint_RR_RL < 0): + if constraint_RR_RL < 0: violation = violation + 1 else: # ub_support_FL_RR = -0.0 # lb_support_FL_RR = -1000 - if (constraint_FL_RR > 0): + if constraint_FL_RR > 0: violation = violation + 1 * 0 - if (RL_contact == 1): - if (FL_contact == 1): + if RL_contact == 1: + if FL_contact == 1: # ub_support_RL_FL = -0.0 # lb_support_RL_FL = -1000 - if (constraint_RL_FL > 0): + if constraint_RL_FL > 0: violation = violation + 1 else: # ub_support_FR_RL = -0.0 # lb_support_FR_RL = -1000 - if (constraint_FR_RL > 0): + if constraint_FR_RL > 0: violation = violation + 1 - if (FL_contact == 1 and FR_contact == 1 and RL_contact == 1 and RR_contact == 1): + if FL_contact == 1 and FR_contact == 1 and RL_contact == 1 and RR_contact == 1: violation = 0 - if (violation >= 1): + if violation >= 1: return True else: return False diff --git a/quadruped_pympc/helpers/srb_inertia_computation.py b/quadruped_pympc/helpers/srb_inertia_computation.py index 231d1fe..026551f 100644 --- a/quadruped_pympc/helpers/srb_inertia_computation.py +++ b/quadruped_pympc/helpers/srb_inertia_computation.py @@ -10,33 +10,45 @@ import sys -sys.path.append(dir_path + '/../') +sys.path.append(dir_path + "/../") # Parameters for both MPC and simulation from quadruped_pympc import config class SrbInertiaComputation: - def __init__(self, ) -> None: - if (config.robot == 'go2'): - urdf_filename = dir_path + '/../simulation/robot_model/go2/go2.urdf' - elif (config.robot == 'aliengo'): - urdf_filename = dir_path + '/../simulation/robot_model/aliengo/aliengo.urdf' - elif (config.robot == 'hyqreal'): - urdf_filename = dir_path + '/../simulation/robot_model/hyqreal/hyqreal.urdf' - elif (config.robot == 'mini_cheetah'): - urdf_filename = dir_path + '/../simulation/robot_model/mini_cheetah/mini_cheetah.urdf' + def __init__( + self, + ) -> None: + if config.robot == "go2": + urdf_filename = dir_path + "/../simulation/robot_model/go2/go2.urdf" + elif config.robot == "aliengo": + urdf_filename = dir_path + "/../simulation/robot_model/aliengo/aliengo.urdf" + elif config.robot == "hyqreal": + urdf_filename = dir_path + "/../simulation/robot_model/hyqreal/hyqreal.urdf" + elif config.robot == "mini_cheetah": + urdf_filename = dir_path + "/../simulation/robot_model/mini_cheetah/mini_cheetah.urdf" self.use_pinocchio = True - if (self.use_pinocchio): + if self.use_pinocchio: self.robot_full = pin.buildModelFromUrdf(urdf_filename) # Create a list of joints to lock - jointsToLock = ['FL_hip_joint', 'FL_thigh_joint', 'FL_calf_joint', - 'FR_hip_joint', 'FR_thigh_joint', 'FR_calf_joint', - 'RL_hip_joint', 'RL_thigh_joint', 'RL_calf_joint', - 'RR_hip_joint', 'RR_thigh_joint', 'RR_calf_joint'] + jointsToLock = [ + "FL_hip_joint", + "FL_thigh_joint", + "FL_calf_joint", + "FR_hip_joint", + "FR_thigh_joint", + "FR_calf_joint", + "RL_hip_joint", + "RL_thigh_joint", + "RL_calf_joint", + "RR_hip_joint", + "RR_thigh_joint", + "RR_calf_joint", + ] # Get the ID of all existing joints self.jointsToLockIDs = [] @@ -45,16 +57,15 @@ def __init__(self, ) -> None: if self.robot_full.existJointName(jn): self.jointsToLockIDs.append(self.robot_full.getJointId(jn)) else: - print('Warning: joint ' + str(jn) + ' does not belong to the model!') + print("Warning: joint " + str(jn) + " does not belong to the model!") else: self.kindyn = KinDynComputations(urdfstring=urdf_filename) self.kindyn.set_frame_velocity_representation(representation=Representations.BODY_FIXED_REPRESENTATION) self.mass_mass_fun = self.kindyn.mass_matrix_fun() def compute_inertia(self, q): - - if (self.use_pinocchio): - if (config.robot == 'aliengo'): + if self.use_pinocchio: + if config.robot == "aliengo": robot_reduced = pin.buildReducedModel(self.robot_full, self.jointsToLockIDs, q[3:]) else: robot_reduced = pin.buildReducedModel(self.robot_full, self.jointsToLockIDs, q[7:]) diff --git a/quadruped_pympc/helpers/swing_generators/explicit_swing_trajectory_generator.py b/quadruped_pympc/helpers/swing_generators/explicit_swing_trajectory_generator.py index 9b492b7..e46a0ca 100644 --- a/quadruped_pympc/helpers/swing_generators/explicit_swing_trajectory_generator.py +++ b/quadruped_pympc/helpers/swing_generators/explicit_swing_trajectory_generator.py @@ -1,34 +1,24 @@ -import numpy as np -import matplotlib.pyplot as plt - -from scipy.interpolate import CubicSpline, Akima1DInterpolator, CubicHermiteSpline -import copy +import matplotlib.pyplot as plt +import numpy as np class SwingTrajectoryGenerator: - def __init__(self, - step_height: float, - swing_period: float) -> None: - + def __init__(self, step_height: float, swing_period: float) -> None: self.step_height = step_height self.swing_period = swing_period self.half_swing_period = swing_period / 2 self.bezier_time_factor = 1 / (swing_period / 2) - - - - def plot_trajectory_3d(self, - curve_points: np.ndarray) -> None: + def plot_trajectory_3d(self, curve_points: np.ndarray) -> None: curve_points = np.array(curve_points) fig = plt.figure() - ax = fig.add_subplot(111, projection='3d') + ax = fig.add_subplot(111, projection="3d") ax.plot(curve_points[:, 0], curve_points[:, 1], curve_points[:, 2]) ax.legend() - plt.title('3D Bézier Curve') + plt.title("3D Bézier Curve") plt.show() def plot_trajectory_references(self, tp, fp, vp, ap): @@ -43,33 +33,31 @@ def plot_trajectory_references(self, tp, fp, vp, ap): # Plot position for i in range(3): - axs[0].plot(time_points, footPosDes_points[:, i], label=f"Position {i+1}") - axs[0].set_xlabel('Time') - axs[0].set_ylabel('Position') + axs[0].plot(time_points, footPosDes_points[:, i], label=f"Position {i + 1}") + axs[0].set_xlabel("Time") + axs[0].set_ylabel("Position") axs[0].legend() # Plot velocity for i in range(3): - axs[1].plot(time_points, footVelDes_points[:, i], label=f"Velocity {i+1}") - axs[1].set_xlabel('Time') - axs[1].set_ylabel('Velocity') + axs[1].plot(time_points, footVelDes_points[:, i], label=f"Velocity {i + 1}") + axs[1].set_xlabel("Time") + axs[1].set_ylabel("Velocity") axs[1].legend() # Plot acceleration for i in range(3): - axs[2].plot(time_points, footAccDes_points[:, i], label=f"Acceleration {i+1}") - axs[2].set_xlabel('Time') - axs[2].set_ylabel('Acceleration') + axs[2].plot(time_points, footAccDes_points[:, i], label=f"Acceleration {i + 1}") + axs[2].set_xlabel("Time") + axs[2].set_ylabel("Acceleration") axs[2].legend() plt.tight_layout() plt.show() - - def compute_control_points(self, - swing_time: float, - lift_off: np.ndarray, - touch_down: np.ndarray) -> (np.array, np.array, np.array, np.array): + def compute_control_points( + self, swing_time: float, lift_off: np.ndarray, touch_down: np.ndarray + ) -> (np.array, np.array, np.array, np.array): cp1, cp2, cp3, cp4 = None, None, None, None middle_point = 0.5 * (lift_off + touch_down) @@ -86,37 +74,54 @@ def compute_control_points(self, return swing_time % self.half_swing_period, cp1, cp2, cp3, cp4 - - def compute_trajectory_references(self, - swing_time: float, - lift_off: np.ndarray, - touch_down: np.ndarray) -> (np.ndarray, np.ndarray, np.ndarray): - bezier_swing_time, cp1, cp2, cp3, cp4 = self.compute_control_points(swing_time, lift_off.squeeze(), touch_down.squeeze()) - - desired_foot_position = (1 - self.bezier_time_factor * bezier_swing_time) ** 3 * cp1 + \ - 3 * (self.bezier_time_factor * bezier_swing_time) * (1 - self.bezier_time_factor * bezier_swing_time) ** 2 * cp2 + \ - 3 * (self.bezier_time_factor * bezier_swing_time) ** 2 * (1 - self.bezier_time_factor * bezier_swing_time) * cp3 + \ - (self.bezier_time_factor * bezier_swing_time) ** 3 * cp4 - - desired_foot_velocity = 3 * (1 - self.bezier_time_factor * bezier_swing_time) ** 2 * (cp2 - cp1) + \ - 6 * (1 - self.bezier_time_factor * bezier_swing_time) * (self.bezier_time_factor * bezier_swing_time) * (cp3 - cp2) + \ - 3 * (self.bezier_time_factor * bezier_swing_time) ** 2 * (cp4 - cp3) - - desired_foot_acceleration = 6 * (1 - self.bezier_time_factor * bezier_swing_time) * (cp3 - 2 * cp2 + cp1) + \ - 6 * (self.bezier_time_factor * bezier_swing_time) * (cp4 - 2 * cp3 + cp2) - - return desired_foot_position.reshape((3,)), desired_foot_velocity.reshape((3,)), desired_foot_acceleration.reshape((3,)) + def compute_trajectory_references( + self, swing_time: float, lift_off: np.ndarray, touch_down: np.ndarray + ) -> (np.ndarray, np.ndarray, np.ndarray): + bezier_swing_time, cp1, cp2, cp3, cp4 = self.compute_control_points( + swing_time, lift_off.squeeze(), touch_down.squeeze() + ) + + desired_foot_position = ( + (1 - self.bezier_time_factor * bezier_swing_time) ** 3 * cp1 + + 3 + * (self.bezier_time_factor * bezier_swing_time) + * (1 - self.bezier_time_factor * bezier_swing_time) ** 2 + * cp2 + + 3 + * (self.bezier_time_factor * bezier_swing_time) ** 2 + * (1 - self.bezier_time_factor * bezier_swing_time) + * cp3 + + (self.bezier_time_factor * bezier_swing_time) ** 3 * cp4 + ) + + desired_foot_velocity = ( + 3 * (1 - self.bezier_time_factor * bezier_swing_time) ** 2 * (cp2 - cp1) + + 6 + * (1 - self.bezier_time_factor * bezier_swing_time) + * (self.bezier_time_factor * bezier_swing_time) + * (cp3 - cp2) + + 3 * (self.bezier_time_factor * bezier_swing_time) ** 2 * (cp4 - cp3) + ) + + desired_foot_acceleration = 6 * (1 - self.bezier_time_factor * bezier_swing_time) * ( + cp3 - 2 * cp2 + cp1 + ) + 6 * (self.bezier_time_factor * bezier_swing_time) * (cp4 - 2 * cp3 + cp2) + + return ( + desired_foot_position.reshape((3,)), + desired_foot_velocity.reshape((3,)), + desired_foot_acceleration.reshape((3,)), + ) # Example: if __name__ == "__main__": step_height = 0.08 swing_period = 0.9 - trajectory_generator = SwingTrajectoryGenerator(step_height=step_height, - swing_period=swing_period) + trajectory_generator = SwingTrajectoryGenerator(step_height=step_height, swing_period=swing_period) - lift_off = np.array([0,0,0]) - touch_down = np.array([0.1,0,0.0]) + lift_off = np.array([0, 0, 0]) + touch_down = np.array([0.1, 0, 0.0]) # Generate trajectory points simulation_dt = 0.002 @@ -126,19 +131,16 @@ def compute_trajectory_references(self, acceleration_points = [] i = 0 for foot_swing_time in np.arange(0.000001, swing_period, 0.002): - desired_foot_position, \ - desired_foot_velocity, \ - desired_foot_acceleration = trajectory_generator.compute_trajectory_references(swing_time=i, - lift_off=lift_off, - touch_down=touch_down) + desired_foot_position, desired_foot_velocity, desired_foot_acceleration = ( + trajectory_generator.compute_trajectory_references(swing_time=i, lift_off=lift_off, touch_down=touch_down) + ) i += simulation_dt - + time_points.append(i) position_points.append(desired_foot_position.squeeze()) velocity_points.append(desired_foot_velocity.squeeze()) acceleration_points.append(desired_foot_acceleration.squeeze()) - + # Plot the generated trajectory trajectory_generator.plot_trajectory_3d(np.array(position_points)) trajectory_generator.plot_trajectory_references(time_points, position_points, velocity_points, acceleration_points) - \ No newline at end of file diff --git a/quadruped_pympc/helpers/swing_generators/ndcurves_swing_trajectory_generator.py b/quadruped_pympc/helpers/swing_generators/ndcurves_swing_trajectory_generator.py index d5bf4d6..f2b427d 100644 --- a/quadruped_pympc/helpers/swing_generators/ndcurves_swing_trajectory_generator.py +++ b/quadruped_pympc/helpers/swing_generators/ndcurves_swing_trajectory_generator.py @@ -1,22 +1,19 @@ -import numpy as np -import matplotlib.pyplot as plt - -import ndcurves # importing tools to plot bezier curves -#from ndcurves.plot import plotBezier +# from ndcurves.plot import plotBezier import copy +import matplotlib.pyplot as plt +import ndcurves +import numpy as np + class SwingTrajectoryGenerator: - def __init__(self, - step_height: float, - swing_period: float) -> None: + def __init__(self, step_height: float, swing_period: float) -> None: self.step_height = step_height self.swing_period = swing_period self.half_swing_period = swing_period / 2 self.bezier_time_factor = 1 / (swing_period / 2) - # Stored swing-trajectory properties dt = 0.002 self._dt = copy.deepcopy(dt) @@ -25,40 +22,36 @@ def __init__(self, # Define the initial and terminal curve constraints self._constraints = ndcurves.curve_constraints(3) - self._constraints.init_vel = np.array([[0., 0., 0.]]).T - self._constraints.end_vel = np.array([[0., 0., 0.]]).T - self._constraints.init_acc = np.array([[0., 0., 0.]]).T - self._constraints.end_acc = np.array([[0., 0., 0.]]).T - + self._constraints.init_vel = np.array([[0.0, 0.0, 0.0]]).T + self._constraints.end_vel = np.array([[0.0, 0.0, 0.0]]).T + self._constraints.init_acc = np.array([[0.0, 0.0, 0.0]]).T + self._constraints.end_acc = np.array([[0.0, 0.0, 0.0]]).T def createBezierCurve(self, x0, xf): - scaling_factor = 0.7105 - p1 = x0 + np.array([0., 0., self._stepHeight / scaling_factor]) - p2 = 0.5 * (x0 + xf) + np.array([0., 0., self._stepHeight / scaling_factor]) - p3 = xf + np.array([0., 0., self._stepHeight / scaling_factor]) + p1 = x0 + np.array([0.0, 0.0, self._stepHeight / scaling_factor]) + p2 = 0.5 * (x0 + xf) + np.array([0.0, 0.0, self._stepHeight / scaling_factor]) + p3 = xf + np.array([0.0, 0.0, self._stepHeight / scaling_factor]) self._curve = ndcurves.bezier( - np.array([x0, p1, p2, p3, xf]).T, self._constraints, 0., (self._N - 1) * self._dt) - + np.array([x0, p1, p2, p3, xf]).T, self._constraints, 0.0, (self._N - 1) * self._dt + ) + def compute_trajectory_references(self, k, lift_off, touch_down): - self.createBezierCurve(lift_off, touch_down) position = self._curve(k * self._dt) velocity = self._curve.derivate(k * self._dt, 1) return position, velocity, 0 - - def plot_trajectory_3d(self, - curve_points: np.array) -> None: + def plot_trajectory_3d(self, curve_points: np.array) -> None: curve_points = np.array(curve_points) fig = plt.figure() - ax = fig.add_subplot(111, projection='3d') + ax = fig.add_subplot(111, projection="3d") ax.plot(curve_points[:, 0], curve_points[:, 1], curve_points[:, 2]) ax.legend() - plt.title('3D Bézier Curve') + plt.title("3D Bézier Curve") plt.show() def plot_trajectory_references(self, tp, fp, vp): @@ -66,58 +59,54 @@ def plot_trajectory_references(self, tp, fp, vp): time_points = np.array(tp) footPosDes_points = np.array(fp) footVelDes_points = np.array(vp) - + # Create subplots for position, velocity, and acceleration fig, axs = plt.subplots(2, 1, figsize=(8, 12)) # Plot position for i in range(3): - axs[0].plot(time_points, footPosDes_points[:, i], label=f"Position {i+1}") - axs[0].set_xlabel('Time') - axs[0].set_ylabel('Position') + axs[0].plot(time_points, footPosDes_points[:, i], label=f"Position {i + 1}") + axs[0].set_xlabel("Time") + axs[0].set_ylabel("Position") axs[0].legend() # Plot velocity for i in range(3): - axs[1].plot(time_points, footVelDes_points[:, i], label=f"Velocity {i+1}") - axs[1].set_xlabel('Time') - axs[1].set_ylabel('Velocity') + axs[1].plot(time_points, footVelDes_points[:, i], label=f"Velocity {i + 1}") + axs[1].set_xlabel("Time") + axs[1].set_ylabel("Velocity") axs[1].legend() - - plt.tight_layout() plt.show() - - - - - # Example: if __name__ == "__main__": step_height = 0.08 swing_period = 0.9 - trajectory_generator = SwingTrajectoryGenerator(step_height=step_height ,swing_period=swing_period, position_gain_fb = 0, velocity_gain_fb = 0) + trajectory_generator = SwingTrajectoryGenerator( + step_height=step_height, swing_period=swing_period, position_gain_fb=0, velocity_gain_fb=0 + ) + + # trajectory_generator.createBezierCurve(np.array([0,0,0]), np.array([0.3,0,0])) - #trajectory_generator.createBezierCurve(np.array([0,0,0]), np.array([0.3,0,0])) - # Generate trajectory points time_points = [] position_points = [] velocity_points = [] - #acceleration_points = [] + # acceleration_points = [] i = 0 for foot_swing_time in np.arange(0.000001, swing_period, 0.002): - desired_foot_position, desired_foot_velocity, _ = trajectory_generator.compute_trajectory_references(i, np.array([0,0,0]), np.array([0.3,0,0])) + desired_foot_position, desired_foot_velocity, _ = trajectory_generator.compute_trajectory_references( + i, np.array([0, 0, 0]), np.array([0.3, 0, 0]) + ) i += 1 - - time_points.append(i*0.002) + + time_points.append(i * 0.002) position_points.append(desired_foot_position.squeeze()) velocity_points.append(desired_foot_velocity.squeeze()) - + # Plot the generated trajectory trajectory_generator.plot_trajectory_3d(np.array(position_points)) trajectory_generator.plot_trajectory_references(time_points, position_points, velocity_points) - \ No newline at end of file diff --git a/quadruped_pympc/helpers/swing_generators/scipy_swing_trajectory_generator.py b/quadruped_pympc/helpers/swing_generators/scipy_swing_trajectory_generator.py index 99044d0..0581290 100644 --- a/quadruped_pympc/helpers/swing_generators/scipy_swing_trajectory_generator.py +++ b/quadruped_pympc/helpers/swing_generators/scipy_swing_trajectory_generator.py @@ -1,15 +1,11 @@ -import numpy as np -import matplotlib.pyplot as plt - -from scipy.interpolate import CubicSpline, Akima1DInterpolator, CubicHermiteSpline -import copy +import matplotlib.pyplot as plt +import numpy as np +from scipy.interpolate import CubicSpline class SwingTrajectoryGenerator: - def __init__(self, - step_height: float, - swing_period: float) -> None: + def __init__(self, step_height: float, swing_period: float) -> None: self.step_height = step_height self.swing_period = swing_period self.half_swing_period = swing_period / 2 @@ -18,32 +14,33 @@ def __init__(self, # Stored swing-trajectory properties self.stepHeight = step_height - - - def createCurve(self, x0, xf): - - #scaling_factor = 0.7105 + # scaling_factor = 0.7105 scaling_factor = 1.5 - - p1 = x0 + np.array([0., 0., self.stepHeight / scaling_factor]) - p2 = 0.5 * (x0 + xf) + np.array([0., 0., self.stepHeight ]) - p3 = xf + np.array([0., 0., self.stepHeight / scaling_factor]) - - + + p1 = x0 + np.array([0.0, 0.0, self.stepHeight / scaling_factor]) + p2 = 0.5 * (x0 + xf) + np.array([0.0, 0.0, self.stepHeight]) + p3 = xf + np.array([0.0, 0.0, self.stepHeight / scaling_factor]) + x = np.array([x0[0], p1[0], p2[0], p3[0], xf[0]]) y = np.array([x0[1], p1[1], p2[1], p3[1], xf[1]]) z = np.array([x0[2], p1[2], p2[2], p3[2], xf[2]]) - t = np.array([0, self.half_swing_period/2, self.half_swing_period, self.half_swing_period*3/2, self.half_swing_period*2]) + t = np.array( + [ + 0, + self.half_swing_period / 2, + self.half_swing_period, + self.half_swing_period * 3 / 2, + self.half_swing_period * 2, + ] + ) self._curve_x = CubicSpline(t, x, bc_type=["clamped", "clamped"]) self._curve_y = CubicSpline(t, y, bc_type=["clamped", "clamped"]) self._curve_z = CubicSpline(t, z, bc_type=["clamped", "clamped"]) - - - #self._curve_x = Akima1DInterpolator(t, x) - #self._curve_y = Akima1DInterpolator(t, y) - #self._curve_z = Akima1DInterpolator(t, z) + # self._curve_x = Akima1DInterpolator(t, x) + # self._curve_y = Akima1DInterpolator(t, y) + # self._curve_z = Akima1DInterpolator(t, z) """dxdt = np.array([0, 0, 0]) dydt = np.array([0, 0, 0]) @@ -68,13 +65,10 @@ def createCurve(self, x0, xf): self._curve_y_acc = self._curve_y_vel.derivative() self._curve_z_acc = self._curve_z_vel.derivative() - - def compute_trajectory_references(self, - swing_time: float, - lift_off: np.array, - touch_down: np.array) -> (np.array, np.array, np.array): - - #if(swing_time == 0): + def compute_trajectory_references( + self, swing_time: float, lift_off: np.array, touch_down: np.array + ) -> (np.array, np.array, np.array): + # if(swing_time == 0): self.createCurve(lift_off, touch_down) position_x = self._curve_x(swing_time) @@ -82,88 +76,75 @@ def compute_trajectory_references(self, position_z = self._curve_z(swing_time) position = np.array([position_x, position_y, position_z]) - + velocity_x = self._curve_x_vel(swing_time) velocity_y = self._curve_y_vel(swing_time) velocity_z = self._curve_z_vel(swing_time) velocity = np.array([velocity_x, velocity_y, velocity_z]) - acceleration_x = self._curve_x_acc(swing_time) acceleration_y = self._curve_y_acc(swing_time) acceleration_z = self._curve_z_acc(swing_time) acceleration = np.array([acceleration_x, acceleration_y, acceleration_z]) - - return position, velocity, acceleration - - def plot_trajectory_3d(self, - curve_points: np.array) -> None: + def plot_trajectory_3d(self, curve_points: np.array) -> None: curve_points = np.array(curve_points) fig = plt.figure() - ax = fig.add_subplot(111, projection='3d') + ax = fig.add_subplot(111, projection="3d") ax.plot(curve_points[:, 0], curve_points[:, 1], curve_points[:, 2]) ax.legend() - plt.title('3D Curve') + plt.title("3D Curve") plt.show() - - def plot_trajectory_references(self, tp, fp, vp, ap): # Convert lists to NumPy arrays for easier plotting time_points = np.array(tp) footPosDes_points = np.array(fp) footVelDes_points = np.array(vp) footAccDes_points = np.array(ap) - + # Create subplots for position, velocity, and acceleration fig, axs = plt.subplots(3, 1, figsize=(8, 12)) # Plot position for i in range(3): - axs[0].plot(time_points, footPosDes_points[:, i], label=f"Position {i+1}") - axs[0].set_xlabel('Time') - axs[0].set_ylabel('Position') + axs[0].plot(time_points, footPosDes_points[:, i], label=f"Position {i + 1}") + axs[0].set_xlabel("Time") + axs[0].set_ylabel("Position") axs[0].legend() # Plot velocity for i in range(3): - axs[1].plot(time_points, footVelDes_points[:, i], label=f"Velocity {i+1}") - axs[1].set_xlabel('Time') - axs[1].set_ylabel('Velocity') + axs[1].plot(time_points, footVelDes_points[:, i], label=f"Velocity {i + 1}") + axs[1].set_xlabel("Time") + axs[1].set_ylabel("Velocity") axs[1].legend() # Plot acceleration for i in range(3): - axs[2].plot(time_points, footAccDes_points[:, i], label=f"Acceleration {i+1}") - axs[2].set_xlabel('Time') - axs[2].set_ylabel('Acceleration') + axs[2].plot(time_points, footAccDes_points[:, i], label=f"Acceleration {i + 1}") + axs[2].set_xlabel("Time") + axs[2].set_ylabel("Acceleration") axs[2].legend() - - plt.tight_layout() plt.show() - - - # Example: if __name__ == "__main__": step_height = 0.08 swing_period = 0.9 - trajectory_generator = SwingTrajectoryGenerator(step_height=step_height, - swing_period=swing_period) + trajectory_generator = SwingTrajectoryGenerator(step_height=step_height, swing_period=swing_period) - lift_off = np.array([0,0,0]) - touch_down = np.array([0.1,-0.2,0.0]) + lift_off = np.array([0, 0, 0]) + touch_down = np.array([0.1, -0.2, 0.0]) # Generate trajectory points simulation_dt = 0.002 @@ -173,19 +154,16 @@ def plot_trajectory_references(self, tp, fp, vp, ap): acceleration_points = [] i = 0 for foot_swing_time in np.arange(0.000001, swing_period, 0.002): - desired_foot_position, \ - desired_foot_velocity, \ - desired_foot_acceleration = trajectory_generator.compute_trajectory_references(swing_time=i, - lift_off=lift_off, - touch_down=touch_down) + desired_foot_position, desired_foot_velocity, desired_foot_acceleration = ( + trajectory_generator.compute_trajectory_references(swing_time=i, lift_off=lift_off, touch_down=touch_down) + ) i += simulation_dt - + time_points.append(i) position_points.append(desired_foot_position.squeeze()) velocity_points.append(desired_foot_velocity.squeeze()) acceleration_points.append(desired_foot_acceleration.squeeze()) - + # Plot the generated trajectory trajectory_generator.plot_trajectory_3d(np.array(position_points)) trajectory_generator.plot_trajectory_references(time_points, position_points, velocity_points, acceleration_points) - \ No newline at end of file diff --git a/quadruped_pympc/helpers/swing_trajectory_controller.py b/quadruped_pympc/helpers/swing_trajectory_controller.py index 6902c69..083a968 100644 --- a/quadruped_pympc/helpers/swing_trajectory_controller.py +++ b/quadruped_pympc/helpers/swing_trajectory_controller.py @@ -1,24 +1,28 @@ -import os import numpy as np -class SwingTrajectoryController: - def __init__(self, - step_height: float, - swing_period: float, - position_gain_fb: np.ndarray, - velocity_gain_fb: np.ndarray, - generator: str) -> None: +class SwingTrajectoryController: + def __init__( + self, + step_height: float, + swing_period: float, + position_gain_fb: np.ndarray, + velocity_gain_fb: np.ndarray, + generator: str, + ) -> None: self.generator = generator - if (self.generator == "ndcurves"): + if self.generator == "ndcurves": from .swing_generators.ndcurves_swing_trajectory_generator import SwingTrajectoryGenerator + self.swing_generator = SwingTrajectoryGenerator(swing_period=swing_period, step_height=step_height) - elif (self.generator == "scipy"): + elif self.generator == "scipy": from .swing_generators.scipy_swing_trajectory_generator import SwingTrajectoryGenerator + self.swing_generator = SwingTrajectoryGenerator(swing_period=swing_period, step_height=step_height) else: from .swing_generators.explicit_swing_trajectory_generator import SwingTrajectoryGenerator + self.swing_generator = SwingTrajectoryGenerator(swing_period=swing_period, step_height=step_height) self.position_gain_fb = position_gain_fb @@ -30,28 +34,23 @@ def __init__(self, self.use_gravity_compensation_only = False def regenerate_swing_trajectory_generator(self, step_height: float, swing_period: float) -> None: - if (self.generator == "ndcurves"): + if self.generator == "ndcurves": from .swing_generators.ndcurves_swing_trajectory_generator import SwingTrajectoryGenerator + self.swing_generator = SwingTrajectoryGenerator(swing_period=swing_period, step_height=step_height) - elif (self.generator == "scipy"): + elif self.generator == "scipy": from .swing_generators.scipy_swing_trajectory_generator import SwingTrajectoryGenerator + self.swing_generator = SwingTrajectoryGenerator(swing_period=swing_period, step_height=step_height) else: from .swing_generators.explicit_swing_trajectory_generator import SwingTrajectoryGenerator + self.swing_generator = SwingTrajectoryGenerator(swing_period=swing_period, step_height=step_height) self.swing_period = swing_period - def compute_swing_control_cartesian_space(self, - leg_id, - q_dot, - J, - J_dot, - lift_off, - touch_down, - foot_pos, - foot_vel, - h, - mass_matrix): + def compute_swing_control_cartesian_space( + self, leg_id, q_dot, J, J_dot, lift_off, touch_down, foot_pos, foot_vel, h, mass_matrix + ): """TODO: Docstring. Args: @@ -75,10 +74,8 @@ def compute_swing_control_cartesian_space(self, """ # Compute trajectory references des_foot_pos, des_foot_vel, des_foot_acc = self.swing_generator.compute_trajectory_references( - self.swing_time[leg_id], - lift_off, - touch_down - ) + self.swing_time[leg_id], lift_off, touch_down + ) err_pos = des_foot_pos - foot_pos err_pos = err_pos.reshape((3,)) @@ -86,48 +83,42 @@ def compute_swing_control_cartesian_space(self, err_vel = des_foot_vel - foot_vel err_vel = err_vel.reshape((3,)) - accelleration = des_foot_acc + \ - self.position_gain_fb * (err_pos) + \ - self.velocity_gain_fb * (err_vel) + accelleration = des_foot_acc + self.position_gain_fb * (err_pos) + self.velocity_gain_fb * (err_vel) accelleration = accelleration.reshape((3,)) # Compute inertia matrix in task space. # Mass Matrix and centrifugal missing tau_swing = J.T @ (self.position_gain_fb * (err_pos) + self.velocity_gain_fb * (err_vel)) - if(self.use_feedback_linearization): + if self.use_feedback_linearization: tau_swing += mass_matrix @ np.linalg.pinv(J) @ (accelleration - J_dot @ q_dot) + h return tau_swing, des_foot_pos, des_foot_vel - - - def compute_swing_control_joint_space(self, - nmpc_joints_pos, - nmpc_joints_vel, - nmpc_joints_acc, - qpos, qvel, - legs_mass_matrix, - legs_qfrc_bias): + def compute_swing_control_joint_space( + self, nmpc_joints_pos, nmpc_joints_vel, nmpc_joints_acc, qpos, qvel, legs_mass_matrix, legs_qfrc_bias + ): error_position = nmpc_joints_pos - qpos - error_position = error_position.reshape((3, )) + error_position = error_position.reshape((3,)) error_velocity = nmpc_joints_vel - qvel - error_velocity = error_velocity.reshape((3, )) + error_velocity = error_velocity.reshape((3,)) accelleration = nmpc_joints_acc accelleration = accelleration.reshape((3,)) - - tau_swing = self.position_gain_fb*error_position + self.velocity_gain_fb*error_velocity + + tau_swing = self.position_gain_fb * error_position + self.velocity_gain_fb * error_velocity # Feedback linearization - if(self.use_feedback_linearization): - tau_swing += legs_mass_matrix@(accelleration + self.position_gain_fb*error_position + self.velocity_gain_fb*error_velocity) + legs_qfrc_bias - elif(self.use_gravity_compensation_only): + if self.use_feedback_linearization: + tau_swing += ( + legs_mass_matrix + @ (accelleration + self.position_gain_fb * error_position + self.velocity_gain_fb * error_velocity) + + legs_qfrc_bias + ) + elif self.use_gravity_compensation_only: tau_swing += legs_qfrc_bias - - return tau_swing, None, None - + return tau_swing, None, None def update_swing_time(self, current_contact, legs_order, dt): for leg_id, leg_name in enumerate(legs_order): @@ -138,25 +129,22 @@ def update_swing_time(self, current_contact, legs_order, dt): else: self.swing_time[leg_id] = 0 - - def check_apex_condition(self, current_contact, interval=0.02): optimize_swing = 0 for leg_id in range(4): # Swing time check - if (current_contact[leg_id] == 0): - if ((self.swing_time[leg_id] > (self.swing_period / 2.) - interval) and \ - (self.swing_time[leg_id] < (self.swing_period / 2.) + interval)): + if current_contact[leg_id] == 0: + if (self.swing_time[leg_id] > (self.swing_period / 2.0) - interval) and ( + self.swing_time[leg_id] < (self.swing_period / 2.0) + interval + ): optimize_swing = 1 return optimize_swing - - def check_full_stance_condition(self, current_contact): stance = 1 # If one leg is not in stance, the robot is not in full stance for leg_id in range(4): - if (current_contact[leg_id] == 0): + if current_contact[leg_id] == 0: stance = 0 return stance @@ -164,11 +152,12 @@ def check_full_stance_condition(self, current_contact): # Example: if __name__ == "__main__": import time + import numpy as np - from tqdm import tqdm # Gym and Simulation related imports from gym_quadruped.quadruped_env import QuadrupedEnv + from gym_quadruped.utils.mujoco.visual import render_sphere, render_vector from gym_quadruped.utils.quadruped_utils import LegsAttr # Config imports @@ -176,9 +165,6 @@ def check_full_stance_condition(self, current_contact): # Helper functions for plotting from quadruped_pympc.helpers.quadruped_utils import plot_swing_mujoco - from gym_quadruped.utils.mujoco.visual import render_vector - from gym_quadruped.utils.mujoco.visual import render_sphere - np.set_printoptions(precision=3, suppress=True) @@ -186,29 +172,39 @@ def check_full_stance_condition(self, current_contact): hip_height = cfg.hip_height robot_leg_joints = cfg.robot_leg_joints robot_feet_geom_names = cfg.robot_feet_geom_names - scene_name = cfg.simulation_params['scene'] - simulation_dt = cfg.simulation_params['dt'] - - state_observables_names = ('base_pos', 'base_lin_vel', 'base_ori_euler_xyz', 'base_ori_quat_wxyz', 'base_ang_vel', - 'qpos_js', 'qvel_js', 'tau_ctrl_setpoint', - 'feet_pos_base', 'feet_vel_base', 'contact_state', 'contact_forces_base',) - + scene_name = cfg.simulation_params["scene"] + simulation_dt = cfg.simulation_params["dt"] + + state_observables_names = ( + "base_pos", + "base_lin_vel", + "base_ori_euler_xyz", + "base_ori_quat_wxyz", + "base_ang_vel", + "qpos_js", + "qvel_js", + "tau_ctrl_setpoint", + "feet_pos_base", + "feet_vel_base", + "contact_state", + "contact_forces_base", + ) # Create the quadruped robot environment ----------------------------------------------------------- - env = QuadrupedEnv(robot=robot_name, - hip_height=hip_height, - legs_joint_names=robot_leg_joints, # Joint names of the legs DoF - feet_geom_name=robot_feet_geom_names, # Geom/Frame id of feet - scene=scene_name, - sim_dt=simulation_dt, - ref_base_lin_vel=0.0, # pass a float for a fixed value - ground_friction_coeff=1.5, # pass a float for a fixed value - base_vel_command_type="human", # "forward", "random", "forward+rotate", "human" - state_obs_names=state_observables_names, # Desired quantities in the 'state' vec - ) + env = QuadrupedEnv( + robot=robot_name, + hip_height=hip_height, + legs_joint_names=robot_leg_joints, # Joint names of the legs DoF + feet_geom_name=robot_feet_geom_names, # Geom/Frame id of feet + scene=scene_name, + sim_dt=simulation_dt, + ref_base_lin_vel=0.0, # pass a float for a fixed value + ground_friction_coeff=1.5, # pass a float for a fixed value + base_vel_command_type="human", # "forward", "random", "forward+rotate", "human" + state_obs_names=state_observables_names, # Desired quantities in the 'state' vec + ) env.reset(random=False) env.render() # Pass in the first render call any mujoco.viewer.KeyCallbackType - feet_traj_geom_ids, feet_GRF_geom_ids = None, LegsAttr(FL=-1, FR=-1, RL=-1, RR=-1) legs_order = ["FL", "FR", "RL", "RR"] @@ -220,67 +216,58 @@ def check_full_stance_condition(self, current_contact): # Torque vector tau = LegsAttr(*[np.zeros((env.mjModel.nv, 1)) for _ in range(4)]) - - # Quadruped PyMPC controller initialization ------------------------------------------------------------- - from quadruped_pympc.interfaces.srbd_controller_interface import SRBDControllerInterface - from quadruped_pympc.interfaces.srbd_batched_controller_interface import SRBDBatchedControllerInterface from quadruped_pympc.interfaces.wb_interface import WBInterface + wb_interface = WBInterface(initial_feet_pos=env.feet_pos(frame="world"), legs_order=legs_order) - wb_interface = WBInterface(initial_feet_pos = env.feet_pos(frame='world'), legs_order = legs_order) - - - nmpc_GRFs = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), - RL=np.zeros(3), RR=np.zeros(3)) - nmpc_footholds = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), - RL=np.zeros(3), RR=np.zeros(3)) + nmpc_GRFs = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), RL=np.zeros(3), RR=np.zeros(3)) + nmpc_footholds = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), RL=np.zeros(3), RR=np.zeros(3)) best_sample_freq = wb_interface.pgg.step_freq import copy - initial_pos=copy.deepcopy(np.array([0,0,0.6])) - initial_angle=copy.deepcopy(env.mjData.qpos[3:7]) + + initial_pos = copy.deepcopy(np.array([0, 0, 0.6])) + initial_angle = copy.deepcopy(env.mjData.qpos[3:7]) while True: env.mjData.qpos[0:3] = initial_pos env.mjData.qpos[3:7] = initial_angle env.mjData.qvel[0:3] = 0 env.mjData.qvel[3:7] = 0 - # Update value from SE or Simulator ---------------------- - feet_pos = env.feet_pos(frame='world') - hip_pos = env.hip_positions(frame='world') - base_lin_vel = env.base_lin_vel(frame='world') - base_ang_vel = env.base_ang_vel(frame='world') + feet_pos = env.feet_pos(frame="world") + hip_pos = env.hip_positions(frame="world") + base_lin_vel = env.base_lin_vel(frame="world") + base_ang_vel = env.base_ang_vel(frame="world") base_ori_euler_xyz = env.base_ori_euler_xyz base_pos = env.base_pos # Get the reference base velocity in the world frame ref_base_lin_vel, ref_base_ang_vel = env.target_base_vel() - + # Get the inertia matrix - if(cfg.simulation_params['use_inertia_recomputation']): + if cfg.simulation_params["use_inertia_recomputation"]: inertia = env.get_base_inertia().flatten() # Reflected inertia of base at qpos, in world frame else: inertia = cfg.inertia.flatten() # Get the qpos and qvel qpos, qvel = env.mjData.qpos, env.mjData.qvel - joints_pos = LegsAttr(FL=qpos[7:10], FR=qpos[10:13], - RL=qpos[13:16], RR=qpos[16:19]) - + joints_pos = LegsAttr(FL=qpos[7:10], FR=qpos[10:13], RL=qpos[13:16], RR=qpos[16:19]) + # Get Centrifugal, Coriolis, Gravity for the swing controller legs_mass_matrix = env.legs_mass_matrix legs_qfrc_bias = env.legs_qfrc_bias # Compute feet jacobian - feet_jac = env.feet_jacobians(frame='world', return_rot_jac=False) - + feet_jac = env.feet_jacobians(frame="world", return_rot_jac=False) + # Compute jacobian derivatives of the contact points jac_feet_dot = (feet_jac - jac_feet_prev) / simulation_dt # Finite difference approximation jac_feet_prev = feet_jac # Update previous Jacobians - + # Compute feet velocities feet_vel = LegsAttr(**{leg_name: feet_jac[leg_name] @ env.mjData.qvel for leg_name in legs_order}) @@ -288,31 +275,31 @@ def check_full_stance_condition(self, current_contact): legs_qvel_idx = env.legs_qvel_idx legs_qpos_idx = env.legs_qpos_idx - - # Update the state and reference ------------------------- - state_current, \ - ref_state, \ - contact_sequence, \ - ref_feet_pos, \ - ref_feet_constraints, \ - contact_sequence_dts, \ - contact_sequence_lenghts, \ - step_height, \ - optimize_swing = wb_interface.update_state_and_reference(base_pos, - base_lin_vel, - base_ori_euler_xyz, - base_ang_vel, - feet_pos, - hip_pos, - joints_pos, - heightmaps, - legs_order, - simulation_dt, - ref_base_lin_vel, - ref_base_ang_vel) - - + ( + state_current, + ref_state, + contact_sequence, + ref_feet_pos, + ref_feet_constraints, + contact_sequence_dts, + contact_sequence_lenghts, + step_height, + optimize_swing, + ) = wb_interface.update_state_and_reference( + base_pos, + base_lin_vel, + base_ori_euler_xyz, + base_ang_vel, + feet_pos, + hip_pos, + joints_pos, + heightmaps, + legs_order, + simulation_dt, + ref_base_lin_vel, + ref_base_ang_vel, + ) ref_feet_pos.FL[2] = 0.3 ref_feet_pos.FR[2] = 0.3 @@ -322,39 +309,35 @@ def check_full_stance_condition(self, current_contact): wb_interface.frg.lift_off_positions.FR[2] = 0.3 wb_interface.frg.lift_off_positions.RL[2] = 0.3 wb_interface.frg.lift_off_positions.RR[2] = 0.3 - nmpc_footholds = LegsAttr(FL=ref_feet_pos['FL'], FR=ref_feet_pos['FR'], - RL=ref_feet_pos['RL'], RR=ref_feet_pos['RR']) - - - - - - + nmpc_footholds = LegsAttr( + FL=ref_feet_pos["FL"], FR=ref_feet_pos["FR"], RL=ref_feet_pos["RL"], RR=ref_feet_pos["RR"] + ) # Compute Swing and Stance Torque --------------------------------------------------------------------------- nmpc_joints_pos = None nmpc_joints_vel = None nmpc_joints_acc = None - tau = wb_interface.compute_stance_and_swing_torque(simulation_dt, - qpos, - qvel, - feet_jac, - jac_feet_dot, - feet_pos, - feet_vel, - legs_qfrc_bias, - legs_mass_matrix, - nmpc_GRFs, - nmpc_footholds, - legs_qpos_idx, - legs_qvel_idx, - tau, - optimize_swing, - best_sample_freq, - nmpc_joints_pos, - nmpc_joints_vel, - nmpc_joints_acc) - + tau = wb_interface.compute_stance_and_swing_torque( + simulation_dt, + qpos, + qvel, + feet_jac, + jac_feet_dot, + feet_pos, + feet_vel, + legs_qfrc_bias, + legs_mass_matrix, + nmpc_GRFs, + nmpc_footholds, + legs_qpos_idx, + legs_qvel_idx, + tau, + optimize_swing, + best_sample_freq, + nmpc_joints_pos, + nmpc_joints_vel, + nmpc_joints_acc, + ) action = np.zeros(env.mjModel.nu) action[env.legs_tau_idx.FL] = tau.FL @@ -362,49 +345,54 @@ def check_full_stance_condition(self, current_contact): action[env.legs_tau_idx.RL] = tau.RL action[env.legs_tau_idx.RR] = tau.RR state, reward, is_terminated, is_truncated, info = env.step(action=action) - - - _, _, feet_GRF = env.feet_contact_state(ground_reaction_forces=True) # Plot the swing trajectory - feet_traj_geom_ids = plot_swing_mujoco(viewer=env.viewer, - swing_traj_controller=wb_interface.stc, - swing_period=wb_interface.stc.swing_period, - swing_time=LegsAttr(FL=wb_interface.stc.swing_time[0], - FR=wb_interface.stc.swing_time[1], - RL=wb_interface.stc.swing_time[2], - RR=wb_interface.stc.swing_time[3]), - lift_off_positions=wb_interface.frg.lift_off_positions, - nmpc_footholds=nmpc_footholds, - ref_feet_pos=ref_feet_pos, - geom_ids=feet_traj_geom_ids) - - + feet_traj_geom_ids = plot_swing_mujoco( + viewer=env.viewer, + swing_traj_controller=wb_interface.stc, + swing_period=wb_interface.stc.swing_period, + swing_time=LegsAttr( + FL=wb_interface.stc.swing_time[0], + FR=wb_interface.stc.swing_time[1], + RL=wb_interface.stc.swing_time[2], + RR=wb_interface.stc.swing_time[3], + ), + lift_off_positions=wb_interface.frg.lift_off_positions, + nmpc_footholds=nmpc_footholds, + ref_feet_pos=ref_feet_pos, + geom_ids=feet_traj_geom_ids, + ) + # Update and Plot the heightmap - if(cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): - #if(stc.check_apex_condition(current_contact, interval=0.01)): + if cfg.simulation_params["visual_foothold_adaptation"] != "blind": + # if(stc.check_apex_condition(current_contact, interval=0.01)): for leg_id, leg_name in enumerate(legs_order): - data = heightmaps[leg_name].data#.update_height_map(ref_feet_pos[leg_name], yaw=env.base_ori_euler_xyz[2]) - if(data is not None): + data = heightmaps[ + leg_name + ].data # .update_height_map(ref_feet_pos[leg_name], yaw=env.base_ori_euler_xyz[2]) + if data is not None: for i in range(data.shape[0]): for j in range(data.shape[1]): - heightmaps[leg_name].geom_ids[i, j] = render_sphere(viewer=env.viewer, - position=([data[i][j][0][0],data[i][j][0][1],data[i][j][0][2]]), - diameter=0.01, - color=[0, 1, 0, .5], - geom_id=heightmaps[leg_name].geom_ids[i, j] - ) - + heightmaps[leg_name].geom_ids[i, j] = render_sphere( + viewer=env.viewer, + position=([data[i][j][0][0], data[i][j][0][1], data[i][j][0][2]]), + diameter=0.01, + color=[0, 1, 0, 0.5], + geom_id=heightmaps[leg_name].geom_ids[i, j], + ) + # Plot the GRF for leg_id, leg_name in enumerate(legs_order): - feet_GRF_geom_ids[leg_name] = render_vector(env.viewer, - vector=feet_GRF[leg_name], - pos=feet_pos[leg_name], - scale=np.linalg.norm(feet_GRF[leg_name]) * 0.005, - color=np.array([0, 1, 0, .5]), - geom_id=feet_GRF_geom_ids[leg_name]) + feet_GRF_geom_ids[leg_name] = render_vector( + env.viewer, + vector=feet_GRF[leg_name], + pos=feet_pos[leg_name], + scale=np.linalg.norm(feet_GRF[leg_name]) * 0.005, + color=np.array([0, 1, 0, 0.5]), + geom_id=feet_GRF_geom_ids[leg_name], + ) env.render() - last_render_time = time.time() \ No newline at end of file + last_render_time = time.time() diff --git a/quadruped_pympc/helpers/terrain_estimator.py b/quadruped_pympc/helpers/terrain_estimator.py index 6073b24..403c7ce 100644 --- a/quadruped_pympc/helpers/terrain_estimator.py +++ b/quadruped_pympc/helpers/terrain_estimator.py @@ -2,17 +2,16 @@ class TerrainEstimator: - def __init__(self, ) -> None: - + def __init__( + self, + ) -> None: self.terrain_roll = 0 self.terrain_pitch = 0 self.terrain_height = 0 - def compute_terrain_estimation(self, - base_position: np.ndarray, - yaw: float, - feet_pos: dict, - current_contact: np.ndarray) -> [float, float]: + def compute_terrain_estimation( + self, base_position: np.ndarray, yaw: float, feet_pos: dict, current_contact: np.ndarray + ) -> [float, float]: """Compute the estimated roll and pitch of the terrain based on the positions of the robot's feet. Parameters @@ -36,17 +35,13 @@ def compute_terrain_estimation(self, """ # Compute roll and pitch for each foot position # Rotation matrix R_yaw - R_W2H = np.array([ - [np.cos(yaw), np.sin(yaw), 0], - [-np.sin(yaw), np.cos(yaw), 0], - [0, 0, 1] - ]) + R_W2H = np.array([[np.cos(yaw), np.sin(yaw), 0], [-np.sin(yaw), np.cos(yaw), 0], [0, 0, 1]]) # Extracting 3-element segments from liftoff_position_z_ and x_op_ - seg0 = feet_pos['FL'] - seg3 = feet_pos['FR'] - seg6 = feet_pos['RL'] - seg9 = feet_pos['RR'] + seg0 = feet_pos["FL"] + seg3 = feet_pos["FR"] + seg6 = feet_pos["RL"] + seg9 = feet_pos["RR"] # Calculating differences # TODO: Feet position in base frame? @@ -57,11 +52,15 @@ def compute_terrain_estimation(self, # Calculating pitch and roll # TODO: Docstring - pitch = (np.arctan(np.abs(left_difference[2]) / np.abs(left_difference[0] + 0.001)) + - np.arctan(np.abs(right_difference[2]) / np.abs(right_difference[0] + 0.001))) * 0.5 + pitch = ( + np.arctan(np.abs(left_difference[2]) / np.abs(left_difference[0] + 0.001)) + + np.arctan(np.abs(right_difference[2]) / np.abs(right_difference[0] + 0.001)) + ) * 0.5 - roll = (np.arctan(np.abs(front_difference[2]) / np.abs(front_difference[1] + 0.001)) + - np.arctan(np.abs(back_difference[2]) / np.abs(back_difference[1] + 0.001))) * 0.5 + roll = ( + np.arctan(np.abs(front_difference[2]) / np.abs(front_difference[1] + 0.001)) + + np.arctan(np.abs(back_difference[2]) / np.abs(back_difference[1] + 0.001)) + ) * 0.5 # Adjusting signs of pitch and roll TODO: Adjusting what and for what? if (front_difference[2] * 0.5 + back_difference[2] * 0.5) < 0: @@ -73,10 +72,10 @@ def compute_terrain_estimation(self, self.terrain_pitch = self.terrain_pitch * 0.8 + pitch * 0.2 # Update the reference height given the foot in contact - z_foot_FL = feet_pos['FL'][2] - z_foot_FR = feet_pos['FR'][2] - z_foot_RL = feet_pos['RL'][2] - z_foot_RR = feet_pos['RR'][2] + z_foot_FL = feet_pos["FL"][2] + z_foot_FR = feet_pos["FR"][2] + z_foot_RL = feet_pos["RL"][2] + z_foot_RR = feet_pos["RR"][2] """number_foot_in_contact = current_contact[0] + \ current_contact[1] + \ current_contact[2] + \ @@ -87,18 +86,14 @@ def compute_terrain_estimation(self, z_foot_RL * current_contact[2] + \ z_foot_RR * current_contact[3]) / number_foot_in_contact self.terrain_height = self.terrain_height * 0.6 + z_foot_mean_temp * 0.4""" - - z_foot_mean_temp = (z_foot_FL + \ - z_foot_FR + \ - z_foot_RL + \ - z_foot_RR) / 4 - self.terrain_height = self.terrain_height * 0.6 + z_foot_mean_temp * 0.4 + z_foot_mean_temp = (z_foot_FL + z_foot_FR + z_foot_RL + z_foot_RR) / 4 + self.terrain_height = self.terrain_height * 0.6 + z_foot_mean_temp * 0.4 - #print("current_contact: ", current_contact) - #print("z_foot_FL: ", z_foot_FL) - #print("z_foot_FR: ", z_foot_FR) - #print("z_foot_RL: ", z_foot_RL) - #print("z_foot_RR: ", z_foot_RR) + # print("current_contact: ", current_contact) + # print("z_foot_FL: ", z_foot_FL) + # print("z_foot_FR: ", z_foot_FR) + # print("z_foot_RL: ", z_foot_RL) + # print("z_foot_RR: ", z_foot_RR) return self.terrain_roll, self.terrain_pitch, self.terrain_height diff --git a/quadruped_pympc/helpers/velocity_modulator.py b/quadruped_pympc/helpers/velocity_modulator.py index a4a780d..b2c1599 100644 --- a/quadruped_pympc/helpers/velocity_modulator.py +++ b/quadruped_pympc/helpers/velocity_modulator.py @@ -1,28 +1,39 @@ import numpy as np + from quadruped_pympc import config as cfg + class VelocityModulator: def __init__(self): self.activated = False - if(cfg.robot=="aliengo"): + if cfg.robot == "aliengo": self.max_distance = 0.15 - elif(cfg.robot=="go1" or cfg.robot=="go2"): + elif cfg.robot == "go1" or cfg.robot == "go2": self.max_distance = 0.12 else: self.max_distance = 0.2 def modulate_velocities(self, ref_base_lin_vel, ref_base_ang_vel, feet_pos, hip_pos): + distance_FL_to_hip_xy = np.sqrt( + np.square(feet_pos.FL[0] - hip_pos.FL[0]) + np.square(feet_pos.FL[1] - hip_pos.FL[1]) + ) + distance_FR_to_hip_xy = np.sqrt( + np.square(feet_pos.FR[0] - hip_pos.FR[0]) + np.square(feet_pos.FR[1] - hip_pos.FR[1]) + ) + distance_RL_to_hip_xy = np.sqrt( + np.square(feet_pos.RL[0] - hip_pos.RL[0]) + np.square(feet_pos.RL[1] - hip_pos.RL[1]) + ) + distance_RR_to_hip_xy = np.sqrt( + np.square(feet_pos.RR[0] - hip_pos.RR[0]) + np.square(feet_pos.RR[1] - hip_pos.RR[1]) + ) - distance_FL_to_hip_xy = np.sqrt(np.square(feet_pos.FL[0] - hip_pos.FL[0]) + np.square(feet_pos.FL[1] - hip_pos.FL[1])) - distance_FR_to_hip_xy = np.sqrt(np.square(feet_pos.FR[0] - hip_pos.FR[0]) + np.square(feet_pos.FR[1] - hip_pos.FR[1])) - distance_RL_to_hip_xy = np.sqrt(np.square(feet_pos.RL[0] - hip_pos.RL[0]) + np.square(feet_pos.RL[1] - hip_pos.RL[1])) - distance_RR_to_hip_xy = np.sqrt(np.square(feet_pos.RR[0] - hip_pos.RR[0]) + np.square(feet_pos.RR[1] - hip_pos.RR[1])) - + if ( + distance_FL_to_hip_xy > self.max_distance + or distance_FR_to_hip_xy > self.max_distance + or distance_RL_to_hip_xy > self.max_distance + or distance_RR_to_hip_xy > self.max_distance + ): + ref_base_lin_vel = ref_base_lin_vel * 0.0 + ref_base_ang_vel = ref_base_ang_vel * 0.0 - if(distance_FL_to_hip_xy > self.max_distance or distance_FR_to_hip_xy > self.max_distance or - distance_RL_to_hip_xy > self.max_distance or distance_RR_to_hip_xy > self.max_distance): - ref_base_lin_vel = ref_base_lin_vel*0.0 - ref_base_ang_vel = ref_base_ang_vel*0.0 - return ref_base_lin_vel, ref_base_ang_vel - diff --git a/quadruped_pympc/helpers/visual_foothold_adaptation.py b/quadruped_pympc/helpers/visual_foothold_adaptation.py index ec1b356..d4382d8 100644 --- a/quadruped_pympc/helpers/visual_foothold_adaptation.py +++ b/quadruped_pympc/helpers/visual_foothold_adaptation.py @@ -6,118 +6,111 @@ except ImportError: print("VFA not installed, not open source yet") -class VisualFootholdAdaptation: - def __init__(self, legs_order, adaptation_strategy='height'): - self.footholds_adaptation = LegsAttr(FL=np.array([0, 0, 0]), - FR=np.array([0, 0, 0]), - RL=np.array([0, 0, 0]), - RR=np.array([0, 0, 0])) +class VisualFootholdAdaptation: + def __init__(self, legs_order, adaptation_strategy="height"): + self.footholds_adaptation = LegsAttr( + FL=np.array([0, 0, 0]), FR=np.array([0, 0, 0]), RL=np.array([0, 0, 0]), RR=np.array([0, 0, 0]) + ) self.footholds_constraints = LegsAttr(FL=None, FR=None, RL=None, RR=None) self.initialized = False self.adaptation_strategy = adaptation_strategy - if(self.adaptation_strategy == 'vfa'): + if self.adaptation_strategy == "vfa": self.vfa_evaluators = LegsAttr(FL=None, FR=None, RL=None, RR=None) for leg_id, leg_name in enumerate(legs_order): self.vfa_evaluators[leg_name] = VFA(leg=leg_name) - def update_footholds_adaptation(self, update_footholds_adaptation): self.footholds_adaptation = update_footholds_adaptation self.initialized = True - def reset(self): self.initialized = False - def get_footholds_adapted(self, reference_footholds): # If the adaptation is not initialized, return the reference footholds - if(self.initialized == False): + if self.initialized == False: self.footholds_adaptation = reference_footholds return reference_footholds, self.footholds_constraints else: return self.footholds_adaptation, self.footholds_constraints - - def get_heightmap_coordinates_foothold_id(self, heightmaps, foothold_id, leg_name): + r = round(foothold_id.item() / heightmaps[leg_name].n) + c = round(foothold_id.item() % heightmaps[leg_name].n) - r = round(foothold_id.item()/heightmaps[leg_name].n) - c = round(foothold_id.item()%heightmaps[leg_name].n) + if r >= heightmaps[leg_name].n: + r = heightmaps[leg_name].n - 1 + if c >= heightmaps[leg_name].n: + c = heightmaps[leg_name].n - 1 - if(r >= heightmaps[leg_name].n): - r = heightmaps[leg_name].n -1 - if(c >= heightmaps[leg_name].n): - c = heightmaps[leg_name].n -1 - return r, c - - def compute_adaptation(self, legs_order, reference_footholds, hip_positions, heightmaps, - forward_vel, base_orientation, base_orientation_rate): - + def compute_adaptation( + self, + legs_order, + reference_footholds, + hip_positions, + heightmaps, + forward_vel, + base_orientation, + base_orientation_rate, + ): for leg_id, leg_name in enumerate(legs_order): - if(heightmaps[leg_name].data is None): + if heightmaps[leg_name].data is None: return False - - - if(self.adaptation_strategy == 'height'): + + if self.adaptation_strategy == "height": for leg_id, leg_name in enumerate(legs_order): height_adjustment = heightmaps[leg_name].get_height(reference_footholds[leg_name]) - if(height_adjustment is not None): + if height_adjustment is not None: reference_footholds[leg_name][2] = height_adjustment - - elif(self.adaptation_strategy == 'vfa'): - gait_phases = 0.0 #for now + elif self.adaptation_strategy == "vfa": + gait_phases = 0.0 # for now for leg_id, leg_name in enumerate(legs_order): # Transform the heightmap in hip frame - + heightmap = heightmaps[leg_name].data hip_position = hip_positions[leg_name] - - heightmap_hip_frame = heightmap[:,:,0,2] - hip_position[2] - - - convex_data, \ - _, \ - safe_map, \ - info = self.vfa_evaluators[leg_name](heightmap_data=heightmap_hip_frame, - base_linear_velocity=forward_vel, - base_orientation=base_orientation, - base_orientation_rate=base_orientation_rate, - gait_phase=gait_phases, - return_info=True - ) - + + heightmap_hip_frame = heightmap[:, :, 0, 2] - hip_position[2] + + convex_data, _, safe_map, info = self.vfa_evaluators[leg_name]( + heightmap_data=heightmap_hip_frame, + base_linear_velocity=forward_vel, + base_orientation=base_orientation, + base_orientation_rate=base_orientation_rate, + gait_phase=gait_phases, + return_info=True, + ) + best_foothold_id = convex_data[0][0] best_convex_area_vertices_id = [convex_data[1][0].x1, convex_data[1][0].y2] - + r, c = self.get_heightmap_coordinates_foothold_id(heightmaps, best_foothold_id, leg_name) - - reference_footholds[leg_name][0:2] = heightmap[r,c,0,:][0:2] - + reference_footholds[leg_name][0:2] = heightmap[r, c, 0, :][0:2] + height_adjustment = heightmaps[leg_name].get_height(reference_footholds[leg_name]) - if(height_adjustment is not None): + if height_adjustment is not None: reference_footholds[leg_name][2] = height_adjustment - - - r_vertex1, c_vertex1 = self.get_heightmap_coordinates_foothold_id(heightmaps, np.array([best_convex_area_vertices_id[0]]), leg_name) - r_vertex2, c_vertex2 = self.get_heightmap_coordinates_foothold_id(heightmaps, np.array([best_convex_area_vertices_id[1]]), leg_name) - vertex1_world_frame = heightmap[r_vertex1,c_vertex1,0,:] - vertex2_world_frame = heightmap[r_vertex2,c_vertex2,0,:] + r_vertex1, c_vertex1 = self.get_heightmap_coordinates_foothold_id( + heightmaps, np.array([best_convex_area_vertices_id[0]]), leg_name + ) + r_vertex2, c_vertex2 = self.get_heightmap_coordinates_foothold_id( + heightmaps, np.array([best_convex_area_vertices_id[1]]), leg_name + ) + vertex1_world_frame = heightmap[r_vertex1, c_vertex1, 0, :] + vertex2_world_frame = heightmap[r_vertex2, c_vertex2, 0, :] self.footholds_constraints[leg_name] = [vertex1_world_frame, vertex2_world_frame] + # print("Safe map: ", safe_map) - #print("Safe map: ", safe_map) - - self.update_footholds_adaptation(reference_footholds) - - return True \ No newline at end of file + + return True diff --git a/quadruped_pympc/interfaces/srbd_batched_controller_interface.py b/quadruped_pympc/interfaces/srbd_batched_controller_interface.py index 0a30cf4..fe12bb2 100644 --- a/quadruped_pympc/interfaces/srbd_batched_controller_interface.py +++ b/quadruped_pympc/interfaces/srbd_batched_controller_interface.py @@ -1,53 +1,49 @@ import numpy as np +from quadruped_pympc import config as cfg from quadruped_pympc.helpers.periodic_gait_generator import PeriodicGaitGenerator -from quadruped_pympc import config as cfg class SRBDBatchedControllerInterface: """This is an interface for a batched controller that uses the SRBD method to optimize the gait""" + def __init__( + self, + ): + """Constructor for the SRBD batched controller interface""" - def __init__(self, ): - """ Constructor for the SRBD batched controller interface """ - - self.type = cfg.mpc_params['type'] - self.mpc_dt = cfg.mpc_params['dt'] - self.horizon = cfg.mpc_params['horizon'] - self.optimize_step_freq = cfg.mpc_params['optimize_step_freq'] - self.step_freq_available = cfg.mpc_params['step_freq_available'] + self.type = cfg.mpc_params["type"] + self.mpc_dt = cfg.mpc_params["dt"] + self.horizon = cfg.mpc_params["horizon"] + self.optimize_step_freq = cfg.mpc_params["optimize_step_freq"] + self.step_freq_available = cfg.mpc_params["step_freq_available"] + from quadruped_pympc.controllers.gradient.nominal.centroidal_nmpc_gait_adaptive import Acados_NMPC_GaitAdaptive - from quadruped_pympc.controllers.gradient.nominal.centroidal_nmpc_gait_adaptive import \ - Acados_NMPC_GaitAdaptive - self.batched_controller = Acados_NMPC_GaitAdaptive() - # in the case of nonuniform discretization, we create a list of dts and horizons for each nonuniform discretization - if (cfg.mpc_params['use_nonuniform_discretization']): - self.contact_sequence_dts = [cfg.mpc_params['dt_fine_grained'], self.mpc_dt] - self.contact_sequence_lenghts = [cfg.mpc_params['horizon_fine_grained'], self.horizon] + if cfg.mpc_params["use_nonuniform_discretization"]: + self.contact_sequence_dts = [cfg.mpc_params["dt_fine_grained"], self.mpc_dt] + self.contact_sequence_lenghts = [cfg.mpc_params["horizon_fine_grained"], self.horizon] else: self.contact_sequence_dts = [self.mpc_dt] self.contact_sequence_lenghts = [self.horizon] - - - - - def optimize_gait(self, - state_current: dict, - ref_state: dict, - inertia: np.ndarray, - pgg_phase_signal: np.ndarray, - pgg_step_freq: float, - pgg_duty_factor: float, - pgg_gait_type: int, - optimize_swing: int) -> float: + def optimize_gait( + self, + state_current: dict, + ref_state: dict, + inertia: np.ndarray, + pgg_phase_signal: np.ndarray, + pgg_step_freq: float, + pgg_duty_factor: float, + pgg_gait_type: int, + optimize_swing: int, + ) -> float: """Optimize the gait using the SRBD method TODO: remove the unused arguments, and not pass pgg but rather its values - + Args: state_current (dict): The current state of the robot ref_state (dict): The reference state of the robot @@ -63,29 +59,25 @@ def optimize_gait(self, Returns: float: The best sample frequency """ - - - best_sample_freq = pgg_step_freq if self.optimize_step_freq and optimize_swing == 1: contact_sequence_temp = np.zeros((len(self.step_freq_available), 4, self.horizon)) for j in range(len(self.step_freq_available)): - pgg_temp = PeriodicGaitGenerator(duty_factor=pgg_duty_factor, - step_freq=self.step_freq_available[j], - gait_type=pgg_gait_type, - horizon=self.horizon) + pgg_temp = PeriodicGaitGenerator( + duty_factor=pgg_duty_factor, + step_freq=self.step_freq_available[j], + gait_type=pgg_gait_type, + horizon=self.horizon, + ) pgg_temp.set_phase_signal(pgg_phase_signal) - contact_sequence_temp[j] = pgg_temp.compute_contact_sequence(contact_sequence_dts=self.contact_sequence_dts, - contact_sequence_lenghts=self.contact_sequence_lenghts) - - - costs, \ - best_sample_freq = self.batched_controller.compute_batch_control(state_current, - ref_state, - contact_sequence_temp) + contact_sequence_temp[j] = pgg_temp.compute_contact_sequence( + contact_sequence_dts=self.contact_sequence_dts, + contact_sequence_lenghts=self.contact_sequence_lenghts, + ) + costs, best_sample_freq = self.batched_controller.compute_batch_control( + state_current, ref_state, contact_sequence_temp + ) - return best_sample_freq - \ No newline at end of file diff --git a/quadruped_pympc/interfaces/srbd_controller_interface.py b/quadruped_pympc/interfaces/srbd_controller_interface.py index 0f972c6..fa31a26 100644 --- a/quadruped_pympc/interfaces/srbd_controller_interface.py +++ b/quadruped_pympc/interfaces/srbd_controller_interface.py @@ -1,91 +1,100 @@ import numpy as np - from gym_quadruped.utils.quadruped_utils import LegsAttr from quadruped_pympc import config as cfg + class SRBDControllerInterface: """This is an interface for a controller that uses the SRBD method to optimize the gait""" + def __init__( + self, + ): + """Constructor for the SRBD controller interface""" - def __init__(self, ): - """ Constructor for the SRBD controller interface """ - - self.type = cfg.mpc_params['type'] - self.mpc_dt = cfg.mpc_params['dt'] - self.horizon = cfg.mpc_params['horizon'] - self.optimize_step_freq = cfg.mpc_params['optimize_step_freq'] - self.step_freq_available = cfg.mpc_params['step_freq_available'] - + self.type = cfg.mpc_params["type"] + self.mpc_dt = cfg.mpc_params["dt"] + self.horizon = cfg.mpc_params["horizon"] + self.optimize_step_freq = cfg.mpc_params["optimize_step_freq"] + self.step_freq_available = cfg.mpc_params["step_freq_available"] self.previous_contact_mpc = np.array([1, 1, 1, 1]) - + # 'nominal' optimized directly the GRF # 'input_rates' optimizes the delta GRF # 'sampling' is a gpu-based mpc that samples the GRF # 'collaborative' optimized directly the GRF and has a passive arm model inside # 'lyapunov' optimized directly the GRF and has a Lyapunov-based stability constraint # 'kynodynamic' sbrd with joints - experimental - if self.type == 'nominal': + if self.type == "nominal": from quadruped_pympc.controllers.gradient.nominal.centroidal_nmpc_nominal import Acados_NMPC_Nominal self.controller = Acados_NMPC_Nominal() if self.optimize_step_freq: - from quadruped_pympc.controllers.gradient.nominal.centroidal_nmpc_gait_adaptive import \ - Acados_NMPC_GaitAdaptive + from quadruped_pympc.controllers.gradient.nominal.centroidal_nmpc_gait_adaptive import ( + Acados_NMPC_GaitAdaptive, + ) self.batched_controller = Acados_NMPC_GaitAdaptive() - elif self.type == 'input_rates': - from quadruped_pympc.controllers.gradient.input_rates.centroidal_nmpc_input_rates import Acados_NMPC_InputRates + elif self.type == "input_rates": + from quadruped_pympc.controllers.gradient.input_rates.centroidal_nmpc_input_rates import ( + Acados_NMPC_InputRates, + ) self.controller = Acados_NMPC_InputRates() if self.optimize_step_freq: - from quadruped_pympc.controllers.gradient.nominal.centroidal_nmpc_gait_adaptive import \ - Acados_NMPC_GaitAdaptive + from quadruped_pympc.controllers.gradient.nominal.centroidal_nmpc_gait_adaptive import ( + Acados_NMPC_GaitAdaptive, + ) self.batched_controller = Acados_NMPC_GaitAdaptive() - elif(cfg.mpc_params['type'] == 'lyapunov'): + elif cfg.mpc_params["type"] == "lyapunov": from quadruped_pympc.controllers.gradient.lyapunov.centroidal_nmpc_lyapunov import Acados_NMPC_Lyapunov self.controller = Acados_NMPC_Lyapunov() if self.optimize_step_freq: - from quadruped_pympc.controllers.gradient.nominal.centroidal_nmpc_gait_adaptive import Acados_NMPC_GaitAdaptive + from quadruped_pympc.controllers.gradient.nominal.centroidal_nmpc_gait_adaptive import ( + Acados_NMPC_GaitAdaptive, + ) + self.batched_controller = Acados_NMPC_GaitAdaptive() - elif(cfg.mpc_params['type'] == 'kinodynamic'): + elif cfg.mpc_params["type"] == "kinodynamic": from quadruped_pympc.controllers.gradient.kinodynamic.kinodynamic_nmpc import Acados_NMPC_KinoDynamic + self.controller = Acados_NMPC_KinoDynamic() if self.optimize_step_freq: - from quadruped_pympc.controllers.gradient.nominal.centroidal_nmpc_gait_adaptive import Acados_NMPC_GaitAdaptive + from quadruped_pympc.controllers.gradient.nominal.centroidal_nmpc_gait_adaptive import ( + Acados_NMPC_GaitAdaptive, + ) + self.batched_controller = Acados_NMPC_GaitAdaptive() - elif self.type == 'sampling': + elif self.type == "sampling": if self.optimize_step_freq: from quadruped_pympc.controllers.sampling.centroidal_nmpc_jax_gait_adaptive import Sampling_MPC else: from quadruped_pympc.controllers.sampling.centroidal_nmpc_jax import Sampling_MPC self.controller = Sampling_MPC() - - - - - def compute_control(self, - state_current: dict, - ref_state: dict, - contact_sequence: np.ndarray, - inertia: np.ndarray, - pgg_phase_signal: np.ndarray, - pgg_step_freq: float, - optimize_swing: int, - external_wrenches: np.ndarray = np.zeros((6,))) -> [LegsAttr, LegsAttr, LegsAttr, LegsAttr, LegsAttr, float]: + def compute_control( + self, + state_current: dict, + ref_state: dict, + contact_sequence: np.ndarray, + inertia: np.ndarray, + pgg_phase_signal: np.ndarray, + pgg_step_freq: float, + optimize_swing: int, + external_wrenches: np.ndarray = np.zeros((6,)), + ) -> [LegsAttr, LegsAttr, LegsAttr, LegsAttr, LegsAttr, float]: """Compute the control using the SRBD method Args: @@ -99,138 +108,143 @@ def compute_control(self, external_wrenches (np.ndarray): The external wrench applied to the robot to compensate Returns: - tuple: The GRFs and the feet positions in world frame, + tuple: The GRFs and the feet positions in world frame, and the best sample frequency (only if the controller is sampling) """ - - current_contact = np.array([contact_sequence[0][0], - contact_sequence[1][0], - contact_sequence[2][0], - contact_sequence[3][0]]) - - # If we use sampling - if (self.type == 'sampling'): + current_contact = np.array( + [contact_sequence[0][0], contact_sequence[1][0], contact_sequence[2][0], contact_sequence[3][0]] + ) + # If we use sampling + if self.type == "sampling": # Convert data to jax and shift previous solution - state_current_jax, \ - reference_state_jax, = self.controller.prepare_state_and_reference(state_current, - ref_state, - current_contact, - self.previous_contact_mpc) + ( + state_current_jax, + reference_state_jax, + ) = self.controller.prepare_state_and_reference( + state_current, ref_state, current_contact, self.previous_contact_mpc + ) self.previous_contact_mpc = current_contact for iter_sampling in range(self.controller.num_sampling_iterations): self.controller = self.controller.with_newkey() - if (self.controller.sampling_method == 'cem_mppi'): - if (iter_sampling == 0): - self.controller = self.controller.with_newsigma(cfg.mpc_params['sigma_cem_mppi']) - - nmpc_GRFs, \ - nmpc_footholds, \ - nmpc_predicted_state,\ - self.controller.best_control_parameters, \ - best_cost, \ - best_sample_freq, \ - costs, \ - sigma_cem_mppi = self.controller.jitted_compute_control(state_current_jax, reference_state_jax, - contact_sequence, self.controller.best_control_parameters, - self.controller.master_key, self.controller.sigma_cem_mppi) + if self.controller.sampling_method == "cem_mppi": + if iter_sampling == 0: + self.controller = self.controller.with_newsigma(cfg.mpc_params["sigma_cem_mppi"]) + + ( + nmpc_GRFs, + nmpc_footholds, + nmpc_predicted_state, + self.controller.best_control_parameters, + best_cost, + best_sample_freq, + costs, + sigma_cem_mppi, + ) = self.controller.jitted_compute_control( + state_current_jax, + reference_state_jax, + contact_sequence, + self.controller.best_control_parameters, + self.controller.master_key, + self.controller.sigma_cem_mppi, + ) self.controller = self.controller.with_newsigma(sigma_cem_mppi) else: nominal_sample_freq = pgg_step_freq - nmpc_GRFs, \ - nmpc_footholds, \ - nmpc_predicted_state,\ - self.controller.best_control_parameters, \ - best_cost, \ - best_sample_freq, \ - costs = self.controller.jitted_compute_control(state_current_jax, reference_state_jax, - contact_sequence, self.controller.best_control_parameters, - self.controller.master_key, pgg_phase_signal, - nominal_sample_freq, optimize_swing) - - nmpc_footholds = LegsAttr(FL=ref_state["ref_foot_FL"][0], - FR=ref_state["ref_foot_FR"][0], - RL=ref_state["ref_foot_RL"][0], - RR=ref_state["ref_foot_RR"][0]) + ( + nmpc_GRFs, + nmpc_footholds, + nmpc_predicted_state, + self.controller.best_control_parameters, + best_cost, + best_sample_freq, + costs, + ) = self.controller.jitted_compute_control( + state_current_jax, + reference_state_jax, + contact_sequence, + self.controller.best_control_parameters, + self.controller.master_key, + pgg_phase_signal, + nominal_sample_freq, + optimize_swing, + ) + + nmpc_footholds = LegsAttr( + FL=ref_state["ref_foot_FL"][0], + FR=ref_state["ref_foot_FR"][0], + RL=ref_state["ref_foot_RL"][0], + RR=ref_state["ref_foot_RR"][0], + ) nmpc_GRFs = np.array(nmpc_GRFs) - + nmpc_joints_pos = None nmpc_joints_vel = None nmpc_joints_acc = None # If we use Gradient-Based MPC else: - if(self.type == 'kinodynamic'): - - nmpc_GRFs, \ - nmpc_footholds, \ - nmpc_joints_pos, \ - nmpc_joints_vel, \ - nmpc_joints_acc, \ - nmpc_predicted_state, \ - status = self.controller.compute_control(state_current, - ref_state, - contact_sequence, - inertia=inertia, - external_wrenches=external_wrenches) - - - - - nmpc_joints_pos = LegsAttr(FL=nmpc_joints_pos[0:3], - FR=nmpc_joints_pos[3:6], - RL=nmpc_joints_pos[6:9], - RR=nmpc_joints_pos[9:12]) - - nmpc_joints_vel = LegsAttr(FL=nmpc_joints_vel[0:3], - FR=nmpc_joints_vel[3:6], - RL=nmpc_joints_vel[6:9], - RR=nmpc_joints_vel[9:12]) - - nmpc_joints_acc = LegsAttr(FL=nmpc_joints_acc[0:3], - FR=nmpc_joints_acc[3:6], - RL=nmpc_joints_acc[6:9], - RR=nmpc_joints_acc[9:12]) - + if self.type == "kinodynamic": + ( + nmpc_GRFs, + nmpc_footholds, + nmpc_joints_pos, + nmpc_joints_vel, + nmpc_joints_acc, + nmpc_predicted_state, + status, + ) = self.controller.compute_control( + state_current, ref_state, contact_sequence, inertia=inertia, external_wrenches=external_wrenches + ) + + nmpc_joints_pos = LegsAttr( + FL=nmpc_joints_pos[0:3], FR=nmpc_joints_pos[3:6], RL=nmpc_joints_pos[6:9], RR=nmpc_joints_pos[9:12] + ) + + nmpc_joints_vel = LegsAttr( + FL=nmpc_joints_vel[0:3], FR=nmpc_joints_vel[3:6], RL=nmpc_joints_vel[6:9], RR=nmpc_joints_vel[9:12] + ) + + nmpc_joints_acc = LegsAttr( + FL=nmpc_joints_acc[0:3], FR=nmpc_joints_acc[3:6], RL=nmpc_joints_acc[6:9], RR=nmpc_joints_acc[9:12] + ) + else: - nmpc_GRFs, \ - nmpc_footholds, \ - nmpc_predicted_state, \ - _ = self.controller.compute_control(state_current, - ref_state, - contact_sequence, - inertia=inertia, - external_wrenches=external_wrenches) - + nmpc_GRFs, nmpc_footholds, nmpc_predicted_state, _ = self.controller.compute_control( + state_current, ref_state, contact_sequence, inertia=inertia, external_wrenches=external_wrenches + ) + nmpc_joints_pos = None nmpc_joints_vel = None nmpc_joints_acc = None - - nmpc_footholds = LegsAttr(FL=nmpc_footholds[0], - FR=nmpc_footholds[1], - RL=nmpc_footholds[2], - RR=nmpc_footholds[3]) - + nmpc_footholds = LegsAttr( + FL=nmpc_footholds[0], FR=nmpc_footholds[1], RL=nmpc_footholds[2], RR=nmpc_footholds[3] + ) best_sample_freq = pgg_step_freq - - # TODO: Indexing should not be hardcoded. Env should provide indexing of leg actuator dimensions. - nmpc_GRFs = LegsAttr(FL=nmpc_GRFs[0:3] * current_contact[0], - FR=nmpc_GRFs[3:6] * current_contact[1], - RL=nmpc_GRFs[6:9] * current_contact[2], - RR=nmpc_GRFs[9:12] * current_contact[3]) - - - - return nmpc_GRFs, nmpc_footholds, nmpc_joints_pos, nmpc_joints_vel, nmpc_joints_acc, best_sample_freq, nmpc_predicted_state - + nmpc_GRFs = LegsAttr( + FL=nmpc_GRFs[0:3] * current_contact[0], + FR=nmpc_GRFs[3:6] * current_contact[1], + RL=nmpc_GRFs[6:9] * current_contact[2], + RR=nmpc_GRFs[9:12] * current_contact[3], + ) + + return ( + nmpc_GRFs, + nmpc_footholds, + nmpc_joints_pos, + nmpc_joints_vel, + nmpc_joints_acc, + best_sample_freq, + nmpc_predicted_state, + ) def compute_RTI(self): - self.controller.acados_ocp_solver.options_set('rti_phase', 1) + self.controller.acados_ocp_solver.options_set("rti_phase", 1) self.controller.acados_ocp_solver.solve() - # print("preparation phase time: ", controller.acados_ocp_solver.get_stats('time_tot')) \ No newline at end of file + # print("preparation phase time: ", controller.acados_ocp_solver.get_stats('time_tot')) diff --git a/quadruped_pympc/interfaces/wb_interface.py b/quadruped_pympc/interfaces/wb_interface.py index 687df33..d41406e 100644 --- a/quadruped_pympc/interfaces/wb_interface.py +++ b/quadruped_pympc/interfaces/wb_interface.py @@ -1,19 +1,19 @@ -import numpy as np -import time -from scipy.spatial.transform import Rotation as R import copy +import time +import numpy as np from gym_quadruped.utils.quadruped_utils import LegsAttr -from quadruped_pympc import config as cfg +from scipy.spatial.transform import Rotation as R +from quadruped_pympc import config as cfg from quadruped_pympc.helpers.foothold_reference_generator import FootholdReferenceGenerator +from quadruped_pympc.helpers.inverse_kinematics.inverse_kinematics_numeric import InverseKinematicsNumeric from quadruped_pympc.helpers.periodic_gait_generator import PeriodicGaitGenerator from quadruped_pympc.helpers.swing_trajectory_controller import SwingTrajectoryController from quadruped_pympc.helpers.terrain_estimator import TerrainEstimator -from quadruped_pympc.helpers.inverse_kinematics.inverse_kinematics_numeric import InverseKinematicsNumeric from quadruped_pympc.helpers.velocity_modulator import VelocityModulator -if(cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): +if cfg.simulation_params["visual_foothold_adaptation"] != "blind": from quadruped_pympc.helpers.visual_foothold_adaptation import VisualFootholdAdaptation @@ -24,88 +24,92 @@ class WBInterface: swing trajectory control, and terrain estimation. """ - def __init__(self, - initial_feet_pos: LegsAttr, - legs_order: tuple[str, str, str, str] = ('FL', 'FR', 'RL', 'RR')): - """ Constructor of the WBInterface class + def __init__(self, initial_feet_pos: LegsAttr, legs_order: tuple[str, str, str, str] = ("FL", "FR", "RL", "RR")): + """Constructor of the WBInterface class Args: initial_feet_pos (LegsAttr): initial feet positions, otherwise they will be all zero legs_order (tuple[str, str, str, str], optional): order of the leg. Defaults to ('FL', 'FR', 'RL', 'RR'). - """ - - - mpc_dt = cfg.mpc_params['dt'] - horizon = cfg.mpc_params['horizon'] + """ + + mpc_dt = cfg.mpc_params["dt"] + horizon = cfg.mpc_params["horizon"] self.legs_order = legs_order - + # Periodic gait generator -------------------------------------------------------------- - gait_name = cfg.simulation_params['gait'] - gait_params = cfg.simulation_params['gait_params'][gait_name] - gait_type, duty_factor, step_frequency = gait_params['type'], gait_params['duty_factor'], gait_params['step_freq'] + gait_name = cfg.simulation_params["gait"] + gait_params = cfg.simulation_params["gait_params"][gait_name] + gait_type, duty_factor, step_frequency = ( + gait_params["type"], + gait_params["duty_factor"], + gait_params["step_freq"], + ) # Given the possibility to use nonuniform discretization, # we generate a contact sequence two times longer and with a dt half of the one of the mpc - self.pgg = PeriodicGaitGenerator(duty_factor=duty_factor, - step_freq=step_frequency, - gait_type=gait_type, - horizon=horizon) + self.pgg = PeriodicGaitGenerator( + duty_factor=duty_factor, step_freq=step_frequency, gait_type=gait_type, horizon=horizon + ) # in the case of nonuniform discretization, we create a list of dts and horizons for each nonuniform discretization - if (cfg.mpc_params['use_nonuniform_discretization']): - self.contact_sequence_dts = [cfg.mpc_params['dt_fine_grained'], mpc_dt] - self.contact_sequence_lenghts = [cfg.mpc_params['horizon_fine_grained'], horizon] + if cfg.mpc_params["use_nonuniform_discretization"]: + self.contact_sequence_dts = [cfg.mpc_params["dt_fine_grained"], mpc_dt] + self.contact_sequence_lenghts = [cfg.mpc_params["horizon_fine_grained"], horizon] else: self.contact_sequence_dts = [mpc_dt] self.contact_sequence_lenghts = [horizon] - # Create the foothold reference generator ------------------------------------------------ stance_time = (1 / self.pgg.step_freq) * self.pgg.duty_factor - self.frg = FootholdReferenceGenerator(stance_time=stance_time, hip_height=cfg.hip_height, lift_off_positions=initial_feet_pos) - + self.frg = FootholdReferenceGenerator( + stance_time=stance_time, hip_height=cfg.hip_height, lift_off_positions=initial_feet_pos + ) # Create swing trajectory generator ------------------------------------------------------ - self.step_height = cfg.simulation_params['step_height'] - swing_period = (1 - self.pgg.duty_factor) * (1 / self.pgg.step_freq) - position_gain_fb = cfg.simulation_params['swing_position_gain_fb'] - velocity_gain_fb = cfg.simulation_params['swing_velocity_gain_fb'] - swing_generator = cfg.simulation_params['swing_generator'] - self.stc = SwingTrajectoryController(step_height=self.step_height, swing_period=swing_period, - position_gain_fb=position_gain_fb, velocity_gain_fb=velocity_gain_fb, - generator=swing_generator) - + self.step_height = cfg.simulation_params["step_height"] + swing_period = (1 - self.pgg.duty_factor) * (1 / self.pgg.step_freq) + position_gain_fb = cfg.simulation_params["swing_position_gain_fb"] + velocity_gain_fb = cfg.simulation_params["swing_velocity_gain_fb"] + swing_generator = cfg.simulation_params["swing_generator"] + self.stc = SwingTrajectoryController( + step_height=self.step_height, + swing_period=swing_period, + position_gain_fb=position_gain_fb, + velocity_gain_fb=velocity_gain_fb, + generator=swing_generator, + ) # Terrain estimator ----------------------------------------------------------------------- self.terrain_computation = TerrainEstimator() - # Inverse Kinematics --------------------------------------------------------------------- - self.ik = InverseKinematicsNumeric() + self.ik = InverseKinematicsNumeric() - if(cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): + if cfg.simulation_params["visual_foothold_adaptation"] != "blind": # Visual foothold adaptation ------------------------------------------------------------- - self.vfa = VisualFootholdAdaptation(legs_order=self.legs_order, adaptation_strategy=cfg.simulation_params['visual_foothold_adaptation']) + self.vfa = VisualFootholdAdaptation( + legs_order=self.legs_order, adaptation_strategy=cfg.simulation_params["visual_foothold_adaptation"] + ) # Velocity modulator --------------------------------------------------------------------- self.vm = VelocityModulator() self.current_contact = np.array([1, 1, 1, 1]) - - - def update_state_and_reference(self, - com_pos: np.ndarray, - base_pos: np.ndarray, - base_lin_vel: np.ndarray, - base_ori_euler_xyz: np.ndarray, - base_ang_vel: np.ndarray, - feet_pos: LegsAttr, - hip_pos: LegsAttr, - joints_pos: LegsAttr, - heightmaps, - legs_order: tuple[str, str, str, str], - simulation_dt: float, - ref_base_lin_vel: np.ndarray, - ref_base_ang_vel: np.ndarray) -> [dict, dict, list, LegsAttr, list, list, float, bool]: + def update_state_and_reference( + self, + com_pos: np.ndarray, + base_pos: np.ndarray, + base_lin_vel: np.ndarray, + base_ori_euler_xyz: np.ndarray, + base_ang_vel: np.ndarray, + feet_pos: LegsAttr, + hip_pos: LegsAttr, + joints_pos: LegsAttr, + heightmaps, + legs_order: tuple[str, str, str, str], + simulation_dt: float, + ref_base_lin_vel: np.ndarray, + ref_base_ang_vel: np.ndarray, + ) -> [dict, dict, list, LegsAttr, list, list, float, bool]: """Update the state and reference for the whole body controller, including the contact sequence, footholds, and terrain estimation. Args: @@ -130,13 +134,12 @@ def update_state_and_reference(self, ref_feet_pos (LegsAttr): where to step in world frame ref_feet_constraints (LegsAttr): constraints for the footholds in the world frame step_height (float): step height - optimize_swing (bool), boolean to inform that the robot is in the apex, hence we can optimize step freq. + optimize_swing (bool), boolean to inform that the robot is in the apex, hence we can optimize step freq. """ - state_current = dict( - position=com_pos + self.frg.com_pos_offset_w, #manual com offset - #position=base_pos, + position=com_pos + self.frg.com_pos_offset_w, # manual com offset + # position=base_pos, linear_velocity=base_lin_vel, orientation=base_ori_euler_xyz, angular_velocity=base_ang_vel, @@ -148,199 +151,223 @@ def update_state_and_reference(self, joint_FR=joints_pos.FR, joint_RL=joints_pos.RL, joint_RR=joints_pos.RR, - ) - + ) + # Modulate the desired velocity if the robot is in strange positions - if(self.vm.activated): - ref_base_lin_vel, ref_base_ang_vel = self.vm.modulate_velocities(ref_base_lin_vel, ref_base_ang_vel, feet_pos, hip_pos) + if self.vm.activated: + ref_base_lin_vel, ref_base_ang_vel = self.vm.modulate_velocities( + ref_base_lin_vel, ref_base_ang_vel, feet_pos, hip_pos + ) - # Update the desired contact sequence --------------------------- - if(self.pgg.start_and_stop_activated): + if self.pgg.start_and_stop_activated: # stop the robot for energy efficency if there is no movement and its safe # only if activated by an an internal flag for now - self.pgg.update_start_and_stop(feet_pos, hip_pos, self.frg.hip_offset, - base_pos, base_ori_euler_xyz, - base_lin_vel, base_ang_vel, - ref_base_lin_vel, ref_base_ang_vel, - self.current_contact) - - self.pgg.run(simulation_dt, self.pgg.step_freq) - contact_sequence = self.pgg.compute_contact_sequence(contact_sequence_dts=self.contact_sequence_dts, - contact_sequence_lenghts=self.contact_sequence_lenghts) + self.pgg.update_start_and_stop( + feet_pos, + hip_pos, + self.frg.hip_offset, + base_pos, + base_ori_euler_xyz, + base_lin_vel, + base_ang_vel, + ref_base_lin_vel, + ref_base_ang_vel, + self.current_contact, + ) + self.pgg.run(simulation_dt, self.pgg.step_freq) + contact_sequence = self.pgg.compute_contact_sequence( + contact_sequence_dts=self.contact_sequence_dts, contact_sequence_lenghts=self.contact_sequence_lenghts + ) previous_contact = self.current_contact - self.current_contact = np.array([contact_sequence[0][0], - contact_sequence[1][0], - contact_sequence[2][0], - contact_sequence[3][0]]) - + self.current_contact = np.array( + [contact_sequence[0][0], contact_sequence[1][0], contact_sequence[2][0], contact_sequence[3][0]] + ) # Compute the reference for the footholds --------------------------------------------------- - self.frg.update_lift_off_positions(previous_contact, self.current_contact, feet_pos, legs_order, self.pgg.gait_type, base_pos, base_ori_euler_xyz) - self.frg.update_touch_down_positions(previous_contact, self.current_contact, feet_pos, legs_order, self.pgg.gait_type, base_pos, base_ori_euler_xyz) + self.frg.update_lift_off_positions( + previous_contact, + self.current_contact, + feet_pos, + legs_order, + self.pgg.gait_type, + base_pos, + base_ori_euler_xyz, + ) + self.frg.update_touch_down_positions( + previous_contact, + self.current_contact, + feet_pos, + legs_order, + self.pgg.gait_type, + base_pos, + base_ori_euler_xyz, + ) ref_feet_pos = self.frg.compute_footholds_reference( base_position=base_pos, base_ori_euler_xyz=base_ori_euler_xyz, base_xy_lin_vel=base_lin_vel[0:2], ref_base_xy_lin_vel=ref_base_lin_vel[0:2], hips_position=hip_pos, - com_height_nominal=cfg.simulation_params['ref_z']) - + com_height_nominal=cfg.simulation_params["ref_z"], + ) # Adjust the footholds given the terrain ----------------------------------------------------- - if(cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): - + if cfg.simulation_params["visual_foothold_adaptation"] != "blind": time_adaptation = time.time() - if(self.stc.check_apex_condition(self.current_contact, interval=0.01) and self.vfa.initialized == False): + if self.stc.check_apex_condition(self.current_contact, interval=0.01) and self.vfa.initialized == False: for leg_id, leg_name in enumerate(legs_order): heightmaps[leg_name].update_height_map(ref_feet_pos[leg_name], yaw=base_ori_euler_xyz[2]) - self.vfa.compute_adaptation(legs_order, ref_feet_pos, hip_pos, heightmaps, base_lin_vel, base_ori_euler_xyz, base_ang_vel) - #print("Adaptation time: ", time.time() - time_adaptation) - - if(self.stc.check_full_stance_condition(self.current_contact)): + self.vfa.compute_adaptation( + legs_order, ref_feet_pos, hip_pos, heightmaps, base_lin_vel, base_ori_euler_xyz, base_ang_vel + ) + # print("Adaptation time: ", time.time() - time_adaptation) + + if self.stc.check_full_stance_condition(self.current_contact): self.vfa.reset() - + ref_feet_pos, ref_feet_constraints = self.vfa.get_footholds_adapted(ref_feet_pos) else: ref_feet_constraints = LegsAttr(FL=None, FR=None, RL=None, RR=None) - # Estimate the terrain slope and elevation ------------------------------------------------------- - terrain_roll, \ - terrain_pitch, \ - terrain_height = self.terrain_computation.compute_terrain_estimation( + terrain_roll, terrain_pitch, terrain_height = self.terrain_computation.compute_terrain_estimation( base_position=base_pos, yaw=base_ori_euler_xyz[2], feet_pos=self.frg.lift_off_positions, - current_contact=self.current_contact) + current_contact=self.current_contact, + ) ref_pos = np.array([0, 0, cfg.hip_height]) - ref_pos[2] = cfg.simulation_params['ref_z'] + terrain_height + ref_pos[2] = cfg.simulation_params["ref_z"] + terrain_height # Rotate the reference base linear velocity to the terrain frame - ref_base_lin_vel = R.from_euler('xyz', [terrain_roll, terrain_pitch, 0]).as_matrix() @ ref_base_lin_vel - + ref_base_lin_vel = R.from_euler("xyz", [terrain_roll, terrain_pitch, 0]).as_matrix() @ ref_base_lin_vel # Update state reference ------------------------------------------------------------------------ - - # Since the MPC close in CoM position, but usually we have desired height for the base, + + # Since the MPC close in CoM position, but usually we have desired height for the base, # we modify the reference to bring the base at the desired height and not the CoM ref_pos[2] -= base_pos[2] - (com_pos[2] + self.frg.com_pos_offset_w[2]) - - if(cfg.mpc_params['type'] != 'kinodynamic'): + + if cfg.mpc_params["type"] != "kinodynamic": ref_state = {} - ref_state |= dict(ref_foot_FL=ref_feet_pos.FL.reshape((1, 3)), - ref_foot_FR=ref_feet_pos.FR.reshape((1, 3)), - ref_foot_RL=ref_feet_pos.RL.reshape((1, 3)), - ref_foot_RR=ref_feet_pos.RR.reshape((1, 3)), - ref_foot_constraints_FL=ref_feet_constraints.FL, - ref_foot_constraints_FR=ref_feet_constraints.FR, - ref_foot_constraints_RL=ref_feet_constraints.RL, - ref_foot_constraints_RR=ref_feet_constraints.RR, - # Also update the reference base linear velocity and - ref_linear_velocity=ref_base_lin_vel, - ref_angular_velocity=ref_base_ang_vel, - ref_orientation=np.array([terrain_roll, terrain_pitch, 0.0]), - ref_position=ref_pos - ) + ref_state |= dict( + ref_foot_FL=ref_feet_pos.FL.reshape((1, 3)), + ref_foot_FR=ref_feet_pos.FR.reshape((1, 3)), + ref_foot_RL=ref_feet_pos.RL.reshape((1, 3)), + ref_foot_RR=ref_feet_pos.RR.reshape((1, 3)), + ref_foot_constraints_FL=ref_feet_constraints.FL, + ref_foot_constraints_FR=ref_feet_constraints.FR, + ref_foot_constraints_RL=ref_feet_constraints.RL, + ref_foot_constraints_RR=ref_feet_constraints.RR, + # Also update the reference base linear velocity and + ref_linear_velocity=ref_base_lin_vel, + ref_angular_velocity=ref_base_ang_vel, + ref_orientation=np.array([terrain_roll, terrain_pitch, 0.0]), + ref_position=ref_pos, + ) else: # In the case of the kinodynamic model, - # we should pass as a reference the X-Y-Z spline of the feet for the horizon, + # we should pass as a reference the X-Y-Z spline of the feet for the horizon, # since in the kynodimic model we are using the feet position as a reference - desired_foot_position_FL = np.zeros((cfg.mpc_params['horizon'], 3)) - desired_foot_position_FR = np.zeros((cfg.mpc_params['horizon'], 3)) - desired_foot_position_RL = np.zeros((cfg.mpc_params['horizon'], 3)) - desired_foot_position_RR = np.zeros((cfg.mpc_params['horizon'], 3)) + desired_foot_position_FL = np.zeros((cfg.mpc_params["horizon"], 3)) + desired_foot_position_FR = np.zeros((cfg.mpc_params["horizon"], 3)) + desired_foot_position_RL = np.zeros((cfg.mpc_params["horizon"], 3)) + desired_foot_position_RR = np.zeros((cfg.mpc_params["horizon"], 3)) for leg_id, leg_name in enumerate(legs_order): lifted_off = [False, False, False, False] - for n in range(cfg.mpc_params['horizon']): - dt_increment_swing = (n)*cfg.mpc_params['dt'] - - if(lifted_off[leg_id] == False and n >= 0): - if(contact_sequence[leg_id][n-1] == 1 and contact_sequence[leg_id][n] == 0): - lifted_off[leg_id] = True - - if(leg_id == 0): - - if(contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == False): + for n in range(cfg.mpc_params["horizon"]): + dt_increment_swing = (n) * cfg.mpc_params["dt"] + + if lifted_off[leg_id] == False and n >= 0: + if contact_sequence[leg_id][n - 1] == 1 and contact_sequence[leg_id][n] == 0: + lifted_off[leg_id] = True + + if leg_id == 0: + if contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == False: desired_foot_position_FL[n] = feet_pos.FL - elif(contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == True): + elif contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == True: desired_foot_position_FL[n] = ref_feet_pos.FL else: - desired_foot_position, \ - desired_foot_velocity, \ - _ = self.stc.swing_generator.compute_trajectory_references(self.stc.swing_time[leg_id] + dt_increment_swing, - self.frg.lift_off_positions[leg_name], - ref_feet_pos.FL) + desired_foot_position, desired_foot_velocity, _ = ( + self.stc.swing_generator.compute_trajectory_references( + self.stc.swing_time[leg_id] + dt_increment_swing, + self.frg.lift_off_positions[leg_name], + ref_feet_pos.FL, + ) + ) desired_foot_position_FL[n] = desired_foot_position - elif(leg_id == 1): - - if(contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == False): + elif leg_id == 1: + if contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == False: desired_foot_position_FR[n] = feet_pos.FR - elif(contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == True): + elif contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == True: desired_foot_position_FR[n] = ref_feet_pos.FR else: - desired_foot_position, \ - desired_foot_velocity, \ - _ = self.stc.swing_generator.compute_trajectory_references(self.stc.swing_time[leg_id] + dt_increment_swing, - self.frg.lift_off_positions[leg_name], - ref_feet_pos.FR) + desired_foot_position, desired_foot_velocity, _ = ( + self.stc.swing_generator.compute_trajectory_references( + self.stc.swing_time[leg_id] + dt_increment_swing, + self.frg.lift_off_positions[leg_name], + ref_feet_pos.FR, + ) + ) desired_foot_position_FR[n] = desired_foot_position - elif(leg_id == 2): - - if(contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == False): + elif leg_id == 2: + if contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == False: desired_foot_position_RL[n] = feet_pos.RL - elif(contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == True): + elif contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == True: desired_foot_position_RL[n] = ref_feet_pos.RL else: - desired_foot_position, \ - desired_foot_velocity, \ - _ = self.stc.swing_generator.compute_trajectory_references(self.stc.swing_time[leg_id] + dt_increment_swing, - self.frg.lift_off_positions[leg_name], - ref_feet_pos.RL) + desired_foot_position, desired_foot_velocity, _ = ( + self.stc.swing_generator.compute_trajectory_references( + self.stc.swing_time[leg_id] + dt_increment_swing, + self.frg.lift_off_positions[leg_name], + ref_feet_pos.RL, + ) + ) desired_foot_position_RL[n] = desired_foot_position - elif(leg_id == 3): - - if(contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == False): + elif leg_id == 3: + if contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == False: desired_foot_position_RR[n] = feet_pos.RR - elif(contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == True): + elif contact_sequence[leg_id][n] == 1 and lifted_off[leg_id] == True: desired_foot_position_RR[n] = ref_feet_pos.RR else: - desired_foot_position, \ - desired_foot_velocity, \ - _ = self.stc.swing_generator.compute_trajectory_references(self.stc.swing_time[leg_id] + dt_increment_swing, - self.frg.lift_off_positions[leg_name], - ref_feet_pos.RR) + desired_foot_position, desired_foot_velocity, _ = ( + self.stc.swing_generator.compute_trajectory_references( + self.stc.swing_time[leg_id] + dt_increment_swing, + self.frg.lift_off_positions[leg_name], + ref_feet_pos.RR, + ) + ) desired_foot_position_RR[n] = desired_foot_position - - #TODO make this more general + + # TODO make this more general ref_state = {} init_qpos = np.array([0, 0.9, -1.8, 0, 0.9, -1.8, 0, 0.9, -1.8, 0, 0.9, -1.8]) - ref_state |= dict(ref_foot_FL=desired_foot_position_FL, - ref_foot_FR=desired_foot_position_FR, - ref_foot_RL=desired_foot_position_RL, - ref_foot_RR=desired_foot_position_RR, - ref_foot_constraints_FL=ref_feet_constraints.FL, - ref_foot_constraints_FR=ref_feet_constraints.FR, - ref_foot_constraints_RL=ref_feet_constraints.RL, - ref_foot_constraints_RR=ref_feet_constraints.RR, - # Also update the reference base linear velocity and - ref_linear_velocity=ref_base_lin_vel, - ref_angular_velocity=ref_base_ang_vel, - ref_orientation=np.array([terrain_roll, terrain_pitch, 0.0]), - ref_position=ref_pos, - ref_joints=init_qpos - ) - - - # ------------------------------------------------------------------------------------------------- + ref_state |= dict( + ref_foot_FL=desired_foot_position_FL, + ref_foot_FR=desired_foot_position_FR, + ref_foot_RL=desired_foot_position_RL, + ref_foot_RR=desired_foot_position_RR, + ref_foot_constraints_FL=ref_feet_constraints.FL, + ref_foot_constraints_FR=ref_feet_constraints.FR, + ref_foot_constraints_RL=ref_feet_constraints.RL, + ref_foot_constraints_RR=ref_feet_constraints.RR, + # Also update the reference base linear velocity and + ref_linear_velocity=ref_base_lin_vel, + ref_angular_velocity=ref_base_ang_vel, + ref_orientation=np.array([terrain_roll, terrain_pitch, 0.0]), + ref_position=ref_pos, + ref_joints=init_qpos, + ) + # ------------------------------------------------------------------------------------------------- - if(cfg.mpc_params['optimize_step_freq']): + if cfg.mpc_params["optimize_step_freq"]: # we can always optimize the step freq, or just at the apex of the swing # to avoid possible jittering in the solution optimize_swing = self.stc.check_apex_condition(self.current_contact) @@ -348,30 +375,30 @@ def update_state_and_reference(self, optimize_swing = 0 return state_current, ref_state, contact_sequence, self.step_height, optimize_swing - - - - def compute_stance_and_swing_torque(self, - simulation_dt: float, - qpos: np.ndarray, - qvel: np.ndarray, - feet_jac: LegsAttr, - jac_feet_dot: LegsAttr, - feet_pos: LegsAttr, - feet_vel: LegsAttr, - legs_qfrc_bias: LegsAttr, - legs_mass_matrix: LegsAttr, - nmpc_GRFs: LegsAttr, - nmpc_footholds: LegsAttr, - legs_qpos_idx: LegsAttr, - legs_qvel_idx: LegsAttr, - tau: LegsAttr, - optimize_swing: int, - best_sample_freq: float, - nmpc_joints_pos, - nmpc_joints_vel, - nmpc_joints_acc, - nmpc_predicted_state) -> LegsAttr: + + def compute_stance_and_swing_torque( + self, + simulation_dt: float, + qpos: np.ndarray, + qvel: np.ndarray, + feet_jac: LegsAttr, + jac_feet_dot: LegsAttr, + feet_pos: LegsAttr, + feet_vel: LegsAttr, + legs_qfrc_bias: LegsAttr, + legs_mass_matrix: LegsAttr, + nmpc_GRFs: LegsAttr, + nmpc_footholds: LegsAttr, + legs_qpos_idx: LegsAttr, + legs_qvel_idx: LegsAttr, + tau: LegsAttr, + optimize_swing: int, + best_sample_freq: float, + nmpc_joints_pos, + nmpc_joints_vel, + nmpc_joints_acc, + nmpc_predicted_state, + ) -> LegsAttr: """Compute the stance and swing torque. Args: @@ -383,102 +410,99 @@ def compute_stance_and_swing_torque(self, feet_vel (LegsAttr): feet velocities in world frame legs_qfrc_bias (LegsAttr): joint forces and torques legs_mass_matrix (LegsAttr): mass matrix of the legs - nmpc_GRFs (LegsAttr): ground reaction forces from the MPC in world frame - nmpc_footholds (LegsAttr): footholds from the MPC in world frame + nmpc_GRFs (LegsAttr): ground reaction forces from the MPC in world frame + nmpc_footholds (LegsAttr): footholds from the MPC in world frame legs_qvel_idx (LegsAttr): joint velocities index tau (LegsAttr): joint torques optimize_swing (int): flag to signal that we need to update the swing trajectory time - best_sample_freq (float): best sample frequency obtained from the + best_sample_freq (float): best sample frequency obtained from the sampling optimization or the batched ocp Returns: LegsAttr: joint torques """ - # If we have optimized the gait, we set all the timing parameters - if (optimize_swing == 1): + if optimize_swing == 1: self.pgg.step_freq = np.array([best_sample_freq])[0] self.frg.stance_time = (1 / self.pgg.step_freq) * self.pgg.duty_factor swing_period = (1 - self.pgg.duty_factor) * (1 / self.pgg.step_freq) self.stc.regenerate_swing_trajectory_generator(step_height=self.step_height, swing_period=swing_period) - - + # Compute Stance Torque --------------------------------------------------------------------------- tau.FL = -np.matmul(feet_jac.FL[:, legs_qvel_idx.FL].T, nmpc_GRFs.FL) tau.FR = -np.matmul(feet_jac.FR[:, legs_qvel_idx.FR].T, nmpc_GRFs.FR) tau.RL = -np.matmul(feet_jac.RL[:, legs_qvel_idx.RL].T, nmpc_GRFs.RL) tau.RR = -np.matmul(feet_jac.RR[:, legs_qvel_idx.RR].T, nmpc_GRFs.RR) - - self.stc.update_swing_time(self.current_contact, self.legs_order, simulation_dt) - # Compute Swing Torque ------------------------------------------------------------------------------ des_foot_pos = LegsAttr(*[np.zeros((3,)) for _ in range(4)]) des_foot_vel = LegsAttr(*[np.zeros((3,)) for _ in range(4)]) - - if(cfg.mpc_params['type'] != 'kinodynamic'): - # The swing controller is in the end-effector space + + if cfg.mpc_params["type"] != "kinodynamic": + # The swing controller is in the end-effector space for leg_id, leg_name in enumerate(self.legs_order): - if self.current_contact[leg_id] == 0: # If in swing phase, compute the swing trajectory tracking control. - tau[leg_name], \ - des_foot_pos[leg_name], \ - des_foot_vel[leg_name] = self.stc.compute_swing_control_cartesian_space( - leg_id=leg_id, - q_dot=qvel[legs_qvel_idx[leg_name]], - J=feet_jac[leg_name][:, legs_qvel_idx[leg_name]], - J_dot=jac_feet_dot[leg_name][:, legs_qvel_idx[leg_name]], - lift_off=self.frg.lift_off_positions[leg_name], - touch_down=nmpc_footholds[leg_name], - foot_pos=feet_pos[leg_name], - foot_vel=feet_vel[leg_name], - h=legs_qfrc_bias[leg_name], - mass_matrix=legs_mass_matrix[leg_name] - ) + if ( + self.current_contact[leg_id] == 0 + ): # If in swing phase, compute the swing trajectory tracking control. + tau[leg_name], des_foot_pos[leg_name], des_foot_vel[leg_name] = ( + self.stc.compute_swing_control_cartesian_space( + leg_id=leg_id, + q_dot=qvel[legs_qvel_idx[leg_name]], + J=feet_jac[leg_name][:, legs_qvel_idx[leg_name]], + J_dot=jac_feet_dot[leg_name][:, legs_qvel_idx[leg_name]], + lift_off=self.frg.lift_off_positions[leg_name], + touch_down=nmpc_footholds[leg_name], + foot_pos=feet_pos[leg_name], + foot_vel=feet_vel[leg_name], + h=legs_qfrc_bias[leg_name], + mass_matrix=legs_mass_matrix[leg_name], + ) + ) else: des_foot_pos[leg_name] = nmpc_footholds[leg_name] - #des_foot_pos[leg_name] = self.frg.touch_down_positions[leg_name] - des_foot_vel[leg_name] = des_foot_vel[leg_name]*0.0 + # des_foot_pos[leg_name] = self.frg.touch_down_positions[leg_name] + des_foot_vel[leg_name] = des_foot_vel[leg_name] * 0.0 else: # The swing controller is in the joint space for leg_id, leg_name in enumerate(self.legs_order): - if self.current_contact[leg_id] == 0: # If in swing phase, compute the swing trajectory tracking control. - tau[leg_name], \ - _, \ - _ = self.stc.compute_swing_control_joint_space(nmpc_joints_pos[leg_name], - nmpc_joints_vel[leg_name], - nmpc_joints_acc[leg_name], - qpos[legs_qpos_idx[leg_name]], - qvel[legs_qvel_idx[leg_name]], - legs_mass_matrix[leg_name], - legs_qfrc_bias[leg_name],) - - + if ( + self.current_contact[leg_id] == 0 + ): # If in swing phase, compute the swing trajectory tracking control. + tau[leg_name], _, _ = self.stc.compute_swing_control_joint_space( + nmpc_joints_pos[leg_name], + nmpc_joints_vel[leg_name], + nmpc_joints_acc[leg_name], + qpos[legs_qpos_idx[leg_name]], + qvel[legs_qvel_idx[leg_name]], + legs_mass_matrix[leg_name], + legs_qfrc_bias[leg_name], + ) + # Compute PD targets for the joints ---------------------------------------------------------------- des_joints_pos = LegsAttr(*[np.zeros((3, 1)) for _ in range(4)]) des_joints_vel = LegsAttr(*[np.zeros((3, 1)) for _ in range(4)]) - if(cfg.mpc_params['type'] != 'kinodynamic'): + if cfg.mpc_params["type"] != "kinodynamic": qpos_predicted = copy.deepcopy(qpos) - #TODO use predicted rotation too - #qpos_predicted[0:3] = nmpc_predicted_state[0:3] - #TODO make the ik explicit and not numerical - temp = self.ik.fun_compute_solution(qpos_predicted, - des_foot_pos.FL, des_foot_pos.FR, - des_foot_pos.RL, des_foot_pos.RR) - - des_joints_pos.FL = np.array(temp[0:3]).reshape((3, )) - des_joints_pos.FR = np.array(temp[3:6]).reshape((3, )) - des_joints_pos.RL = np.array(temp[6:9]).reshape((3, )) - des_joints_pos.RR = np.array(temp[9:12]).reshape((3, )) - - - #des_joints_vel.FL = (des_joints_pos.FL - qpos[legs_qpos_idx.FL])/self.contact_sequence_dts[0] - #des_joints_vel.FR = (des_joints_pos.FR - qpos[legs_qpos_idx.FR])/self.contact_sequence_dts[0] - #des_joints_vel.RL = (des_joints_pos.RL - qpos[legs_qpos_idx.RL])/self.contact_sequence_dts[0] - #des_joints_vel.RR = (des_joints_pos.RR - qpos[legs_qpos_idx.RR])/self.contact_sequence_dts[0] - #TODO This should be done over the the desired joint positions jacobian + # TODO use predicted rotation too + # qpos_predicted[0:3] = nmpc_predicted_state[0:3] + # TODO make the ik explicit and not numerical + temp = self.ik.fun_compute_solution( + qpos_predicted, des_foot_pos.FL, des_foot_pos.FR, des_foot_pos.RL, des_foot_pos.RR + ) + + des_joints_pos.FL = np.array(temp[0:3]).reshape((3,)) + des_joints_pos.FR = np.array(temp[3:6]).reshape((3,)) + des_joints_pos.RL = np.array(temp[6:9]).reshape((3,)) + des_joints_pos.RR = np.array(temp[9:12]).reshape((3,)) + + # des_joints_vel.FL = (des_joints_pos.FL - qpos[legs_qpos_idx.FL])/self.contact_sequence_dts[0] + # des_joints_vel.FR = (des_joints_pos.FR - qpos[legs_qpos_idx.FR])/self.contact_sequence_dts[0] + # des_joints_vel.RL = (des_joints_pos.RL - qpos[legs_qpos_idx.RL])/self.contact_sequence_dts[0] + # des_joints_vel.RR = (des_joints_pos.RR - qpos[legs_qpos_idx.RR])/self.contact_sequence_dts[0] + # TODO This should be done over the the desired joint positions jacobian des_joints_vel.FL = np.linalg.pinv(feet_jac.FL[:, legs_qvel_idx.FL]) @ des_foot_vel.FL des_joints_vel.FR = np.linalg.pinv(feet_jac.FR[:, legs_qvel_idx.FR]) @ des_foot_vel.FR des_joints_vel.RL = np.linalg.pinv(feet_jac.RL[:, legs_qvel_idx.RL]) @ des_foot_vel.RL @@ -492,41 +516,40 @@ def compute_stance_and_swing_torque(self, # Saturate of desired joint positions and velocities max_joints_pos_difference = 3 max_joints_vel_difference = 10.0 - + # Calculate the difference - actual_joints_pos = LegsAttr(**{leg_name:qpos[legs_qpos_idx[leg_name]] for leg_name in self.legs_order}) - actual_joints_vel = LegsAttr(**{leg_name:qvel[legs_qvel_idx[leg_name]] for leg_name in self.legs_order}) + actual_joints_pos = LegsAttr(**{leg_name: qpos[legs_qpos_idx[leg_name]] for leg_name in self.legs_order}) + actual_joints_vel = LegsAttr(**{leg_name: qvel[legs_qvel_idx[leg_name]] for leg_name in self.legs_order}) # Saturate the difference for each leg for leg in ["FL", "FR", "RL", "RR"]: joints_pos_difference = des_joints_pos[leg] - actual_joints_pos[leg] - saturated_joints_pos_difference = np.clip(joints_pos_difference, -max_joints_pos_difference, max_joints_pos_difference) + saturated_joints_pos_difference = np.clip( + joints_pos_difference, -max_joints_pos_difference, max_joints_pos_difference + ) des_joints_pos[leg] = actual_joints_pos[leg] + saturated_joints_pos_difference joints_vel_difference = des_joints_vel[leg] - actual_joints_vel[leg] - saturated_joints_vel_difference = np.clip(joints_vel_difference, -max_joints_vel_difference, max_joints_vel_difference) + saturated_joints_vel_difference = np.clip( + joints_vel_difference, -max_joints_vel_difference, max_joints_vel_difference + ) des_joints_vel[leg] = actual_joints_vel[leg] + saturated_joints_vel_difference - - return tau, des_joints_pos, des_joints_vel - - - def reset(self, - initial_feet_pos: LegsAttr): + def reset(self, initial_feet_pos: LegsAttr): """Reset the whole body interface Args: initial_feet_pos (LegsAttr): initial feet positions """ - + self.pgg.reset() - #self.frg.reset() - #self.stc.reset() - #self.terrain_computation.reset() + # self.frg.reset() + # self.stc.reset() + # self.terrain_computation.reset() self.frg.lift_off_positions = initial_feet_pos - if(cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): + if cfg.simulation_params["visual_foothold_adaptation"] != "blind": self.vfa.reset() self.current_contact = np.array([1, 1, 1, 1]) - return \ No newline at end of file + return diff --git a/quadruped_pympc/quadruped_pympc_wrapper.py b/quadruped_pympc/quadruped_pympc_wrapper.py index 3e788c6..0844bca 100644 --- a/quadruped_pympc/quadruped_pympc_wrapper.py +++ b/quadruped_pympc/quadruped_pympc_wrapper.py @@ -1,88 +1,80 @@ -from quadruped_pympc.interfaces.srbd_controller_interface import SRBDControllerInterface -from quadruped_pympc.interfaces.srbd_batched_controller_interface import SRBDBatchedControllerInterface -from quadruped_pympc.interfaces.wb_interface import WBInterface - -from gym_quadruped.utils.quadruped_utils import LegsAttr -from quadruped_pympc import config as cfg - import numpy as np +from gym_quadruped.utils.quadruped_utils import LegsAttr +from quadruped_pympc import config as cfg +from quadruped_pympc.interfaces.srbd_batched_controller_interface import SRBDBatchedControllerInterface +from quadruped_pympc.interfaces.srbd_controller_interface import SRBDControllerInterface +from quadruped_pympc.interfaces.wb_interface import WBInterface -_DEFAULT_OBS = ('ref_base_height', 'ref_base_angles', 'nmpc_GRFs', 'nmpc_footholds', 'swing_time') +_DEFAULT_OBS = ("ref_base_height", "ref_base_angles", "nmpc_GRFs", "nmpc_footholds", "swing_time") class QuadrupedPyMPC_Wrapper: """A simple class wrapper of all the mpc submodules (swing, contact generator, mpc itself).""" - def __init__(self, - initial_feet_pos: LegsAttr, - legs_order: tuple[str, str, str, str] = ('FL', 'FR', 'RL', 'RR'), - quadrupedpympc_observables_names: tuple[str, ...] = _DEFAULT_OBS): - """ Constructor of the QuadrupedPyMPC_Wrapper class. + def __init__( + self, + initial_feet_pos: LegsAttr, + legs_order: tuple[str, str, str, str] = ("FL", "FR", "RL", "RR"), + quadrupedpympc_observables_names: tuple[str, ...] = _DEFAULT_OBS, + ): + """Constructor of the QuadrupedPyMPC_Wrapper class. Args: initial_feet_pos (LegsAttr): initial feet positions, otherwise they will be all zero. legs_order (tuple[str, str, str, str], optional): order of the leg. Defaults to ('FL', 'FR', 'RL', 'RR'). quadrupedpympc_observables_names (tuple[str, ...], optional): list of observable to save. Defaults to _DEFAULT_OBS. - """ + """ - self.mpc_frequency = cfg.simulation_params['mpc_frequency'] + self.mpc_frequency = cfg.simulation_params["mpc_frequency"] self.srbd_controller_interface = SRBDControllerInterface() - if(cfg.mpc_params['type'] != 'sampling' and cfg.mpc_params['optimize_step_freq']): + if cfg.mpc_params["type"] != "sampling" and cfg.mpc_params["optimize_step_freq"]: self.srbd_batched_controller_interface = SRBDBatchedControllerInterface() - self.wb_interface = WBInterface(initial_feet_pos = initial_feet_pos(frame='world'), - legs_order = legs_order) - + self.wb_interface = WBInterface(initial_feet_pos=initial_feet_pos(frame="world"), legs_order=legs_order) - self.nmpc_GRFs = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), - RL=np.zeros(3), RR=np.zeros(3)) - self.nmpc_footholds = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), - RL=np.zeros(3), RR=np.zeros(3)) - self.nmpc_joints_pos = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), - RL=np.zeros(3), RR=np.zeros(3)) - self.nmpc_joints_vel = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), - RL=np.zeros(3), RR=np.zeros(3)) - self.nmpc_joints_acc = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), - RL=np.zeros(3), RR=np.zeros(3)) + self.nmpc_GRFs = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), RL=np.zeros(3), RR=np.zeros(3)) + self.nmpc_footholds = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), RL=np.zeros(3), RR=np.zeros(3)) + self.nmpc_joints_pos = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), RL=np.zeros(3), RR=np.zeros(3)) + self.nmpc_joints_vel = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), RL=np.zeros(3), RR=np.zeros(3)) + self.nmpc_joints_acc = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), RL=np.zeros(3), RR=np.zeros(3)) self.nmpc_predicted_state = np.zeros(12) self.best_sample_freq = self.wb_interface.pgg.step_freq - self.quadrupedpympc_observables_names = quadrupedpympc_observables_names self.quadrupedpympc_observables = {} - - - def compute_actions(self, - com_pos: np.ndarray, - base_pos: np.ndarray, - base_lin_vel: np.ndarray, - base_ori_euler_xyz: np.ndarray, - base_ang_vel: np.ndarray, - feet_pos: LegsAttr, - hip_pos: LegsAttr, - joints_pos: LegsAttr, - heightmaps, - legs_order: tuple[str, str, str, str], - simulation_dt: float, - ref_base_lin_vel: np.ndarray, - ref_base_ang_vel: np.ndarray, - step_num: int, - qpos: np.ndarray, - qvel: np.ndarray, - feet_jac: LegsAttr, - jac_feet_dot: LegsAttr, - feet_vel: LegsAttr, - legs_qfrc_bias: LegsAttr, - legs_mass_matrix: LegsAttr, - legs_qpos_idx: LegsAttr, - legs_qvel_idx: LegsAttr, - tau: LegsAttr, - inertia: np.ndarray) -> LegsAttr: - """ Given the current state of the robot (and the reference), + def compute_actions( + self, + com_pos: np.ndarray, + base_pos: np.ndarray, + base_lin_vel: np.ndarray, + base_ori_euler_xyz: np.ndarray, + base_ang_vel: np.ndarray, + feet_pos: LegsAttr, + hip_pos: LegsAttr, + joints_pos: LegsAttr, + heightmaps, + legs_order: tuple[str, str, str, str], + simulation_dt: float, + ref_base_lin_vel: np.ndarray, + ref_base_ang_vel: np.ndarray, + step_num: int, + qpos: np.ndarray, + qvel: np.ndarray, + feet_jac: LegsAttr, + jac_feet_dot: LegsAttr, + feet_vel: LegsAttr, + legs_qfrc_bias: LegsAttr, + legs_mass_matrix: LegsAttr, + legs_qpos_idx: LegsAttr, + legs_qvel_idx: LegsAttr, + tau: LegsAttr, + inertia: np.ndarray, + ) -> LegsAttr: + """Given the current state of the robot (and the reference), compute the torques to be applied to the motors. Args: @@ -112,159 +104,151 @@ def compute_actions(self, Returns: LegsAttr: torques to be applied to the motors - """ + """ - # Update the state and reference ------------------------- - state_current, \ - ref_state, \ - contact_sequence, \ - step_height, \ - optimize_swing = self.wb_interface.update_state_and_reference(com_pos, - base_pos, - base_lin_vel, - base_ori_euler_xyz, - base_ang_vel, - feet_pos, - hip_pos, - joints_pos, - heightmaps, - legs_order, - simulation_dt, - ref_base_lin_vel, - ref_base_ang_vel) - - - + state_current, ref_state, contact_sequence, step_height, optimize_swing = ( + self.wb_interface.update_state_and_reference( + com_pos, + base_pos, + base_lin_vel, + base_ori_euler_xyz, + base_ang_vel, + feet_pos, + hip_pos, + joints_pos, + heightmaps, + legs_order, + simulation_dt, + ref_base_lin_vel, + ref_base_ang_vel, + ) + ) # Solve OCP --------------------------------------------------------------------------------------- if step_num % round(1 / (self.mpc_frequency * simulation_dt)) == 0: - - self.nmpc_GRFs, \ - self.nmpc_footholds, \ - self.nmpc_joints_pos, \ - self.nmpc_joints_vel, \ - self.nmpc_joints_acc, \ - self.best_sample_freq, \ - self.nmpc_predicted_state = self.srbd_controller_interface.compute_control(state_current, - ref_state, - contact_sequence, - inertia, - self.wb_interface.pgg.phase_signal, - self.wb_interface.pgg.step_freq, - optimize_swing) - - if(cfg.mpc_params['type'] != 'sampling' and cfg.mpc_params['use_RTI']): + ( + self.nmpc_GRFs, + self.nmpc_footholds, + self.nmpc_joints_pos, + self.nmpc_joints_vel, + self.nmpc_joints_acc, + self.best_sample_freq, + self.nmpc_predicted_state, + ) = self.srbd_controller_interface.compute_control( + state_current, + ref_state, + contact_sequence, + inertia, + self.wb_interface.pgg.phase_signal, + self.wb_interface.pgg.step_freq, + optimize_swing, + ) + + if cfg.mpc_params["type"] != "sampling" and cfg.mpc_params["use_RTI"]: # If the controller is gradient and is using RTI, we need to linearize the mpc after its computation # this helps to minize the delay between new state->control in a real case scenario. self.srbd_controller_interface.compute_RTI() - # Update the gait - if(cfg.mpc_params['type'] != 'sampling' and cfg.mpc_params['optimize_step_freq']): - self.best_sample_freq = self.srbd_batched_controller_interface.optimize_gait(state_current, - ref_state, - inertia, - self.wb_interface.pgg.phase_signal, - self.wb_interface.pgg.step_freq, - self.wb_interface.pgg.duty_factor, - self.wb_interface.pgg.gait_type, - optimize_swing) + if cfg.mpc_params["type"] != "sampling" and cfg.mpc_params["optimize_step_freq"]: + self.best_sample_freq = self.srbd_batched_controller_interface.optimize_gait( + state_current, + ref_state, + inertia, + self.wb_interface.pgg.phase_signal, + self.wb_interface.pgg.step_freq, + self.wb_interface.pgg.duty_factor, + self.wb_interface.pgg.gait_type, + optimize_swing, + ) - - - # Compute Swing and Stance Torque --------------------------------------------------------------------------- - tau, \ - des_joints_pos, \ - des_joints_vel = self.wb_interface.compute_stance_and_swing_torque(simulation_dt, - qpos, - qvel, - feet_jac, - jac_feet_dot, - feet_pos, - feet_vel, - legs_qfrc_bias, - legs_mass_matrix, - self.nmpc_GRFs, - self.nmpc_footholds, - legs_qpos_idx, - legs_qvel_idx, - tau, - optimize_swing, - self.best_sample_freq, - self.nmpc_joints_pos, - self.nmpc_joints_vel, - self.nmpc_joints_acc, - self.nmpc_predicted_state) - + tau, des_joints_pos, des_joints_vel = self.wb_interface.compute_stance_and_swing_torque( + simulation_dt, + qpos, + qvel, + feet_jac, + jac_feet_dot, + feet_pos, + feet_vel, + legs_qfrc_bias, + legs_mass_matrix, + self.nmpc_GRFs, + self.nmpc_footholds, + legs_qpos_idx, + legs_qvel_idx, + tau, + optimize_swing, + self.best_sample_freq, + self.nmpc_joints_pos, + self.nmpc_joints_vel, + self.nmpc_joints_acc, + self.nmpc_predicted_state, + ) # Do some PD control over the joints (these values are normally passed # to a low-level motor controller, here we can try to simulate it) - kp_joint_motor = cfg.simulation_params['impedence_joint_position_gain'] - kd_joint_motor = cfg.simulation_params['impedence_joint_velocity_gain'] - #for leg in legs_order: + kp_joint_motor = cfg.simulation_params["impedence_joint_position_gain"] + kd_joint_motor = cfg.simulation_params["impedence_joint_velocity_gain"] + # for leg in legs_order: # tau[leg] += kp_joint_motor * (des_joints_pos[leg] - qpos[legs_qpos_idx[leg]]) + \ # kd_joint_motor * (des_joints_vel[leg] - qvel[legs_qvel_idx[leg]]) - # Save some observables ------------------------------------------------------------------------------------- self.quadrupedpympc_observables = {} for obs_name in self.quadrupedpympc_observables_names: - if obs_name == 'ref_base_height': - data = {'ref_base_height': ref_state['ref_position'][2]} - elif obs_name == 'ref_base_angles': - data = {'ref_base_angles': ref_state['ref_orientation']} - elif obs_name == 'ref_feet_pos': - ref_feet_pos = LegsAttr(FL=ref_state['ref_foot_FL'].reshape(3,1), - FR=ref_state['ref_foot_FR'].reshape(3,1), - RL=ref_state['ref_foot_RL'].reshape(3,1), - RR=ref_state['ref_foot_RR'].reshape(3,1)) - data = {'ref_feet_pos': ref_feet_pos} - elif obs_name == 'ref_feet_constraints': - ref_feet_constraints = LegsAttr(FL=ref_state['ref_foot_FL_constraints'], - FR=ref_state['ref_foot_FR_constraints'], - RL=ref_state['ref_foot_RL_constraints'], - RR=ref_state['ref_foot_RR_constraints']) - data = {'ref_feet_constraints': ref_feet_constraints} - elif obs_name == 'nmpc_GRFs': - data = {'nmpc_GRFs': self.nmpc_GRFs} - elif obs_name == 'nmpc_footholds': - data = {'nmpc_footholds': self.nmpc_footholds} - elif obs_name == 'swing_time': - data = {'swing_time': self.wb_interface.stc.swing_time} - elif obs_name == 'phase_signal': - data = {'phase_signal': self.wb_interface.pgg._phase_signal} - elif obs_name == 'lift_off_positions': - data = {'lift_off_positions': self.wb_interface.frg.lift_off_positions} - + if obs_name == "ref_base_height": + data = {"ref_base_height": ref_state["ref_position"][2]} + elif obs_name == "ref_base_angles": + data = {"ref_base_angles": ref_state["ref_orientation"]} + elif obs_name == "ref_feet_pos": + ref_feet_pos = LegsAttr( + FL=ref_state["ref_foot_FL"].reshape(3, 1), + FR=ref_state["ref_foot_FR"].reshape(3, 1), + RL=ref_state["ref_foot_RL"].reshape(3, 1), + RR=ref_state["ref_foot_RR"].reshape(3, 1), + ) + data = {"ref_feet_pos": ref_feet_pos} + elif obs_name == "ref_feet_constraints": + ref_feet_constraints = LegsAttr( + FL=ref_state["ref_foot_FL_constraints"], + FR=ref_state["ref_foot_FR_constraints"], + RL=ref_state["ref_foot_RL_constraints"], + RR=ref_state["ref_foot_RR_constraints"], + ) + data = {"ref_feet_constraints": ref_feet_constraints} + elif obs_name == "nmpc_GRFs": + data = {"nmpc_GRFs": self.nmpc_GRFs} + elif obs_name == "nmpc_footholds": + data = {"nmpc_footholds": self.nmpc_footholds} + elif obs_name == "swing_time": + data = {"swing_time": self.wb_interface.stc.swing_time} + elif obs_name == "phase_signal": + data = {"phase_signal": self.wb_interface.pgg._phase_signal} + elif obs_name == "lift_off_positions": + data = {"lift_off_positions": self.wb_interface.frg.lift_off_positions} + else: data = {} raise ValueError(f"Unknown observable name: {obs_name}") - + self.quadrupedpympc_observables.update(data) - return tau - - - def get_obs(self,) -> dict: - """ Get some user-defined observables from withing the control loop. + def get_obs( + self, + ) -> dict: + """Get some user-defined observables from withing the control loop. Returns: Dict: dictionary of observables """ return self.quadrupedpympc_observables - - - def reset(self, - initial_feet_pos: LegsAttr): - """ Reset the controller.""" + def reset(self, initial_feet_pos: LegsAttr): + """Reset the controller.""" self.wb_interface.reset(initial_feet_pos) self.srbd_controller_interface.controller.reset() - - - \ No newline at end of file diff --git a/simulation/batched_simulations.py b/simulation/batched_simulations.py index bf35f29..8410942 100644 --- a/simulation/batched_simulations.py +++ b/simulation/batched_simulations.py @@ -1,14 +1,15 @@ -# Description: This script is used to simulate multiple MPC runs in parallel, +# Description: This script is used to simulate multiple MPC runs in parallel, # mainly to test the performance of the MPC controller over time. # Authors: # Giulio Turrisi import multiprocessing -from multiprocessing import Process import time +from multiprocessing import Process + import numpy as np -import matplotlib.pyplot as plt + # Import the simulation module that will run mujoco import simulation @@ -17,24 +18,26 @@ NUM_EPISODES = 20 -if __name__ == '__main__': +if __name__ == "__main__": render = False render_only_first = True output_simulations_tracking = None output_simulations_success_rate = None base_vel_command_type = "forward+rotate" - - + # Run the simulation in parallel processes = [] manager = multiprocessing.Manager() return_dict = manager.dict() for i in range(NUM_PROCESSES): - if(render_only_first and i == 0): + if render_only_first and i == 0: render = True else: render = False - p = Process(target=simulation.run_simulation, args=(i, NUM_EPISODES, return_dict, i*NUM_PROCESSES, base_vel_command_type, render)) + p = Process( + target=simulation.run_simulation, + args=(i, NUM_EPISODES, return_dict, i * NUM_PROCESSES, base_vel_command_type, render), + ) p.start() time.sleep(2) processes.append(p) @@ -42,24 +45,33 @@ for p in processes: p.join() - # Concatenate the data from the processes for proc_j in range(NUM_PROCESSES): for ep_i in range(NUM_EPISODES): - if(output_simulations_tracking is None): - output_simulations_tracking = return_dict['process'+str(proc_j)+'_ctrl_state_history_ep'+str(ep_i)] + if output_simulations_tracking is None: + output_simulations_tracking = return_dict[ + "process" + str(proc_j) + "_ctrl_state_history_ep" + str(ep_i) + ] else: - output_simulations_tracking = np.concatenate((output_simulations_tracking, return_dict['process'+str(proc_j)+'_ctrl_state_history_ep'+str(ep_i)])) + output_simulations_tracking = np.concatenate( + ( + output_simulations_tracking, + return_dict["process" + str(proc_j) + "_ctrl_state_history_ep" + str(ep_i)], + ) + ) - if(output_simulations_success_rate is None): - output_simulations_success_rate = np.array([return_dict['process'+str(proc_j)+'_success_rate_ep'+str(ep_i)]]) + if output_simulations_success_rate is None: + output_simulations_success_rate = np.array( + [return_dict["process" + str(proc_j) + "_success_rate_ep" + str(ep_i)]] + ) else: - output_simulations_success_rate = np.concatenate((output_simulations_success_rate, np.array([return_dict['process'+str(proc_j)+'_success_rate_ep'+str(ep_i)]]))) + output_simulations_success_rate = np.concatenate( + ( + output_simulations_success_rate, + np.array([return_dict["process" + str(proc_j) + "_success_rate_ep" + str(ep_i)]]), + ) + ) - print("Tracking data mean: ", np.mean(output_simulations_tracking, axis=0)) print("Tracking data std: ", np.std(output_simulations_tracking, axis=0)) - print("Success rate: ", np.sum(output_simulations_success_rate)/len(output_simulations_success_rate), "%") - - - \ No newline at end of file + print("Success rate: ", np.sum(output_simulations_success_rate) / len(output_simulations_success_rate), "%") diff --git a/simulation/simulation.py b/simulation/simulation.py index 2c56d9a..ff2de11 100644 --- a/simulation/simulation.py +++ b/simulation/simulation.py @@ -4,63 +4,71 @@ # Giulio Turrisi, Daniel Ordonez import time + import numpy as np -from tqdm import tqdm # Gym and Simulation related imports from gym_quadruped.quadruped_env import QuadrupedEnv +from gym_quadruped.utils.mujoco.visual import render_sphere, render_vector from gym_quadruped.utils.quadruped_utils import LegsAttr +from tqdm import tqdm # Config imports from quadruped_pympc import config as cfg # Helper functions for plotting from quadruped_pympc.helpers.quadruped_utils import plot_swing_mujoco -from gym_quadruped.utils.mujoco.visual import render_vector -from gym_quadruped.utils.mujoco.visual import render_sphere # PyMPC controller imports from quadruped_pympc.quadruped_pympc_wrapper import QuadrupedPyMPC_Wrapper # HeightMap import -if(cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): +if cfg.simulation_params["visual_foothold_adaptation"] != "blind": from gym_quadruped.sensors.heightmap import HeightMap - - -def run_simulation(process=0, num_episodes=500, return_dict=None, seed_number=0, base_vel_command_type="human" ,render=True): - +def run_simulation( + process=0, num_episodes=500, return_dict=None, seed_number=0, base_vel_command_type="human", render=True +): np.set_printoptions(precision=3, suppress=True) - np.random.seed(seed_number) + np.random.seed(seed_number) robot_name = cfg.robot hip_height = cfg.hip_height robot_leg_joints = cfg.robot_leg_joints robot_feet_geom_names = cfg.robot_feet_geom_names - scene_name = cfg.simulation_params['scene'] - simulation_dt = cfg.simulation_params['dt'] - - state_observables_names = ('base_pos', 'base_lin_vel', 'base_ori_euler_xyz', 'base_ori_quat_wxyz', 'base_ang_vel', - 'qpos_js', 'qvel_js', 'tau_ctrl_setpoint', - 'feet_pos_base', 'feet_vel_base', 'contact_state', 'contact_forces_base',) - + scene_name = cfg.simulation_params["scene"] + simulation_dt = cfg.simulation_params["dt"] + + state_observables_names = ( + "base_pos", + "base_lin_vel", + "base_ori_euler_xyz", + "base_ori_quat_wxyz", + "base_ang_vel", + "qpos_js", + "qvel_js", + "tau_ctrl_setpoint", + "feet_pos_base", + "feet_vel_base", + "contact_state", + "contact_forces_base", + ) # Create the quadruped robot environment ----------------------------------------------------------- - env = QuadrupedEnv(robot=robot_name, - hip_height=hip_height, - legs_joint_names=robot_leg_joints, # Joint names of the legs DoF - feet_geom_name=robot_feet_geom_names, # Geom/Frame id of feet - scene=scene_name, - sim_dt=simulation_dt, - ref_base_lin_vel=(0, 0.6), # pass a float for a fixed value - ref_base_ang_vel=(-0.2, 0.2), # pass a float for a fixed value - ground_friction_coeff=1.5, # pass a float for a fixed value - base_vel_command_type=base_vel_command_type, # "forward", "random", "forward+rotate", "human" - state_obs_names=state_observables_names, # Desired quantities in the 'state' vec - ) - - + env = QuadrupedEnv( + robot=robot_name, + hip_height=hip_height, + legs_joint_names=robot_leg_joints, # Joint names of the legs DoF + feet_geom_name=robot_feet_geom_names, # Geom/Frame id of feet + scene=scene_name, + sim_dt=simulation_dt, + ref_base_lin_vel=(0, 0.6), # pass a float for a fixed value + ref_base_ang_vel=(-0.2, 0.2), # pass a float for a fixed value + ground_friction_coeff=1.5, # pass a float for a fixed value + base_vel_command_type=base_vel_command_type, # "forward", "random", "forward+rotate", "human" + state_obs_names=state_observables_names, # Desired quantities in the 'state' vec + ) # Some robots require a change in the zero joint-space configuration. If provided apply it if cfg.qpos0_js is not None: @@ -70,7 +78,6 @@ def run_simulation(process=0, num_episodes=500, return_dict=None, seed_number=0, if render: env.render() # Pass in the first render call any mujoco.viewer.KeyCallbackType - # Initialization of variables used in the main control loop -------------------------------- # Jacobian matrices @@ -81,42 +88,57 @@ def run_simulation(process=0, num_episodes=500, return_dict=None, seed_number=0, # Torque limits tau_soft_limits_scalar = 0.9 tau_limits = LegsAttr( - FL=env.mjModel.actuator_ctrlrange[env.legs_tau_idx.FL]*tau_soft_limits_scalar, - FR=env.mjModel.actuator_ctrlrange[env.legs_tau_idx.FR]*tau_soft_limits_scalar, - RL=env.mjModel.actuator_ctrlrange[env.legs_tau_idx.RL]*tau_soft_limits_scalar, - RR=env.mjModel.actuator_ctrlrange[env.legs_tau_idx.RR]*tau_soft_limits_scalar) + FL=env.mjModel.actuator_ctrlrange[env.legs_tau_idx.FL] * tau_soft_limits_scalar, + FR=env.mjModel.actuator_ctrlrange[env.legs_tau_idx.FR] * tau_soft_limits_scalar, + RL=env.mjModel.actuator_ctrlrange[env.legs_tau_idx.RL] * tau_soft_limits_scalar, + RR=env.mjModel.actuator_ctrlrange[env.legs_tau_idx.RR] * tau_soft_limits_scalar, + ) - - # Feet positions and Legs order feet_pos = None feet_traj_geom_ids, feet_GRF_geom_ids = None, LegsAttr(FL=-1, FR=-1, RL=-1, RR=-1) legs_order = ["FL", "FR", "RL", "RR"] - # Create HeightMap ----------------------------------------------------------------------- - if(cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): + if cfg.simulation_params["visual_foothold_adaptation"] != "blind": resolution_vfa = 0.04 dimension_vfa = 7 - heightmaps = LegsAttr(FL=HeightMap(n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mj_model=env.mjModel, mj_data=env.mjData), - FR=HeightMap(n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mj_model=env.mjModel, mj_data=env.mjData), - RL=HeightMap(n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mj_model=env.mjModel, mj_data=env.mjData), - RR=HeightMap(n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mj_model=env.mjModel, mj_data=env.mjData)) + heightmaps = LegsAttr( + FL=HeightMap( + n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mj_model=env.mjModel, mj_data=env.mjData + ), + FR=HeightMap( + n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mj_model=env.mjModel, mj_data=env.mjData + ), + RL=HeightMap( + n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mj_model=env.mjModel, mj_data=env.mjData + ), + RR=HeightMap( + n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mj_model=env.mjModel, mj_data=env.mjData + ), + ) else: heightmaps = None - # Quadruped PyMPC controller initialization ------------------------------------------------------------- - mpc_frequency = cfg.simulation_params['mpc_frequency'] - - - quadrupedpympc_observables_names = ('ref_base_height', 'ref_base_angles', 'ref_feet_pos', - 'nmpc_GRFs', 'nmpc_footholds', - 'swing_time', 'phase_signal', 'lift_off_positions') - - quadrupedpympc_wrapper = QuadrupedPyMPC_Wrapper(initial_feet_pos=env.feet_pos, legs_order=legs_order, - quadrupedpympc_observables_names = quadrupedpympc_observables_names) - + mpc_frequency = cfg.simulation_params["mpc_frequency"] + + quadrupedpympc_observables_names = ( + "ref_base_height", + "ref_base_angles", + "ref_feet_pos", + "nmpc_GRFs", + "nmpc_footholds", + "swing_time", + "phase_signal", + "lift_off_positions", + ) + + quadrupedpympc_wrapper = QuadrupedPyMPC_Wrapper( + initial_feet_pos=env.feet_pos, + legs_order=legs_order, + quadrupedpympc_observables_names=quadrupedpympc_observables_names, + ) # -------------------------------------------------------------- RENDER_FREQ = 30 # Hz @@ -126,46 +148,43 @@ def run_simulation(process=0, num_episodes=500, return_dict=None, seed_number=0, state_obs_history, ctrl_state_history = [], [] for episode_num in tqdm(range(N_EPISODES), desc="Episodes"): - ep_state_obs_history, ep_ctrl_state_history = [], [] for _ in range(N_STEPS_PER_EPISODE): step_start = time.time() - # Update value from SE or Simulator ---------------------- - feet_pos = env.feet_pos(frame='world') - hip_pos = env.hip_positions(frame='world') - base_lin_vel = env.base_lin_vel(frame='world') - base_ang_vel = env.base_ang_vel(frame='world') + feet_pos = env.feet_pos(frame="world") + hip_pos = env.hip_positions(frame="world") + base_lin_vel = env.base_lin_vel(frame="world") + base_ang_vel = env.base_ang_vel(frame="world") base_ori_euler_xyz = env.base_ori_euler_xyz base_pos = env.base_pos com_pos = env.com - + # Get the reference base velocity in the world frame ref_base_lin_vel, ref_base_ang_vel = env.target_base_vel() - + # Get the inertia matrix - if(cfg.simulation_params['use_inertia_recomputation']): + if cfg.simulation_params["use_inertia_recomputation"]: inertia = env.get_base_inertia().flatten() # Reflected inertia of base at qpos, in world frame else: inertia = cfg.inertia.flatten() # Get the qpos and qvel qpos, qvel = env.mjData.qpos, env.mjData.qvel - joints_pos = LegsAttr(FL=qpos[7:10], FR=qpos[10:13], - RL=qpos[13:16], RR=qpos[16:19]) - + joints_pos = LegsAttr(FL=qpos[7:10], FR=qpos[10:13], RL=qpos[13:16], RR=qpos[16:19]) + # Get Centrifugal, Coriolis, Gravity for the swing controller legs_mass_matrix = env.legs_mass_matrix legs_qfrc_bias = env.legs_qfrc_bias # Compute feet jacobian - feet_jac = env.feet_jacobians(frame='world', return_rot_jac=False) - + feet_jac = env.feet_jacobians(frame="world", return_rot_jac=False) + # Compute jacobian derivatives of the contact points jac_feet_dot = (feet_jac - jac_feet_prev) / simulation_dt # Finite difference approximation jac_feet_prev = feet_jac # Update previous Jacobians - + # Compute feet velocities feet_vel = LegsAttr(**{leg_name: feet_jac[leg_name] @ env.mjData.qvel for leg_name in legs_order}) @@ -173,15 +192,34 @@ def run_simulation(process=0, num_episodes=500, return_dict=None, seed_number=0, legs_qvel_idx = env.legs_qvel_idx legs_qpos_idx = env.legs_qpos_idx - - # Quadruped PyMPC controller -------------------------------------------------------------- - tau = quadrupedpympc_wrapper.compute_actions(com_pos, base_pos, base_lin_vel, base_ori_euler_xyz, base_ang_vel, - feet_pos, hip_pos, joints_pos, - heightmaps, - legs_order, simulation_dt, ref_base_lin_vel, ref_base_ang_vel, - env.step_num, qpos, qvel, feet_jac, jac_feet_dot, feet_vel, legs_qfrc_bias, - legs_mass_matrix, legs_qpos_idx, legs_qvel_idx, tau, inertia) + tau = quadrupedpympc_wrapper.compute_actions( + com_pos, + base_pos, + base_lin_vel, + base_ori_euler_xyz, + base_ang_vel, + feet_pos, + hip_pos, + joints_pos, + heightmaps, + legs_order, + simulation_dt, + ref_base_lin_vel, + ref_base_ang_vel, + env.step_num, + qpos, + qvel, + feet_jac, + jac_feet_dot, + feet_vel, + legs_qfrc_bias, + legs_mass_matrix, + legs_qpos_idx, + legs_qvel_idx, + tau, + inertia, + ) # Limit tau between tau_limits for leg in ["FL", "FR", "RL", "RR"]: tau_min, tau_max = tau_limits[leg][:, 0], tau_limits[leg][:, 1] @@ -189,9 +227,6 @@ def run_simulation(process=0, num_episodes=500, return_dict=None, seed_number=0, quadrupedpympc_observables = quadrupedpympc_wrapper.get_obs() - - - # Set control and mujoco step ---------------------------------------------------------------------- action = np.zeros(env.mjModel.nu) action[env.legs_tau_idx.FL] = tau.FL @@ -199,67 +234,74 @@ def run_simulation(process=0, num_episodes=500, return_dict=None, seed_number=0, action[env.legs_tau_idx.RL] = tau.RL action[env.legs_tau_idx.RR] = tau.RR - #action_noise = np.random.normal(0, 2, size=env.mjModel.nu) - #action += action_noise + # action_noise = np.random.normal(0, 2, size=env.mjModel.nu) + # action += action_noise state, reward, is_terminated, is_truncated, info = env.step(action=action) - # Store the history of observations and control ------------------------------------------------------- ep_state_obs_history.append(state) base_lin_vel_err = ref_base_lin_vel - base_lin_vel base_ang_vel_err = ref_base_ang_vel - base_ang_vel base_poz_z_err = quadrupedpympc_observables["ref_base_height"] - base_pos[2] - ctrl_state = np.concatenate((base_lin_vel_err, base_ang_vel_err, [base_poz_z_err], quadrupedpympc_observables["phase_signal"])) + ctrl_state = np.concatenate( + (base_lin_vel_err, base_ang_vel_err, [base_poz_z_err], quadrupedpympc_observables["phase_signal"]) + ) ep_ctrl_state_history.append(ctrl_state) - # Render only at a certain frequency ----------------------------------------------------------------- if render and (time.time() - last_render_time > 1.0 / RENDER_FREQ or env.step_num == 1): _, _, feet_GRF = env.feet_contact_state(ground_reaction_forces=True) # Plot the swing trajectory - feet_traj_geom_ids = plot_swing_mujoco(viewer=env.viewer, - swing_traj_controller=quadrupedpympc_wrapper.wb_interface.stc, - swing_period=quadrupedpympc_wrapper.wb_interface.stc.swing_period, - swing_time=LegsAttr(FL=quadrupedpympc_observables["swing_time"][0], - FR=quadrupedpympc_observables["swing_time"][1], - RL=quadrupedpympc_observables["swing_time"][2], - RR=quadrupedpympc_observables["swing_time"][3]), - lift_off_positions=quadrupedpympc_observables["lift_off_positions"], - nmpc_footholds=quadrupedpympc_observables["nmpc_footholds"], - ref_feet_pos=quadrupedpympc_observables["ref_feet_pos"], - geom_ids=feet_traj_geom_ids) - - + feet_traj_geom_ids = plot_swing_mujoco( + viewer=env.viewer, + swing_traj_controller=quadrupedpympc_wrapper.wb_interface.stc, + swing_period=quadrupedpympc_wrapper.wb_interface.stc.swing_period, + swing_time=LegsAttr( + FL=quadrupedpympc_observables["swing_time"][0], + FR=quadrupedpympc_observables["swing_time"][1], + RL=quadrupedpympc_observables["swing_time"][2], + RR=quadrupedpympc_observables["swing_time"][3], + ), + lift_off_positions=quadrupedpympc_observables["lift_off_positions"], + nmpc_footholds=quadrupedpympc_observables["nmpc_footholds"], + ref_feet_pos=quadrupedpympc_observables["ref_feet_pos"], + geom_ids=feet_traj_geom_ids, + ) + # Update and Plot the heightmap - if(cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): - #if(stc.check_apex_condition(current_contact, interval=0.01)): + if cfg.simulation_params["visual_foothold_adaptation"] != "blind": + # if(stc.check_apex_condition(current_contact, interval=0.01)): for leg_id, leg_name in enumerate(legs_order): - data = heightmaps[leg_name].data#.update_height_map(ref_feet_pos[leg_name], yaw=env.base_ori_euler_xyz[2]) - if(data is not None): + data = heightmaps[ + leg_name + ].data # .update_height_map(ref_feet_pos[leg_name], yaw=env.base_ori_euler_xyz[2]) + if data is not None: for i in range(data.shape[0]): for j in range(data.shape[1]): - heightmaps[leg_name].geom_ids[i, j] = render_sphere(viewer=env.viewer, - position=([data[i][j][0][0],data[i][j][0][1],data[i][j][0][2]]), - diameter=0.01, - color=[0, 1, 0, .5], - geom_id=heightmaps[leg_name].geom_ids[i, j] - ) - + heightmaps[leg_name].geom_ids[i, j] = render_sphere( + viewer=env.viewer, + position=([data[i][j][0][0], data[i][j][0][1], data[i][j][0][2]]), + diameter=0.01, + color=[0, 1, 0, 0.5], + geom_id=heightmaps[leg_name].geom_ids[i, j], + ) + # Plot the GRF for leg_id, leg_name in enumerate(legs_order): - feet_GRF_geom_ids[leg_name] = render_vector(env.viewer, - vector=feet_GRF[leg_name], - pos=feet_pos[leg_name], - scale=np.linalg.norm(feet_GRF[leg_name]) * 0.005, - color=np.array([0, 1, 0, .5]), - geom_id=feet_GRF_geom_ids[leg_name]) + feet_GRF_geom_ids[leg_name] = render_vector( + env.viewer, + vector=feet_GRF[leg_name], + pos=feet_pos[leg_name], + scale=np.linalg.norm(feet_GRF[leg_name]) * 0.005, + color=np.array([0, 1, 0, 0.5]), + geom_id=feet_GRF_geom_ids[leg_name], + ) env.render() last_render_time = time.time() - # Reset the environment if the episode is terminated ------------------------------------------------ if env.step_num >= N_STEPS_PER_EPISODE or is_terminated or is_truncated: if is_terminated: @@ -267,26 +309,25 @@ def run_simulation(process=0, num_episodes=500, return_dict=None, seed_number=0, else: state_obs_history.append(ep_state_obs_history) ctrl_state_history.append(ep_ctrl_state_history) - - env.reset(random=True) - quadrupedpympc_wrapper.reset(initial_feet_pos = env.feet_pos(frame='world')) - - if(return_dict is not None): - return_dict['process'+str(process)+'_ctrl_state_history_ep'+str(episode_num)] = np.array(ctrl_state_history).reshape(-1, len(ctrl_state)) - if(is_terminated or is_truncated): - return_dict['process'+str(process)+'_success_rate_ep'+str(episode_num)] = 0 + + env.reset(random=True) + quadrupedpympc_wrapper.reset(initial_feet_pos=env.feet_pos(frame="world")) + + if return_dict is not None: + return_dict["process" + str(process) + "_ctrl_state_history_ep" + str(episode_num)] = np.array( + ctrl_state_history + ).reshape(-1, len(ctrl_state)) + if is_terminated or is_truncated: + return_dict["process" + str(process) + "_success_rate_ep" + str(episode_num)] = 0 else: - return_dict['process'+str(process)+'_success_rate_ep'+str(episode_num)] = 1 - - break + return_dict["process" + str(process) + "_success_rate_ep" + str(episode_num)] = 1 + break env.close() return return_dict - - -if __name__ == '__main__': +if __name__ == "__main__": run_simulation() - #run_simulation(num_episodes=1, render=False) \ No newline at end of file + # run_simulation(num_episodes=1, render=False) From 39862e4588881e4d4fd39a02bfc3c19b4e7f49ec Mon Sep 17 00:00:00 2001 From: Daniel Ordonez Date: Wed, 26 Mar 2025 16:51:02 -0500 Subject: [PATCH 2/2] Dataset generation pipeline using the simulation script. This commit introduces changes to the simulation script so it can be used for generating the datasets of sensor observations from controlled episodes of the robots in `gym_quadruped` using controllers from this repo. An example script for generating and saving the data is also provided. --- .gitignore | 15 +- pyproject.toml | 2 +- .../gradient/kinodynamic/kinodynamic_nmpc.py | 1 - .../explicit_swing_trajectory_generator.py | 1 - .../scipy_swing_trajectory_generator.py | 1 - simulation/generate_dataset.py | 55 +++++ simulation/simulation.py | 220 ++++++++++-------- 7 files changed, 189 insertions(+), 106 deletions(-) create mode 100644 simulation/generate_dataset.py diff --git a/.gitignore b/.gitignore index a760536..5d0fafd 100644 --- a/.gitignore +++ b/.gitignore @@ -303,17 +303,14 @@ pyrightconfig.json -quadruped_pympc/controllers/gradient/nominal/c_generated_code_gait_adaptive/ - - - .idea/ -MUJOCO_LOG.TXT - -MUJOCO_LOG.TXT - -cfg/ +# Custom ignore parts. +*MUJOCO_LOG.TXT +quadruped_pympc/controllers/gradient/nominal/c_generated_code_gait_adaptive/ +# Datasets are heavy and shouls be saved to huggingface instead of handled here. +/datasets/ *.pkl +*.h5 diff --git a/pyproject.toml b/pyproject.toml index 248363b..f3944e2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,7 @@ build-backend = "hatchling.build" [project] name = "quadruped_pympc" -version = "0.1.1" +version = "0.1.2" description = "A model predictive controller for quadruped robots based on the single rigid body model and written in python. Gradient-based (acados) or Sampling-based (jax)." authors = [ { name="Giulio Turrisi", email="giulio.turrisi@iit.it" }, diff --git a/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_nmpc.py b/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_nmpc.py index 7df705f..f2c65f9 100644 --- a/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_nmpc.py +++ b/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_nmpc.py @@ -20,7 +20,6 @@ sys.path.append(dir_path + "/../../") - from quadruped_pympc import config diff --git a/quadruped_pympc/helpers/swing_generators/explicit_swing_trajectory_generator.py b/quadruped_pympc/helpers/swing_generators/explicit_swing_trajectory_generator.py index e46a0ca..5b737ee 100644 --- a/quadruped_pympc/helpers/swing_generators/explicit_swing_trajectory_generator.py +++ b/quadruped_pympc/helpers/swing_generators/explicit_swing_trajectory_generator.py @@ -1,4 +1,3 @@ - import matplotlib.pyplot as plt import numpy as np diff --git a/quadruped_pympc/helpers/swing_generators/scipy_swing_trajectory_generator.py b/quadruped_pympc/helpers/swing_generators/scipy_swing_trajectory_generator.py index 0581290..f0abbda 100644 --- a/quadruped_pympc/helpers/swing_generators/scipy_swing_trajectory_generator.py +++ b/quadruped_pympc/helpers/swing_generators/scipy_swing_trajectory_generator.py @@ -1,4 +1,3 @@ - import matplotlib.pyplot as plt import numpy as np from scipy.interpolate import CubicSpline diff --git a/simulation/generate_dataset.py b/simulation/generate_dataset.py new file mode 100644 index 0000000..182813c --- /dev/null +++ b/simulation/generate_dataset.py @@ -0,0 +1,55 @@ +# Authors: +# Giulio Turrisi, Daniel Ordonez +import itertools +import os +import pathlib +import sys +from pprint import pprint + +from gym_quadruped.quadruped_env import QuadrupedEnv +from gym_quadruped.utils.data.h5py import H5Reader + +# TODO: Remove horrible hack +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) +from simulation import run_simulation + +if __name__ == "__main__": + from quadruped_pympc import config as cfg + + qpympc_cfg = cfg + # Custom changes to the config here: + + storage_path = pathlib.Path(__file__).parent.parent / "datasets" + data_path = run_simulation( + qpympc_cfg=qpympc_cfg, + num_seconds_per_episode=1, + num_episodes=5, + render=True, + recording_path=storage_path, + ) + + dataset = H5Reader(file_path=data_path) + + print(f"Number of trajectories: {dataset.len()}") + print("Parameters used to generate the data") + pprint(dataset.env_hparams) + # We can reconstruct the environment used to reproduce the dataset using + reproduced_env = QuadrupedEnv(**dataset.env_hparams) + + # We can also list the observations in the dataset + obs_names = dataset.env_hparams["state_obs_names"] + + print(f"\n\n{str(list(itertools.chain(['-'] * 100)))}\n\n") + print( + f"Dataset stored at: {data_path} \n" + f"Number of trajectories: {dataset.len()} \n" + f"Dataset observations names: \n{obs_names}" + ) + + # And since we use the same names as QuadrupedEnv, we can get the group representations for free + from gym_quadruped.utils.quadruped_utils import configure_observation_space_representations + + obs_reps = configure_observation_space_representations(robot_name=dataset.env_hparams["robot"], obs_names=obs_names) + + for obs_name in obs_names: + print(f"Obs: {obs_name} \n Representation: {obs_reps[obs_name]}") diff --git a/simulation/simulation.py b/simulation/simulation.py index ff2de11..43bea23 100644 --- a/simulation/simulation.py +++ b/simulation/simulation.py @@ -1,9 +1,11 @@ # Description: This script is used to simulate the full model of the robot in mujoco +import pathlib # Authors: # Giulio Turrisi, Daniel Ordonez - import time +from os import PathLike +from pprint import pprint import numpy as np @@ -13,47 +15,38 @@ from gym_quadruped.utils.quadruped_utils import LegsAttr from tqdm import tqdm -# Config imports -from quadruped_pympc import config as cfg - # Helper functions for plotting from quadruped_pympc.helpers.quadruped_utils import plot_swing_mujoco # PyMPC controller imports from quadruped_pympc.quadruped_pympc_wrapper import QuadrupedPyMPC_Wrapper -# HeightMap import -if cfg.simulation_params["visual_foothold_adaptation"] != "blind": - from gym_quadruped.sensors.heightmap import HeightMap - def run_simulation( - process=0, num_episodes=500, return_dict=None, seed_number=0, base_vel_command_type="human", render=True + qpympc_cfg, + process=0, + num_episodes=500, + num_seconds_per_episode=60, + ref_base_lin_vel=(0.0, 4.0), + ref_base_ang_vel=(-0.4, 0.4), + friction_coeff=(0.5, 1.0), + base_vel_command_type="human", + seed=0, + render=True, + recording_path: PathLike = None, ): np.set_printoptions(precision=3, suppress=True) - np.random.seed(seed_number) - - robot_name = cfg.robot - hip_height = cfg.hip_height - robot_leg_joints = cfg.robot_leg_joints - robot_feet_geom_names = cfg.robot_feet_geom_names - scene_name = cfg.simulation_params["scene"] - simulation_dt = cfg.simulation_params["dt"] - - state_observables_names = ( - "base_pos", - "base_lin_vel", - "base_ori_euler_xyz", - "base_ori_quat_wxyz", - "base_ang_vel", - "qpos_js", - "qvel_js", - "tau_ctrl_setpoint", - "feet_pos_base", - "feet_vel_base", - "contact_state", - "contact_forces_base", - ) + np.random.seed(seed) + + robot_name = qpympc_cfg.robot + hip_height = qpympc_cfg.hip_height + robot_leg_joints = qpympc_cfg.robot_leg_joints + robot_feet_geom_names = qpympc_cfg.robot_feet_geom_names + scene_name = qpympc_cfg.simulation_params["scene"] + simulation_dt = qpympc_cfg.simulation_params["dt"] + + # Save all observables available. + state_obs_names = list(QuadrupedEnv.ALL_OBS) # + list(IMU.ALL_OBS) # Create the quadruped robot environment ----------------------------------------------------------- env = QuadrupedEnv( @@ -63,26 +56,26 @@ def run_simulation( feet_geom_name=robot_feet_geom_names, # Geom/Frame id of feet scene=scene_name, sim_dt=simulation_dt, - ref_base_lin_vel=(0, 0.6), # pass a float for a fixed value - ref_base_ang_vel=(-0.2, 0.2), # pass a float for a fixed value - ground_friction_coeff=1.5, # pass a float for a fixed value + ref_base_lin_vel=np.asarray(ref_base_lin_vel) * hip_height, # pass a float for a fixed value + ref_base_ang_vel=ref_base_ang_vel, # pass a float for a fixed value + ground_friction_coeff=friction_coeff, # pass a float for a fixed value base_vel_command_type=base_vel_command_type, # "forward", "random", "forward+rotate", "human" - state_obs_names=state_observables_names, # Desired quantities in the 'state' vec + state_obs_names=tuple(state_obs_names), # Desired quantities in the 'state' vec ) + pprint(env.get_hyperparameters()) # Some robots require a change in the zero joint-space configuration. If provided apply it - if cfg.qpos0_js is not None: - env.mjModel.qpos0 = np.concatenate((env.mjModel.qpos0[:7], cfg.qpos0_js)) + if qpympc_cfg.qpos0_js is not None: + env.mjModel.qpos0 = np.concatenate((env.mjModel.qpos0[:7], qpympc_cfg.qpos0_js)) env.reset(random=False) if render: env.render() # Pass in the first render call any mujoco.viewer.KeyCallbackType # Initialization of variables used in the main control loop -------------------------------- - # Jacobian matrices jac_feet_prev = LegsAttr(*[np.zeros((3, env.mjModel.nv)) for _ in range(4)]) - jac_feet_dot = LegsAttr(*[np.zeros((3, env.mjModel.nv)) for _ in range(4)]) + # jac_feet_dot = LegsAttr(*[np.zeros((3, env.mjModel.nv)) for _ in range(4)]) # Torque vector tau = LegsAttr(*[np.zeros((env.mjModel.nv, 1)) for _ in range(4)]) # Torque limits @@ -95,12 +88,13 @@ def run_simulation( ) # Feet positions and Legs order - feet_pos = None feet_traj_geom_ids, feet_GRF_geom_ids = None, LegsAttr(FL=-1, FR=-1, RL=-1, RR=-1) legs_order = ["FL", "FR", "RL", "RR"] # Create HeightMap ----------------------------------------------------------------------- - if cfg.simulation_params["visual_foothold_adaptation"] != "blind": + if qpympc_cfg.simulation_params["visual_foothold_adaptation"] != "blind": + from gym_quadruped.sensors.heightmap import HeightMap + resolution_vfa = 0.04 dimension_vfa = 7 heightmaps = LegsAttr( @@ -121,8 +115,7 @@ def run_simulation( heightmaps = None # Quadruped PyMPC controller initialization ------------------------------------------------------------- - mpc_frequency = cfg.simulation_params["mpc_frequency"] - + # mpc_frequency = qpympc_cfg.simulation_params["mpc_frequency"] quadrupedpympc_observables_names = ( "ref_base_height", "ref_base_angles", @@ -132,26 +125,48 @@ def run_simulation( "swing_time", "phase_signal", "lift_off_positions", + # "base_lin_vel_err", + # "base_ang_vel_err", + # "base_poz_z_err", ) quadrupedpympc_wrapper = QuadrupedPyMPC_Wrapper( initial_feet_pos=env.feet_pos, - legs_order=legs_order, + legs_order=tuple(legs_order), quadrupedpympc_observables_names=quadrupedpympc_observables_names, ) - # -------------------------------------------------------------- + # Data recording ------------------------------------------------------------------------------------------- + if recording_path is not None: + from gym_quadruped.utils.data.h5py import H5Writer + + root_path = pathlib.Path(recording_path) + assert root_path.exists(), f"Recording path does not exist: {root_path.absolute()}" + dataset_path = ( + root_path + / f"{robot_name}/{scene_name}" + / f"lin_vel={ref_base_lin_vel} ang_vel={ref_base_ang_vel} friction={friction_coeff}" + / f"ep={num_episodes}_steps={int(num_seconds_per_episode // simulation_dt):d}.h5" + ) + h5py_writer = H5Writer( + file_path=dataset_path, + env=env, + extra_obs=None, # TODO: Make this automatically configured. Not hardcoded + ) + print(f"\n Recording data to: {dataset_path.absolute()}") + else: + h5py_writer = None + + # ----------------------------------------------------------------------------------------------------------- RENDER_FREQ = 30 # Hz N_EPISODES = num_episodes - N_STEPS_PER_EPISODE = 2000 if env.base_vel_command_type != "human" else 20000 + N_STEPS_PER_EPISODE = int(num_seconds_per_episode // simulation_dt) last_render_time = time.time() state_obs_history, ctrl_state_history = [], [] - for episode_num in tqdm(range(N_EPISODES), desc="Episodes"): - ep_state_obs_history, ep_ctrl_state_history = [], [] - for _ in range(N_STEPS_PER_EPISODE): - step_start = time.time() - + for episode_num in range(N_EPISODES): + ep_state_history, ep_ctrl_state_history, ep_time = [], [], [] + for _ in tqdm(range(N_STEPS_PER_EPISODE), desc=f"Ep:{episode_num:d}-steps:", total=N_STEPS_PER_EPISODE): # Update value from SE or Simulator ---------------------- feet_pos = env.feet_pos(frame="world") hip_pos = env.hip_positions(frame="world") @@ -165,33 +180,29 @@ def run_simulation( ref_base_lin_vel, ref_base_ang_vel = env.target_base_vel() # Get the inertia matrix - if cfg.simulation_params["use_inertia_recomputation"]: + if qpympc_cfg.simulation_params["use_inertia_recomputation"]: inertia = env.get_base_inertia().flatten() # Reflected inertia of base at qpos, in world frame else: - inertia = cfg.inertia.flatten() + inertia = qpympc_cfg.inertia.flatten() # Get the qpos and qvel qpos, qvel = env.mjData.qpos, env.mjData.qvel - joints_pos = LegsAttr(FL=qpos[7:10], FR=qpos[10:13], RL=qpos[13:16], RR=qpos[16:19]) + # Idx of the leg + legs_qvel_idx = env.legs_qvel_idx # leg_name: [idx1, idx2, idx3] ... + legs_qpos_idx = env.legs_qpos_idx # leg_name: [idx1, idx2, idx3] ... + joints_pos = LegsAttr(FL=legs_qvel_idx.FL, FR=legs_qvel_idx.FR, RL=legs_qvel_idx.RL, RR=legs_qvel_idx.RR) # Get Centrifugal, Coriolis, Gravity for the swing controller legs_mass_matrix = env.legs_mass_matrix legs_qfrc_bias = env.legs_qfrc_bias - # Compute feet jacobian feet_jac = env.feet_jacobians(frame="world", return_rot_jac=False) - # Compute jacobian derivatives of the contact points jac_feet_dot = (feet_jac - jac_feet_prev) / simulation_dt # Finite difference approximation jac_feet_prev = feet_jac # Update previous Jacobians - # Compute feet velocities feet_vel = LegsAttr(**{leg_name: feet_jac[leg_name] @ env.mjData.qvel for leg_name in legs_order}) - # Idx of the leg - legs_qvel_idx = env.legs_qvel_idx - legs_qpos_idx = env.legs_qpos_idx - # Quadruped PyMPC controller -------------------------------------------------------------- tau = quadrupedpympc_wrapper.compute_actions( com_pos, @@ -225,9 +236,7 @@ def run_simulation( tau_min, tau_max = tau_limits[leg][:, 0], tau_limits[leg][:, 1] tau[leg] = np.clip(tau[leg], tau_min, tau_max) - quadrupedpympc_observables = quadrupedpympc_wrapper.get_obs() - - # Set control and mujoco step ---------------------------------------------------------------------- + # Set control and mujoco step ------------------------------------------------------------------------- action = np.zeros(env.mjModel.nu) action[env.legs_tau_idx.FL] = tau.FL action[env.legs_tau_idx.FR] = tau.FR @@ -237,16 +246,22 @@ def run_simulation( # action_noise = np.random.normal(0, 2, size=env.mjModel.nu) # action += action_noise + # Apply the action to the environment and evolve sim -------------------------------------------------- state, reward, is_terminated, is_truncated, info = env.step(action=action) + # Get Controller state observables + ctrl_state = quadrupedpympc_wrapper.get_obs() + # Store the history of observations and control ------------------------------------------------------- - ep_state_obs_history.append(state) - base_lin_vel_err = ref_base_lin_vel - base_lin_vel - base_ang_vel_err = ref_base_ang_vel - base_ang_vel - base_poz_z_err = quadrupedpympc_observables["ref_base_height"] - base_pos[2] - ctrl_state = np.concatenate( - (base_lin_vel_err, base_ang_vel_err, [base_poz_z_err], quadrupedpympc_observables["phase_signal"]) - ) + # base_lin_vel_err = ref_base_lin_vel - base_lin_vel + # base_ang_vel_err = ref_base_ang_vel - base_ang_vel + base_poz_z_err = ctrl_state["ref_base_height"] - base_pos[2] + # ctrl_state["base_lin_vel_err"] = base_lin_vel_err + # ctrl_state["base_ang_vel_err"] = base_ang_vel_err + ctrl_state["base_poz_z_err"] = base_poz_z_err + + ep_state_history.append(state) + ep_time.append(env.simulation_time) ep_ctrl_state_history.append(ctrl_state) # Render only at a certain frequency ----------------------------------------------------------------- @@ -259,19 +274,19 @@ def run_simulation( swing_traj_controller=quadrupedpympc_wrapper.wb_interface.stc, swing_period=quadrupedpympc_wrapper.wb_interface.stc.swing_period, swing_time=LegsAttr( - FL=quadrupedpympc_observables["swing_time"][0], - FR=quadrupedpympc_observables["swing_time"][1], - RL=quadrupedpympc_observables["swing_time"][2], - RR=quadrupedpympc_observables["swing_time"][3], + FL=ctrl_state["swing_time"][0], + FR=ctrl_state["swing_time"][1], + RL=ctrl_state["swing_time"][2], + RR=ctrl_state["swing_time"][3], ), - lift_off_positions=quadrupedpympc_observables["lift_off_positions"], - nmpc_footholds=quadrupedpympc_observables["nmpc_footholds"], - ref_feet_pos=quadrupedpympc_observables["ref_feet_pos"], + lift_off_positions=ctrl_state["lift_off_positions"], + nmpc_footholds=ctrl_state["nmpc_footholds"], + ref_feet_pos=ctrl_state["ref_feet_pos"], geom_ids=feet_traj_geom_ids, ) # Update and Plot the heightmap - if cfg.simulation_params["visual_foothold_adaptation"] != "blind": + if qpympc_cfg.simulation_params["visual_foothold_adaptation"] != "blind": # if(stc.check_apex_condition(current_contact, interval=0.01)): for leg_id, leg_name in enumerate(legs_order): data = heightmaps[ @@ -307,27 +322,46 @@ def run_simulation( if is_terminated: print("Environment terminated") else: - state_obs_history.append(ep_state_obs_history) + state_obs_history.append(ep_state_history) ctrl_state_history.append(ep_ctrl_state_history) env.reset(random=True) quadrupedpympc_wrapper.reset(initial_feet_pos=env.feet_pos(frame="world")) - if return_dict is not None: - return_dict["process" + str(process) + "_ctrl_state_history_ep" + str(episode_num)] = np.array( - ctrl_state_history - ).reshape(-1, len(ctrl_state)) - if is_terminated or is_truncated: - return_dict["process" + str(process) + "_success_rate_ep" + str(episode_num)] = 0 - else: - return_dict["process" + str(process) + "_success_rate_ep" + str(episode_num)] = 1 - - break + if h5py_writer is not None: # Save episode trajectory data to disk. + ep_obs_history = collate_obs(ep_state_history) # | collate_obs(ep_ctrl_state_history) + ep_traj_time = np.asarray(ep_time)[:, np.newaxis] + h5py_writer.append_trajectory(state_obs_traj=ep_obs_history, time=ep_traj_time) env.close() - return return_dict + if h5py_writer is not None: + return h5py_writer.file_path + + +def collate_obs(list_of_dicts) -> dict[str, np.ndarray]: + """Collates a list of dictionaries containing observation names and numpy arrays + into a single dictionary of stacked numpy arrays. + """ + if not list_of_dicts: + raise ValueError("Input list is empty.") + + # Get all keys (assumes all dicts have the same keys) + keys = list_of_dicts[0].keys() + + # Stack the values per key + collated = {key: np.stack([d[key] for d in list_of_dicts], axis=0) for key in keys} + collated = {key: v[:, None] if v.ndim == 1 else v for key, v in collated.items()} + return collated if __name__ == "__main__": - run_simulation() + from quadruped_pympc import config as cfg + + qpympc_cfg = cfg + # Custom changes to the config here: + pass + + # Run the simulation with the desired configuration..... + run_simulation(qpympc_cfg=qpympc_cfg) + # run_simulation(num_episodes=1, render=False)