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 0916cfb..63882dd 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/config.py b/quadruped_pympc/config.py index 4afd49f..42b9d87 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'): 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'): 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,119 +117,92 @@ # '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 = { @@ -198,32 +215,25 @@ '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 7b44e82..979cad4 100644 --- a/quadruped_pympc/controllers/gradient/collaborative/centroidal_model_collaborative.py +++ b/quadruped_pympc/controllers/gradient/collaborative/centroidal_model_collaborative.py @@ -3,10 +3,9 @@ # 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 @@ -17,7 +16,7 @@ import sys sys.path.append(dir_path) -sys.path.append(dir_path + '/../../') +sys.path.append(dir_path + "/../../") from quadruped_pympc import config @@ -165,7 +164,7 @@ def __init__(self) -> None: 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: """ diff --git a/quadruped_pympc/controllers/gradient/collaborative/centroidal_nmpc_collaborative.py b/quadruped_pympc/controllers/gradient/collaborative/centroidal_nmpc_collaborative.py index 192d2eb..48dc351 100644 --- a/quadruped_pympc/controllers/gradient/collaborative/centroidal_nmpc_collaborative.py +++ b/quadruped_pympc/controllers/gradient/collaborative/centroidal_nmpc_collaborative.py @@ -5,13 +5,14 @@ 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 @@ -20,7 +21,7 @@ import sys sys.path.append(dir_path) -sys.path.append(dir_path + '/../../') +sys.path.append(dir_path + "/../../") from quadruped_pympc import config @@ -485,8 +486,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", @@ -1220,22 +1221,22 @@ def compute_control( 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'] 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'] 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]]) @@ -1311,7 +1312,7 @@ def compute_control( # 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 @@ -1541,13 +1542,13 @@ def compute_control( # 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 control = self.acados_ocp_solver.get(0, "u") 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 93d0cf4..4ef0acf 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,9 +3,10 @@ # 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 @@ -157,7 +158,7 @@ def __init__(self) -> None: 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: """ 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 f567eeb..390c804 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 @@ -2,7 +2,10 @@ # Authors: Giulio Turrisi - +import copy import os + +import casadi as cs import numpy as np import scipy.linalg import casadi as cs @@ -14,6 +17,7 @@ ACADOS_INFTY = ACADOS_INFTY = 1000 from quadruped_pympc import config + from .centroidal_model_input_rates import Centroidal_Model_InputRates @@ -47,7 +51,7 @@ def __init__(self): 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)) @@ -249,8 +253,8 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: # 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 @@ -483,8 +487,8 @@ 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", @@ -1233,14 +1237,14 @@ def compute_control( 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!!! @@ -1309,7 +1313,7 @@ def compute_control( # 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 @@ -1519,7 +1523,7 @@ def compute_control( # 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')) diff --git a/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_model.py b/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_model.py index 44b363b..fd8872d 100644 --- a/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_model.py +++ b/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_model.py @@ -3,10 +3,9 @@ # 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 @@ -17,7 +16,7 @@ import sys sys.path.append(dir_path) -sys.path.append(dir_path + '/../../') +sys.path.append(dir_path + "/../../") from liecasadi import SO3 @@ -29,8 +28,8 @@ 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 diff --git a/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_nmpc.py b/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_nmpc.py index 304ac42..bcb9673 100644 --- a/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_nmpc.py +++ b/quadruped_pympc/controllers/gradient/kinodynamic/kinodynamic_nmpc.py @@ -2,15 +2,15 @@ # 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 @@ -19,7 +19,7 @@ import sys sys.path.append(dir_path) -sys.path.append(dir_path + '/../../') +sys.path.append(dir_path + "/../../") from quadruped_pympc import config @@ -41,7 +41,7 @@ def __init__(self): 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)) @@ -506,8 +506,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", @@ -1207,14 +1207,14 @@ def compute_control( # 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!! @@ -1266,7 +1266,7 @@ def compute_control( # 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 @@ -1420,13 +1420,13 @@ def compute_control( # 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 control = self.acados_ocp_solver.get(0, "u") diff --git a/quadruped_pympc/controllers/gradient/lyapunov/centroidal_model_lyapunov.py b/quadruped_pympc/controllers/gradient/lyapunov/centroidal_model_lyapunov.py index 684fed1..e4f9747 100644 --- a/quadruped_pympc/controllers/gradient/lyapunov/centroidal_model_lyapunov.py +++ b/quadruped_pympc/controllers/gradient/lyapunov/centroidal_model_lyapunov.py @@ -3,10 +3,9 @@ # 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 @@ -17,7 +16,7 @@ import sys sys.path.append(dir_path) -sys.path.append(dir_path + '/../../../') +sys.path.append(dir_path + "/../../../") import config @@ -170,7 +169,7 @@ def __init__(self) -> None: 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: """ diff --git a/quadruped_pympc/controllers/gradient/lyapunov/centroidal_nmpc_lyapunov.py b/quadruped_pympc/controllers/gradient/lyapunov/centroidal_nmpc_lyapunov.py index 39e00bb..e4562bf 100644 --- a/quadruped_pympc/controllers/gradient/lyapunov/centroidal_nmpc_lyapunov.py +++ b/quadruped_pympc/controllers/gradient/lyapunov/centroidal_nmpc_lyapunov.py @@ -2,22 +2,22 @@ import pathlib # Authors: Giulio Turrisi - +from acados_template import AcadosOcp, AcadosOcpSolver from acados_template import AcadosOcp, AcadosOcpSolver ACADOS_INFTY = 1000 +import copy + +import casadi as cs +import config import numpy as np import scipy.linalg -import casadi as cs -import copy 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): @@ -33,9 +33,9 @@ def __init__(self): 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)) @@ -60,7 +60,7 @@ 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( @@ -75,7 +75,7 @@ def __init__(self): 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 @@ -231,8 +231,8 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: # 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 @@ -506,8 +506,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", @@ -1297,13 +1297,13 @@ def compute_control( 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!!! @@ -1336,19 +1336,19 @@ def compute_control( # Fill last step horizon reference (self.states_dim - no control action!!) 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 @@ -1568,7 +1568,7 @@ def compute_control( # 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')) diff --git a/quadruped_pympc/controllers/gradient/nominal/centroidal_model_nominal.py b/quadruped_pympc/controllers/gradient/nominal/centroidal_model_nominal.py index 2f2846a..52803c6 100644 --- a/quadruped_pympc/controllers/gradient/nominal/centroidal_model_nominal.py +++ b/quadruped_pympc/controllers/gradient/nominal/centroidal_model_nominal.py @@ -3,10 +3,9 @@ # 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 @@ -17,7 +16,7 @@ import sys sys.path.append(dir_path) -sys.path.append(dir_path + '/../../') +sys.path.append(dir_path + "/../../") from quadruped_pympc import config @@ -157,7 +156,7 @@ def __init__(self) -> None: 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: """ 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 a6a1ede..34db3b5 100644 --- a/quadruped_pympc/controllers/gradient/nominal/centroidal_nmpc_gait_adaptive.py +++ b/quadruped_pympc/controllers/gradient/nominal/centroidal_nmpc_gait_adaptive.py @@ -2,18 +2,19 @@ # 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 @@ -23,16 +24,16 @@ 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_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)) @@ -56,8 +57,8 @@ def __init__(self): ) # 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 @@ -214,8 +215,8 @@ def create_ocp_solver_description(self, acados_model, num_threads_in_batch_solve # 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 @@ -241,7 +242,7 @@ def create_ocp_solver_description(self, acados_model, num_threads_in_batch_solve 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": @@ -446,8 +447,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", @@ -935,7 +936,7 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h 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: @@ -1085,20 +1086,20 @@ def compute_batch_control( # 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 @@ -1122,7 +1123,7 @@ def compute_batch_control( # 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"] state_acados = np.concatenate( ( @@ -1167,10 +1168,10 @@ def compute_batch_control( 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) diff --git a/quadruped_pympc/controllers/gradient/nominal/centroidal_nmpc_nominal.py b/quadruped_pympc/controllers/gradient/nominal/centroidal_nmpc_nominal.py index 00f228f..4732d10 100644 --- a/quadruped_pympc/controllers/gradient/nominal/centroidal_nmpc_nominal.py +++ b/quadruped_pympc/controllers/gradient/nominal/centroidal_nmpc_nominal.py @@ -6,14 +6,16 @@ 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: @@ -21,18 +23,18 @@ 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_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,7 +55,7 @@ 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( @@ -68,7 +70,7 @@ def __init__(self): 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 @@ -207,8 +209,8 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: # 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 @@ -234,7 +236,7 @@ def create_ocp_solver_description(self, acados_model) -> AcadosOcp: 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": @@ -438,8 +440,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", @@ -933,7 +935,7 @@ def set_stage_constraint(self, constraint, state, reference, contact_sequence, h 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: @@ -1162,14 +1164,14 @@ def compute_control( 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!!! @@ -1221,20 +1223,20 @@ def compute_control( # 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 @@ -1439,7 +1441,7 @@ def compute_control( # 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')) @@ -1453,7 +1455,7 @@ def compute_control( 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!) diff --git a/quadruped_pympc/controllers/sampling/centroidal_model_jax.py b/quadruped_pympc/controllers/sampling/centroidal_model_jax.py index 8a443cc..ac197c4 100644 --- a/quadruped_pympc/controllers/sampling/centroidal_model_jax.py +++ b/quadruped_pympc/controllers/sampling/centroidal_model_jax.py @@ -3,6 +3,7 @@ # Authors: Giulio Turrisi - +import os import time import os @@ -12,15 +13,13 @@ import sys sys.path.append(dir_path) -sys.path.append(dir_path + '/../') - -from quadruped_pympc import config +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 +from quadruped_pympc import config dtype_general = 'float32' @@ -37,12 +36,12 @@ def __init__(self, dt, device) -> None: 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 @@ -60,7 +59,7 @@ def __init__(self, dt, device) -> None: ) ) 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) diff --git a/quadruped_pympc/controllers/sampling/centroidal_nmpc_jax.py b/quadruped_pympc/controllers/sampling/centroidal_nmpc_jax.py index f8873cb..fb23072 100644 --- a/quadruped_pympc/controllers/sampling/centroidal_nmpc_jax.py +++ b/quadruped_pympc/controllers/sampling/centroidal_nmpc_jax.py @@ -9,22 +9,20 @@ os.environ['XLA_FLAGS'] = '--xla_gpu_triton_gemm_any=True' 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 dtype_general = 'float32' +dtype_general = "float32" class Sampling_MPC: """This is a small class that implements a sampling based control law""" @@ -44,12 +42,12 @@ def __init__( 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 @@ -59,9 +57,9 @@ def __init__( 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] @@ -229,7 +227,7 @@ def __init__( 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.mu = config.mpc_params["mu"] # maximum allowed z contact forces self.f_z_max = config.mpc_params['grf_max'] 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 2dcb832..e91ddbe 100644 --- a/quadruped_pympc/controllers/sampling/centroidal_nmpc_jax_gait_adaptive.py +++ b/quadruped_pympc/controllers/sampling/centroidal_nmpc_jax_gait_adaptive.py @@ -9,21 +9,17 @@ os.environ['XLA_FLAGS'] = '--xla_gpu_triton_gemm_any=True' 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 dtype_general = 'float32' @@ -38,12 +34,12 @@ def __init__(self, device="gpu"): 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 @@ -53,9 +49,9 @@ def __init__(self, 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] @@ -223,7 +219,7 @@ def __init__(self, device="gpu"): 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.mu = config.mpc_params["mu"] # maximum allowed z contact forces self.f_z_max = config.mpc_params['grf_max'] diff --git a/quadruped_pympc/helpers/foothold_reference_generator.py b/quadruped_pympc/helpers/foothold_reference_generator.py index 05f617d..c6519ca 100644 --- a/quadruped_pympc/helpers/foothold_reference_generator.py +++ b/quadruped_pympc/helpers/foothold_reference_generator.py @@ -1,11 +1,11 @@ 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 @@ -131,7 +131,7 @@ def compute_footholds_reference( 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] @@ -189,14 +189,14 @@ def update_touch_down_positions( 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]) diff --git a/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_numeric.py b/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_numeric.py index 84c0224..1931b23 100644 --- a/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_numeric.py +++ b/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_numeric.py @@ -3,6 +3,7 @@ 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 @@ -11,10 +12,10 @@ # 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 @@ -210,10 +211,10 @@ def compute_solution( 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] 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 bf65d77..067ad4d 100644 --- a/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_numeric_mujoco.py +++ b/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_numeric_mujoco.py @@ -7,15 +7,16 @@ # 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 @@ -28,6 +29,7 @@ from gym_quadruped.quadruped_env import QuadrupedEnv +from quadruped_pympc import config as cfg IT_MAX = 5 DT = 1e-2 @@ -154,10 +156,10 @@ def compute_solution( 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] @@ -173,7 +175,7 @@ def compute_solution( 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) @@ -203,7 +205,7 @@ def compute_solution( 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) diff --git a/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_qp.py b/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_qp.py index c22a791..473a385 100644 --- a/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_qp.py +++ b/quadruped_pympc/helpers/inverse_kinematics/inverse_kinematics_qp.py @@ -3,7 +3,7 @@ np.set_printoptions(precision=3, suppress=True) from numpy.linalg import norm import time -import unittest + import casadi as cs # import example_robot_data as robex @@ -14,11 +14,12 @@ 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 # Mujoco magic + import mujoco import mujoco.viewer @@ -26,7 +27,6 @@ import pinocchio as pin from pinocchio import casadi as cpin -import copy # Class for solving a generic inverse kinematics problem @@ -191,7 +191,7 @@ def callback(self, q: np.ndarray) -> None: 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:] diff --git a/quadruped_pympc/helpers/periodic_gait_generator.py b/quadruped_pympc/helpers/periodic_gait_generator.py index 50e5f17..0258a86 100644 --- a/quadruped_pympc/helpers/periodic_gait_generator.py +++ b/quadruped_pympc/helpers/periodic_gait_generator.py @@ -1,8 +1,8 @@ -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 +from quadruped_pympc.helpers.quadruped_utils import GaitType class PeriodicGaitGenerator: diff --git a/quadruped_pympc/helpers/periodic_gait_generator_jax.py b/quadruped_pympc/helpers/periodic_gait_generator_jax.py index 8fc5f65..615d112 100644 --- a/quadruped_pympc/helpers/periodic_gait_generator_jax.py +++ b/quadruped_pympc/helpers/periodic_gait_generator_jax.py @@ -1,5 +1,5 @@ -import copy -import numpy as np +import os + import jax import jax.numpy as jnp @@ -11,6 +11,8 @@ sys.path.append(dir_path + '/../') +sys.path.append(dir_path + "/../") + # Parameters for both MPC and simulation from quadruped_pympc import config diff --git a/quadruped_pympc/helpers/quadruped_utils.py b/quadruped_pympc/helpers/quadruped_utils.py index 25fd368..9eeafb3 100644 --- a/quadruped_pympc/helpers/quadruped_utils.py +++ b/quadruped_pympc/helpers/quadruped_utils.py @@ -59,14 +59,14 @@ def plot_swing_mujoco( 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 @@ -101,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) diff --git a/quadruped_pympc/helpers/srb_inertia_computation.py b/quadruped_pympc/helpers/srb_inertia_computation.py index 21d3e6a..644e8eb 100644 --- a/quadruped_pympc/helpers/srb_inertia_computation.py +++ b/quadruped_pympc/helpers/srb_inertia_computation.py @@ -10,7 +10,7 @@ import sys -sys.path.append(dir_path + '/../') +sys.path.append(dir_path + "/../") # Parameters for both MPC and simulation from quadruped_pympc import config @@ -55,7 +55,7 @@ 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) 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 4f269ab..1638f3a 100644 --- a/quadruped_pympc/helpers/swing_generators/explicit_swing_trajectory_generator.py +++ b/quadruped_pympc/helpers/swing_generators/explicit_swing_trajectory_generator.py @@ -1,9 +1,5 @@ -import numpy as np import matplotlib.pyplot as plt - - -from scipy.interpolate import CubicSpline, Akima1DInterpolator, CubicHermiteSpline -import copy +import numpy as np class SwingTrajectoryGenerator: @@ -17,11 +13,11 @@ 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): 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 ad0f2d1..69ef7f5 100644 --- a/quadruped_pympc/helpers/swing_generators/ndcurves_swing_trajectory_generator.py +++ b/quadruped_pympc/helpers/swing_generators/ndcurves_swing_trajectory_generator.py @@ -7,6 +7,10 @@ # 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: @@ -48,11 +52,11 @@ 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): 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 3800842..7d1ecf4 100644 --- a/quadruped_pympc/helpers/swing_generators/scipy_swing_trajectory_generator.py +++ b/quadruped_pympc/helpers/swing_generators/scipy_swing_trajectory_generator.py @@ -1,9 +1,6 @@ -import numpy as np import matplotlib.pyplot as plt - - -from scipy.interpolate import CubicSpline, Akima1DInterpolator, CubicHermiteSpline -import copy +import numpy as np +from scipy.interpolate import CubicSpline class SwingTrajectoryGenerator: @@ -97,11 +94,11 @@ 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): diff --git a/quadruped_pympc/helpers/swing_trajectory_controller.py b/quadruped_pympc/helpers/swing_trajectory_controller.py index 438a5e1..82c9609 100644 --- a/quadruped_pympc/helpers/swing_trajectory_controller.py +++ b/quadruped_pympc/helpers/swing_trajectory_controller.py @@ -1,4 +1,3 @@ -import os import numpy as np @@ -153,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 @@ -165,8 +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) @@ -174,8 +172,8 @@ 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'] + scene_name = cfg.simulation_params["scene"] + simulation_dt = cfg.simulation_params["dt"] state_observables_names = ( 'base_pos', @@ -219,8 +217,6 @@ def check_full_stance_condition(self, current_contact): 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) @@ -241,10 +237,10 @@ def check_full_stance_condition(self, current_contact): 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 diff --git a/quadruped_pympc/helpers/terrain_estimator.py b/quadruped_pympc/helpers/terrain_estimator.py index a20c956..a96faa8 100644 --- a/quadruped_pympc/helpers/terrain_estimator.py +++ b/quadruped_pympc/helpers/terrain_estimator.py @@ -36,10 +36,10 @@ def compute_terrain_estimation( 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? @@ -70,10 +70,10 @@ def compute_terrain_estimation( 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] + \ diff --git a/quadruped_pympc/helpers/velocity_modulator.py b/quadruped_pympc/helpers/velocity_modulator.py index a62cc7a..b2c1599 100644 --- a/quadruped_pympc/helpers/velocity_modulator.py +++ b/quadruped_pympc/helpers/velocity_modulator.py @@ -1,4 +1,5 @@ import numpy as np + from quadruped_pympc import config as cfg diff --git a/quadruped_pympc/interfaces/srbd_batched_controller_interface.py b/quadruped_pympc/interfaces/srbd_batched_controller_interface.py index c412b02..7bc8d23 100644 --- a/quadruped_pympc/interfaces/srbd_batched_controller_interface.py +++ b/quadruped_pympc/interfaces/srbd_batched_controller_interface.py @@ -1,8 +1,8 @@ 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: diff --git a/quadruped_pympc/interfaces/srbd_controller_interface.py b/quadruped_pympc/interfaces/srbd_controller_interface.py index c5e2880..f9e779c 100644 --- a/quadruped_pympc/interfaces/srbd_controller_interface.py +++ b/quadruped_pympc/interfaces/srbd_controller_interface.py @@ -1,5 +1,4 @@ import numpy as np - from gym_quadruped.utils.quadruped_utils import LegsAttr from quadruped_pympc import config as cfg @@ -25,7 +24,7 @@ def __init__(self): # '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() @@ -75,7 +74,7 @@ def __init__(self): 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: @@ -241,6 +240,6 @@ def compute_control( ) 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')) diff --git a/quadruped_pympc/interfaces/wb_interface.py b/quadruped_pympc/interfaces/wb_interface.py index 5598292..c2271b5 100644 --- a/quadruped_pympc/interfaces/wb_interface.py +++ b/quadruped_pympc/interfaces/wb_interface.py @@ -1,16 +1,16 @@ -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': @@ -241,9 +241,9 @@ def update_state_and_reference( ) 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 ------------------------------------------------------------------------ @@ -272,10 +272,10 @@ def update_state_and_reference( # In the case of the kinodynamic model, # 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']): diff --git a/quadruped_pympc/quadruped_pympc_wrapper.py b/quadruped_pympc/quadruped_pympc_wrapper.py index 3d74b82..54faba5 100644 --- a/quadruped_pympc/quadruped_pympc_wrapper.py +++ b/quadruped_pympc/quadruped_pympc_wrapper.py @@ -1,14 +1,12 @@ -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: @@ -28,7 +26,7 @@ def __init__( 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() diff --git a/simulation/batched_simulations.py b/simulation/batched_simulations.py index 1243060..8410942 100644 --- a/simulation/batched_simulations.py +++ b/simulation/batched_simulations.py @@ -5,10 +5,10 @@ # 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 @@ -18,7 +18,7 @@ NUM_EPISODES = 20 -if __name__ == '__main__': +if __name__ == "__main__": render = False render_only_first = True output_simulations_tracking = None @@ -50,25 +50,25 @@ 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) + "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)], + 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)]] + [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)]]), + np.array([return_dict["process" + str(proc_j) + "_success_rate_ep" + str(ep_i)]]), ) ) diff --git a/simulation/generate_dataset.py b/simulation/generate_dataset.py new file mode 100644 index 0000000..0beac2e --- /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=False, + 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 23187e2..43bea23 100644 --- a/simulation/simulation.py +++ b/simulation/simulation.py @@ -1,59 +1,52 @@ # 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 -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 -from quadruped_pympc import config as cfg +from tqdm import tqdm # 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': - 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,42 +115,63 @@ 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', - 'ref_feet_pos', - 'nmpc_GRFs', - 'nmpc_footholds', - 'swing_time', - 'phase_signal', - 'lift_off_positions', + "ref_base_height", + "ref_base_angles", + "ref_feet_pos", + "nmpc_GRFs", + "nmpc_footholds", + "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') - 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 @@ -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) - + 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')) + 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__": + 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) -if __name__ == '__main__': - run_simulation() # run_simulation(num_episodes=1, render=False)