diff --git a/quadruped_pympc/interfaces/srbd_batched_controller_interface.py b/quadruped_pympc/interfaces/srbd_batched_controller_interface.py index f8f6bee..5f8e6ba 100644 --- a/quadruped_pympc/interfaces/srbd_batched_controller_interface.py +++ b/quadruped_pympc/interfaces/srbd_batched_controller_interface.py @@ -4,25 +4,26 @@ from gym_quadruped.utils.quadruped_utils import LegsAttr -from quadruped_pympc import config as cfg +# from quadruped_pympc import config as cfg class SRBDBatchedControllerInterface: """This is an interface for a batched controller that uses the SRBD method to optimize the gait""" - def __init__(self, ): + def __init__(self, mpc_cfg): """ Constructor for the SRBD batched controller interface """ - - self.type = cfg.mpc_params['type'] - self.mpc_dt = cfg.mpc_params['dt'] - self.horizon = cfg.mpc_params['horizon'] - self.optimize_step_freq = cfg.mpc_params['optimize_step_freq'] - self.num_parallel_computations = cfg.mpc_params['num_parallel_computations'] - self.sampling_method = cfg.mpc_params['sampling_method'] - self.control_parametrization = cfg.mpc_params['control_parametrization'] - self.num_sampling_iterations = cfg.mpc_params['num_sampling_iterations'] - self.sigma_cem_mppi = cfg.mpc_params['sigma_cem_mppi'] - self.step_freq_available = cfg.mpc_params['step_freq_available'] + self._mpc_cfg = mpc_cfg + + self.type = self._mpc_cfg['type'] + self.mpc_dt = self._mpc_cfg['dt'] + self.horizon = self._mpc_cfg['horizon'] + self.optimize_step_freq = self._mpc_cfg['optimize_step_freq'] + self.num_parallel_computations = self._mpc_cfg['num_parallel_computations'] + self.sampling_method = self._mpc_cfg['sampling_method'] + self.control_parametrization = self._mpc_cfg['control_parametrization'] + self.num_sampling_iterations = self._mpc_cfg['num_sampling_iterations'] + self.sigma_cem_mppi = self._mpc_cfg['sigma_cem_mppi'] + self.step_freq_available = self._mpc_cfg['step_freq_available'] from quadruped_pympc.controllers.gradient.nominal.centroidal_nmpc_gait_adaptive import \ diff --git a/quadruped_pympc/interfaces/srbd_controller_interface.py b/quadruped_pympc/interfaces/srbd_controller_interface.py index c98e18e..9064fa4 100644 --- a/quadruped_pympc/interfaces/srbd_controller_interface.py +++ b/quadruped_pympc/interfaces/srbd_controller_interface.py @@ -2,20 +2,19 @@ from gym_quadruped.utils.quadruped_utils import LegsAttr -from quadruped_pympc import config as cfg class SRBDControllerInterface: """This is an interface for a controller that uses the SRBD method to optimize the gait""" - def __init__(self, ): + def __init__(self, mpc_cfg): """ Constructor for the SRBD controller interface """ - - self.type = cfg.mpc_params['type'] - self.mpc_dt = cfg.mpc_params['dt'] - self.horizon = cfg.mpc_params['horizon'] - self.optimize_step_freq = cfg.mpc_params['optimize_step_freq'] - self.step_freq_available = cfg.mpc_params['step_freq_available'] + self._mpc_cfg = mpc_cfg + self.type = self._mpc_cfg['type'] + self.mpc_dt = self._mpc_cfg['dt'] + self.horizon = self._mpc_cfg['horizon'] + self.optimize_step_freq = self._mpc_cfg['optimize_step_freq'] + self.step_freq_available = self._mpc_cfg['step_freq_available'] self.previous_contact_mpc = np.array([1, 1, 1, 1]) @@ -46,12 +45,12 @@ def __init__(self, ): self.batched_controller = Acados_NMPC_GaitAdaptive() - elif(cfg.mpc_params['type'] == 'lyapunov'): + elif(self._mpc_cfg['type'] == 'lyapunov'): from quadruped_pympc.controllers.gradient.lyapunov.centroidal_nmpc_lyapunov import Acados_NMPC_Lyapunov self.controller = Acados_NMPC_Lyapunov() - if cfg.mpc_params['optimize_step_freq']: + if self._mpc_cfg['optimize_step_freq']: from quadruped_pympc.controllers.gradient.nominal.centroidal_nmpc_gait_adaptive import Acados_NMPC_GaitAdaptive self.batched_controller = Acados_NMPC_GaitAdaptive() @@ -120,7 +119,7 @@ def compute_control(self, self.controller = self.controller.with_newkey() if (self.controller.sampling_method == 'cem_mppi'): if (iter_sampling == 0): - self.controller = self.controller.with_newsigma(cfg.mpc_params['sigma_cem_mppi']) + self.controller = self.controller.with_newsigma(self._mpc_cfg['sigma_cem_mppi']) nmpc_GRFs, \ nmpc_footholds, \ diff --git a/quadruped_pympc/interfaces/wb_interface.py b/quadruped_pympc/interfaces/wb_interface.py index 9fdfe06..8c7858e 100644 --- a/quadruped_pympc/interfaces/wb_interface.py +++ b/quadruped_pympc/interfaces/wb_interface.py @@ -2,15 +2,12 @@ import time from gym_quadruped.utils.quadruped_utils import LegsAttr -from quadruped_pympc import config as cfg from quadruped_pympc.helpers.foothold_reference_generator import FootholdReferenceGenerator 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 -if(cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): - from quadruped_pympc.helpers.visual_foothold_adaptation import VisualFootholdAdaptation class WBInterface: @@ -21,6 +18,8 @@ class WBInterface: """ def __init__(self, + mpc_cfg, + sim_cfg, initial_feet_pos: LegsAttr, legs_order: tuple[str, str, str, str] = ('FL', 'FR', 'RL', 'RR')): """ Constructor of the WBInterface class @@ -30,14 +29,16 @@ def __init__(self, legs_order (tuple[str, str, str, str], optional): order of the leg. Defaults to ('FL', 'FR', 'RL', 'RR'). """ + self._mpc_cfg = mpc_cfg + self._sim_cfg = sim_cfg - mpc_dt = cfg.mpc_params['dt'] - horizon = cfg.mpc_params['horizon'] + mpc_dt = self._mpc_cfg['dt'] + horizon = self._mpc_cfg['horizon'] self.legs_order = legs_order # Periodic gait generator -------------------------------------------------------------- - gait_name = cfg.simulation_params['gait'] - gait_params = cfg.simulation_params['gait_params'][gait_name] + gait_name = self._sim_cfg['gait'] + gait_params = self._sim_cfg['gait_params'][gait_name] gait_type, duty_factor, step_frequency = gait_params['type'], gait_params['duty_factor'], gait_params['step_freq'] # Given the possibility to use nonuniform discretization, # we generate a contact sequence two times longer and with a dt half of the one of the mpc @@ -46,9 +47,9 @@ def __init__(self, gait_type=gait_type, horizon=horizon) # in the case of nonuniform discretization, we create a list of dts and horizons for each nonuniform discretization - if (cfg.mpc_params['use_nonuniform_discretization']): - self.contact_sequence_dts = [cfg.mpc_params['dt_fine_grained'], mpc_dt] - self.contact_sequence_lenghts = [cfg.mpc_params['horizon_fine_grained'], horizon] + if (self._mpc_cfg['use_nonuniform_discretization']): + self.contact_sequence_dts = [self._mpc_cfg['dt_fine_grained'], mpc_dt] + self.contact_sequence_lenghts = [self._mpc_cfg['horizon_fine_grained'], horizon] else: self.contact_sequence_dts = [mpc_dt] self.contact_sequence_lenghts = [horizon] @@ -56,15 +57,15 @@ def __init__(self, # Create the foothold reference generator ------------------------------------------------ stance_time = (1 / self.pgg.step_freq) * self.pgg.duty_factor - self.frg = FootholdReferenceGenerator(stance_time=stance_time, hip_height=cfg.hip_height, lift_off_positions=initial_feet_pos) + self.frg = FootholdReferenceGenerator(stance_time=stance_time, hip_height=self._sim_cfg['ref_z'], lift_off_positions=initial_feet_pos) # Create swing trajectory generator ------------------------------------------------------ - self.step_height = cfg.simulation_params['step_height'] + self.step_height = self._sim_cfg['step_height'] swing_period = (1 - self.pgg.duty_factor) * (1 / self.pgg.step_freq) - position_gain_fb = cfg.simulation_params['swing_position_gain_fb'] - velocity_gain_fb = cfg.simulation_params['swing_velocity_gain_fb'] - swing_generator = cfg.simulation_params['swing_generator'] + position_gain_fb = self._sim_cfg['swing_position_gain_fb'] + velocity_gain_fb = self._sim_cfg['swing_velocity_gain_fb'] + swing_generator = self._sim_cfg['swing_generator'] self.stc = SwingTrajectoryController(step_height=self.step_height, swing_period=swing_period, position_gain_fb=position_gain_fb, velocity_gain_fb=velocity_gain_fb, generator=swing_generator) @@ -73,9 +74,12 @@ def __init__(self, # Terrain estimator ----------------------------------------------------------------------- self.terrain_computation = TerrainEstimator() - if(cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): - # Visual foothold adaptation ------------------------------------------------------------- - self.vfa = VisualFootholdAdaptation(legs_order=self.legs_order, adaptation_strategy=cfg.simulation_params['visual_foothold_adaptation']) + # Visual foothold adaptation ------------------------------------------------------------- + if(self._sim_cfg['visual_foothold_adaptation'] != 'blind'): + from quadruped_pympc.helpers.visual_foothold_adaptation import VisualFootholdAdaptation + self.vfa = VisualFootholdAdaptation( + legs_order=self.legs_order, adaptation_strategy=self._sim_cfg['visual_foothold_adaptation'] + ) self.current_contact = np.array([1, 1, 1, 1]) @@ -154,11 +158,11 @@ def update_state_and_reference(self, base_xy_lin_vel=base_lin_vel[0:2], ref_base_xy_lin_vel=ref_base_lin_vel[0:2], hips_position=hip_pos, - com_height_nominal=cfg.simulation_params['ref_z']) + com_height_nominal=self._sim_cfg['ref_z']) # Adjust the footholds given the terrain ----------------------------------------------------- - if(cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): + if(self._sim_cfg['visual_foothold_adaptation'] != 'blind'): time_adaptation = time.time() if(self.stc.check_apex_condition(self.current_contact, interval=0.01) and self.vfa.initialized == False): @@ -177,12 +181,12 @@ def update_state_and_reference(self, for _, leg_name in enumerate(legs_order): #TODO this should be in the terrain frame ref_feet_pos_height_diffence = ref_feet_pos[leg_name][2] - self.frg.lift_off_positions[leg_name][2] - if(ref_feet_pos_height_diffence > (cfg.simulation_params['step_height']) and ref_feet_pos_height_diffence > 0.0): + if(ref_feet_pos_height_diffence > (self._sim_cfg['step_height']) and ref_feet_pos_height_diffence > 0.0): step_height = ref_feet_pos_height_diffence + 0.05 step_height_modification = True self.stc.regenerate_swing_trajectory_generator(step_height=step_height, swing_period=self.stc.swing_period) if(step_height_modification == False): - step_height = cfg.simulation_params['step_height'] + step_height = self._sim_cfg['step_height'] self.stc.regenerate_swing_trajectory_generator(step_height=step_height, swing_period=self.stc.swing_period)""" else: @@ -198,8 +202,8 @@ def update_state_and_reference(self, feet_pos=self.frg.lift_off_positions, current_contact=self.current_contact) - ref_pos = np.array([0, 0, cfg.hip_height]) - ref_pos[2] = cfg.simulation_params['ref_z'] + terrain_height + ref_pos = np.array([0, 0, self._sim_cfg['ref_z']]) + ref_pos[2] = self._sim_cfg['ref_z'] + terrain_height # Update state reference ------------------------------------------------------------------------ @@ -217,7 +221,7 @@ def update_state_and_reference(self, # ------------------------------------------------------------------------------------------------- - if(cfg.mpc_params['optimize_step_freq']): + if(self._mpc_cfg['optimize_step_freq']): # we can always optimize the step freq, or just at the apex of the swing # to avoid possible jittering in the solution optimize_swing = self.stc.check_apex_condition(self.current_contact) @@ -322,7 +326,7 @@ def reset(self, #self.stc.reset() #self.terrain_computation.reset() self.frg.lift_off_positions = initial_feet_pos - if(cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): + if(self._sim_cfg['visual_foothold_adaptation'] != 'blind'): self.vfa.reset() self.current_contact = np.array([1, 1, 1, 1]) return \ No newline at end of file diff --git a/quadruped_pympc/quadruped_pympc_wrapper.py b/quadruped_pympc/quadruped_pympc_wrapper.py index fea50fd..41f5b14 100644 --- a/quadruped_pympc/quadruped_pympc_wrapper.py +++ b/quadruped_pympc/quadruped_pympc_wrapper.py @@ -3,73 +3,73 @@ from quadruped_pympc.interfaces.wb_interface import WBInterface from gym_quadruped.utils.quadruped_utils import LegsAttr -from quadruped_pympc import config as cfg +# from quadruped_pympc import config as cfg # TODO: bad modularity import numpy as np - _DEFAULT_OBS = ('ref_base_height', 'ref_base_angles', 'nmpc_GRFs', 'nmpc_footholds', 'swing_time') class QuadrupedPyMPC_Wrapper: """A simple class wrapper of all the mpc submodules (swing, contact generator, mpc itself).""" - def __init__(self, - initial_feet_pos: LegsAttr, - legs_order: tuple[str, str, str, str] = ('FL', 'FR', 'RL', 'RR'), - quadrupedpympc_observables_names: tuple[str, ...] = _DEFAULT_OBS): + def __init__(self, + mpc_cfg, + sim_cfg, + initial_feet_pos: LegsAttr, + legs_order: tuple[str, str, str, str] = ('FL', 'FR', 'RL', 'RR'), + mpc_observables_names: tuple[str, ...] = _DEFAULT_OBS): """ Constructor of the QuadrupedPyMPC_Wrapper class. Args: initial_feet_pos (LegsAttr): initial feet positions, otherwise they will be all zero. legs_order (tuple[str, str, str, str], optional): order of the leg. Defaults to ('FL', 'FR', 'RL', 'RR'). - quadrupedpympc_observables_names (tuple[str, ...], optional): list of observable to save. Defaults to _DEFAULT_OBS. - """ + mpc_observables_names (tuple[str, ...], optional): list of observable to save. Defaults to + _DEFAULT_OBS. + """ + self._mpc_cfg = mpc_cfg + self._sim_cfg = sim_cfg - self.mpc_frequency = cfg.simulation_params['mpc_frequency'] + self.mpc_frequency = self._sim_cfg['mpc_frequency'] - self.srbd_controller_interface = SRBDControllerInterface() + self.srbd_controller_interface = SRBDControllerInterface(mpc_cfg=self._mpc_cfg) - if(cfg.mpc_params['type'] != 'sampling' and cfg.mpc_params['optimize_step_freq']): - self.srbd_batched_controller_interface = SRBDBatchedControllerInterface() + if self._mpc_cfg['type'] != 'sampling' and self._mpc_cfg['optimize_step_freq']: + self.srbd_batched_controller_interface = SRBDBatchedControllerInterface(mpc_cfg=self._mpc_cfg) - self.wb_interface = WBInterface(initial_feet_pos = initial_feet_pos(frame='world'), - legs_order = legs_order) - + self.wb_interface = WBInterface(mpc_cfg=self._mpc_cfg, + sim_cfg=self._sim_cfg, + initial_feet_pos=initial_feet_pos(frame='world'), + legs_order=legs_order) - self.nmpc_GRFs = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), - RL=np.zeros(3), RR=np.zeros(3)) - self.nmpc_footholds = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), - RL=np.zeros(3), RR=np.zeros(3)) + self.nmpc_GRFs = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), RL=np.zeros(3), RR=np.zeros(3)) + self.nmpc_footholds = LegsAttr(FL=np.zeros(3), FR=np.zeros(3), RL=np.zeros(3), RR=np.zeros(3)) self.best_sample_freq = self.wb_interface.pgg.step_freq - - self.quadrupedpympc_observables_names = quadrupedpympc_observables_names + self.quadrupedpympc_observables_names = mpc_observables_names self.quadrupedpympc_observables = {} - - - - def compute_actions(self, - base_pos: np.ndarray, - base_lin_vel: np.ndarray, - base_ori_euler_xyz: np.ndarray, - base_ang_vel: np.ndarray, - feet_pos: LegsAttr, - hip_pos: LegsAttr, - heightmaps, - legs_order: tuple[str, str, str, str], - simulation_dt: float, - ref_base_lin_vel: np.ndarray, - ref_base_ang_vel: np.ndarray, - step_num: int, - qvel: np.ndarray, - feet_jac: LegsAttr, + + def compute_actions(self, + base_pos: np.ndarray, + base_lin_vel: np.ndarray, + base_ori_euler_xyz: np.ndarray, + base_ang_vel: np.ndarray, + feet_pos: LegsAttr, + hip_pos: LegsAttr, + heightmaps, + legs_order: tuple[str, str, str, str], + simulation_dt: float, + ref_base_lin_vel: np.ndarray, + ref_base_ang_vel: np.ndarray, + step_num: int, + qvel: np.ndarray, + feet_jac: LegsAttr, jac_feet_dot: LegsAttr, - feet_vel: LegsAttr, - legs_qfrc_bias: LegsAttr, - legs_mass_matrix: LegsAttr, - legs_qvel_idx: LegsAttr, - tau: LegsAttr, + feet_vel: LegsAttr, + legs_qfrc_bias: LegsAttr, + legs_mass_matrix: LegsAttr, + legs_qvel_idx: LegsAttr, + tau: LegsAttr, inertia: np.ndarray) -> LegsAttr: """ Given the current state of the robot (and the reference), compute the torques to be applied to the motors. @@ -99,85 +99,73 @@ def compute_actions(self, Returns: LegsAttr: torques to be applied to the motors - """ + """ - # Update the state and reference ------------------------- state_current, \ - ref_state, \ - contact_sequence, \ - ref_feet_pos, \ - ref_feet_constraints, \ - contact_sequence_dts, \ - contact_sequence_lenghts, \ - step_height, \ - optimize_swing = self.wb_interface.update_state_and_reference(base_pos, - base_lin_vel, - base_ori_euler_xyz, - base_ang_vel, - feet_pos, - hip_pos, - heightmaps, - legs_order, - simulation_dt, - ref_base_lin_vel, - ref_base_ang_vel) - - - + ref_state, \ + contact_sequence, \ + ref_feet_pos, \ + ref_feet_constraints, \ + contact_sequence_dts, \ + contact_sequence_lenghts, \ + step_height, \ + optimize_swing = self.wb_interface.update_state_and_reference(base_pos, + base_lin_vel, + base_ori_euler_xyz, + base_ang_vel, + feet_pos, + hip_pos, + heightmaps, + legs_order, + simulation_dt, + ref_base_lin_vel, + ref_base_ang_vel) # Solve OCP --------------------------------------------------------------------------------------- if step_num % round(1 / (self.mpc_frequency * simulation_dt)) == 0: - self.nmpc_GRFs, \ - self.nmpc_footholds, \ - self.best_sample_freq = self.srbd_controller_interface.compute_control(state_current, - ref_state, - contact_sequence, - inertia, - self.wb_interface.pgg, - ref_feet_pos, - contact_sequence_dts, - contact_sequence_lenghts, - step_height, - optimize_swing) - - + self.nmpc_GRFs, \ + self.nmpc_footholds, \ + self.best_sample_freq = self.srbd_controller_interface.compute_control(state_current, + ref_state, + contact_sequence, + inertia, + self.wb_interface.pgg, + ref_feet_pos, + contact_sequence_dts, + contact_sequence_lenghts, + step_height, + optimize_swing) # Update the gait - if(cfg.mpc_params['type'] != 'sampling' and cfg.mpc_params['optimize_step_freq']): + if (self._mpc_cfg['type'] != 'sampling' and self._mpc_cfg['optimize_step_freq']): self.best_sample_freq = self.srbd_batched_controller_interface.optimize_gait(state_current, - ref_state, - contact_sequence, - inertia, - self.wb_interface.pgg, - ref_feet_pos, - contact_sequence_dts, - contact_sequence_lenghts, - step_height, - optimize_swing) - - - - + ref_state, + contact_sequence, + inertia, + self.wb_interface.pgg, + ref_feet_pos, + contact_sequence_dts, + contact_sequence_lenghts, + step_height, + optimize_swing) + # Compute Swing and Stance Torque --------------------------------------------------------------------------- tau = self.wb_interface.compute_stance_and_swing_torque(simulation_dt, - qvel, - feet_jac, - jac_feet_dot, - feet_pos, - feet_vel, - legs_qfrc_bias, - legs_mass_matrix, - self.nmpc_GRFs, - self.nmpc_footholds, - legs_qvel_idx, - tau, - optimize_swing, - self.best_sample_freq) - - - + qvel, + feet_jac, + jac_feet_dot, + feet_pos, + feet_vel, + legs_qfrc_bias, + legs_mass_matrix, + self.nmpc_GRFs, + self.nmpc_footholds, + legs_qvel_idx, + tau, + optimize_swing, + self.best_sample_freq) # Save some observables ------------------------------------------------------------------------------------- self.quadrupedpympc_observables = {} @@ -200,19 +188,16 @@ def compute_actions(self, data = {'phase_signal': self.wb_interface.pgg._phase_signal} elif obs_name == 'lift_off_positions': data = {'lift_off_positions': self.wb_interface.frg.lift_off_positions} - + else: data = {} raise ValueError(f"Unknown observable name: {obs_name}") - + self.quadrupedpympc_observables.update(data) - return tau - - - def get_obs(self,) -> dict: + def get_obs(self, ) -> dict: """ Get some user-defined observables from withing the control loop. Returns: @@ -220,13 +205,8 @@ def get_obs(self,) -> dict: """ return self.quadrupedpympc_observables - - - def reset(self, + def reset(self, initial_feet_pos: LegsAttr): """ Reset the controller.""" self.wb_interface.reset(initial_feet_pos) - - - \ No newline at end of file diff --git a/simulation/simulation.py b/simulation/simulation.py index b9ce0ce..2fa1c04 100644 --- a/simulation/simulation.py +++ b/simulation/simulation.py @@ -22,14 +22,6 @@ # 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 - - - - - if __name__ == '__main__': np.set_printoptions(precision=3, suppress=True) @@ -44,7 +36,6 @@ 'qpos_js', 'qvel_js', 'tau_ctrl_setpoint', 'feet_pos_base', 'feet_vel_base', 'contact_state', 'contact_forces_base',) - # Create the quadruped robot environment ----------------------------------------------------------- env = QuadrupedEnv(robot=robot_name, hip_height=hip_height, @@ -58,8 +49,6 @@ state_obs_names=state_observables_names, # Desired quantities in the 'state' vec ) - - # Some robots require a change in the zero joint-space configuration. If provided apply it if cfg.qpos0_js is not None: env.mjModel.qpos0 = np.concatenate((env.mjModel.qpos0[:7], cfg.qpos0_js)) @@ -67,7 +56,6 @@ env.reset(random=False) env.render() # Pass in the first render call any mujoco.viewer.KeyCallbackType - # Initialization of variables used in the main control loop -------------------------------- # Jacobian matrices @@ -75,36 +63,42 @@ 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)]) - + # Feet positions and Legs order feet_pos = None feet_traj_geom_ids, feet_GRF_geom_ids = None, LegsAttr(FL=-1, FR=-1, RL=-1, RR=-1) legs_order = ["FL", "FR", "RL", "RR"] - # Create HeightMap ----------------------------------------------------------------------- - if(cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): + if (cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): + from gym_quadruped.sensors.heightmap import HeightMap + resolution_vfa = 0.04 dimension_vfa = 7 - heightmaps = LegsAttr(FL=HeightMap(n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mjModel=env.mjModel, mjData=env.mjData), - FR=HeightMap(n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mjModel=env.mjModel, mjData=env.mjData), - RL=HeightMap(n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mjModel=env.mjModel, mjData=env.mjData), - RR=HeightMap(n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mjModel=env.mjModel, mjData=env.mjData)) + heightmaps = LegsAttr( + FL=HeightMap(n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mjModel=env.mjModel, + mjData=env.mjData), + FR=HeightMap(n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mjModel=env.mjModel, + mjData=env.mjData), + RL=HeightMap(n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mjModel=env.mjModel, + mjData=env.mjData), + RR=HeightMap(n=dimension_vfa, dist_x=resolution_vfa, dist_y=resolution_vfa, mjModel=env.mjModel, + mjData=env.mjData)) else: heightmaps = None - # Quadruped PyMPC controller initialization ------------------------------------------------------------- mpc_frequency = cfg.simulation_params['mpc_frequency'] + mpc_obs_names = ('ref_base_height', 'ref_base_angles', 'ref_feet_pos', + 'nmpc_GRFs', 'nmpc_footholds', + 'swing_time', 'phase_signal', 'lift_off_positions') - quadrupedpympc_observables_names = ('ref_base_height', 'ref_base_angles', 'ref_feet_pos', - 'nmpc_GRFs', 'nmpc_footholds', - 'swing_time', 'phase_signal', 'lift_off_positions') - - quadrupedpympc_wrapper = QuadrupedPyMPC_Wrapper(initial_feet_pos=env.feet_pos, legs_order=legs_order, - quadrupedpympc_observables_names = quadrupedpympc_observables_names) - + mpc_wrapper = QuadrupedPyMPC_Wrapper(mpc_cfg=cfg.mpc_params, + sim_cfg=cfg.simulation_params, + initial_feet_pos=env.feet_pos, + legs_order=legs_order, + mpc_observables_names=mpc_obs_names) # -------------------------------------------------------------- RENDER_FREQ = 30 # Hz @@ -119,7 +113,6 @@ for _ in range(N_STEPS_PER_EPISODE): step_start = time.time() - # Update value from SE or Simulator ---------------------- feet_pos = env.feet_pos(frame='world') hip_pos = env.hip_positions(frame='world') @@ -127,50 +120,45 @@ base_ang_vel = env.base_ang_vel(frame='world') base_ori_euler_xyz = env.base_ori_euler_xyz base_pos = env.base_pos - + # Get the reference base velocity in the world frame ref_base_lin_vel, ref_base_ang_vel = env.target_base_vel() - + # Get the inertia matrix - if(cfg.simulation_params['use_inertia_recomputation']): + if (cfg.simulation_params['use_inertia_recomputation']): inertia = env.get_base_inertia().flatten() # Reflected inertia of base at qpos, in world frame else: inertia = cfg.inertia.flatten() # Get the qpos and qvel qpos, qvel = env.mjData.qpos, env.mjData.qvel - + # Get Centrifugal, Coriolis, Gravity for the swing controller legs_mass_matrix = env.legs_mass_matrix legs_qfrc_bias = env.legs_qfrc_bias # Compute feet jacobian feet_jac = env.feet_jacobians(frame='world', return_rot_jac=False) - + # Compute jacobian derivatives of the contact points jac_feet_dot = (feet_jac - jac_feet_prev) / simulation_dt # Finite difference approximation jac_feet_prev = feet_jac # Update previous Jacobians - + # Compute feet velocities feet_vel = LegsAttr(**{leg_name: feet_jac[leg_name] @ env.mjData.qvel for leg_name in legs_order}) # Idx of the leg legs_qvel_idx = env.legs_qvel_idx - - # Quadruped PyMPC controller -------------------------------------------------------------- - tau = quadrupedpympc_wrapper.compute_actions(base_pos, base_lin_vel, base_ori_euler_xyz, base_ang_vel, - feet_pos, hip_pos, heightmaps, - legs_order, simulation_dt, ref_base_lin_vel, ref_base_ang_vel, - env.step_num, qvel, feet_jac, jac_feet_dot, feet_vel, legs_qfrc_bias, - legs_mass_matrix, legs_qvel_idx, tau, inertia) - - quadrupedpympc_observables = quadrupedpympc_wrapper.get_obs() + tau = mpc_wrapper.compute_actions(base_pos, base_lin_vel, base_ori_euler_xyz, base_ang_vel, + feet_pos, hip_pos, heightmaps, + legs_order, simulation_dt, ref_base_lin_vel, ref_base_ang_vel, + env.step_num, qvel, feet_jac, jac_feet_dot, feet_vel, legs_qfrc_bias, + legs_mass_matrix, legs_qvel_idx, tau, inertia) - + quadrupedpympc_observables = mpc_wrapper.get_obs() - # Set control and mujoco step ---------------------------------------------------------------------- action = np.zeros(env.mjModel.nu) action[env.legs_tau_idx.FL] = tau.FL @@ -183,49 +171,54 @@ state, reward, is_terminated, is_truncated, info = env.step(action=action) - # Store the history of observations and control ------------------------------------------------------- ep_state_obs_history.append(state) base_lin_vel_err = ref_base_lin_vel - base_lin_vel base_ang_vel_err = ref_base_ang_vel - base_ang_vel base_poz_z_err = quadrupedpympc_observables["ref_base_height"] - base_pos[2] - ctrl_state = np.concatenate((base_lin_vel_err, base_ang_vel_err, [base_poz_z_err], quadrupedpympc_observables["phase_signal"])) + ctrl_state = np.concatenate( + (base_lin_vel_err, base_ang_vel_err, [base_poz_z_err], quadrupedpympc_observables["phase_signal"])) ep_ctrl_state_history.append(ctrl_state) - # Render only at a certain frequency ----------------------------------------------------------------- if time.time() - last_render_time > 1.0 / RENDER_FREQ or env.step_num == 1: _, _, feet_GRF = env.feet_contact_state(ground_reaction_forces=True) # Plot the swing trajectory feet_traj_geom_ids = plot_swing_mujoco(viewer=env.viewer, - swing_traj_controller=quadrupedpympc_wrapper.wb_interface.stc, - swing_period=quadrupedpympc_wrapper.wb_interface.stc.swing_period, - swing_time=LegsAttr(FL=quadrupedpympc_observables["swing_time"][0], - FR=quadrupedpympc_observables["swing_time"][1], - RL=quadrupedpympc_observables["swing_time"][2], - RR=quadrupedpympc_observables["swing_time"][3]), - lift_off_positions=quadrupedpympc_observables["lift_off_positions"], + swing_traj_controller=mpc_wrapper.wb_interface.stc, + swing_period=mpc_wrapper.wb_interface.stc.swing_period, + swing_time=LegsAttr( + FL=quadrupedpympc_observables["swing_time"][0], + FR=quadrupedpympc_observables["swing_time"][1], + RL=quadrupedpympc_observables["swing_time"][2], + RR=quadrupedpympc_observables["swing_time"][3]), + lift_off_positions=quadrupedpympc_observables[ + "lift_off_positions"], nmpc_footholds=quadrupedpympc_observables["nmpc_footholds"], ref_feet_pos=quadrupedpympc_observables["ref_feet_pos"], geom_ids=feet_traj_geom_ids) - - + # Update and Plot the heightmap - if(cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): - #if(stc.check_apex_condition(current_contact, interval=0.01)): + if (cfg.simulation_params['visual_foothold_adaptation'] != 'blind'): + # if(stc.check_apex_condition(current_contact, interval=0.01)): for leg_id, leg_name in enumerate(legs_order): - data = heightmaps[leg_name].data#.update_height_map(ref_feet_pos[leg_name], yaw=env.base_ori_euler_xyz[2]) - if(data is not None): + data = heightmaps[ + leg_name].data # .update_height_map(ref_feet_pos[leg_name], yaw=env.base_ori_euler_xyz[2]) + if (data is not None): for i in range(data.shape[0]): for j in range(data.shape[1]): - heightmaps[leg_name].geom_ids[i, j] = render_sphere(viewer=env.viewer, - position=([data[i][j][0][0],data[i][j][0][1],data[i][j][0][2]]), - diameter=0.01, - color=[0, 1, 0, .5], - geom_id=heightmaps[leg_name].geom_ids[i, j] - ) - + heightmaps[leg_name].geom_ids[i, j] = render_sphere(viewer=env.viewer, + position=([data[i][j][0][0], + data[i][j][0][1], + data[i][j][0][2]]), + diameter=0.01, + color=[0, 1, 0, .5], + geom_id= + heightmaps[leg_name].geom_ids[ + i, j] + ) + # Plot the GRF for leg_id, leg_name in enumerate(legs_order): feet_GRF_geom_ids[leg_name] = render_vector(env.viewer, @@ -238,7 +231,6 @@ env.render() last_render_time = time.time() - # Reset the environment if the episode is terminated ------------------------------------------------ if env.step_num > N_STEPS_PER_EPISODE or is_terminated or is_truncated: if is_terminated: @@ -247,12 +239,7 @@ state_obs_history.append(ep_state_obs_history) ctrl_state_history.append(ep_ctrl_state_history) env.reset(random=False) - - quadrupedpympc_wrapper.reset(initial_feet_pos = env.feet_pos(frame='world')) - - + mpc_wrapper.reset(initial_feet_pos=env.feet_pos(frame='world')) env.close() - -