diff --git a/Go2Py/__init__.py b/Go2Py/__init__.py index 4aab8b1..f398eaf 100644 --- a/Go2Py/__init__.py +++ b/Go2Py/__init__.py @@ -6,4 +6,4 @@ ASSETS_PATH = os.path.join(os.path.dirname(__file__), "assets") GO2_USD_PATH = os.path.join(os.path.dirname(__file__), "assets/usd/go2.usd") GO2_ISAACSIM_CFG_PATH = os.path.join( os.path.dirname(__file__), "sim/isaacsim/sim_config.yaml" -) \ No newline at end of file +) diff --git a/Go2Py/control/walk_these_ways.py b/Go2Py/control/walk_these_ways.py index 47b7403..b09dcef 100755 --- a/Go2Py/control/walk_these_ways.py +++ b/Go2Py/control/walk_these_ways.py @@ -5,6 +5,7 @@ import time import numpy as np import torch + def loadParameters(path): with open(path + "/parameters.pkl", "rb") as file: pkl_cfg = pkl.load(file) @@ -53,7 +54,7 @@ class HistoryWrapper: privileged_obs = info["privileged_obs"] self.obs_history = torch.cat( - (self.obs_history[:, self.env.num_obs :], obs), dim=-1 + (self.obs_history[:, self.env.num_obs:], obs), dim=-1 ) return ( { @@ -70,7 +71,7 @@ class HistoryWrapper: obs = self.env.get_observations() privileged_obs = self.env.get_privileged_observations() self.obs_history = torch.cat( - (self.obs_history[:, self.env.num_obs :], obs), dim=-1 + (self.obs_history[:, self.env.num_obs:], obs), dim=-1 ) return { "obs": obs, @@ -82,7 +83,7 @@ class HistoryWrapper: obs = self.env.get_obs() privileged_obs = self.env.get_privileged_observations() self.obs_history = torch.cat( - (self.obs_history[:, self.env.num_obs :], obs), dim=-1 + (self.obs_history[:, self.env.num_obs:], obs), dim=-1 ) return { "obs": obs, @@ -130,7 +131,8 @@ class CommandInterface: self.stance_width_cmd = 0.25 def setGaitType(self, gait_type): - assert gait_type in [key for key in self.gaits.keys()], f'The gain type should be in {[key for key in self.gaits.keys()]}' + assert gait_type in [key for key in self.gaits.keys()], f'The gain type should be in { + [key for key in self.gaits.keys()]}' self.gait = torch.tensor(self.gaits[gait_type]) def get_command(self): @@ -249,18 +251,18 @@ class WalkTheseWaysAgent: "RL_calf_joint", ] - policy_to_unitree_map=[] - unitree_to_policy_map=[] - for i,policy_joint_name in enumerate(policy_joint_names): - id = np.where([name==policy_joint_name for name in unitree_joint_names])[0][0] - policy_to_unitree_map.append((i,id)) - self.policy_to_unitree_map=np.array(policy_to_unitree_map) + policy_to_unitree_map = [] + unitree_to_policy_map = [] + for i, policy_joint_name in enumerate(policy_joint_names): + id = np.where([name == policy_joint_name for name in unitree_joint_names])[0][0] + policy_to_unitree_map.append((i, id)) + self.policy_to_unitree_map = np.array(policy_to_unitree_map) + + for i, unitree_joint_name in enumerate(unitree_joint_names): + id = np.where([name == unitree_joint_name for name in policy_joint_names])[0][0] + unitree_to_policy_map.append((i, id)) + self.unitree_to_policy_map = np.array(unitree_to_policy_map) - for i,unitree_joint_name in enumerate(unitree_joint_names): - id = np.where([name==unitree_joint_name for name in policy_joint_names])[0][0] - unitree_to_policy_map.append((i,id)) - self.unitree_to_policy_map=np.array(unitree_to_policy_map) - self.default_dof_pos = np.array( [ self.cfg["init_state"]["default_joint_angles"][name] @@ -342,8 +344,8 @@ class WalkTheseWaysAgent: joint_state = self.robot.getJointStates() if joint_state is not None: self.gravity_vector = self.robot.getGravityInBody() - self.dof_pos = joint_state['q'][self.unitree_to_policy_map[:,1]] - self.dof_vel = joint_state['dq'][self.unitree_to_policy_map[:,1]] + self.dof_pos = joint_state['q'][self.unitree_to_policy_map[:, 1]] + self.dof_vel = joint_state['dq'][self.unitree_to_policy_map[:, 1]] if reset_timer: self.reset_gait_indices() @@ -403,10 +405,10 @@ class WalkTheseWaysAgent: self.joint_vel_target - self.dof_vel ) * self.d_gains # self.lcm_bridge.sendCommands(command_for_robot) - self.robot.setCommands(self.joint_pos_target[self.policy_to_unitree_map[:,1]],\ - self.joint_vel_target[self.policy_to_unitree_map[:,1]],\ - self.p_gains[self.policy_to_unitree_map[:,1]],\ - self.d_gains[self.policy_to_unitree_map[:,1]], + self.robot.setCommands(self.joint_pos_target[self.policy_to_unitree_map[:, 1]], + self.joint_vel_target[self.policy_to_unitree_map[:, 1]], + self.p_gains[self.policy_to_unitree_map[:, 1]], + self.d_gains[self.policy_to_unitree_map[:, 1]], np.zeros(12)) def reset(self): diff --git a/Go2Py/estimation/contact.py b/Go2Py/estimation/contact.py index bad1480..5ff0503 100644 --- a/Go2Py/estimation/contact.py +++ b/Go2Py/estimation/contact.py @@ -1,14 +1,15 @@ import numpy as np + class HysteresisContactDetector: def __init__(self, upper_limit, lower_limit): self.upper_limit = upper_limit self.lower_limit = lower_limit self.contact_state = np.zeros(4) - + def update(self, contact_forces): - self.contact_state[np.where(contact_forces>self.upper_limit)[0]]=1 - self.contact_state[np.where(contact_forces self.upper_limit)[0]] = 1 + self.contact_state[np.where(contact_forces < self.lower_limit)[0]] = 0 def getContactStates(self): - return self.contact_state \ No newline at end of file + return self.contact_state diff --git a/Go2Py/joy.py b/Go2Py/joy.py index c7f3fd9..cb8de9d 100644 --- a/Go2Py/joy.py +++ b/Go2Py/joy.py @@ -15,6 +15,7 @@ from pygame.locals import ( K_q, ) + @dataclass class xKeySwitch: R1: int @@ -45,12 +46,14 @@ class xRockerBtn: L2: float ly: float + class PyGameJoyManager: def __init__(self, user_callback=None): self.dt = 0.01 # Sampling frequence of the joystick self.running = False self.joy_thread_handle = threading.Thread(target=self.joy_thread) - # an optional callback that can be used to send the reading values to a user provided function + # an optional callback that can be used to send the reading values to a + # user provided function self.user_callback = user_callback pygame.init() self.offset = None @@ -222,4 +225,4 @@ class Logitech3DPro: "11": self.digital_raw[10], "12": self.digital_raw[11], } - return data \ No newline at end of file + return data diff --git a/Go2Py/robot/fsm.py b/Go2Py/robot/fsm.py index 27a885b..a65e61d 100644 --- a/Go2Py/robot/fsm.py +++ b/Go2Py/robot/fsm.py @@ -13,25 +13,29 @@ class FSM: self.user_controller_callback = user_controller_callback self.state = "damping" - self.tracking_kp = np.array(4*[150, 150, 150.]).reshape(12) - self.tracking_kv = np.array(12*[3.]) - self.damping_kv = np.array(12*[2.]) - + self.tracking_kp = np.array(4 * [150, 150, 150.]).reshape(12) + self.tracking_kv = np.array(12 * [3.]) + self.damping_kv = np.array(12 * [2.]) self.tracking_complete = True - self.robot.setCommands(np.zeros(12), np.zeros(12), np.zeros(12), self.damping_kv, np.zeros(12)) - self.fsm_dT = 1./50. - self.control_dT = 1./50. + self.robot.setCommands( + np.zeros(12), + np.zeros(12), + np.zeros(12), + self.damping_kv, + np.zeros(12)) + self.fsm_dT = 1. / 50. + self.control_dT = 1. / 50. self.dT = self.fsm_dT - self.modes = {"tracking":self.trackingControlUpdate, - "damping" :self.dampingControlUpdate, + self.modes = {"tracking": self.trackingControlUpdate, + "damping": self.dampingControlUpdate, "user": self.userControlUpdate, - } + } self.setMode("damping") self.running = True - self.fsm_thread = threading.Thread(target = self.update) + self.fsm_thread = threading.Thread(target=self.update) self.fsm_thread.start() # if the robot is a simulation, create a thread for stepping it if self.robot.simulated: @@ -42,26 +46,40 @@ class FSM: assert mode in self.modes.keys(), 'the requested control update mode is not implemented' self.updateCommands = self.modes[mode] # print(f'setting mode to {mode}') - + def moveTo(self, target, duration=0.5): # assert self.tracking_complete, 'The previous moveTo command is not completed yet!' self.q_start = self.robot.getJointStates()['q'] self.q_target = target self.time = 0.0 self.move_duration = duration - self.q_des = lambda t: [self.q_start + np.clip((t)/self.move_duration,0, 1)*(self.q_target - self.q_start), \ - True if np.clip((t)/self.move_duration,0, 1)==1 else False] # q_des(t), Movement finished + self.q_des = lambda t: [self.q_start + + np.clip((t) / + self.move_duration, 0, 1) * + (self.q_target - + self.q_start), True if np.clip((t) / + self.move_duration, 0, 1) == 1 else False] # q_des(t), Movement finished self.tracking_complete = False self.setMode("tracking") def trackingControlUpdate(self): - self.time +=self.dT + self.time += self.dT q_des, done = self.q_des(self.time) - self.robot.setCommands(q_des, np.zeros(12), self.tracking_kp, self.tracking_kv, np.zeros(12)) + self.robot.setCommands( + q_des, + np.zeros(12), + self.tracking_kp, + self.tracking_kv, + np.zeros(12)) self.tracking_complete = done def dampingControlUpdate(self): - self.robot.setCommands(np.zeros(12), np.zeros(12), np.zeros(12), self.damping_kv, np.zeros(12)) + self.robot.setCommands( + np.zeros(12), + np.zeros(12), + np.zeros(12), + self.damping_kv, + np.zeros(12)) def userControlUpdate(self): if self.user_controller_callback is not None: @@ -71,13 +89,13 @@ class FSM: while self.running: self.robot.step() time.sleep(self.robot.dt) - + def update(self): while self.running: getattr(self, self.state)() time.sleep(self.dT) self.updateCommands() - + def close(self): self.running = False @@ -98,7 +116,7 @@ class FSM: self.state = 'standing' else: self.state = "pre_standing" - + def standing(self): # print("standing") if self.tracking_complete: @@ -113,9 +131,14 @@ class FSM: self.setMode("user") self.dT = self.control_dT self.state = "user_loop" - self.robot.setCommands(np.zeros(12), np.zeros(12), np.zeros(12), np.zeros(12), np.zeros(12)) + self.robot.setCommands( + np.zeros(12), + np.zeros(12), + np.zeros(12), + np.zeros(12), + np.zeros(12)) elif self.remote.standUpDownSeq() or self.robot.overheat(): - self.moveTo(self.robot.sitting_q, duration = 1.5) + self.moveTo(self.robot.sitting_q, duration=1.5) self.state = "sitting" def user_loop(self): @@ -125,7 +148,7 @@ class FSM: self.setMode("damping") elif self.remote.standUpDownSeq() or self.safety.controlTimeout(): self.dT = self.fsm_dT - self.moveTo(self.robot.standing_q, duration = 1) + self.moveTo(self.robot.standing_q, duration=1) self.timer = time.time() self.state = "switch_back_to_locked_stance" else: @@ -140,6 +163,6 @@ class FSM: self.state = "sitting" def switch_back_to_locked_stance(self): - if time.time()-self.timer > 0.5: + if time.time() - self.timer > 0.5: # print("going back to locked stance") - self.state = "locked_stance" \ No newline at end of file + self.state = "locked_stance" diff --git a/Go2Py/robot/interface.py b/Go2Py/robot/interface.py index 1cd7934..1c04064 100644 --- a/Go2Py/robot/interface.py +++ b/Go2Py/robot/interface.py @@ -21,11 +21,10 @@ from scipy.spatial.transform import Rotation from Go2Py.joy import xKeySwitch, xRockerBtn - class GO2Real(): def __init__( self, - mode = 'lowlevel', # 'highlevel' or 'lowlevel' + mode='lowlevel', # 'highlevel' or 'lowlevel' vx_max=0.5, vy_max=0.4, ωz_max=0.5, @@ -33,20 +32,20 @@ class GO2Real(): assert mode in ['highlevel', 'lowlevel'], "mode should be either 'highlevel' or 'lowlevel'" self.mode = mode self.simulated = False - self.prestanding_q = np.array([ 0.0, 1.26186061, -2.5, - 0.0, 1.25883281, -2.5, - 0.0, 1.27193761, -2.6, - 0.0, 1.27148342, -2.6]) + self.prestanding_q = np.array([0.0, 1.26186061, -2.5, + 0.0, 1.25883281, -2.5, + 0.0, 1.27193761, -2.6, + 0.0, 1.27148342, -2.6]) - self.sitting_q = np.array([-0.02495611, 1.26249647, -2.82826662, - 0.04563564, 1.2505368 , -2.7933557 , - -0.30623949, 1.28283751, -2.82314873, - 0.26400229, 1.29355574, -2.84276843]) + self.sitting_q = np.array([-0.02495611, 1.26249647, -2.82826662, + 0.04563564, 1.2505368, -2.7933557, + -0.30623949, 1.28283751, -2.82314873, + 0.26400229, 1.29355574, -2.84276843]) - self.standing_q = np.array([ 0.0, 0.77832842, -1.56065452, - 0.0, 0.76754963, -1.56634164, - 0.0, 0.76681757, -1.53601146, - 0.0, 0.75422204, -1.53229916]) + self.standing_q = np.array([0.0, 0.77832842, -1.56065452, + 0.0, 0.76754963, -1.56634164, + 0.0, 0.76681757, -1.53601146, + 0.0, 0.75422204, -1.53229916]) self.latest_command_stamp = time.time() self.highcmd_topic_name = "rt/go2/twist_cmd" self.lowcmd_topic_name = "rt/go2py/low_cmd" @@ -56,7 +55,7 @@ class GO2Real(): self.participant = DomainParticipant() self.lowstate_topic = Topic(self.participant, self.lowstate_topic_name, Go2pyState_) self.state_reader = DataReader(self.participant, self.lowstate_topic) - + self.lowcmd_topic = Topic(self.participant, self.lowcmd_topic_name, Go2pyLowCmd_) self.lowcmd_writer = DataWriter(self.participant, self.lowcmd_topic) @@ -67,11 +66,11 @@ class GO2Real(): self.P_v_max = np.diag([1 / self.vx_max**2, 1 / self.vy_max**2]) self.ωz_max = ωz_max self.ωz_min = -ωz_max - + self.state = None - self.setCommands = {'lowlevel':self.setCommandsLow, - 'highlevel':self.setCommandsHigh}[self.mode] - self.state_thread = Thread(target = self.state_update) + self.setCommands = {'lowlevel': self.setCommandsLow, + 'highlevel': self.setCommandsHigh}[self.mode] + self.state_thread = Thread(target=self.state_update) self.running = True self.state_thread.start() @@ -88,23 +87,23 @@ class GO2Real(): gyro = self.state.gyro quat = self.state.quat temp = self.state.imu_temp - return {'accel':accel, 'gyro':gyro, 'quat':quat, 'temp':temp} + return {'accel': accel, 'gyro': gyro, 'quat': quat, 'temp': temp} def getFootContacts(self): """Returns the raw foot contact forces""" - footContacts = self.state.contact + footContacts = self.state.contact return footContacts def getJointStates(self): """Returns the joint angles (q) and velocities (dq) of the robot""" - return {'q':self.state.q, - 'dq':self.state.dq, - 'tau_est':self.state.tau, - 'temperature':self.state.motor_temp} + return {'q': self.state.q, + 'dq': self.state.dq, + 'tau_est': self.state.tau, + 'temperature': self.state.motor_temp} def getRemoteState(self): - """A method to get the state of the wireless remote control. - Returns a xRockerBtn object: + """A method to get the state of the wireless remote control. + Returns a xRockerBtn object: - head: [head1, head2] - keySwitch: xKeySwitch object - lx: float @@ -145,7 +144,7 @@ class GO2Real(): v_x = ly * self.vx_max v_y = lx * self.vy_max ω = rx * self.ωz_max - + return v_x, v_y, ω def getBatteryState(self): @@ -166,7 +165,7 @@ class GO2Real(): assert q_des.size == dq_des.size == kp.size == kd.size == tau_ff.size == 12, "q, dq, kp, kd, tau_ff should have size 12" lowcmd = Go2pyLowCmd_( q_des, - dq_des, + dq_des, kp, kd, tau_ff @@ -194,5 +193,5 @@ class GO2Real(): def getGravityInBody(self): q = self.getIMU()['quat'] R = Rotation.from_quat(q).as_matrix() - g_in_body = R.T@np.array([0.0, 0.0, -1.0]).reshape(3, 1) - return g_in_body \ No newline at end of file + g_in_body = R.T @ np.array([0.0, 0.0, -1.0]).reshape(3, 1) + return g_in_body diff --git a/Go2Py/robot/model.py b/Go2Py/robot/model.py index a1f94ed..9bee225 100644 --- a/Go2Py/robot/model.py +++ b/Go2Py/robot/model.py @@ -18,17 +18,19 @@ class Go2Model: q_reordering_idx (np.ndarray): Index array for reordering the joint position vector. ef_J_ (dict): A dictionary storing the Jacobians for the end-effector frames. """ + def __init__(self): """ Initializes the Go2Model class by loading the URDF, setting up the robot model, and calculating initial dimensions. """ - self.robot = pin.RobotWrapper.BuildFromURDF(urdf_path, urdf_root_path, pin.JointModelFreeFlyer()) + self.robot = pin.RobotWrapper.BuildFromURDF( + urdf_path, urdf_root_path, pin.JointModelFreeFlyer()) self.data = self.robot.data # Standing joint configuration in Unitree Joint order self.ef_frames = ['FR_foot', 'FL_foot', 'RR_foot', 'RL_foot'] - self.dq_reordering_idx = np.array([0, 1, 2, 3, 4, 5,\ + self.dq_reordering_idx = np.array([0, 1, 2, 3, 4, 5, 9, 10, 11, 6, 7, 8, 15, 16, 17, 12, 13, 14]) - self.q_reordering_idx = np.array([9, 10, 11, 6, 7, 8, 15, 16, 17, 12, 13, 14])-6 + self.q_reordering_idx = np.array([9, 10, 11, 6, 7, 8, 15, 16, 17, 12, 13, 14]) - 6 self.ef_Jb_ = {} self.ef_Jw_ = {} @@ -49,11 +51,21 @@ class Go2Model: q_neutral = np.asarray([0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) self.robot.framesForwardKinematics(q_neutral) - self.h = np.linalg.norm(self.robot.data.oMf[ID_FR_HAA].translation - self.robot.data.oMf[ID_RR_HAA].translation) - self.b = np.linalg.norm(self.robot.data.oMf[ID_FR_HAA].translation - self.robot.data.oMf[ID_FL_HAA].translation) - self.l1 = np.linalg.norm(self.robot.data.oMf[ID_FR_HAA].translation - self.robot.data.oMf[ID_FR_HFE].translation) - self.l2 = np.linalg.norm(self.robot.data.oMf[ID_FR_HFE].translation - self.robot.data.oMf[ID_FR_KFE].translation) - self.l3 = np.linalg.norm(self.robot.data.oMf[ID_FR_KFE].translation - self.robot.data.oMf[ID_FR_FOOT].translation) + self.h = np.linalg.norm( + self.robot.data.oMf[ID_FR_HAA].translation - + self.robot.data.oMf[ID_RR_HAA].translation) + self.b = np.linalg.norm( + self.robot.data.oMf[ID_FR_HAA].translation - + self.robot.data.oMf[ID_FL_HAA].translation) + self.l1 = np.linalg.norm( + self.robot.data.oMf[ID_FR_HAA].translation - + self.robot.data.oMf[ID_FR_HFE].translation) + self.l2 = np.linalg.norm( + self.robot.data.oMf[ID_FR_HFE].translation - + self.robot.data.oMf[ID_FR_KFE].translation) + self.l3 = np.linalg.norm( + self.robot.data.oMf[ID_FR_KFE].translation - + self.robot.data.oMf[ID_FR_FOOT].translation) self.M_ = None self.Minv_ = None self.nle_ = None @@ -64,7 +76,7 @@ class Go2Model: # print(self.l1) # print(self.l2) # print(self.l3) - + def inverseKinematics(self, T, feet_pos): """ Calculates the inverse kinematics for the robot given a desired state. @@ -77,7 +89,7 @@ class Go2Model: np.ndarray: A numpy array of size 12 representing the joint angles of the legs. """ rB = np.asarray(T[0:3, -1]) # Base position (3D vector) - R = T[:3,:3] # Body orientation (quaternion converted to rotation matrix) + R = T[:3, :3] # Body orientation (quaternion converted to rotation matrix) sx = [1, 1, -1, -1] sy = [-1, 1, -1, 1] @@ -86,7 +98,7 @@ class Go2Model: for i in range(4): r_HB = np.array([sx[i] * self.h / 2, sy[i] * self.b / 2, 0]) # Hip offset (3D vector) - rf = np.asarray(feet_pos[3*i:3*i+3]) # Foot position (3D vector) + rf = np.asarray(feet_pos[3 * i:3 * i + 3]) # Foot position (3D vector) r_fH = R.T @ (rf - rB) - r_HB # Foot relative to hip in body frame (3D vector) x = r_fH[0] @@ -103,21 +115,21 @@ class Go2Model: k1 = self.l2 + self.l3 * c3 k2 = self.l3 * s3 r1 = np.sqrt(k1**2 + k2**2) - t2 = np.arctan2(-x/r1, np.sqrt(et) / r1) - np.arctan2(k2/r1, k1/r1) + t2 = np.arctan2(-x / r1, np.sqrt(et) / r1) - np.arctan2(k2 / r1, k1 / r1) # Theta 1 calculation zv = self.l2 * np.cos(t2) + self.l3 * np.cos(t2 + t3) m1 = sy[i] * self.l1 m2 = -zv r2 = np.sqrt(m1**2 + m2**2) - t1 = np.arctan2(z/r2, y/r2) - np.arctan2(m2/r2, m1/r2) + t1 = np.arctan2(z / r2, y / r2) - np.arctan2(m2 / r2, m1 / r2) - joint_angles[3*i:3*i+3] = np.array([t1, t2, t3]) + joint_angles[3 * i:3 * i + 3] = np.array([t1, t2, t3]) # TODO: Implement joint axis direction multiplication from URDF return joint_angles - + def forwardKinematics(self, T, q): """ Computes the forward kinematics for the robot given a transformation matrix and joint configuration. @@ -131,14 +143,14 @@ class Go2Model: """ q_ = np.hstack([pin.SE3ToXYZQUATtuple(pin.SE3(T)), q[self.q_reordering_idx]]) self.robot.framesForwardKinematics(q_) - ef_frames = ['base_link','FR_foot', 'FL_foot', 'RR_foot', 'RL_foot'] + ef_frames = ['base_link', 'FR_foot', 'FL_foot', 'RR_foot', 'RL_foot'] data = {} - return {frame:self.robot.data.oMf[self.robot.model.getFrameId(frame)].homogeneous \ - for frame in ef_frames} + return {frame: self.robot.data.oMf[self.robot.model.getFrameId(frame)].homogeneous + for frame in ef_frames} def updateKinematics(self, q): """ - Updates the kinematic states. + Updates the kinematic states. Args: q (np.ndarray): A numpy array of size 19 representing the [x, y, z, qx, qy, qz, qw] and joint configurations in FR, FL, RR, RL order. @@ -147,14 +159,18 @@ class Go2Model: self.robot.framesForwardKinematics(q)\ for ef_frame in self.ef_frames: - Jw = self.robot.getFrameJacobian(self.robot.model.getFrameId(ef_frame), pin.ReferenceFrame.LOCAL_WORLD_ALIGNED) - Jb = self.robot.getFrameJacobian(self.robot.model.getFrameId(ef_frame), pin.ReferenceFrame.LOCAL) - self.ef_Jw_[ef_frame]=Jw[:, self.dq_reordering_idx] - self.ef_Jb_[ef_frame]=Jb[:, self.dq_reordering_idx] - + Jw = self.robot.getFrameJacobian( + self.robot.model.getFrameId(ef_frame), + pin.ReferenceFrame.LOCAL_WORLD_ALIGNED) + Jb = self.robot.getFrameJacobian( + self.robot.model.getFrameId(ef_frame), + pin.ReferenceFrame.LOCAL) + self.ef_Jw_[ef_frame] = Jw[:, self.dq_reordering_idx] + self.ef_Jb_[ef_frame] = Jb[:, self.dq_reordering_idx] + def updateKinematicsPose(self, q, T): """ - Updates the kinematic states. + Updates the kinematic states. Args: q (np.ndarray): A numpy array of size 12 representing the joint configurations in FR, FL, RR, RL order. @@ -165,27 +181,32 @@ class Go2Model: self.robot.framesForwardKinematics(q_)\ for ef_frame in self.ef_frames: - Jw = self.robot.getFrameJacobian(self.robot.model.getFrameId(ef_frame), pin.ReferenceFrame.LOCAL_WORLD_ALIGNED) - Jb = self.robot.getFrameJacobian(self.robot.model.getFrameId(ef_frame), pin.ReferenceFrame.LOCAL) - self.ef_Jw_[ef_frame]=Jw[:, self.dq_reordering_idx] - self.ef_Jb_[ef_frame]=Jb[:, self.dq_reordering_idx] + Jw = self.robot.getFrameJacobian( + self.robot.model.getFrameId(ef_frame), + pin.ReferenceFrame.LOCAL_WORLD_ALIGNED) + Jb = self.robot.getFrameJacobian( + self.robot.model.getFrameId(ef_frame), + pin.ReferenceFrame.LOCAL) + self.ef_Jw_[ef_frame] = Jw[:, self.dq_reordering_idx] + self.ef_Jb_[ef_frame] = Jb[:, self.dq_reordering_idx] def updateDynamics(self, q, dq): """ - Updates the dynamical states. + Updates the dynamical states. Args: q (np.ndarray): A numpy array of size 19 representing the [x, y, z, qx, qy, qz, qw] and joint configurations in FR, FL, RR, RL order. dq (np.ndarray): A numpy array of size 18 representing the [vx, vy, vz, wx, wy, wz] and joint configurations in FR, FL, RR, RL order. """ - self.robot.centroidalMomentum(q,dq) + self.robot.centroidalMomentum(q, dq) self.nle_ = self.robot.nle(q, dq)[self.dq_reordering_idx] self.g_ = self.robot.gravity(q)[self.dq_reordering_idx] - self.M_ = self.robot.mass(q)[self.dq_reordering_idx,:] - self.M_ = self.M_[:,self.dq_reordering_idx] - self.Minv_ = pin.computeMinverse(self.robot.model, self.robot.data, q)[self.dq_reordering_idx,:] - self.Minv_ = self.Minv_[:,self.dq_reordering_idx] - + self.M_ = self.robot.mass(q)[self.dq_reordering_idx, :] + self.M_ = self.M_[:, self.dq_reordering_idx] + self.Minv_ = pin.computeMinverse( + self.robot.model, self.robot.data, q)[ + self.dq_reordering_idx, :] + self.Minv_ = self.Minv_[:, self.dq_reordering_idx] def updateAll(q, dq): """ @@ -212,7 +233,7 @@ class Go2Model: dq_ = np.hstack([v, dq[self.q_reordering_idx]]) self.updateKinematics(q_) self.updateDynamics(q_, dq_) - + def getInfo(self): """ Retrieves the current dynamics and kinematic information of the robot. @@ -221,17 +242,19 @@ class Go2Model: dict: A dictionary containing the robot's mass matrix, inverse mass matrix, non-linear effects, gravity vector, and Jacobians for the end-effectors. """ return { - 'M':self.M_, - 'Minv':self.Minv_, - 'nle':self.nle_, - 'g':self.g_, - 'J_w':self.ef_Jw_, - 'J_b':self.ef_Jb_, + 'M': self.M_, + 'Minv': self.Minv_, + 'nle': self.nle_, + 'g': self.g_, + 'J_w': self.ef_Jw_, + 'J_b': self.ef_Jb_, } - + def getGroundReactionForce(self, tau_est, body_acceleration=None): - if body_acceleration is None: - grf = {key:np.linalg.pinv(self.ef_Jw_[key][:3,6:].T)@(tau_est.squeeze() - self.nle_[6:]) for key in self.ef_Jw_.keys()} + if body_acceleration is None: + grf = {key: np.linalg.pinv( + self.ef_Jw_[key][:3, 6:].T) @ (tau_est.squeeze() - self.nle_[6:]) for key in self.ef_Jw_.keys()} else: - raise NotImplementedError("Ground reaction force with body dynamics is not implemented") - return grf \ No newline at end of file + raise NotImplementedError( + "Ground reaction force with body dynamics is not implemented") + return grf diff --git a/Go2Py/robot/remote.py b/Go2Py/robot/remote.py index 419d840..7c006c3 100644 --- a/Go2Py/robot/remote.py +++ b/Go2Py/robot/remote.py @@ -1,19 +1,24 @@ import threading from pynput import keyboard + class BaseRemote: def __init__(self): pass + def startSeq(self): return False + def standUpDownSeq(self): return False def flushStates(self): pass + remote = BaseRemote() + class KeyboardRemote(BaseRemote): def __init__(self): super().__init__() @@ -32,7 +37,6 @@ class KeyboardRemote(BaseRemote): except AttributeError: pass # Special keys (like space) will be handled here - def _on_release(self, key): try: if key.char == 's': # Start sequence @@ -62,6 +66,7 @@ class KeyboardRemote(BaseRemote): self.stand_up_down_seq_flag = False self.start_seq_flag = False + class UnitreeRemote(BaseRemote): def __init__(self, robot): self.robot = robot @@ -86,4 +91,3 @@ class UnitreeRemote(BaseRemote): return True else: return False - \ No newline at end of file diff --git a/Go2Py/robot/safety.py b/Go2Py/robot/safety.py index 4e2c2fe..37a5fe4 100644 --- a/Go2Py/robot/safety.py +++ b/Go2Py/robot/safety.py @@ -1,14 +1,16 @@ import time + + class SafetyHypervisor(): def __init__(self, robot): self.robot = robot def unsafe(self): return False - + def controlTimeout(self): if time.time() - self.robot.latest_command_stamp > 0.1: print('controller timeout') return True else: - return False \ No newline at end of file + return False diff --git a/Go2Py/sim/gym/__init__.py b/Go2Py/sim/gym/__init__.py index 79872ed..23e4aaa 100755 --- a/Go2Py/sim/gym/__init__.py +++ b/Go2Py/sim/gym/__init__.py @@ -2,6 +2,10 @@ import os -MINI_GYM_ROOT_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.realpath(__file__))), '..') +MINI_GYM_ROOT_DIR = os.path.join( + os.path.dirname( + os.path.dirname( + os.path.realpath(__file__))), + '..') print(MINI_GYM_ROOT_DIR) MINI_GYM_ENVS_DIR = os.path.join(MINI_GYM_ROOT_DIR, 'go_gym', 'envs') diff --git a/Go2Py/sim/gym/envs/base/base_task.py b/Go2Py/sim/gym/envs/base/base_task.py index ba68575..5843d40 100755 --- a/Go2Py/sim/gym/envs/base/base_task.py +++ b/Go2Py/sim/gym/envs/base/base_task.py @@ -25,7 +25,8 @@ class BaseTask(gym.Env): sim_device_type, self.sim_device_id = gymutil.parse_device_str(self.sim_device) self.headless = headless - # env device is GPU only if sim is on GPU and use_gpu_pipeline=True, otherwise returned tensors are copied to CPU by physX. + # env device is GPU only if sim is on GPU and use_gpu_pipeline=True, + # otherwise returned tensors are copied to CPU by physX. if sim_device_type == 'cuda' and sim_params.use_gpu_pipeline: self.device = self.sim_device else: @@ -33,7 +34,7 @@ class BaseTask(gym.Env): # graphics device for rendering, -1 for no rendering self.graphics_device_id = self.sim_device_id - if self.headless == True: + if self.headless: self.graphics_device_id = self.sim_device_id self.num_obs = cfg.env.num_observations @@ -54,15 +55,22 @@ class BaseTask(gym.Env): torch._C._jit_set_profiling_executor(False) # allocate buffers - self.obs_buf = torch.zeros(self.num_envs, self.num_obs, device=self.device, dtype=torch.float) + self.obs_buf = torch.zeros( + self.num_envs, + self.num_obs, + device=self.device, + dtype=torch.float) self.rew_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.float) self.rew_buf_pos = torch.zeros(self.num_envs, device=self.device, dtype=torch.float) self.rew_buf_neg = torch.zeros(self.num_envs, device=self.device, dtype=torch.float) self.reset_buf = torch.ones(self.num_envs, device=self.device, dtype=torch.long) self.episode_length_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.long) self.time_out_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) - self.privileged_obs_buf = torch.zeros(self.num_envs, self.num_privileged_obs, device=self.device, - dtype=torch.float) + self.privileged_obs_buf = torch.zeros( + self.num_envs, + self.num_privileged_obs, + device=self.device, + dtype=torch.float) # self.num_privileged_obs = self.num_obs self.extras = {} @@ -76,7 +84,7 @@ class BaseTask(gym.Env): self.viewer = None # if running with a viewer, set up keyboard shortcuts and camera - if self.headless == False: + if not self.headless: # subscribe to keyboard shortcuts self.viewer = self.gym.create_viewer( self.sim, gymapi.CameraProperties()) @@ -132,6 +140,6 @@ class BaseTask(gym.Env): self.gym.poll_viewer_events(self.viewer) def close(self): - if self.headless == False: + if not self.headless: self.gym.destroy_viewer(self.viewer) self.gym.destroy_sim(self.sim) diff --git a/Go2Py/sim/gym/envs/base/curriculum.py b/Go2Py/sim/gym/envs/base/curriculum.py index 5ea70a8..2970a3b 100644 --- a/Go2Py/sim/gym/envs/base/curriculum.py +++ b/Go2Py/sim/gym/envs/base/curriculum.py @@ -32,14 +32,18 @@ class Curriculum: self.indices = indices = {} for key, v_range in key_ranges.items(): bin_size = (v_range[1] - v_range[0]) / v_range[2] - cfg[key] = np.linspace(v_range[0] + bin_size / 2, v_range[1] - bin_size / 2, v_range[2]) - indices[key] = np.linspace(0, v_range[2]-1, v_range[2]) + cfg[key] = np.linspace( + v_range[0] + bin_size / 2, + v_range[1] - bin_size / 2, + v_range[2]) + indices[key] = np.linspace(0, v_range[2] - 1, v_range[2]) self.lows = np.array([range[0] for range in key_ranges.values()]) self.highs = np.array([range[1] for range in key_ranges.values()]) # self.bin_sizes = {key: arr[1] - arr[0] for key, arr in cfg.items()} - self.bin_sizes = {key: (v_range[1] - v_range[0]) / v_range[2] for key, v_range in key_ranges.items()} + self.bin_sizes = {key: (v_range[1] - v_range[0]) / v_range[2] + for key, v_range in key_ranges.items()} self._raw_grid = np.stack(np.meshgrid(*cfg.values(), indexing='ij')) self._idx_grid = np.stack(np.meshgrid(*indices.values(), indexing='ij')) @@ -66,7 +70,7 @@ class Curriculum: def sample_bins(self, batch_size, low=None, high=None): """default to uniform""" - if low is not None and high is not None: # if bounds given + if low is not None and high is not None: # if bounds given valid_inds = np.logical_and( self.grid >= low[:, None], self.grid <= high[:, None] @@ -74,7 +78,7 @@ class Curriculum: temp_weights = np.zeros_like(self.weights) temp_weights[valid_inds] = self.weights[valid_inds] inds = self.rng.choice(self.indices, batch_size, p=temp_weights / temp_weights.sum()) - else: # if no bounds given + else: # if no bounds given inds = self.rng.choice(self.indices, batch_size, p=self.weights / self.weights.sum()) return self.grid.T[inds], inds @@ -82,7 +86,7 @@ class Curriculum: def sample_uniform_from_cell(self, centroids): bin_sizes = np.array([*self.bin_sizes.values()]) low, high = centroids + bin_sizes / 2, centroids - bin_sizes / 2 - return self.rng.uniform(low, high)#.clip(self.lows, self.highs) + return self.rng.uniform(low, high) # .clip(self.lows, self.highs) def sample(self, batch_size, low=None, high=None): cgf_centroid, inds = self.sample_bins(batch_size, low=low, high=high) @@ -145,11 +149,12 @@ class RewardThresholdCurriculum(Curriculum): # if len(is_success) > 0 and is_success.any(): # print("successes") - self.weights[bin_inds[is_success]] = np.clip(self.weights[bin_inds[is_success]] + 0.2, 0, 1) + self.weights[bin_inds[is_success]] = np.clip( + self.weights[bin_inds[is_success]] + 0.2, 0, 1) adjacents = self.get_local_bins(bin_inds[is_success], ranges=local_range) for adjacent in adjacents: - #print(adjacent) - #print(self.grid[:, adjacent]) + # print(adjacent) + # print(self.grid[:, adjacent]) adjacent_inds = np.array(adjacent.nonzero()[0]) self.weights[adjacent_inds] = np.clip(self.weights[adjacent_inds] + 0.2, 0, 1) @@ -158,6 +163,7 @@ class RewardThresholdCurriculum(Curriculum): self.episode_ang_vel_raw[bin_inds] = ang_vel_raw.cpu().numpy() self.episode_duration[bin_inds] = episode_duration.cpu().numpy() + if __name__ == '__main__': r = RewardThresholdCurriculum(100, x=(-1, 1, 5), y=(-1, 1, 2), z=(-1, 1, 11)) @@ -171,9 +177,13 @@ if __name__ == '__main__': for adjacent in adjacents: adjacent_inds = np.array(adjacent.nonzero()[0]) print(adjacent_inds) - r.update(bin_inds=adjacent_inds, lin_vel_rewards=np.ones_like(adjacent_inds), - ang_vel_rewards=np.ones_like(adjacent_inds), lin_vel_threshold=0.0, ang_vel_threshold=0.0, - local_range=0.5) + r.update( + bin_inds=adjacent_inds, + lin_vel_rewards=np.ones_like(adjacent_inds), + ang_vel_rewards=np.ones_like(adjacent_inds), + lin_vel_threshold=0.0, + ang_vel_threshold=0.0, + local_range=0.5) samples, bins = r.sample(10_000) diff --git a/Go2Py/sim/gym/envs/base/legged_robot.py b/Go2Py/sim/gym/envs/base/legged_robot.py index bd8e4cc..5f60f58 100755 --- a/Go2Py/sim/gym/envs/base/legged_robot.py +++ b/Go2Py/sim/gym/envs/base/legged_robot.py @@ -1,5 +1,11 @@ # License: see [LICENSE, LICENSES/legged_gym/LICENSE] +from .legged_robot_config import Cfg +from Go2Py.sim.gym.utils.terrain import Terrain +from Go2Py.sim.gym.utils.math_utils import quat_apply_yaw, wrap_to_pi, get_scale_shift +from Go2Py.sim.gym.envs.base.base_task import BaseTask +from Go2Py.sim.gym import MINI_GYM_ROOT_DIR +import torch import os from typing import Dict @@ -7,13 +13,6 @@ from isaacgym import gymtorch, gymapi, gymutil from isaacgym.torch_utils import * assert gymtorch -import torch - -from Go2Py.sim.gym import MINI_GYM_ROOT_DIR -from Go2Py.sim.gym.envs.base.base_task import BaseTask -from Go2Py.sim.gym.utils.math_utils import quat_apply_yaw, wrap_to_pi, get_scale_shift -from Go2Py.sim.gym.utils.terrain import Terrain -from .legged_robot_config import Cfg class LeggedRobot(BaseTask): @@ -38,7 +37,8 @@ class LeggedRobot(BaseTask): self.debug_viz = False self.init_done = False self.initial_dynamics_dict = initial_dynamics_dict - if eval_cfg is not None: self._parse_cfg(eval_cfg) + if eval_cfg is not None: + self._parse_cfg(eval_cfg) self._parse_cfg(self.cfg) super().__init__(self.cfg, sim_params, physics_engine, sim_device, headless, self.eval_cfg) @@ -89,7 +89,7 @@ class LeggedRobot(BaseTask): def post_physics_step(self): """ check terminations, compute observations and rewards - calls self._post_physics_step_callback() for common computations + calls self._post_physics_step_callback() for common computations calls self._draw_debug_vis() if needed """ self.gym.refresh_actor_root_state_tensor(self.sim) @@ -105,14 +105,16 @@ class LeggedRobot(BaseTask): # prepare quantities self.base_pos[:] = self.root_states[:self.num_envs, 0:3] self.base_quat[:] = self.root_states[:self.num_envs, 3:7] - self.base_lin_vel[:] = quat_rotate_inverse(self.base_quat, self.root_states[:self.num_envs, 7:10]) - self.base_ang_vel[:] = quat_rotate_inverse(self.base_quat, self.root_states[:self.num_envs, 10:13]) + self.base_lin_vel[:] = quat_rotate_inverse( + self.base_quat, self.root_states[:self.num_envs, 7:10]) + self.base_ang_vel[:] = quat_rotate_inverse( + self.base_quat, self.root_states[:self.num_envs, 10:13]) self.projected_gravity[:] = quat_rotate_inverse(self.base_quat, self.gravity_vec) self.foot_velocities = self.rigid_body_state.view(self.num_envs, self.num_bodies, 13 )[:, self.feet_indices, 7:10] - self.foot_positions = self.rigid_body_state.view(self.num_envs, self.num_bodies, 13)[:, self.feet_indices, - 0:3] + self.foot_positions = self.rigid_body_state.view(self.num_envs, self.num_bodies, 13)[ + :, self.feet_indices, 0:3] self._post_physics_step_callback() @@ -138,13 +140,14 @@ class LeggedRobot(BaseTask): def check_termination(self): """ Check if environments need to be reset """ - self.reset_buf = torch.any(torch.norm(self.contact_forces[:, self.termination_contact_indices, :], dim=-1) > 1., - dim=1) - self.time_out_buf = self.episode_length_buf > self.cfg.env.max_episode_length # no terminal reward for time-outs + self.reset_buf = torch.any(torch.norm( + self.contact_forces[:, self.termination_contact_indices, :], dim=-1) > 1., dim=1) + # no terminal reward for time-outs + self.time_out_buf = self.episode_length_buf > self.cfg.env.max_episode_length self.reset_buf |= self.time_out_buf if self.cfg.rewards.use_terminal_body_height: - self.body_height_buf = torch.mean(self.root_states[:, 2].unsqueeze(1) - self.measured_heights, dim=1) \ - < self.cfg.rewards.terminal_body_height + self.body_height_buf = torch.mean(self.root_states[:, 2].unsqueeze( + 1) - self.measured_heights, dim=1) < self.cfg.rewards.terminal_body_height self.reset_buf = torch.logical_or(self.body_height_buf, self.reset_buf) def reset_idx(self, env_ids): @@ -217,11 +220,13 @@ class LeggedRobot(BaseTask): self.extras["train/episode"]["min_command_yaw_vel"] = torch.min(self.commands[:, 2]) self.extras["train/episode"]["max_command_yaw_vel"] = torch.max(self.commands[:, 2]) if self.cfg.commands.num_commands > 9: - self.extras["train/episode"]["min_command_swing_height"] = torch.min(self.commands[:, 9]) - self.extras["train/episode"]["max_command_swing_height"] = torch.max(self.commands[:, 9]) + self.extras["train/episode"]["min_command_swing_height"] = torch.min( + self.commands[:, 9]) + self.extras["train/episode"]["max_command_swing_height"] = torch.max( + self.commands[:, 9]) for curriculum, category in zip(self.curricula, self.category_names): - self.extras["train/episode"][f"command_area_{category}"] = np.sum(curriculum.weights) / \ - curriculum.weights.shape[0] + self.extras["train/episode"][f"command_area_{category}"] = np.sum( + curriculum.weights) / curriculum.weights.shape[0] self.extras["train/episode"]["min_action"] = torch.min(self.actions) self.extras["train/episode"]["max_action"] = torch.max(self.actions) @@ -249,16 +254,14 @@ class LeggedRobot(BaseTask): self.dof_pos[env_ids] = dof_pos self.dof_vel[env_ids] = 0. - self.gym.set_dof_state_tensor_indexed(self.sim, - gymtorch.unwrap_tensor(self.dof_state), - gymtorch.unwrap_tensor(env_ids_int32), len(env_ids_int32)) + self.gym.set_dof_state_tensor_indexed(self.sim, gymtorch.unwrap_tensor( + self.dof_state), gymtorch.unwrap_tensor(env_ids_int32), len(env_ids_int32)) # base position self.root_states[env_ids] = base_state.to(self.device) - self.gym.set_actor_root_state_tensor_indexed(self.sim, - gymtorch.unwrap_tensor(self.root_states), - gymtorch.unwrap_tensor(env_ids_int32), len(env_ids_int32)) + self.gym.set_actor_root_state_tensor_indexed(self.sim, gymtorch.unwrap_tensor( + self.root_states), gymtorch.unwrap_tensor(env_ids_int32), len(env_ids_int32)) def compute_reward(self): """ Compute rewards @@ -283,8 +286,9 @@ class LeggedRobot(BaseTask): self.command_sums[name] += rew if self.cfg.rewards.only_positive_rewards: self.rew_buf[:] = torch.clip(self.rew_buf[:], min=0.) - elif self.cfg.rewards.only_positive_rewards_ji22_style: #TODO: update - self.rew_buf[:] = self.rew_buf_pos[:] * torch.exp(self.rew_buf_neg[:] / self.cfg.rewards.sigma_rew_neg) + elif self.cfg.rewards.only_positive_rewards_ji22_style: # TODO: update + self.rew_buf[:] = self.rew_buf_pos[:] * \ + torch.exp(self.rew_buf_neg[:] / self.cfg.rewards.sigma_rew_neg) self.episode_sums["total"] += self.rew_buf # add termination reward after clipping if "termination" in self.reward_scales: @@ -295,19 +299,23 @@ class LeggedRobot(BaseTask): self.command_sums["lin_vel_raw"] += self.base_lin_vel[:, 0] self.command_sums["ang_vel_raw"] += self.base_ang_vel[:, 2] - self.command_sums["lin_vel_residual"] += (self.base_lin_vel[:, 0] - self.commands[:, 0]) ** 2 - self.command_sums["ang_vel_residual"] += (self.base_ang_vel[:, 2] - self.commands[:, 2]) ** 2 + self.command_sums["lin_vel_residual"] += (self.base_lin_vel[:, + 0] - self.commands[:, 0]) ** 2 + self.command_sums["ang_vel_residual"] += (self.base_ang_vel[:, + 2] - self.commands[:, 2]) ** 2 self.command_sums["ep_timesteps"] += 1 def compute_observations(self): """ Computes observations """ self.obs_buf = torch.cat((self.projected_gravity, - (self.dof_pos[:, :self.num_actuated_dof] - self.default_dof_pos[:, - :self.num_actuated_dof]) * self.obs_scales.dof_pos, - self.dof_vel[:, :self.num_actuated_dof] * self.obs_scales.dof_vel, - self.actions - ), dim=-1) + (self.dof_pos[:, + :self.num_actuated_dof] - self.default_dof_pos[:, + :self.num_actuated_dof]) * self.obs_scales.dof_pos, + self.dof_vel[:, + :self.num_actuated_dof] * self.obs_scales.dof_vel, + self.actions), + dim=-1) # if self.cfg.env.observe_command and not self.cfg.env.observe_height_command: # self.obs_buf = torch.cat((self.projected_gravity, # self.commands[:, :3] * self.commands_scale, @@ -317,13 +325,12 @@ class LeggedRobot(BaseTask): # ), dim=-1) if self.cfg.env.observe_command: - self.obs_buf = torch.cat((self.projected_gravity, - self.commands * self.commands_scale, - (self.dof_pos[:, :self.num_actuated_dof] - self.default_dof_pos[:, - :self.num_actuated_dof]) * self.obs_scales.dof_pos, - self.dof_vel[:, :self.num_actuated_dof] * self.obs_scales.dof_vel, - self.actions - ), dim=-1) + self.obs_buf = torch.cat((self.projected_gravity, self.commands * + self.commands_scale, (self.dof_pos[:, :self.num_actuated_dof] - + self.default_dof_pos[:, :self.num_actuated_dof]) * + self.obs_scales.dof_pos, self.dof_vel[:, :self.num_actuated_dof] * + self.obs_scales.dof_vel, self.actions), dim=- + 1) if self.cfg.env.observe_two_prev_actions: self.obs_buf = torch.cat((self.obs_buf, @@ -343,9 +350,11 @@ class LeggedRobot(BaseTask): if self.cfg.env.observe_vel: if self.cfg.commands.global_reference: - self.obs_buf = torch.cat((self.root_states[:self.num_envs, 7:10] * self.obs_scales.lin_vel, + self.obs_buf = torch.cat((self.root_states[:self.num_envs, + 7:10] * self.obs_scales.lin_vel, self.base_ang_vel * self.obs_scales.ang_vel, - self.obs_buf), dim=-1) + self.obs_buf), + dim=-1) else: self.obs_buf = torch.cat((self.base_lin_vel * self.obs_scales.lin_vel, self.base_ang_vel * self.obs_scales.ang_vel, @@ -368,8 +377,7 @@ class LeggedRobot(BaseTask): if self.cfg.env.observe_contact_states: self.obs_buf = torch.cat((self.obs_buf, (self.contact_forces[:, self.feet_indices, 2] > 1.).view( - self.num_envs, - -1) * 1.0), dim=1) + self.num_envs, -1) * 1.0), dim=1) # add noise if needed if self.add_noise: @@ -381,7 +389,8 @@ class LeggedRobot(BaseTask): self.next_privileged_obs_buf = torch.empty(self.num_envs, 0).to(self.device) if self.cfg.env.priv_observe_friction: - friction_coeffs_scale, friction_coeffs_shift = get_scale_shift(self.cfg.normalization.friction_range) + friction_coeffs_scale, friction_coeffs_shift = get_scale_shift( + self.cfg.normalization.friction_range) self.privileged_obs_buf = torch.cat((self.privileged_obs_buf, (self.friction_coeffs[:, 0].unsqueeze( 1) - friction_coeffs_shift) * friction_coeffs_scale), @@ -394,16 +403,21 @@ class LeggedRobot(BaseTask): self.ground_friction_coeffs = self._get_ground_frictions(range(self.num_envs)) ground_friction_coeffs_scale, ground_friction_coeffs_shift = get_scale_shift( self.cfg.normalization.ground_friction_range) - self.privileged_obs_buf = torch.cat((self.privileged_obs_buf, - (self.ground_friction_coeffs.unsqueeze( - 1) - ground_friction_coeffs_shift) * ground_friction_coeffs_scale), - dim=1) - self.next_privileged_obs_buf = torch.cat((self.next_privileged_obs_buf, - (self.ground_friction_coeffs.unsqueeze( - 1) - friction_coeffs_shift) * friction_coeffs_scale), - dim=1) + self.privileged_obs_buf = torch.cat( + (self.privileged_obs_buf, + (self.ground_friction_coeffs.unsqueeze(1) - + ground_friction_coeffs_shift) * + ground_friction_coeffs_scale), + dim=1) + self.next_privileged_obs_buf = torch.cat( + (self.next_privileged_obs_buf, + (self.ground_friction_coeffs.unsqueeze(1) - + friction_coeffs_shift) * + friction_coeffs_scale), + dim=1) if self.cfg.env.priv_observe_restitution: - restitutions_scale, restitutions_shift = get_scale_shift(self.cfg.normalization.restitution_range) + restitutions_scale, restitutions_shift = get_scale_shift( + self.cfg.normalization.restitution_range) self.privileged_obs_buf = torch.cat((self.privileged_obs_buf, (self.restitutions[:, 0].unsqueeze( 1) - restitutions_shift) * restitutions_scale), @@ -413,46 +427,54 @@ class LeggedRobot(BaseTask): 1) - restitutions_shift) * restitutions_scale), dim=1) if self.cfg.env.priv_observe_base_mass: - payloads_scale, payloads_shift = get_scale_shift(self.cfg.normalization.added_mass_range) - self.privileged_obs_buf = torch.cat((self.privileged_obs_buf, - (self.payloads.unsqueeze(1) - payloads_shift) * payloads_scale), - dim=1) - self.next_privileged_obs_buf = torch.cat((self.next_privileged_obs_buf, - (self.payloads.unsqueeze(1) - payloads_shift) * payloads_scale), - dim=1) + payloads_scale, payloads_shift = get_scale_shift( + self.cfg.normalization.added_mass_range) + self.privileged_obs_buf = torch.cat( + (self.privileged_obs_buf, (self.payloads.unsqueeze(1) - payloads_shift) * payloads_scale), dim=1) + self.next_privileged_obs_buf = torch.cat( + (self.next_privileged_obs_buf, + (self.payloads.unsqueeze(1) - payloads_shift) * payloads_scale), + dim=1) if self.cfg.env.priv_observe_com_displacement: com_displacements_scale, com_displacements_shift = get_scale_shift( self.cfg.normalization.com_displacement_range) - self.privileged_obs_buf = torch.cat((self.privileged_obs_buf, - ( - self.com_displacements - com_displacements_shift) * com_displacements_scale), - dim=1) - self.next_privileged_obs_buf = torch.cat((self.next_privileged_obs_buf, - ( - self.com_displacements - com_displacements_shift) * com_displacements_scale), - dim=1) + self.privileged_obs_buf = torch.cat( + (self.privileged_obs_buf, + (self.com_displacements - + com_displacements_shift) * + com_displacements_scale), + dim=1) + self.next_privileged_obs_buf = torch.cat( + (self.next_privileged_obs_buf, + (self.com_displacements - + com_displacements_shift) * + com_displacements_scale), + dim=1) if self.cfg.env.priv_observe_motor_strength: - motor_strengths_scale, motor_strengths_shift = get_scale_shift(self.cfg.normalization.motor_strength_range) - self.privileged_obs_buf = torch.cat((self.privileged_obs_buf, - ( - self.motor_strengths - motor_strengths_shift) * motor_strengths_scale), - dim=1) - self.next_privileged_obs_buf = torch.cat((self.next_privileged_obs_buf, - ( - self.motor_strengths - motor_strengths_shift) * motor_strengths_scale), - dim=1) + motor_strengths_scale, motor_strengths_shift = get_scale_shift( + self.cfg.normalization.motor_strength_range) + self.privileged_obs_buf = torch.cat( + (self.privileged_obs_buf, + (self.motor_strengths - + motor_strengths_shift) * + motor_strengths_scale), + dim=1) + self.next_privileged_obs_buf = torch.cat( + (self.next_privileged_obs_buf, + (self.motor_strengths - + motor_strengths_shift) * + motor_strengths_scale), + dim=1) if self.cfg.env.priv_observe_motor_offset: - motor_offset_scale, motor_offset_shift = get_scale_shift(self.cfg.normalization.motor_offset_range) - self.privileged_obs_buf = torch.cat((self.privileged_obs_buf, - ( - self.motor_offsets - motor_offset_shift) * motor_offset_scale), - dim=1) - self.next_privileged_obs_buf = torch.cat((self.privileged_obs_buf, - ( - self.motor_offsets - motor_offset_shift) * motor_offset_scale), - dim=1) + motor_offset_scale, motor_offset_shift = get_scale_shift( + self.cfg.normalization.motor_offset_range) + self.privileged_obs_buf = torch.cat( + (self.privileged_obs_buf, (self.motor_offsets - motor_offset_shift) * motor_offset_scale), dim=1) + self.next_privileged_obs_buf = torch.cat( + (self.privileged_obs_buf, (self.motor_offsets - motor_offset_shift) * motor_offset_scale), dim=1) if self.cfg.env.priv_observe_body_height: - body_height_scale, body_height_shift = get_scale_shift(self.cfg.normalization.body_height_range) + body_height_scale, body_height_shift = get_scale_shift( + self.cfg.normalization.body_height_range) self.privileged_obs_buf = torch.cat((self.privileged_obs_buf, ((self.root_states[:self.num_envs, 2]).view( self.num_envs, -1) - body_height_shift) * body_height_scale), @@ -462,22 +484,21 @@ class LeggedRobot(BaseTask): self.num_envs, -1) - body_height_shift) * body_height_scale), dim=1) if self.cfg.env.priv_observe_body_velocity: - body_velocity_scale, body_velocity_shift = get_scale_shift(self.cfg.normalization.body_velocity_range) - self.privileged_obs_buf = torch.cat((self.privileged_obs_buf, - ((self.base_lin_vel).view(self.num_envs, - -1) - body_velocity_shift) * body_velocity_scale), - dim=1) - self.next_privileged_obs_buf = torch.cat((self.next_privileged_obs_buf, - ((self.base_lin_vel).view(self.num_envs, - -1) - body_velocity_shift) * body_velocity_scale), - dim=1) + body_velocity_scale, body_velocity_shift = get_scale_shift( + self.cfg.normalization.body_velocity_range) + self.privileged_obs_buf = torch.cat( + (self.privileged_obs_buf, ((self.base_lin_vel).view( + self.num_envs, -1) - body_velocity_shift) * body_velocity_scale), dim=1) + self.next_privileged_obs_buf = torch.cat( + (self.next_privileged_obs_buf, ((self.base_lin_vel).view( + self.num_envs, -1) - body_velocity_shift) * body_velocity_scale), dim=1) if self.cfg.env.priv_observe_gravity: gravity_scale, gravity_shift = get_scale_shift(self.cfg.normalization.gravity_range) self.privileged_obs_buf = torch.cat((self.privileged_obs_buf, (self.gravities - gravity_shift) / gravity_scale), dim=1) - self.next_privileged_obs_buf = torch.cat((self.next_privileged_obs_buf, - (self.gravities - gravity_shift) / gravity_scale), dim=1) + self.next_privileged_obs_buf = torch.cat( + (self.next_privileged_obs_buf, (self.gravities - gravity_shift) / gravity_scale), dim=1) if self.cfg.env.priv_observe_clock_inputs: self.privileged_obs_buf = torch.cat((self.privileged_obs_buf, @@ -487,20 +508,28 @@ class LeggedRobot(BaseTask): self.privileged_obs_buf = torch.cat((self.privileged_obs_buf, self.desired_contact_states), dim=-1) - assert self.privileged_obs_buf.shape[ - 1] == self.cfg.env.num_privileged_obs, f"num_privileged_obs ({self.cfg.env.num_privileged_obs}) != the number of privileged observations ({self.privileged_obs_buf.shape[1]}), you will discard data from the student!" + assert self.privileged_obs_buf.shape[1] == self.cfg.env.num_privileged_obs, f"num_privileged_obs ({ + self.cfg.env.num_privileged_obs}) != the number of privileged observations ({ + self.privileged_obs_buf.shape[1]}), you will discard data from the student!" def create_sim(self): """ Creates simulation, terrain and evironments """ self.up_axis_idx = 2 # 2 for z, 1 for y -> adapt gravity accordingly - self.sim = self.gym.create_sim(self.sim_device_id, self.graphics_device_id, self.physics_engine, - self.sim_params) + self.sim = self.gym.create_sim( + self.sim_device_id, + self.graphics_device_id, + self.physics_engine, + self.sim_params) mesh_type = self.cfg.terrain.mesh_type if mesh_type in ['heightfield', 'trimesh']: if self.eval_cfg is not None: - self.terrain = Terrain(self.cfg.terrain, self.num_train_envs, self.eval_cfg.terrain, self.num_eval_envs) + self.terrain = Terrain( + self.cfg.terrain, + self.num_train_envs, + self.eval_cfg.terrain, + self.num_eval_envs) else: self.terrain = Terrain(self.cfg.terrain, self.num_train_envs) if mesh_type == 'plane': @@ -510,11 +539,11 @@ class LeggedRobot(BaseTask): elif mesh_type == 'trimesh': self._create_trimesh() elif mesh_type is not None: - raise ValueError("Terrain mesh type not recognised. Allowed types are [None, plane, heightfield, trimesh]") + raise ValueError( + "Terrain mesh type not recognised. Allowed types are [None, plane, heightfield, trimesh]") self._create_envs() - def set_camera(self, position, lookat): """ Set camera position and direction """ @@ -522,6 +551,7 @@ class LeggedRobot(BaseTask): cam_target = gymapi.Vec3(lookat[0], lookat[1], lookat[2]) self.gym.viewer_camera_look_at(self.viewer, None, cam_pos, cam_target) \ + def set_main_agent_pose(self, loc, quat): self.root_states[0, 0:3] = torch.Tensor(loc) self.root_states[0, 3:7] = torch.Tensor(quat) @@ -539,17 +569,20 @@ class LeggedRobot(BaseTask): ret = func(env_ids_train, self.cfg) if len(env_ids_eval) > 0: ret_eval = func(env_ids_eval, self.eval_cfg) - if ret is not None and ret_eval is not None: ret = torch.cat((ret, ret_eval), axis=-1) + if ret is not None and ret_eval is not None: + ret = torch.cat((ret, ret_eval), axis=-1) return ret - def _randomize_gravity(self, external_force = None): + def _randomize_gravity(self, external_force=None): if external_force is not None: self.gravities[:, :] = external_force.unsqueeze(0) elif self.cfg.domain_rand.randomize_gravity: min_gravity, max_gravity = self.cfg.domain_rand.gravity_range - external_force = torch.rand(3, dtype=torch.float, device=self.device, + external_force = torch.rand(3, + dtype=torch.float, + device=self.device, requires_grad=False) * (max_gravity - min_gravity) + min_gravity self.gravities[:, :] = external_force.unsqueeze(0) @@ -591,10 +624,22 @@ class LeggedRobot(BaseTask): [numpy.array]: Modified DOF properties """ if env_id == 0: - self.dof_pos_limits = torch.zeros(self.num_dof, 2, dtype=torch.float, device=self.device, - requires_grad=False) - self.dof_vel_limits = torch.zeros(self.num_dof, dtype=torch.float, device=self.device, requires_grad=False) - self.torque_limits = torch.zeros(self.num_dof, dtype=torch.float, device=self.device, requires_grad=False) + self.dof_pos_limits = torch.zeros( + self.num_dof, + 2, + dtype=torch.float, + device=self.device, + requires_grad=False) + self.dof_vel_limits = torch.zeros( + self.num_dof, + dtype=torch.float, + device=self.device, + requires_grad=False) + self.torque_limits = torch.zeros( + self.num_dof, + dtype=torch.float, + device=self.device, + requires_grad=False) for i in range(len(props)): self.dof_pos_limits[i, 0] = props["lower"][i].item() self.dof_pos_limits[i, 1] = props["upper"][i].item() @@ -612,25 +657,35 @@ class LeggedRobot(BaseTask): if cfg.domain_rand.randomize_base_mass: min_payload, max_payload = cfg.domain_rand.added_mass_range # self.payloads[env_ids] = -1.0 - self.payloads[env_ids] = torch.rand(len(env_ids), dtype=torch.float, device=self.device, - requires_grad=False) * (max_payload - min_payload) + min_payload + self.payloads[env_ids] = torch.rand( + len(env_ids), + dtype=torch.float, + device=self.device, + requires_grad=False) * ( + max_payload - min_payload) + min_payload if cfg.domain_rand.randomize_com_displacement: min_com_displacement, max_com_displacement = cfg.domain_rand.com_displacement_range - self.com_displacements[env_ids, :] = torch.rand(len(env_ids), 3, dtype=torch.float, device=self.device, - requires_grad=False) * ( - max_com_displacement - min_com_displacement) + min_com_displacement + self.com_displacements[env_ids, + :] = torch.rand(len(env_ids), + 3, + dtype=torch.float, + device=self.device, + requires_grad=False) * (max_com_displacement - min_com_displacement) + min_com_displacement if cfg.domain_rand.randomize_friction: min_friction, max_friction = cfg.domain_rand.friction_range - self.friction_coeffs[env_ids, :] = torch.rand(len(env_ids), 1, dtype=torch.float, device=self.device, - requires_grad=False) * ( - max_friction - min_friction) + min_friction + self.friction_coeffs[env_ids, + :] = torch.rand(len(env_ids), + 1, + dtype=torch.float, + device=self.device, + requires_grad=False) * (max_friction - min_friction) + min_friction if cfg.domain_rand.randomize_restitution: min_restitution, max_restitution = cfg.domain_rand.restitution_range - self.restitutions[env_ids] = torch.rand(len(env_ids), 1, dtype=torch.float, device=self.device, - requires_grad=False) * ( - max_restitution - min_restitution) + min_restitution + self.restitutions[env_ids] = torch.rand( + len(env_ids), 1, dtype=torch.float, device=self.device, requires_grad=False) * ( + max_restitution - min_restitution) + min_restitution def refresh_actor_rigid_shape_props(self, env_ids, cfg): for env_id in env_ids: @@ -645,31 +700,40 @@ class LeggedRobot(BaseTask): def _randomize_dof_props(self, env_ids, cfg): if cfg.domain_rand.randomize_motor_strength: min_strength, max_strength = cfg.domain_rand.motor_strength_range - self.motor_strengths[env_ids, :] = torch.rand(len(env_ids), dtype=torch.float, device=self.device, - requires_grad=False).unsqueeze(1) * ( - max_strength - min_strength) + min_strength + self.motor_strengths[env_ids, + :] = torch.rand(len(env_ids), + dtype=torch.float, + device=self.device, + requires_grad=False).unsqueeze(1) * (max_strength - min_strength) + min_strength if cfg.domain_rand.randomize_motor_offset: min_offset, max_offset = cfg.domain_rand.motor_offset_range - self.motor_offsets[env_ids, :] = torch.rand(len(env_ids), self.num_dof, dtype=torch.float, - device=self.device, requires_grad=False) * ( - max_offset - min_offset) + min_offset + self.motor_offsets[env_ids, + :] = torch.rand(len(env_ids), + self.num_dof, + dtype=torch.float, + device=self.device, + requires_grad=False) * (max_offset - min_offset) + min_offset if cfg.domain_rand.randomize_Kp_factor: min_Kp_factor, max_Kp_factor = cfg.domain_rand.Kp_factor_range - self.Kp_factors[env_ids, :] = torch.rand(len(env_ids), dtype=torch.float, device=self.device, - requires_grad=False).unsqueeze(1) * ( - max_Kp_factor - min_Kp_factor) + min_Kp_factor + self.Kp_factors[env_ids, + :] = torch.rand(len(env_ids), + dtype=torch.float, + device=self.device, + requires_grad=False).unsqueeze(1) * (max_Kp_factor - min_Kp_factor) + min_Kp_factor if cfg.domain_rand.randomize_Kd_factor: min_Kd_factor, max_Kd_factor = cfg.domain_rand.Kd_factor_range - self.Kd_factors[env_ids, :] = torch.rand(len(env_ids), dtype=torch.float, device=self.device, - requires_grad=False).unsqueeze(1) * ( - max_Kd_factor - min_Kd_factor) + min_Kd_factor + self.Kd_factors[env_ids, + :] = torch.rand(len(env_ids), + dtype=torch.float, + device=self.device, + requires_grad=False).unsqueeze(1) * (max_Kd_factor - min_Kd_factor) + min_Kd_factor def _process_rigid_body_props(self, props, env_id): self.default_body_mass = props[0].mass props[0].mass = self.default_body_mass + self.payloads[env_id] - props[0].com = gymapi.Vec3(self.com_displacements[env_id, 0], self.com_displacements[env_id, 1], - self.com_displacements[env_id, 2]) + props[0].com = gymapi.Vec3(self.com_displacements[env_id, 0], + self.com_displacements[env_id, 1], self.com_displacements[env_id, 2]) return props def _post_physics_step_callback(self): @@ -678,17 +742,23 @@ class LeggedRobot(BaseTask): """ # teleport robots to prevent falling off the edge - self._call_train_eval(self._teleport_robots, torch.arange(self.num_envs, device=self.device)) + self._call_train_eval( + self._teleport_robots, torch.arange( + self.num_envs, device=self.device)) # resample commands sample_interval = int(self.cfg.commands.resampling_time / self.dt) - env_ids = (self.episode_length_buf % sample_interval == 0).nonzero(as_tuple=False).flatten() + env_ids = ( + self.episode_length_buf % + sample_interval == 0).nonzero( + as_tuple=False).flatten() self._resample_commands(env_ids) self._step_contact_targets() # measure terrain heights if self.cfg.terrain.measure_heights: - self.measured_heights = self._get_heights(torch.arange(self.num_envs, device=self.device), self.cfg) + self.measured_heights = self._get_heights( + torch.arange(self.num_envs, device=self.device), self.cfg) # push robots self._call_train_eval(self._push_robots, torch.arange(self.num_envs, device=self.device)) @@ -709,7 +779,8 @@ class LeggedRobot(BaseTask): def _resample_commands(self, env_ids): - if len(env_ids) == 0: return + if len(env_ids) == 0: + return timesteps = int(self.cfg.commands.resampling_time / self.dt) ep_len = min(self.cfg.env.max_episode_length, timesteps) @@ -729,61 +800,71 @@ class LeggedRobot(BaseTask): "tracking_contacts_shaped_vel"]: if key in self.command_sums.keys(): task_rewards.append(self.command_sums[key][env_ids_in_category] / ep_len) - success_thresholds.append(self.curriculum_thresholds[key] * self.reward_scales[key]) + success_thresholds.append( + self.curriculum_thresholds[key] * + self.reward_scales[key]) old_bins = self.env_command_bins[env_ids_in_category.cpu().numpy()] if len(success_thresholds) > 0: - curriculum.update(old_bins, task_rewards, success_thresholds, - local_range=np.array( - [0.55, 0.55, 0.55, 0.55, 0.35, 0.25, 0.25, 0.25, 0.25, 1.0, 1.0, 1.0, 1.0, 1.0, - 1.0])) + curriculum.update(old_bins, task_rewards, success_thresholds, local_range=np.array( + [0.55, 0.55, 0.55, 0.55, 0.35, 0.25, 0.25, 0.25, 0.25, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0])) # assign resampled environments to new categories random_env_floats = torch.rand(len(env_ids), device=self.device) probability_per_category = 1. / len(self.category_names) - category_env_ids = [env_ids[torch.logical_and(probability_per_category * i <= random_env_floats, - random_env_floats < probability_per_category * (i + 1))] for i in - range(len(self.category_names))] + category_env_ids = [env_ids[torch.logical_and(probability_per_category * + i <= random_env_floats, random_env_floats < probability_per_category * + (i + + 1))] for i in range(len(self.category_names))] # sample from new category curricula for i, (category, env_ids_in_category, curriculum) in enumerate( zip(self.category_names, category_env_ids, self.curricula)): batch_size = len(env_ids_in_category) - if batch_size == 0: continue + if batch_size == 0: + continue new_commands, new_bin_inds = curriculum.sample(batch_size=batch_size) self.env_command_bins[env_ids_in_category.cpu().numpy()] = new_bin_inds self.env_command_categories[env_ids_in_category.cpu().numpy()] = i - self.commands[env_ids_in_category, :] = torch.Tensor(new_commands[:, :self.cfg.commands.num_commands]).to( - self.device) + self.commands[env_ids_in_category, :] = torch.Tensor( + new_commands[:, :self.cfg.commands.num_commands]).to(self.device) if self.cfg.commands.num_commands > 5: if self.cfg.commands.gaitwise_curricula: - for i, (category, env_ids_in_category) in enumerate(zip(self.category_names, category_env_ids)): + for i, (category, env_ids_in_category) in enumerate( + zip(self.category_names, category_env_ids)): if category == "pronk": # pronking - self.commands[env_ids_in_category, 5] = (self.commands[env_ids_in_category, 5] / 2 - 0.25) % 1 - self.commands[env_ids_in_category, 6] = (self.commands[env_ids_in_category, 6] / 2 - 0.25) % 1 - self.commands[env_ids_in_category, 7] = (self.commands[env_ids_in_category, 7] / 2 - 0.25) % 1 + self.commands[env_ids_in_category, 5] = ( + self.commands[env_ids_in_category, 5] / 2 - 0.25) % 1 + self.commands[env_ids_in_category, 6] = ( + self.commands[env_ids_in_category, 6] / 2 - 0.25) % 1 + self.commands[env_ids_in_category, 7] = ( + self.commands[env_ids_in_category, 7] / 2 - 0.25) % 1 elif category == "trot": # trotting - self.commands[env_ids_in_category, 5] = self.commands[env_ids_in_category, 5] / 2 + 0.25 + self.commands[env_ids_in_category, + 5] = self.commands[env_ids_in_category, 5] / 2 + 0.25 self.commands[env_ids_in_category, 6] = 0 self.commands[env_ids_in_category, 7] = 0 elif category == "pace": # pacing self.commands[env_ids_in_category, 5] = 0 - self.commands[env_ids_in_category, 6] = self.commands[env_ids_in_category, 6] / 2 + 0.25 + self.commands[env_ids_in_category, + 6] = self.commands[env_ids_in_category, 6] / 2 + 0.25 self.commands[env_ids_in_category, 7] = 0 elif category == "bound": # bounding self.commands[env_ids_in_category, 5] = 0 self.commands[env_ids_in_category, 6] = 0 - self.commands[env_ids_in_category, 7] = self.commands[env_ids_in_category, 7] / 2 + 0.25 + self.commands[env_ids_in_category, + 7] = self.commands[env_ids_in_category, 7] / 2 + 0.25 elif self.cfg.commands.exclusive_phase_offset: random_env_floats = torch.rand(len(env_ids), device=self.device) trotting_envs = env_ids[random_env_floats < 0.34] - pacing_envs = env_ids[torch.logical_and(0.34 <= random_env_floats, random_env_floats < 0.67)] + pacing_envs = env_ids[torch.logical_and( + 0.34 <= random_env_floats, random_env_floats < 0.67)] bounding_envs = env_ids[0.67 <= random_env_floats] self.commands[pacing_envs, 5] = 0 self.commands[bounding_envs, 5] = 0 @@ -795,8 +876,10 @@ class LeggedRobot(BaseTask): elif self.cfg.commands.balance_gait_distribution: random_env_floats = torch.rand(len(env_ids), device=self.device) pronking_envs = env_ids[random_env_floats <= 0.25] - trotting_envs = env_ids[torch.logical_and(0.25 <= random_env_floats, random_env_floats < 0.50)] - pacing_envs = env_ids[torch.logical_and(0.50 <= random_env_floats, random_env_floats < 0.75)] + trotting_envs = env_ids[torch.logical_and( + 0.25 <= random_env_floats, random_env_floats < 0.50)] + pacing_envs = env_ids[torch.logical_and( + 0.50 <= random_env_floats, random_env_floats < 0.75)] bounding_envs = env_ids[0.75 <= random_env_floats] self.commands[pronking_envs, 5] = (self.commands[pronking_envs, 5] / 2 - 0.25) % 1 self.commands[pronking_envs, 6] = (self.commands[pronking_envs, 6] / 2 - 0.25) % 1 @@ -817,7 +900,8 @@ class LeggedRobot(BaseTask): self.commands[env_ids, 7] = (torch.round(2 * self.commands[env_ids, 7])) / 2.0 % 1 # setting the smaller commands to zero - self.commands[env_ids, :2] *= (torch.norm(self.commands[env_ids, :2], dim=1) > 0.2).unsqueeze(1) + self.commands[env_ids, + :2] *= (torch.norm(self.commands[env_ids, :2], dim=1) > 0.2).unsqueeze(1) # reset command sums for key in self.command_sums.keys(): @@ -843,15 +927,17 @@ class LeggedRobot(BaseTask): self.gait_indices + bounds, self.gait_indices + phases] - self.foot_indices = torch.remainder(torch.cat([foot_indices[i].unsqueeze(1) for i in range(4)], dim=1), 1.0) + self.foot_indices = torch.remainder( + torch.cat([foot_indices[i].unsqueeze(1) for i in range(4)], dim=1), 1.0) for idxs in foot_indices: stance_idxs = torch.remainder(idxs, 1) < durations swing_idxs = torch.remainder(idxs, 1) > durations - idxs[stance_idxs] = torch.remainder(idxs[stance_idxs], 1) * (0.5 / durations[stance_idxs]) - idxs[swing_idxs] = 0.5 + (torch.remainder(idxs[swing_idxs], 1) - durations[swing_idxs]) * ( - 0.5 / (1 - durations[swing_idxs])) + idxs[stance_idxs] = torch.remainder( + idxs[stance_idxs], 1) * (0.5 / durations[stance_idxs]) + idxs[swing_idxs] = 0.5 + (torch.remainder(idxs[swing_idxs], 1) - + durations[swing_idxs]) * (0.5 / (1 - durations[swing_idxs])) # if self.cfg.commands.durations_warp_clock_inputs: @@ -872,29 +958,109 @@ class LeggedRobot(BaseTask): # von mises distribution kappa = self.cfg.rewards.kappa_gait_probs - smoothing_cdf_start = torch.distributions.normal.Normal(0, - kappa).cdf # (x) + torch.distributions.normal.Normal(1, kappa).cdf(x)) / 2 + # (x) + torch.distributions.normal.Normal(1, kappa).cdf(x)) / 2 + smoothing_cdf_start = torch.distributions.normal.Normal(0, kappa).cdf - smoothing_multiplier_FL = (smoothing_cdf_start(torch.remainder(foot_indices[0], 1.0)) * ( - 1 - smoothing_cdf_start(torch.remainder(foot_indices[0], 1.0) - 0.5)) + - smoothing_cdf_start(torch.remainder(foot_indices[0], 1.0) - 1) * ( - 1 - smoothing_cdf_start( - torch.remainder(foot_indices[0], 1.0) - 0.5 - 1))) - smoothing_multiplier_FR = (smoothing_cdf_start(torch.remainder(foot_indices[1], 1.0)) * ( - 1 - smoothing_cdf_start(torch.remainder(foot_indices[1], 1.0) - 0.5)) + - smoothing_cdf_start(torch.remainder(foot_indices[1], 1.0) - 1) * ( - 1 - smoothing_cdf_start( - torch.remainder(foot_indices[1], 1.0) - 0.5 - 1))) - smoothing_multiplier_RL = (smoothing_cdf_start(torch.remainder(foot_indices[2], 1.0)) * ( - 1 - smoothing_cdf_start(torch.remainder(foot_indices[2], 1.0) - 0.5)) + - smoothing_cdf_start(torch.remainder(foot_indices[2], 1.0) - 1) * ( - 1 - smoothing_cdf_start( - torch.remainder(foot_indices[2], 1.0) - 0.5 - 1))) - smoothing_multiplier_RR = (smoothing_cdf_start(torch.remainder(foot_indices[3], 1.0)) * ( - 1 - smoothing_cdf_start(torch.remainder(foot_indices[3], 1.0) - 0.5)) + - smoothing_cdf_start(torch.remainder(foot_indices[3], 1.0) - 1) * ( - 1 - smoothing_cdf_start( - torch.remainder(foot_indices[3], 1.0) - 0.5 - 1))) + smoothing_multiplier_FL = ( + smoothing_cdf_start( + torch.remainder( + foot_indices[0], + 1.0)) * + ( + 1 - + smoothing_cdf_start( + torch.remainder( + foot_indices[0], + 1.0) - + 0.5)) + + smoothing_cdf_start( + torch.remainder( + foot_indices[0], + 1.0) - + 1) * + ( + 1 - + smoothing_cdf_start( + torch.remainder( + foot_indices[0], + 1.0) - + 0.5 - + 1))) + smoothing_multiplier_FR = ( + smoothing_cdf_start( + torch.remainder( + foot_indices[1], + 1.0)) * + ( + 1 - + smoothing_cdf_start( + torch.remainder( + foot_indices[1], + 1.0) - + 0.5)) + + smoothing_cdf_start( + torch.remainder( + foot_indices[1], + 1.0) - + 1) * + ( + 1 - + smoothing_cdf_start( + torch.remainder( + foot_indices[1], + 1.0) - + 0.5 - + 1))) + smoothing_multiplier_RL = ( + smoothing_cdf_start( + torch.remainder( + foot_indices[2], + 1.0)) * + ( + 1 - + smoothing_cdf_start( + torch.remainder( + foot_indices[2], + 1.0) - + 0.5)) + + smoothing_cdf_start( + torch.remainder( + foot_indices[2], + 1.0) - + 1) * + ( + 1 - + smoothing_cdf_start( + torch.remainder( + foot_indices[2], + 1.0) - + 0.5 - + 1))) + smoothing_multiplier_RR = ( + smoothing_cdf_start( + torch.remainder( + foot_indices[3], + 1.0)) * + ( + 1 - + smoothing_cdf_start( + torch.remainder( + foot_indices[3], + 1.0) - + 0.5)) + + smoothing_cdf_start( + torch.remainder( + foot_indices[3], + 1.0) - + 1) * + ( + 1 - + smoothing_cdf_start( + torch.remainder( + foot_indices[3], + 1.0) - + 0.5 - + 1))) self.desired_contact_states[:, 0] = smoothing_multiplier_FL self.desired_contact_states[:, 1] = smoothing_multiplier_FR @@ -917,7 +1083,8 @@ class LeggedRobot(BaseTask): """ # pd controller actions_scaled = actions[:, :12] * self.cfg.control.action_scale - actions_scaled[:, [0, 3, 6, 9]] *= self.cfg.control.hip_scale_reduction # scale down hip flexion range + # scale down hip flexion range + actions_scaled[:, [0, 3, 6, 9]] *= self.cfg.control.hip_scale_reduction if self.cfg.domain_rand.randomize_lag_timesteps: self.lag_buffer = self.lag_buffer[1:] + [actions_scaled.clone()] @@ -930,15 +1097,20 @@ class LeggedRobot(BaseTask): if control_type == "actuator_net": self.joint_pos_err = self.dof_pos - self.joint_pos_target + self.motor_offsets self.joint_vel = self.dof_vel - torques = self.actuator_network(self.joint_pos_err, self.joint_pos_err_last, self.joint_pos_err_last_last, - self.joint_vel, self.joint_vel_last, self.joint_vel_last_last) + torques = self.actuator_network( + self.joint_pos_err, + self.joint_pos_err_last, + self.joint_pos_err_last_last, + self.joint_vel, + self.joint_vel_last, + self.joint_vel_last_last) self.joint_pos_err_last_last = torch.clone(self.joint_pos_err_last) self.joint_pos_err_last = torch.clone(self.joint_pos_err) self.joint_vel_last_last = torch.clone(self.joint_vel_last) self.joint_vel_last = torch.clone(self.joint_vel) elif control_type == "P": torques = self.p_gains * self.Kp_factors * ( - self.joint_pos_target - self.dof_pos + self.motor_offsets) - self.d_gains * self.Kd_factors * self.dof_vel + self.joint_pos_target - self.dof_pos + self.motor_offsets) - self.d_gains * self.Kd_factors * self.dof_vel else: raise NameError(f"Unknown controller type: {control_type}") @@ -953,14 +1125,17 @@ class LeggedRobot(BaseTask): Args: env_ids (List[int]): Environemnt ids """ - self.dof_pos[env_ids] = self.default_dof_pos * torch_rand_float(0.5, 1.5, (len(env_ids), self.num_dof), - device=self.device) + self.dof_pos[env_ids] = self.default_dof_pos * \ + torch_rand_float(0.5, 1.5, (len(env_ids), self.num_dof), device=self.device) self.dof_vel[env_ids] = 0. env_ids_int32 = env_ids.to(dtype=torch.int32) - self.gym.set_dof_state_tensor_indexed(self.sim, - gymtorch.unwrap_tensor(self.dof_state), - gymtorch.unwrap_tensor(env_ids_int32), len(env_ids_int32)) + self.gym.set_dof_state_tensor_indexed( + self.sim, + gymtorch.unwrap_tensor( + self.dof_state), + gymtorch.unwrap_tensor(env_ids_int32), + len(env_ids_int32)) def _reset_root_states(self, env_ids, cfg): """ Resets ROOT states position and velocities of selected environmments @@ -973,12 +1148,18 @@ class LeggedRobot(BaseTask): if self.custom_origins: self.root_states[env_ids] = self.base_init_state self.root_states[env_ids, :3] += self.env_origins[env_ids] - self.root_states[env_ids, 0:1] += torch_rand_float(-cfg.terrain.x_init_range, - cfg.terrain.x_init_range, (len(env_ids), 1), - device=self.device) - self.root_states[env_ids, 1:2] += torch_rand_float(-cfg.terrain.y_init_range, - cfg.terrain.y_init_range, (len(env_ids), 1), - device=self.device) + self.root_states[env_ids, + 0:1] += torch_rand_float(-cfg.terrain.x_init_range, + cfg.terrain.x_init_range, + (len(env_ids), + 1), + device=self.device) + self.root_states[env_ids, + 1:2] += torch_rand_float(-cfg.terrain.y_init_range, + cfg.terrain.y_init_range, + (len(env_ids), + 1), + device=self.device) self.root_states[env_ids, 0] += cfg.terrain.x_init_offset self.root_states[env_ids, 1] += cfg.terrain.y_init_offset else: @@ -993,12 +1174,12 @@ class LeggedRobot(BaseTask): self.root_states[env_ids, 3:7] = quat # base velocities - self.root_states[env_ids, 7:13] = torch_rand_float(-0.5, 0.5, (len(env_ids), 6), - device=self.device) # [7:10]: lin vel, [10:13]: ang vel + # [7:10]: lin vel, [10:13]: ang vel + self.root_states[env_ids, + 7:13] = torch_rand_float(-0.5, 0.5, (len(env_ids), 6), device=self.device) env_ids_int32 = env_ids.to(dtype=torch.int32) - self.gym.set_actor_root_state_tensor_indexed(self.sim, - gymtorch.unwrap_tensor(self.root_states), - gymtorch.unwrap_tensor(env_ids_int32), len(env_ids_int32)) + self.gym.set_actor_root_state_tensor_indexed(self.sim, gymtorch.unwrap_tensor( + self.root_states), gymtorch.unwrap_tensor(env_ids_int32), len(env_ids_int32)) if cfg.env.record_video and 0 in env_ids: if self.complete_video_frames is None: @@ -1018,12 +1199,14 @@ class LeggedRobot(BaseTask): """ Random pushes the robots. Emulates an impulse by setting a randomized base velocity. """ if cfg.domain_rand.push_robots: - env_ids = env_ids[self.episode_length_buf[env_ids] % int(cfg.domain_rand.push_interval) == 0] + env_ids = env_ids[self.episode_length_buf[env_ids] % + int(cfg.domain_rand.push_interval) == 0] max_vel = cfg.domain_rand.max_push_vel_xy self.root_states[env_ids, 7:9] = torch_rand_float(-max_vel, max_vel, (len(env_ids), 2), device=self.device) # lin vel x/y - self.gym.set_actor_root_state_tensor(self.sim, gymtorch.unwrap_tensor(self.root_states)) + self.gym.set_actor_root_state_tensor( + self.sim, gymtorch.unwrap_tensor(self.root_states)) def _teleport_robots(self, env_ids, cfg): """ Teleports any robots that are too close to the edge to the other side @@ -1034,20 +1217,25 @@ class LeggedRobot(BaseTask): x_offset = int(cfg.terrain.x_offset * cfg.terrain.horizontal_scale) low_x_ids = env_ids[self.root_states[env_ids, 0] < thresh + x_offset] - self.root_states[low_x_ids, 0] += cfg.terrain.terrain_length * (cfg.terrain.num_rows - 1) + self.root_states[low_x_ids, 0] += cfg.terrain.terrain_length * \ + (cfg.terrain.num_rows - 1) - high_x_ids = env_ids[ - self.root_states[env_ids, 0] > cfg.terrain.terrain_length * cfg.terrain.num_rows - thresh + x_offset] - self.root_states[high_x_ids, 0] -= cfg.terrain.terrain_length * (cfg.terrain.num_rows - 1) + high_x_ids = env_ids[self.root_states[env_ids, 0] > + cfg.terrain.terrain_length * cfg.terrain.num_rows - thresh + x_offset] + self.root_states[high_x_ids, 0] -= cfg.terrain.terrain_length * \ + (cfg.terrain.num_rows - 1) low_y_ids = env_ids[self.root_states[env_ids, 1] < thresh] - self.root_states[low_y_ids, 1] += cfg.terrain.terrain_width * (cfg.terrain.num_cols - 1) + self.root_states[low_y_ids, 1] += cfg.terrain.terrain_width * \ + (cfg.terrain.num_cols - 1) - high_y_ids = env_ids[ - self.root_states[env_ids, 1] > cfg.terrain.terrain_width * cfg.terrain.num_cols - thresh] - self.root_states[high_y_ids, 1] -= cfg.terrain.terrain_width * (cfg.terrain.num_cols - 1) + high_y_ids = env_ids[self.root_states[env_ids, 1] > + cfg.terrain.terrain_width * cfg.terrain.num_cols - thresh] + self.root_states[high_y_ids, 1] -= cfg.terrain.terrain_width * \ + (cfg.terrain.num_cols - 1) - self.gym.set_actor_root_state_tensor(self.sim, gymtorch.unwrap_tensor(self.root_states)) + self.gym.set_actor_root_state_tensor( + self.sim, gymtorch.unwrap_tensor(self.root_states)) self.gym.refresh_actor_root_state_tensor(self.sim) def _get_noise_scale_vec(self, cfg): @@ -1064,23 +1252,46 @@ class LeggedRobot(BaseTask): self.add_noise = self.cfg.noise.add_noise noise_scales = self.cfg.noise_scales noise_level = self.cfg.noise.noise_level - noise_vec = torch.cat((torch.ones(3) * noise_scales.gravity * noise_level, - torch.ones( - self.num_actuated_dof) * noise_scales.dof_pos * noise_level * self.obs_scales.dof_pos, - torch.ones( - self.num_actuated_dof) * noise_scales.dof_vel * noise_level * self.obs_scales.dof_vel, - torch.zeros(self.num_actions), - ), dim=0) + noise_vec = torch.cat( + (torch.ones(3) * + noise_scales.gravity * + noise_level, + torch.ones( + self.num_actuated_dof) * + noise_scales.dof_pos * + noise_level * + self.obs_scales.dof_pos, + torch.ones( + self.num_actuated_dof) * + noise_scales.dof_vel * + noise_level * + self.obs_scales.dof_vel, + torch.zeros( + self.num_actions), + ), + dim=0) if self.cfg.env.observe_command: - noise_vec = torch.cat((torch.ones(3) * noise_scales.gravity * noise_level, - torch.zeros(self.cfg.commands.num_commands), - torch.ones( - self.num_actuated_dof) * noise_scales.dof_pos * noise_level * self.obs_scales.dof_pos, - torch.ones( - self.num_actuated_dof) * noise_scales.dof_vel * noise_level * self.obs_scales.dof_vel, - torch.zeros(self.num_actions), - ), dim=0) + noise_vec = torch.cat( + (torch.ones(3) * + noise_scales.gravity * + noise_level, + torch.zeros( + self.cfg.commands.num_commands), + torch.ones( + self.num_actuated_dof) * + noise_scales.dof_pos * + noise_level * + self.obs_scales.dof_pos, + torch.ones( + self.num_actuated_dof) * + noise_scales.dof_vel * + noise_level * + self.obs_scales.dof_vel, + torch.zeros( + self.num_actions), + ), + dim=0) if self.cfg.env.observe_two_prev_actions: noise_vec = torch.cat((noise_vec, torch.zeros(self.num_actions) @@ -1094,15 +1305,26 @@ class LeggedRobot(BaseTask): torch.zeros(4) ), dim=0) if self.cfg.env.observe_vel: - noise_vec = torch.cat((torch.ones(3) * noise_scales.lin_vel * noise_level * self.obs_scales.lin_vel, - torch.ones(3) * noise_scales.ang_vel * noise_level * self.obs_scales.ang_vel, - noise_vec - ), dim=0) + noise_vec = torch.cat( + (torch.ones(3) * + noise_scales.lin_vel * + noise_level * + self.obs_scales.lin_vel, + torch.ones(3) * + noise_scales.ang_vel * + noise_level * + self.obs_scales.ang_vel, + noise_vec), + dim=0) if self.cfg.env.observe_only_lin_vel: - noise_vec = torch.cat((torch.ones(3) * noise_scales.lin_vel * noise_level * self.obs_scales.lin_vel, - noise_vec - ), dim=0) + noise_vec = torch.cat( + (torch.ones(3) * + noise_scales.lin_vel * + noise_level * + self.obs_scales.lin_vel, + noise_vec), + dim=0) if self.cfg.env.observe_yaw: noise_vec = torch.cat((noise_vec, @@ -1114,7 +1336,6 @@ class LeggedRobot(BaseTask): torch.ones(4) * noise_scales.contact_states * noise_level, ), dim=0) - noise_vec = noise_vec.to(self.device) return noise_vec @@ -1137,87 +1358,152 @@ class LeggedRobot(BaseTask): # create some wrapper tensors for different slices self.root_states = gymtorch.wrap_tensor(actor_root_state) self.dof_state = gymtorch.wrap_tensor(dof_state_tensor) - self.net_contact_forces = gymtorch.wrap_tensor(net_contact_forces)[:self.num_envs * self.num_bodies, :] + self.net_contact_forces = gymtorch.wrap_tensor( + net_contact_forces)[:self.num_envs * self.num_bodies, :] self.dof_pos = self.dof_state.view(self.num_envs, self.num_dof, 2)[..., 0] self.base_pos = self.root_states[:self.num_envs, 0:3] self.dof_vel = self.dof_state.view(self.num_envs, self.num_dof, 2)[..., 1] self.base_quat = self.root_states[:self.num_envs, 3:7] - self.rigid_body_state = gymtorch.wrap_tensor(rigid_body_state)[:self.num_envs * self.num_bodies, :] - self.foot_velocities = self.rigid_body_state.view(self.num_envs, self.num_bodies, 13)[:, - self.feet_indices, - 7:10] - self.foot_positions = self.rigid_body_state.view(self.num_envs, self.num_bodies, 13)[:, self.feet_indices, - 0:3] + self.rigid_body_state = gymtorch.wrap_tensor( + rigid_body_state)[:self.num_envs * self.num_bodies, :] + self.foot_velocities = self.rigid_body_state.view(self.num_envs, self.num_bodies, 13)[ + :, self.feet_indices, 7:10] + self.foot_positions = self.rigid_body_state.view(self.num_envs, self.num_bodies, 13)[ + :, self.feet_indices, 0:3] self.prev_base_pos = self.base_pos.clone() self.prev_foot_velocities = self.foot_velocities.clone() - self.lag_buffer = [torch.zeros_like(self.dof_pos) for i in range(self.cfg.domain_rand.lag_timesteps+1)] + self.lag_buffer = [ + torch.zeros_like( + self.dof_pos) for i in range( + self.cfg.domain_rand.lag_timesteps + 1)] - self.contact_forces = gymtorch.wrap_tensor(net_contact_forces)[:self.num_envs * self.num_bodies, :].view(self.num_envs, -1, - 3) # shape: num_envs, num_bodies, xyz axis + self.contact_forces = gymtorch.wrap_tensor(net_contact_forces)[ + :self.num_envs * self.num_bodies, :].view(self.num_envs, -1, 3) # shape: num_envs, num_bodies, xyz axis # initialize some data used later on self.common_step_counter = 0 self.extras = {} if self.cfg.terrain.measure_heights: - self.height_points = self._init_height_points(torch.arange(self.num_envs, device=self.device), self.cfg) + self.height_points = self._init_height_points( + torch.arange(self.num_envs, device=self.device), self.cfg) self.measured_heights = 0 self.noise_scale_vec = self._get_noise_scale_vec(self.cfg) # , self.eval_cfg) - self.gravity_vec = to_torch(get_axis_params(-1., self.up_axis_idx), device=self.device).repeat( - (self.num_envs, 1)) + self.gravity_vec = to_torch(get_axis_params(-1., self.up_axis_idx), + device=self.device).repeat((self.num_envs, 1)) self.forward_vec = to_torch([1., 0., 0.], device=self.device).repeat((self.num_envs, 1)) - self.torques = torch.zeros(self.num_envs, self.num_dof, dtype=torch.float, device=self.device, - requires_grad=False) - self.p_gains = torch.zeros(self.num_dof, dtype=torch.float, device=self.device, requires_grad=False) - self.d_gains = torch.zeros(self.num_dof, dtype=torch.float, device=self.device, requires_grad=False) - self.actions = torch.zeros(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, - requires_grad=False) - self.last_actions = torch.zeros(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, - requires_grad=False) - self.last_last_actions = torch.zeros(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, - requires_grad=False) + self.torques = torch.zeros( + self.num_envs, + self.num_dof, + dtype=torch.float, + device=self.device, + requires_grad=False) + self.p_gains = torch.zeros( + self.num_dof, + dtype=torch.float, + device=self.device, + requires_grad=False) + self.d_gains = torch.zeros( + self.num_dof, + dtype=torch.float, + device=self.device, + requires_grad=False) + self.actions = torch.zeros( + self.num_envs, + self.num_actions, + dtype=torch.float, + device=self.device, + requires_grad=False) + self.last_actions = torch.zeros( + self.num_envs, + self.num_actions, + dtype=torch.float, + device=self.device, + requires_grad=False) + self.last_last_actions = torch.zeros( + self.num_envs, + self.num_actions, + dtype=torch.float, + device=self.device, + requires_grad=False) self.joint_pos_target = torch.zeros(self.num_envs, self.num_dof, dtype=torch.float, device=self.device, requires_grad=False) - self.last_joint_pos_target = torch.zeros(self.num_envs, self.num_dof, dtype=torch.float, device=self.device, - requires_grad=False) - self.last_last_joint_pos_target = torch.zeros(self.num_envs, self.num_dof, dtype=torch.float, - device=self.device, - requires_grad=False) + self.last_joint_pos_target = torch.zeros( + self.num_envs, + self.num_dof, + dtype=torch.float, + device=self.device, + requires_grad=False) + self.last_last_joint_pos_target = torch.zeros( + self.num_envs, + self.num_dof, + dtype=torch.float, + device=self.device, + requires_grad=False) self.last_dof_vel = torch.zeros_like(self.dof_vel) self.last_root_vel = torch.zeros_like(self.root_states[:, 7:13]) - - self.commands_value = torch.zeros(self.num_envs, self.cfg.commands.num_commands, dtype=torch.float, - device=self.device, requires_grad=False) + self.commands_value = torch.zeros( + self.num_envs, + self.cfg.commands.num_commands, + dtype=torch.float, + device=self.device, + requires_grad=False) self.commands = torch.zeros_like(self.commands_value) # x vel, y vel, yaw vel, heading - self.commands_scale = torch.tensor([self.obs_scales.lin_vel, self.obs_scales.lin_vel, self.obs_scales.ang_vel, - self.obs_scales.body_height_cmd, self.obs_scales.gait_freq_cmd, - self.obs_scales.gait_phase_cmd, self.obs_scales.gait_phase_cmd, - self.obs_scales.gait_phase_cmd, self.obs_scales.gait_phase_cmd, - self.obs_scales.footswing_height_cmd, self.obs_scales.body_pitch_cmd, - self.obs_scales.body_roll_cmd, self.obs_scales.stance_width_cmd, - self.obs_scales.stance_length_cmd, self.obs_scales.aux_reward_cmd], - device=self.device, requires_grad=False, )[:self.cfg.commands.num_commands] - self.desired_contact_states = torch.zeros(self.num_envs, 4, dtype=torch.float, device=self.device, - requires_grad=False, ) + self.commands_scale = torch.tensor([self.obs_scales.lin_vel, + self.obs_scales.lin_vel, + self.obs_scales.ang_vel, + self.obs_scales.body_height_cmd, + self.obs_scales.gait_freq_cmd, + self.obs_scales.gait_phase_cmd, + self.obs_scales.gait_phase_cmd, + self.obs_scales.gait_phase_cmd, + self.obs_scales.gait_phase_cmd, + self.obs_scales.footswing_height_cmd, + self.obs_scales.body_pitch_cmd, + self.obs_scales.body_roll_cmd, + self.obs_scales.stance_width_cmd, + self.obs_scales.stance_length_cmd, + self.obs_scales.aux_reward_cmd], + device=self.device, + requires_grad=False, + )[:self.cfg.commands.num_commands] + self.desired_contact_states = torch.zeros( + self.num_envs, + 4, + dtype=torch.float, + device=self.device, + requires_grad=False, + ) - - self.feet_air_time = torch.zeros(self.num_envs, self.feet_indices.shape[0], dtype=torch.float, - device=self.device, requires_grad=False) - self.last_contacts = torch.zeros(self.num_envs, len(self.feet_indices), dtype=torch.bool, device=self.device, + self.feet_air_time = torch.zeros( + self.num_envs, + self.feet_indices.shape[0], + dtype=torch.float, + device=self.device, + requires_grad=False) + self.last_contacts = torch.zeros(self.num_envs, + len(self.feet_indices), + dtype=torch.bool, + device=self.device, requires_grad=False) - self.last_contact_filt = torch.zeros(self.num_envs, len(self.feet_indices), dtype=torch.bool, - device=self.device, - requires_grad=False) - self.base_lin_vel = quat_rotate_inverse(self.base_quat, self.root_states[:self.num_envs, 7:10]) - self.base_ang_vel = quat_rotate_inverse(self.base_quat, self.root_states[:self.num_envs, 10:13]) + self.last_contact_filt = torch.zeros(self.num_envs, len( + self.feet_indices), dtype=torch.bool, device=self.device, requires_grad=False) + self.base_lin_vel = quat_rotate_inverse( + self.base_quat, self.root_states[:self.num_envs, 7:10]) + self.base_ang_vel = quat_rotate_inverse( + self.base_quat, self.root_states[:self.num_envs, 10:13]) self.projected_gravity = quat_rotate_inverse(self.base_quat, self.gravity_vec) # joint positions offsets and PD gains - self.default_dof_pos = torch.zeros(self.num_dof, dtype=torch.float, device=self.device, requires_grad=False) + self.default_dof_pos = torch.zeros( + self.num_dof, + dtype=torch.float, + device=self.device, + requires_grad=False) for i in range(self.num_dofs): name = self.dof_names[i] angle = self.cfg.init_state.default_joint_angles[name] @@ -1236,11 +1522,19 @@ class LeggedRobot(BaseTask): self.default_dof_pos = self.default_dof_pos.unsqueeze(0) if self.cfg.control.control_type == "actuator_net": - actuator_path = f'{os.path.dirname(os.path.dirname(os.path.realpath(__file__)))}/../../resources/actuator_nets/unitree_go2.pt' + actuator_path = f'{ + os.path.dirname( + os.path.dirname( + os.path.realpath(__file__)))}/../../resources/actuator_nets/unitree_go2.pt' actuator_network = torch.jit.load(actuator_path).to(self.device) - def eval_actuator_network(joint_pos, joint_pos_last, joint_pos_last_last, joint_vel, joint_vel_last, - joint_vel_last_last): + def eval_actuator_network( + joint_pos, + joint_pos_last, + joint_pos_last_last, + joint_vel, + joint_vel_last, + joint_vel_last_last): xs = torch.cat((joint_pos.unsqueeze(-1), joint_pos_last.unsqueeze(-1), joint_pos_last_last.unsqueeze(-1), @@ -1259,29 +1553,59 @@ class LeggedRobot(BaseTask): def _init_custom_buffers__(self): # domain randomization properties - self.friction_coeffs = self.default_friction * torch.ones(self.num_envs, 4, dtype=torch.float, device=self.device, - requires_grad=False) - self.restitutions = self.default_restitution * torch.ones(self.num_envs, 4, dtype=torch.float, device=self.device, - requires_grad=False) - self.payloads = torch.zeros(self.num_envs, dtype=torch.float, device=self.device, requires_grad=False) - self.com_displacements = torch.zeros(self.num_envs, 3, dtype=torch.float, device=self.device, - requires_grad=False) - self.motor_strengths = torch.ones(self.num_envs, self.num_dof, dtype=torch.float, device=self.device, - requires_grad=False) - self.motor_offsets = torch.zeros(self.num_envs, self.num_dof, dtype=torch.float, device=self.device, - requires_grad=False) - self.Kp_factors = torch.ones(self.num_envs, self.num_dof, dtype=torch.float, device=self.device, - requires_grad=False) - self.Kd_factors = torch.ones(self.num_envs, self.num_dof, dtype=torch.float, device=self.device, - requires_grad=False) + self.friction_coeffs = self.default_friction * \ + torch.ones(self.num_envs, 4, dtype=torch.float, device=self.device, requires_grad=False) + self.restitutions = self.default_restitution * \ + torch.ones(self.num_envs, 4, dtype=torch.float, device=self.device, requires_grad=False) + self.payloads = torch.zeros( + self.num_envs, + dtype=torch.float, + device=self.device, + requires_grad=False) + self.com_displacements = torch.zeros( + self.num_envs, + 3, + dtype=torch.float, + device=self.device, + requires_grad=False) + self.motor_strengths = torch.ones( + self.num_envs, + self.num_dof, + dtype=torch.float, + device=self.device, + requires_grad=False) + self.motor_offsets = torch.zeros( + self.num_envs, + self.num_dof, + dtype=torch.float, + device=self.device, + requires_grad=False) + self.Kp_factors = torch.ones( + self.num_envs, + self.num_dof, + dtype=torch.float, + device=self.device, + requires_grad=False) + self.Kd_factors = torch.ones( + self.num_envs, + self.num_dof, + dtype=torch.float, + device=self.device, + requires_grad=False) self.gravities = torch.zeros(self.num_envs, 3, dtype=torch.float, device=self.device, requires_grad=False) - self.gravity_vec = to_torch(get_axis_params(-1., self.up_axis_idx), device=self.device).repeat( - (self.num_envs, 1)) + self.gravity_vec = to_torch(get_axis_params(-1., self.up_axis_idx), + device=self.device).repeat((self.num_envs, 1)) # if custom initialization values were passed in, set them here - dynamics_params = ["friction_coeffs", "restitutions", "payloads", "com_displacements", "motor_strengths", - "Kp_factors", "Kd_factors"] + dynamics_params = [ + "friction_coeffs", + "restitutions", + "payloads", + "com_displacements", + "motor_strengths", + "Kp_factors", + "Kd_factors"] if self.initial_dynamics_dict is not None: for k, v in self.initial_dynamics_dict.items(): if k in dynamics_params: @@ -1291,10 +1615,14 @@ class LeggedRobot(BaseTask): requires_grad=False) self.clock_inputs = torch.zeros(self.num_envs, 4, dtype=torch.float, device=self.device, requires_grad=False) - self.doubletime_clock_inputs = torch.zeros(self.num_envs, 4, dtype=torch.float, device=self.device, - requires_grad=False) - self.halftime_clock_inputs = torch.zeros(self.num_envs, 4, dtype=torch.float, device=self.device, - requires_grad=False) + self.doubletime_clock_inputs = torch.zeros( + self.num_envs, 4, dtype=torch.float, device=self.device, requires_grad=False) + self.halftime_clock_inputs = torch.zeros( + self.num_envs, + 4, + dtype=torch.float, + device=self.device, + requires_grad=False) def _init_command_distribution(self, env_ids): # new style curriculum @@ -1348,11 +1676,11 @@ class LeggedRobot(BaseTask): self.cfg.commands.limit_stance_width[1], self.cfg.commands.num_bins_stance_width), stance_length=(self.cfg.commands.limit_stance_length[0], - self.cfg.commands.limit_stance_length[1], - self.cfg.commands.num_bins_stance_length), + self.cfg.commands.limit_stance_length[1], + self.cfg.commands.num_bins_stance_length), aux_reward_coef=(self.cfg.commands.limit_aux_reward_coef[0], - self.cfg.commands.limit_aux_reward_coef[1], - self.cfg.commands.num_bins_aux_reward_coef), + self.cfg.commands.limit_aux_reward_coef[1], + self.cfg.commands.num_bins_aux_reward_coef), )] if self.cfg.commands.curriculum_type == "LipschitzCurriculum": @@ -1361,24 +1689,38 @@ class LeggedRobot(BaseTask): binary_phases=self.cfg.commands.binary_phases) self.env_command_bins = np.zeros(len(env_ids), dtype=np.int) self.env_command_categories = np.zeros(len(env_ids), dtype=np.int) - low = np.array( - [self.cfg.commands.lin_vel_x[0], self.cfg.commands.lin_vel_y[0], - self.cfg.commands.ang_vel_yaw[0], self.cfg.commands.body_height_cmd[0], - self.cfg.commands.gait_frequency_cmd_range[0], - self.cfg.commands.gait_phase_cmd_range[0], self.cfg.commands.gait_offset_cmd_range[0], - self.cfg.commands.gait_bound_cmd_range[0], self.cfg.commands.gait_duration_cmd_range[0], - self.cfg.commands.footswing_height_range[0], self.cfg.commands.body_pitch_range[0], - self.cfg.commands.body_roll_range[0],self.cfg.commands.stance_width_range[0], - self.cfg.commands.stance_length_range[0], self.cfg.commands.aux_reward_coef_range[0], ]) - high = np.array( - [self.cfg.commands.lin_vel_x[1], self.cfg.commands.lin_vel_y[1], - self.cfg.commands.ang_vel_yaw[1], self.cfg.commands.body_height_cmd[1], - self.cfg.commands.gait_frequency_cmd_range[1], - self.cfg.commands.gait_phase_cmd_range[1], self.cfg.commands.gait_offset_cmd_range[1], - self.cfg.commands.gait_bound_cmd_range[1], self.cfg.commands.gait_duration_cmd_range[1], - self.cfg.commands.footswing_height_range[1], self.cfg.commands.body_pitch_range[1], - self.cfg.commands.body_roll_range[1],self.cfg.commands.stance_width_range[1], - self.cfg.commands.stance_length_range[1], self.cfg.commands.aux_reward_coef_range[1], ]) + low = np.array([self.cfg.commands.lin_vel_x[0], + self.cfg.commands.lin_vel_y[0], + self.cfg.commands.ang_vel_yaw[0], + self.cfg.commands.body_height_cmd[0], + self.cfg.commands.gait_frequency_cmd_range[0], + self.cfg.commands.gait_phase_cmd_range[0], + self.cfg.commands.gait_offset_cmd_range[0], + self.cfg.commands.gait_bound_cmd_range[0], + self.cfg.commands.gait_duration_cmd_range[0], + self.cfg.commands.footswing_height_range[0], + self.cfg.commands.body_pitch_range[0], + self.cfg.commands.body_roll_range[0], + self.cfg.commands.stance_width_range[0], + self.cfg.commands.stance_length_range[0], + self.cfg.commands.aux_reward_coef_range[0], + ]) + high = np.array([self.cfg.commands.lin_vel_x[1], + self.cfg.commands.lin_vel_y[1], + self.cfg.commands.ang_vel_yaw[1], + self.cfg.commands.body_height_cmd[1], + self.cfg.commands.gait_frequency_cmd_range[1], + self.cfg.commands.gait_phase_cmd_range[1], + self.cfg.commands.gait_offset_cmd_range[1], + self.cfg.commands.gait_bound_cmd_range[1], + self.cfg.commands.gait_duration_cmd_range[1], + self.cfg.commands.footswing_height_range[1], + self.cfg.commands.body_pitch_range[1], + self.cfg.commands.body_roll_range[1], + self.cfg.commands.stance_width_range[1], + self.cfg.commands.stance_length_range[1], + self.cfg.commands.aux_reward_coef_range[1], + ]) for curriculum in self.curricula: curriculum.set_to(low=low, high=high) @@ -1406,7 +1748,10 @@ class LeggedRobot(BaseTask): continue if not hasattr(self.reward_container, '_reward_' + name): - print(f"Warning: reward {'_reward_' + name} has nonzero coefficient but was not found!") + print( + f"Warning: reward { + '_reward_' + + name} has nonzero coefficient but was not found!") else: self.reward_names.append(name) self.reward_functions.append(getattr(self.reward_container, '_reward_' + name)) @@ -1415,13 +1760,16 @@ class LeggedRobot(BaseTask): self.episode_sums = { name: torch.zeros(self.num_envs, dtype=torch.float, device=self.device, requires_grad=False) for name in self.reward_scales.keys()} - self.episode_sums["total"] = torch.zeros(self.num_envs, dtype=torch.float, device=self.device, - requires_grad=False) + self.episode_sums["total"] = torch.zeros( + self.num_envs, + dtype=torch.float, + device=self.device, + requires_grad=False) self.episode_sums_eval = { name: -1 * torch.ones(self.num_envs, dtype=torch.float, device=self.device, requires_grad=False) for name in self.reward_scales.keys()} - self.episode_sums_eval["total"] = torch.zeros(self.num_envs, dtype=torch.float, device=self.device, - requires_grad=False) + self.episode_sums_eval["total"] = torch.zeros( + self.num_envs, dtype=torch.float, device=self.device, requires_grad=False) self.command_sums = { name: torch.zeros(self.num_envs, dtype=torch.float, device=self.device, requires_grad=False) for name in @@ -1457,8 +1805,11 @@ class LeggedRobot(BaseTask): print(self.terrain.heightsamples.shape, hf_params.nbRows, hf_params.nbColumns) self.gym.add_heightfield(self.sim, self.terrain.heightsamples.T, hf_params) - self.height_samples = torch.tensor(self.terrain.heightsamples).view(self.terrain.tot_rows, - self.terrain.tot_cols).to(self.device) + self.height_samples = torch.tensor( + self.terrain.heightsamples).view( + self.terrain.tot_rows, + self.terrain.tot_cols).to( + self.device) def _create_trimesh(self): """ Adds a triangle mesh terrain to the simulation, sets parameters based on the cfg. @@ -1475,14 +1826,17 @@ class LeggedRobot(BaseTask): tm_params.restitution = self.cfg.terrain.restitution self.gym.add_triangle_mesh(self.sim, self.terrain.vertices.flatten(order='C'), self.terrain.triangles.flatten(order='C'), tm_params) - self.height_samples = torch.tensor(self.terrain.heightsamples).view(self.terrain.tot_rows, - self.terrain.tot_cols).to(self.device) + self.height_samples = torch.tensor( + self.terrain.heightsamples).view( + self.terrain.tot_rows, + self.terrain.tot_cols).to( + self.device) def _create_envs(self): """ Creates environments: 1. loads the robot URDF/MJCF asset, 2. For each environment - 2.1 creates the environment, + 2.1 creates the environment, 2.2 calls DOF and Rigid shape properties callbacks, 2.3 create actor with these properties and add them to the env 3. Store indices of different bodies of the robot @@ -1526,16 +1880,31 @@ class LeggedRobot(BaseTask): for name in self.cfg.asset.terminate_after_contacts_on: termination_contact_names.extend([s for s in body_names if name in s]) - base_init_state_list = self.cfg.init_state.pos + self.cfg.init_state.rot + self.cfg.init_state.lin_vel + self.cfg.init_state.ang_vel - self.base_init_state = to_torch(base_init_state_list, device=self.device, requires_grad=False) + base_init_state_list = self.cfg.init_state.pos + self.cfg.init_state.rot + \ + self.cfg.init_state.lin_vel + self.cfg.init_state.ang_vel + self.base_init_state = to_torch( + base_init_state_list, + device=self.device, + requires_grad=False) start_pose = gymapi.Transform() start_pose.p = gymapi.Vec3(*self.base_init_state[:3]) self.env_origins = torch.zeros(self.num_envs, 3, device=self.device, requires_grad=False) - self.terrain_levels = torch.zeros(self.num_envs, device=self.device, requires_grad=False, dtype=torch.long) - self.terrain_origins = torch.zeros(self.num_envs, 3, device=self.device, requires_grad=False) - self.terrain_types = torch.zeros(self.num_envs, device=self.device, requires_grad=False, dtype=torch.long) - self._call_train_eval(self._get_env_origins, torch.arange(self.num_envs, device=self.device)) + self.terrain_levels = torch.zeros( + self.num_envs, + device=self.device, + requires_grad=False, + dtype=torch.long) + self.terrain_origins = torch.zeros( + self.num_envs, 3, device=self.device, requires_grad=False) + self.terrain_types = torch.zeros( + self.num_envs, + device=self.device, + requires_grad=False, + dtype=torch.long) + self._call_train_eval( + self._get_env_origins, torch.arange( + self.num_envs, device=self.device)) env_lower = gymapi.Vec3(0., 0., 0.) env_upper = gymapi.Vec3(0., 0., 0.) self.actor_handles = [] @@ -1545,60 +1914,90 @@ class LeggedRobot(BaseTask): self.default_friction = rigid_shape_props_asset[1].friction self.default_restitution = rigid_shape_props_asset[1].restitution self._init_custom_buffers__() - self._call_train_eval(self._randomize_rigid_body_props, torch.arange(self.num_envs, device=self.device)) + self._call_train_eval( + self._randomize_rigid_body_props, + torch.arange( + self.num_envs, + device=self.device)) self._randomize_gravity() for i in range(self.num_envs): # create env instance - env_handle = self.gym.create_env(self.sim, env_lower, env_upper, int(np.sqrt(self.num_envs))) + env_handle = self.gym.create_env( + self.sim, env_lower, env_upper, int( + np.sqrt( + self.num_envs))) pos = self.env_origins[i].clone() - pos[0:1] += torch_rand_float(-self.cfg.terrain.x_init_range, self.cfg.terrain.x_init_range, (1, 1), + pos[0:1] += torch_rand_float(-self.cfg.terrain.x_init_range, + self.cfg.terrain.x_init_range, + (1, + 1), device=self.device).squeeze(1) - pos[1:2] += torch_rand_float(-self.cfg.terrain.y_init_range, self.cfg.terrain.y_init_range, (1, 1), + pos[1:2] += torch_rand_float(-self.cfg.terrain.y_init_range, + self.cfg.terrain.y_init_range, + (1, + 1), device=self.device).squeeze(1) start_pose.p = gymapi.Vec3(*pos) rigid_shape_props = self._process_rigid_shape_props(rigid_shape_props_asset, i) self.gym.set_asset_rigid_shape_properties(self.robot_asset, rigid_shape_props) - anymal_handle = self.gym.create_actor(env_handle, self.robot_asset, start_pose, "anymal", i, - self.cfg.asset.self_collisions, 0) + anymal_handle = self.gym.create_actor( + env_handle, + self.robot_asset, + start_pose, + "anymal", + i, + self.cfg.asset.self_collisions, + 0) dof_props = self._process_dof_props(dof_props_asset, i) self.gym.set_actor_dof_properties(env_handle, anymal_handle, dof_props) body_props = self.gym.get_actor_rigid_body_properties(env_handle, anymal_handle) body_props = self._process_rigid_body_props(body_props, i) - self.gym.set_actor_rigid_body_properties(env_handle, anymal_handle, body_props, recomputeInertia=True) + self.gym.set_actor_rigid_body_properties( + env_handle, anymal_handle, body_props, recomputeInertia=True) self.envs.append(env_handle) self.actor_handles.append(anymal_handle) - self.feet_indices = torch.zeros(len(feet_names), dtype=torch.long, device=self.device, requires_grad=False) + self.feet_indices = torch.zeros( + len(feet_names), + dtype=torch.long, + device=self.device, + requires_grad=False) for i in range(len(feet_names)): - self.feet_indices[i] = self.gym.find_actor_rigid_body_handle(self.envs[0], self.actor_handles[0], - feet_names[i]) + self.feet_indices[i] = self.gym.find_actor_rigid_body_handle( + self.envs[0], self.actor_handles[0], feet_names[i]) - self.penalised_contact_indices = torch.zeros(len(penalized_contact_names), dtype=torch.long, device=self.device, - requires_grad=False) + self.penalised_contact_indices = torch.zeros( + len(penalized_contact_names), + dtype=torch.long, + device=self.device, + requires_grad=False) for i in range(len(penalized_contact_names)): - self.penalised_contact_indices[i] = self.gym.find_actor_rigid_body_handle(self.envs[0], - self.actor_handles[0], - penalized_contact_names[i]) + self.penalised_contact_indices[i] = self.gym.find_actor_rigid_body_handle( + self.envs[0], self.actor_handles[0], penalized_contact_names[i]) - self.termination_contact_indices = torch.zeros(len(termination_contact_names), dtype=torch.long, - device=self.device, requires_grad=False) + self.termination_contact_indices = torch.zeros( + len(termination_contact_names), + dtype=torch.long, + device=self.device, + requires_grad=False) for i in range(len(termination_contact_names)): - self.termination_contact_indices[i] = self.gym.find_actor_rigid_body_handle(self.envs[0], - self.actor_handles[0], - termination_contact_names[i]) + self.termination_contact_indices[i] = self.gym.find_actor_rigid_body_handle( + self.envs[0], self.actor_handles[0], termination_contact_names[i]) # if recording video, set up camera if self.cfg.env.record_video: self.camera_props = gymapi.CameraProperties() self.camera_props.width = 360 self.camera_props.height = 240 self.rendering_camera = self.gym.create_camera_sensor(self.envs[0], self.camera_props) - self.gym.set_camera_location(self.rendering_camera, self.envs[0], gymapi.Vec3(1.5, 1, 3.0), - gymapi.Vec3(0, 0, 0)) + self.gym.set_camera_location( + self.rendering_camera, self.envs[0], gymapi.Vec3( + 1.5, 1, 3.0), gymapi.Vec3( + 0, 0, 0)) if self.eval_cfg is not None: - self.rendering_camera_eval = self.gym.create_camera_sensor(self.envs[self.num_train_envs], - self.camera_props) + self.rendering_camera_eval = self.gym.create_camera_sensor( + self.envs[self.num_train_envs], self.camera_props) self.gym.set_camera_location(self.rendering_camera_eval, self.envs[self.num_train_envs], gymapi.Vec3(1.5, 1, 3.0), gymapi.Vec3(0, 0, 0)) @@ -1611,29 +2010,53 @@ class LeggedRobot(BaseTask): def render(self, mode="rgb_array"): assert mode == "rgb_array" bx, by, bz = self.root_states[0, 0], self.root_states[0, 1], self.root_states[0, 2] - self.gym.set_camera_location(self.rendering_camera, self.envs[0], gymapi.Vec3(bx, by - 1.0, bz + 1.0), - gymapi.Vec3(bx, by, bz)) + self.gym.set_camera_location( + self.rendering_camera, + self.envs[0], + gymapi.Vec3( + bx, + by - 1.0, + bz + 1.0), + gymapi.Vec3( + bx, + by, + bz)) self.gym.step_graphics(self.sim) self.gym.render_all_camera_sensors(self.sim) - img = self.gym.get_camera_image(self.sim, self.envs[0], self.rendering_camera, gymapi.IMAGE_COLOR) + img = self.gym.get_camera_image( + self.sim, + self.envs[0], + self.rendering_camera, + gymapi.IMAGE_COLOR) w, h = img.shape return img.reshape([w, h // 4, 4]) def _render_headless(self): - if self.record_now and self.complete_video_frames is not None and len(self.complete_video_frames) == 0: + if self.record_now and self.complete_video_frames is not None and len( + self.complete_video_frames) == 0: bx, by, bz = self.root_states[0, 0], self.root_states[0, 1], self.root_states[0, 2] - self.gym.set_camera_location(self.rendering_camera, self.envs[0], gymapi.Vec3(bx, by - 1.0, bz + 1.0), - gymapi.Vec3(bx, by, bz)) - self.video_frame = self.gym.get_camera_image(self.sim, self.envs[0], self.rendering_camera, - gymapi.IMAGE_COLOR) - self.video_frame = self.video_frame.reshape((self.camera_props.height, self.camera_props.width, 4)) + self.gym.set_camera_location( + self.rendering_camera, + self.envs[0], + gymapi.Vec3( + bx, + by - 1.0, + bz + 1.0), + gymapi.Vec3( + bx, + by, + bz)) + self.video_frame = self.gym.get_camera_image( + self.sim, self.envs[0], self.rendering_camera, gymapi.IMAGE_COLOR) + self.video_frame = self.video_frame.reshape( + (self.camera_props.height, self.camera_props.width, 4)) self.video_frames.append(self.video_frame) if self.record_eval_now and self.complete_video_frames_eval is not None and len( self.complete_video_frames_eval) == 0: if self.eval_cfg is not None: - bx, by, bz = self.root_states[self.num_train_envs, 0], self.root_states[self.num_train_envs, 1], \ - self.root_states[self.num_train_envs, 2] + bx, by, bz = self.root_states[self.num_train_envs, + 0], self.root_states[self.num_train_envs, 1], self.root_states[self.num_train_envs, 2] self.gym.set_camera_location(self.rendering_camera_eval, self.envs[self.num_train_envs], gymapi.Vec3(bx, by - 1.0, bz + 1.0), gymapi.Vec3(bx, by, bz)) @@ -1681,25 +2104,34 @@ class LeggedRobot(BaseTask): # put robots at the origins defined by the terrain max_init_level = cfg.terrain.max_init_terrain_level min_init_level = cfg.terrain.min_init_terrain_level - if not cfg.terrain.curriculum: max_init_level = cfg.terrain.num_rows - 1 - if not cfg.terrain.curriculum: min_init_level = 0 + if not cfg.terrain.curriculum: + max_init_level = cfg.terrain.num_rows - 1 + if not cfg.terrain.curriculum: + min_init_level = 0 if cfg.terrain.center_robots: min_terrain_level = cfg.terrain.num_rows // 2 - cfg.terrain.center_span max_terrain_level = cfg.terrain.num_rows // 2 + cfg.terrain.center_span - 1 min_terrain_type = cfg.terrain.num_cols // 2 - cfg.terrain.center_span max_terrain_type = cfg.terrain.num_cols // 2 + cfg.terrain.center_span - 1 - self.terrain_levels[env_ids] = torch.randint(min_terrain_level, max_terrain_level + 1, (len(env_ids),), - device=self.device) - self.terrain_types[env_ids] = torch.randint(min_terrain_type, max_terrain_type + 1, (len(env_ids),), - device=self.device) + self.terrain_levels[env_ids] = torch.randint( + min_terrain_level, max_terrain_level + 1, (len(env_ids),), device=self.device) + self.terrain_types[env_ids] = torch.randint( + min_terrain_type, max_terrain_type + 1, (len(env_ids),), device=self.device) else: - self.terrain_levels[env_ids] = torch.randint(min_init_level, max_init_level + 1, (len(env_ids),), - device=self.device) - self.terrain_types[env_ids] = torch.div(torch.arange(len(env_ids), device=self.device), - (len(env_ids) / cfg.terrain.num_cols), rounding_mode='floor').to( + self.terrain_levels[env_ids] = torch.randint( + min_init_level, max_init_level + 1, (len(env_ids),), device=self.device) + self.terrain_types[env_ids] = torch.div( + torch.arange( + len(env_ids), + device=self.device), + (len(env_ids) / cfg.terrain.num_cols), + rounding_mode='floor').to( torch.long) cfg.terrain.max_terrain_level = cfg.terrain.num_rows - cfg.terrain.terrain_origins = torch.from_numpy(cfg.terrain.env_origins).to(self.device).to(torch.float) + cfg.terrain.terrain_origins = torch.from_numpy( + cfg.terrain.env_origins).to( + self.device).to( + torch.float) self.env_origins[env_ids] = cfg.terrain.terrain_origins[ self.terrain_levels[env_ids], self.terrain_types[env_ids]] else: @@ -1727,7 +2159,8 @@ class LeggedRobot(BaseTask): cfg.domain_rand.push_interval = np.ceil(cfg.domain_rand.push_interval_s / self.dt) cfg.domain_rand.rand_interval = np.ceil(cfg.domain_rand.rand_interval_s / self.dt) - cfg.domain_rand.gravity_rand_interval = np.ceil(cfg.domain_rand.gravity_rand_interval_s / self.dt) + cfg.domain_rand.gravity_rand_interval = np.ceil( + cfg.domain_rand.gravity_rand_interval_s / self.dt) cfg.domain_rand.gravity_rand_duration = np.ceil( cfg.domain_rand.gravity_rand_interval * cfg.domain_rand.gravity_impulse_duration) @@ -1764,7 +2197,12 @@ class LeggedRobot(BaseTask): grid_x, grid_y = torch.meshgrid(x, y) cfg.env.num_height_points = grid_x.numel() - points = torch.zeros(len(env_ids), cfg.env.num_height_points, 3, device=self.device, requires_grad=False) + points = torch.zeros( + len(env_ids), + cfg.env.num_height_points, + 3, + device=self.device, + requires_grad=False) points[:, :, 0] = grid_x.flatten() points[:, :, 1] = grid_y.flatten() return points @@ -1783,7 +2221,11 @@ class LeggedRobot(BaseTask): [type]: [description] """ if cfg.terrain.mesh_type == 'plane': - return torch.zeros(len(env_ids), cfg.env.num_height_points, device=self.device, requires_grad=False) + return torch.zeros( + len(env_ids), + cfg.env.num_height_points, + device=self.device, + requires_grad=False) elif cfg.terrain.mesh_type == 'none': raise NameError("Can't measure height with terrain mesh type 'none'") diff --git a/Go2Py/sim/gym/envs/base/legged_robot_config.py b/Go2Py/sim/gym/envs/base/legged_robot_config.py index f3a06bd..0ef22ce 100755 --- a/Go2Py/sim/gym/envs/base/legged_robot_config.py +++ b/Go2Py/sim/gym/envs/base/legged_robot_config.py @@ -8,7 +8,8 @@ class Cfg(PrefixProto, cli=False): num_envs = 4096 num_observations = 235 num_scalar_observations = 42 - # if not None a privilige_obs_buf will be returned by step() (critic obs for assymetric training). None is returned otherwise + # if not None a privilige_obs_buf will be returned by step() (critic obs + # for assymetric training). None is returned otherwise num_privileged_obs = 18 privileged_future_horizon = 1 num_actions = 12 @@ -75,7 +76,8 @@ class Cfg(PrefixProto, cli=False): terrain_smoothness = 0.005 measure_heights = True # 1mx1.6m rectangle (without center line) - measured_points_x = [-0.8, -0.7, -0.6, -0.5, -0.4, -0.3, -0.2, -0.1, 0., 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8] + measured_points_x = [-0.8, -0.7, -0.6, -0.5, -0.4, -0.3, - + 0.2, -0.1, 0., 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8] measured_points_y = [-0.5, -0.4, -0.3, -0.2, -0.1, 0., 0.1, 0.2, 0.3, 0.4, 0.5] selected = False # select a unique terrain type and pass all arguments terrain_kwargs = None # Dict of arguments for selected terrain @@ -207,7 +209,7 @@ class Cfg(PrefixProto, cli=False): default_joint_angles = {"joint_a": 0., "joint_b": 0.} class control(PrefixProto, cli=False): - control_type = 'actuator_net' #'P' # P: position, V: velocity, T: torques + control_type = 'actuator_net' # 'P' # P: position, V: velocity, T: torques # PD Drive parameters: stiffness = {'joint_a': 10.0, 'joint_b': 15.} # [N*m/rad] damping = {'joint_a': 1.0, 'joint_b': 1.5} # [N*m*s/rad] @@ -223,10 +225,12 @@ class Cfg(PrefixProto, cli=False): penalize_contacts_on = [] terminate_after_contacts_on = [] disable_gravity = False - # merge bodies connected by fixed joints. Specific fixed joints can be kept by adding " <... dont_collapse="true"> + # merge bodies connected by fixed joints. Specific fixed joints can be + # kept by adding " <... dont_collapse="true"> collapse_fixed_joints = True fix_base_link = False # fixe the base of the robot - default_dof_drive_mode = 3 # see GymDofDriveModeFlags (0 is none, 1 is pos tgt, 2 is vel tgt, 3 effort) + # see GymDofDriveModeFlags (0 is none, 1 is pos tgt, 2 is vel tgt, 3 effort) + default_dof_drive_mode = 3 self_collisions = 0 # 1 to disable, 0 to enable...bitwise filter # replace collision cylinders with capsules, leads to faster/more stable simulation replace_cylinder_with_capsule = True @@ -270,7 +274,8 @@ class Cfg(PrefixProto, cli=False): lag_timesteps = 6 class rewards(PrefixProto, cli=False): - only_positive_rewards = True # if true negative total rewards are clipped at zero (avoids early termination problems) + # if true negative total rewards are clipped at zero (avoids early termination problems) + only_positive_rewards = True only_positive_rewards_ji22_style = False sigma_rew_neg = 5 reward_container_name = "CoRLRewards" diff --git a/Go2Py/sim/gym/envs/go2/velocity_tracking/__init__.py b/Go2Py/sim/gym/envs/go2/velocity_tracking/__init__.py index 64943f2..1742708 100644 --- a/Go2Py/sim/gym/envs/go2/velocity_tracking/__init__.py +++ b/Go2Py/sim/gym/envs/go2/velocity_tracking/__init__.py @@ -8,22 +8,37 @@ from Go2Py.sim.gym.envs.base.legged_robot_config import Cfg class VelocityTrackingEasyEnv(LeggedRobot): - def __init__(self, sim_device, headless, num_envs=None, prone=False, deploy=False, - cfg: Cfg = None, eval_cfg: Cfg = None, initial_dynamics_dict=None, physics_engine="SIM_PHYSX"): + def __init__( + self, + sim_device, + headless, + num_envs=None, + prone=False, + deploy=False, + cfg: Cfg = None, + eval_cfg: Cfg = None, + initial_dynamics_dict=None, + physics_engine="SIM_PHYSX"): if num_envs is not None: cfg.env.num_envs = num_envs - + sim_params = gymapi.SimParams() gymutil.parse_sim_config(vars(cfg.sim), sim_params) - super().__init__(cfg, sim_params, physics_engine, sim_device, headless, eval_cfg, initial_dynamics_dict) - + super().__init__( + cfg, + sim_params, + physics_engine, + sim_device, + headless, + eval_cfg, + initial_dynamics_dict) def step(self, actions): self.obs_buf, self.privileged_obs_buf, self.rew_buf, self.reset_buf, self.extras = super().step(actions) - self.foot_positions = self.rigid_body_state.view(self.num_envs, self.num_bodies, 13)[:, self.feet_indices, - 0:3] + self.foot_positions = self.rigid_body_state.view(self.num_envs, self.num_bodies, 13)[ + :, self.feet_indices, 0:3] self.extras.update({ "privileged_obs": self.privileged_obs_buf, @@ -45,6 +60,6 @@ class VelocityTrackingEasyEnv(LeggedRobot): def reset(self): self.reset_idx(torch.arange(self.num_envs, device=self.device)) - obs, _, _, _ = self.step(torch.zeros(self.num_envs, self.num_actions, device=self.device, requires_grad=False)) + obs, _, _, _ = self.step(torch.zeros(self.num_envs, self.num_actions, + device=self.device, requires_grad=False)) return obs - diff --git a/Go2Py/sim/gym/envs/rewards/corl_rewards.py b/Go2Py/sim/gym/envs/rewards/corl_rewards.py index 93b24e9..f4af703 100644 --- a/Go2Py/sim/gym/envs/rewards/corl_rewards.py +++ b/Go2Py/sim/gym/envs/rewards/corl_rewards.py @@ -4,6 +4,7 @@ from Go2Py.sim.gym.utils.math_utils import quat_apply_yaw, wrap_to_pi, get_scale from isaacgym.torch_utils import * from isaacgym import gymapi + class CoRLRewards: def __init__(self, env): self.env = env @@ -14,7 +15,8 @@ class CoRLRewards: # ------------ reward functions---------------- def _reward_tracking_lin_vel(self): # Tracking of linear velocity commands (xy axes) - lin_vel_error = torch.sum(torch.square(self.env.commands[:, :2] - self.env.base_lin_vel[:, :2]), dim=1) + lin_vel_error = torch.sum(torch.square( + self.env.commands[:, :2] - self.env.base_lin_vel[:, :2]), dim=1) return torch.exp(-lin_vel_error / self.env.cfg.rewards.tracking_sigma) def _reward_tracking_ang_vel(self): @@ -40,7 +42,12 @@ class CoRLRewards: def _reward_dof_acc(self): # Penalize dof accelerations - return torch.sum(torch.square((self.env.last_dof_vel - self.env.dof_vel) / self.env.dt), dim=1) + return torch.sum( + torch.square( + (self.env.last_dof_vel - + self.env.dof_vel) / + self.env.dt), + dim=1) def _reward_action_rate(self): # Penalize changes in actions @@ -48,12 +55,13 @@ class CoRLRewards: def _reward_collision(self): # Penalize collisions on selected bodies - return torch.sum(1. * (torch.norm(self.env.contact_forces[:, self.env.penalised_contact_indices, :], dim=-1) > 0.1), - dim=1) + return torch.sum( + 1. * (torch.norm(self.env.contact_forces[:, self.env.penalised_contact_indices, :], dim=-1) > 0.1), dim=1) def _reward_dof_pos_limits(self): # Penalize dof positions too close to the limit - out_of_limits = -(self.env.dof_pos - self.env.dof_pos_limits[:, 0]).clip(max=0.) # lower limit + out_of_limits = -(self.env.dof_pos - + self.env.dof_pos_limits[:, 0]).clip(max=0.) # lower limit out_of_limits += (self.env.dof_pos - self.env.dof_pos_limits[:, 1]).clip(min=0.) return torch.sum(out_of_limits, dim=1) @@ -71,7 +79,7 @@ class CoRLRewards: reward = 0 for i in range(4): reward += - (1 - desired_contact[:, i]) * ( - 1 - torch.exp(-1 * foot_forces[:, i] ** 2 / self.env.cfg.rewards.gait_force_sigma)) + 1 - torch.exp(-1 * foot_forces[:, i] ** 2 / self.env.cfg.rewards.gait_force_sigma)) return reward / 4 def _reward_tracking_contacts_shaped_vel(self): @@ -79,8 +87,9 @@ class CoRLRewards: desired_contact = self.env.desired_contact_states reward = 0 for i in range(4): - reward += - (desired_contact[:, i] * ( - 1 - torch.exp(-1 * foot_velocities[:, i] ** 2 / self.env.cfg.rewards.gait_vel_sigma))) + reward += - (desired_contact[:, + i] * (1 - torch.exp(-1 * foot_velocities[:, + i] ** 2 / self.env.cfg.rewards.gait_vel_sigma))) return reward / 4 def _reward_dof_pos(self): @@ -93,13 +102,18 @@ class CoRLRewards: def _reward_action_smoothness_1(self): # Penalize changes in actions - diff = torch.square(self.env.joint_pos_target[:, :self.env.num_actuated_dof] - self.env.last_joint_pos_target[:, :self.env.num_actuated_dof]) + diff = torch.square(self.env.joint_pos_target[:, + :self.env.num_actuated_dof] - self.env.last_joint_pos_target[:, + :self.env.num_actuated_dof]) diff = diff * (self.env.last_actions[:, :self.env.num_dof] != 0) # ignore first step return torch.sum(diff, dim=1) def _reward_action_smoothness_2(self): # Penalize changes in actions - diff = torch.square(self.env.joint_pos_target[:, :self.env.num_actuated_dof] - 2 * self.env.last_joint_pos_target[:, :self.env.num_actuated_dof] + self.env.last_last_joint_pos_target[:, :self.env.num_actuated_dof]) + diff = torch.square(self.env.joint_pos_target[:, + :self.env.num_actuated_dof] - 2 * self.env.last_joint_pos_target[:, + :self.env.num_actuated_dof] + self.env.last_last_joint_pos_target[:, + :self.env.num_actuated_dof]) diff = diff * (self.env.last_actions[:, :self.env.num_dof] != 0) # ignore first step diff = diff * (self.env.last_last_actions[:, :self.env.num_dof] != 0) # ignore second step return torch.sum(diff, dim=1) @@ -108,55 +122,66 @@ class CoRLRewards: contact = self.env.contact_forces[:, self.env.feet_indices, 2] > 1. contact_filt = torch.logical_or(contact, self.env.last_contacts) self.env.last_contacts = contact - foot_velocities = torch.square(torch.norm(self.env.foot_velocities[:, :, 0:2], dim=2).view(self.env.num_envs, -1)) + foot_velocities = torch.square(torch.norm( + self.env.foot_velocities[:, :, 0:2], dim=2).view(self.env.num_envs, -1)) rew_slip = torch.sum(contact_filt * foot_velocities, dim=1) return rew_slip def _reward_feet_contact_vel(self): reference_heights = 0 near_ground = self.env.foot_positions[:, :, 2] - reference_heights < 0.03 - foot_velocities = torch.square(torch.norm(self.env.foot_velocities[:, :, 0:3], dim=2).view(self.env.num_envs, -1)) + foot_velocities = torch.square(torch.norm( + self.env.foot_velocities[:, :, 0:3], dim=2).view(self.env.num_envs, -1)) rew_contact_vel = torch.sum(near_ground * foot_velocities, dim=1) return rew_contact_vel def _reward_feet_contact_forces(self): # penalize high contact forces - return torch.sum((torch.norm(self.env.contact_forces[:, self.env.feet_indices, :], - dim=-1) - self.env.cfg.rewards.max_contact_force).clip(min=0.), dim=1) + return torch.sum((torch.norm(self.env.contact_forces[:, + self.env.feet_indices, + :], + dim=-1) - self.env.cfg.rewards.max_contact_force).clip(min=0.), + dim=1) def _reward_feet_clearance_cmd_linear(self): - phases = 1 - torch.abs(1.0 - torch.clip((self.env.foot_indices * 2.0) - 1.0, 0.0, 1.0) * 2.0) - foot_height = (self.env.foot_positions[:, :, 2]).view(self.env.num_envs, -1)# - reference_heights - target_height = self.env.commands[:, 9].unsqueeze(1) * phases + 0.02 # offset for foot radius 2cm - rew_foot_clearance = torch.square(target_height - foot_height) * (1 - self.env.desired_contact_states) + phases = 1 - \ + torch.abs(1.0 - torch.clip((self.env.foot_indices * 2.0) - 1.0, 0.0, 1.0) * 2.0) + foot_height = (self.env.foot_positions[:, :, 2]).view( + self.env.num_envs, -1) # - reference_heights + target_height = self.env.commands[:, 9].unsqueeze( + 1) * phases + 0.02 # offset for foot radius 2cm + rew_foot_clearance = torch.square( + target_height - foot_height) * (1 - self.env.desired_contact_states) return torch.sum(rew_foot_clearance, dim=1) def _reward_feet_impact_vel(self): prev_foot_velocities = self.env.prev_foot_velocities[:, :, 2].view(self.env.num_envs, -1) - contact_states = torch.norm(self.env.contact_forces[:, self.env.feet_indices, :], dim=-1) > 1.0 + contact_states = torch.norm( + self.env.contact_forces[:, self.env.feet_indices, :], dim=-1) > 1.0 - rew_foot_impact_vel = contact_states * torch.square(torch.clip(prev_foot_velocities, -100, 0)) + rew_foot_impact_vel = contact_states * \ + torch.square(torch.clip(prev_foot_velocities, -100, 0)) return torch.sum(rew_foot_impact_vel, dim=1) - def _reward_collision(self): # Penalize collisions on selected bodies - return torch.sum(1. * (torch.norm(self.env.contact_forces[:, self.env.penalised_contact_indices, :], dim=-1) > 0.1), - dim=1) + return torch.sum( + 1. * (torch.norm(self.env.contact_forces[:, self.env.penalised_contact_indices, :], dim=-1) > 0.1), dim=1) def _reward_orientation_control(self): # Penalize non flat base orientation roll_pitch_commands = self.env.commands[:, 10:12] - quat_roll = quat_from_angle_axis(-roll_pitch_commands[:, 1], - torch.tensor([1, 0, 0], device=self.env.device, dtype=torch.float)) - quat_pitch = quat_from_angle_axis(-roll_pitch_commands[:, 0], - torch.tensor([0, 1, 0], device=self.env.device, dtype=torch.float)) + quat_roll = quat_from_angle_axis(-roll_pitch_commands[:, 1], torch.tensor( + [1, 0, 0], device=self.env.device, dtype=torch.float)) + quat_pitch = quat_from_angle_axis(-roll_pitch_commands[:, 0], torch.tensor( + [0, 1, 0], device=self.env.device, dtype=torch.float)) desired_base_quat = quat_mul(quat_roll, quat_pitch) desired_projected_gravity = quat_rotate_inverse(desired_base_quat, self.env.gravity_vec) - return torch.sum(torch.square(self.env.projected_gravity[:, :2] - desired_projected_gravity[:, :2]), dim=1) + return torch.sum(torch.square( + self.env.projected_gravity[:, :2] - desired_projected_gravity[:, :2]), dim=1) def _reward_raibert_heuristic(self): cur_footsteps_translated = self.env.foot_positions - self.env.base_pos.unsqueeze(1) @@ -168,17 +193,41 @@ class CoRLRewards: # nominal positions: [FR, FL, RR, RL] if self.env.cfg.commands.num_commands >= 13: desired_stance_width = self.env.commands[:, 12:13] - desired_ys_nom = torch.cat([desired_stance_width / 2, -desired_stance_width / 2, desired_stance_width / 2, -desired_stance_width / 2], dim=1) + desired_ys_nom = torch.cat([desired_stance_width / + 2, - + desired_stance_width / + 2, desired_stance_width / + 2, - + desired_stance_width / + 2], dim=1) else: desired_stance_width = 0.3 - desired_ys_nom = torch.tensor([desired_stance_width / 2, -desired_stance_width / 2, desired_stance_width / 2, -desired_stance_width / 2], device=self.env.device).unsqueeze(0) + desired_ys_nom = torch.tensor([desired_stance_width / + 2, - + desired_stance_width / + 2, desired_stance_width / + 2, - + desired_stance_width / + 2], device=self.env.device).unsqueeze(0) if self.env.cfg.commands.num_commands >= 14: desired_stance_length = self.env.commands[:, 13:14] - desired_xs_nom = torch.cat([desired_stance_length / 2, desired_stance_length / 2, -desired_stance_length / 2, -desired_stance_length / 2], dim=1) + desired_xs_nom = torch.cat([desired_stance_length / + 2, desired_stance_length / + 2, - + desired_stance_length / + 2, - + desired_stance_length / + 2], dim=1) else: desired_stance_length = 0.45 - desired_xs_nom = torch.tensor([desired_stance_length / 2, desired_stance_length / 2, -desired_stance_length / 2, -desired_stance_length / 2], device=self.env.device).unsqueeze(0) + desired_xs_nom = torch.tensor([desired_stance_length / + 2, desired_stance_length / + 2, - + desired_stance_length / + 2, - + desired_stance_length / + 2], device=self.env.device).unsqueeze(0) # raibert offsets phases = torch.abs(1.0 - (self.env.foot_indices * 2.0)) * 1.0 - 0.5 @@ -193,10 +242,12 @@ class CoRLRewards: desired_ys_nom = desired_ys_nom + desired_ys_offset desired_xs_nom = desired_xs_nom + desired_xs_offset - desired_footsteps_body_frame = torch.cat((desired_xs_nom.unsqueeze(2), desired_ys_nom.unsqueeze(2)), dim=2) + desired_footsteps_body_frame = torch.cat( + (desired_xs_nom.unsqueeze(2), desired_ys_nom.unsqueeze(2)), dim=2) - err_raibert_heuristic = torch.abs(desired_footsteps_body_frame - footsteps_in_body_frame[:, :, 0:2]) + err_raibert_heuristic = torch.abs( + desired_footsteps_body_frame - footsteps_in_body_frame[:, :, 0:2]) reward = torch.sum(torch.square(err_raibert_heuristic), dim=(1, 2)) - return reward \ No newline at end of file + return reward diff --git a/Go2Py/sim/gym/envs/wrappers/history_wrapper.py b/Go2Py/sim/gym/envs/wrappers/history_wrapper.py index a2f6976..ced20b0 100755 --- a/Go2Py/sim/gym/envs/wrappers/history_wrapper.py +++ b/Go2Py/sim/gym/envs/wrappers/history_wrapper.py @@ -1,7 +1,8 @@ +import gym +import torch import isaacgym assert isaacgym -import torch -import gym + class HistoryWrapper(gym.Wrapper): def __init__(self, env): @@ -21,7 +22,8 @@ class HistoryWrapper(gym.Wrapper): privileged_obs = info["privileged_obs"] self.obs_history = torch.cat((self.obs_history[:, self.env.num_obs:], obs), dim=-1) - return {'obs': obs, 'privileged_obs': privileged_obs, 'obs_history': self.obs_history}, rew, done, info + return {'obs': obs, 'privileged_obs': privileged_obs, + 'obs_history': self.obs_history}, rew, done, info def get_observations(self): obs = self.env.get_observations() diff --git a/Go2Py/sim/gym/utils/terrain.py b/Go2Py/sim/gym/utils/terrain.py index 6d1f9c1..340af5f 100755 --- a/Go2Py/sim/gym/utils/terrain.py +++ b/Go2Py/sim/gym/utils/terrain.py @@ -30,10 +30,8 @@ class Terrain: self.heightsamples = self.height_field_raw if self.type == "trimesh": - self.vertices, self.triangles = terrain_utils.convert_heightfield_to_trimesh(self.height_field_raw, - self.cfg.horizontal_scale, - self.cfg.vertical_scale, - self.cfg.slope_treshold) + self.vertices, self.triangles = terrain_utils.convert_heightfield_to_trimesh( + self.height_field_raw, self.cfg.horizontal_scale, self.cfg.vertical_scale, self.cfg.slope_treshold) def load_cfgs(self): self._load_cfg(self.cfg) @@ -45,14 +43,16 @@ class Terrain: return self.cfg.row_indices, self.cfg.col_indices, [], [] else: self._load_cfg(self.eval_cfg) - self.eval_cfg.row_indices = np.arange(self.cfg.tot_rows, self.cfg.tot_rows + self.eval_cfg.tot_rows) + self.eval_cfg.row_indices = np.arange( + self.cfg.tot_rows, self.cfg.tot_rows + self.eval_cfg.tot_rows) self.eval_cfg.col_indices = np.arange(0, self.eval_cfg.tot_cols) self.eval_cfg.x_offset = self.cfg.tot_rows self.eval_cfg.rows_offset = self.cfg.num_rows return self.cfg.row_indices, self.cfg.col_indices, self.eval_cfg.row_indices, self.eval_cfg.col_indices def _load_cfg(self, cfg): - cfg.proportions = [np.sum(cfg.terrain_proportions[:i + 1]) for i in range(len(cfg.terrain_proportions))] + cfg.proportions = [np.sum(cfg.terrain_proportions[:i + 1]) + for i in range(len(cfg.terrain_proportions))] cfg.num_sub_terrains = cfg.num_rows * cfg.num_cols cfg.env_origins = np.zeros((cfg.num_rows, cfg.num_cols, 3)) @@ -128,32 +128,54 @@ class Terrain: terrain_utils.pyramid_sloped_terrain(terrain, slope=slope, platform_size=3.) elif choice < proportions[1]: terrain_utils.pyramid_sloped_terrain(terrain, slope=slope, platform_size=3.) - terrain_utils.random_uniform_terrain(terrain, min_height=-0.05, max_height=0.05, - step=self.cfg.terrain_smoothness, downsampled_scale=0.2) + terrain_utils.random_uniform_terrain( + terrain, + min_height=-0.05, + max_height=0.05, + step=self.cfg.terrain_smoothness, + downsampled_scale=0.2) elif choice < proportions[3]: if choice < proportions[2]: step_height *= -1 - terrain_utils.pyramid_stairs_terrain(terrain, step_width=0.31, step_height=step_height, platform_size=3.) + terrain_utils.pyramid_stairs_terrain( + terrain, step_width=0.31, step_height=step_height, platform_size=3.) elif choice < proportions[4]: num_rectangles = 20 rectangle_min_size = 1. rectangle_max_size = 2. - terrain_utils.discrete_obstacles_terrain(terrain, discrete_obstacles_height, rectangle_min_size, - rectangle_max_size, num_rectangles, platform_size=3.) + terrain_utils.discrete_obstacles_terrain( + terrain, + discrete_obstacles_height, + rectangle_min_size, + rectangle_max_size, + num_rectangles, + platform_size=3.) elif choice < proportions[5]: - terrain_utils.stepping_stones_terrain(terrain, stone_size=stepping_stones_size, - stone_distance=stone_distance, max_height=0., platform_size=4.) + terrain_utils.stepping_stones_terrain( + terrain, + stone_size=stepping_stones_size, + stone_distance=stone_distance, + max_height=0., + platform_size=4.) elif choice < proportions[6]: pass elif choice < proportions[7]: pass elif choice < proportions[8]: - terrain_utils.random_uniform_terrain(terrain, min_height=-cfg.terrain_noise_magnitude, - max_height=cfg.terrain_noise_magnitude, step=0.005, - downsampled_scale=0.2) + terrain_utils.random_uniform_terrain( + terrain, + min_height=- + cfg.terrain_noise_magnitude, + max_height=cfg.terrain_noise_magnitude, + step=0.005, + downsampled_scale=0.2) elif choice < proportions[9]: - terrain_utils.random_uniform_terrain(terrain, min_height=-0.05, max_height=0.05, - step=self.cfg.terrain_smoothness, downsampled_scale=0.2) + terrain_utils.random_uniform_terrain( + terrain, + min_height=-0.05, + max_height=0.05, + step=self.cfg.terrain_smoothness, + downsampled_scale=0.2) terrain.height_field_raw[0:terrain.length // 2, :] = 0 return terrain @@ -174,7 +196,7 @@ class Terrain: x2 = int((cfg.terrain_length / 2. + 1) / terrain.horizontal_scale) + cfg.x_offset y1 = int((cfg.terrain_width / 2. - 1) / terrain.horizontal_scale) y2 = int((cfg.terrain_width / 2. + 1) / terrain.horizontal_scale) - env_origin_z = np.max(self.height_field_raw[start_x: end_x, start_y:end_y]) * terrain.vertical_scale + env_origin_z = np.max( + self.height_field_raw[start_x: end_x, start_y:end_y]) * terrain.vertical_scale cfg.env_origins[i, j] = [env_origin_x, env_origin_y, env_origin_z] - diff --git a/Go2Py/sim/isaacsim/go2.py b/Go2Py/sim/isaacsim/go2.py index c51e004..b18e0e7 100644 --- a/Go2Py/sim/isaacsim/go2.py +++ b/Go2Py/sim/isaacsim/go2.py @@ -42,7 +42,7 @@ class UnitreeGo2(Articulation): prim.GetReferences().AddReference(self.usd_path) breakpoint() super().__init__( - prim_path=self._prim_path+'/Go2/base', + prim_path=self._prim_path + '/Go2/base', name=name, position=position, orientation=orientation, diff --git a/Go2Py/sim/isaacsim/isaacsim_node.py b/Go2Py/sim/isaacsim/isaacsim_node.py index 24adaea..12b338b 100644 --- a/Go2Py/sim/isaacsim/isaacsim_node.py +++ b/Go2Py/sim/isaacsim/isaacsim_node.py @@ -1,5 +1,18 @@ # launch Isaac Sim before any other imports # default first two lines in any standalone application +from Go2Py.sim.utils import ( + NumpyMemMapDataPipe, + load_config, + simulationManager, +) +from Go2Py.sim.isaacsim.utils import AnnotatorManager +from Go2Py.sim.isaacsim.go2 import UnitreeGo2 +from Go2Py.sim.isaacsim.lcm_types.unitree_lowlevel import UnitreeLowCommand +import Go2Py +from omni.isaac.sensor import RotatingLidarPhysX +from omni.isaac.core.utils.prims import define_prim, get_prim_at_path +from omni.isaac.core.utils.nucleus import get_assets_root_path +from omni.isaac.core import World import pickle import sys import time @@ -9,20 +22,7 @@ from omni.isaac.kit import SimulationApp simulation_app = SimulationApp({"headless": False}) # we can also run as headless. # import cv2 -from omni.isaac.core import World -from omni.isaac.core.utils.nucleus import get_assets_root_path -from omni.isaac.core.utils.prims import define_prim, get_prim_at_path -from omni.isaac.sensor import RotatingLidarPhysX -import Go2Py -from Go2Py.sim.isaacsim.lcm_types.unitree_lowlevel import UnitreeLowCommand -from Go2Py.sim.isaacsim.go2 import UnitreeGo2 -from Go2Py.sim.isaacsim.utils import AnnotatorManager -from Go2Py.sim.utils import ( - NumpyMemMapDataPipe, - load_config, - simulationManager, -) cfg = load_config(Go2Py.GO2_ISAACSIM_CFG_PATH) robots = cfg["robots"] @@ -154,7 +154,7 @@ world.reset() go2.initialize() # while simulation_app.is_running(): # world.step(render=True) - + # breakpoint() # sim_manager = simulationManager( diff --git a/Go2Py/sim/mujoco.py b/Go2Py/sim/mujoco.py index 478ecd7..954115e 100644 --- a/Go2Py/sim/mujoco.py +++ b/Go2Py/sim/mujoco.py @@ -9,17 +9,19 @@ from pinocchio.robot_wrapper import RobotWrapper from Go2Py import ASSETS_PATH import os from scipy.spatial.transform import Rotation + + class Go2Sim: def __init__(self, render=True, dt=0.002): - + self.model = mujoco.MjModel.from_xml_path( os.path.join(ASSETS_PATH, 'mujoco/go2.xml') ) self.simulated = True self.data = mujoco.MjData(self.model) self.dt = dt - _render_dt = 1/60 - self.render_ds_ratio = max(1, _render_dt//dt) + _render_dt = 1 / 60 + self.render_ds_ratio = max(1, _render_dt // dt) if render: self.viewer = mujoco.viewer.launch_passive(self.model, self.data) @@ -36,21 +38,21 @@ class Go2Sim: self.renderer = None self.render = render self.step_counter = 0 - - self.prestanding_q = np.array([ 0.0, 1.26186061, -2.5, - 0.0, 1.25883281, -2.5, - 0.0, 1.27193761, -2.6, - 0.0, 1.27148342, -2.6]) - self.sitting_q = np.array([-0.02495611, 1.26249647, -2.82826662, - 0.04563564, 1.2505368 , -2.7933557 , - -0.30623949, 1.28283751, -2.82314873, - 0.26400229, 1.29355574, -2.84276843]) + self.prestanding_q = np.array([0.0, 1.26186061, -2.5, + 0.0, 1.25883281, -2.5, + 0.0, 1.27193761, -2.6, + 0.0, 1.27148342, -2.6]) - self.standing_q = np.array([ 0.0, 0.77832842, -1.56065452, - 0.0, 0.76754963, -1.56634164, - 0.0, 0.76681757, -1.53601146, - 0.0, 0.75422204, -1.53229916]) + self.sitting_q = np.array([-0.02495611, 1.26249647, -2.82826662, + 0.04563564, 1.2505368, -2.7933557, + -0.30623949, 1.28283751, -2.82314873, + 0.26400229, 1.29355574, -2.84276843]) + + self.standing_q = np.array([0.0, 0.77832842, -1.56065452, + 0.0, 0.76754963, -1.56634164, + 0.0, 0.76681757, -1.53601146, + 0.0, 0.75422204, -1.53229916]) self.q0 = self.sitting_q self.pos0 = np.array([0., 0., 0.1]) @@ -77,7 +79,7 @@ class Go2Sim: ) self.data.qpos = self.q_nominal self.data.qvel = np.zeros(18) - + def standUpReset(self): self.q0 = self.standing_q self.pos0 = np.array([0., 0., 0.33]) @@ -95,18 +97,18 @@ class Go2Sim: self.viewer.sync() def getJointStates(self): - return {"q":self.data.qpos[7:], - "dq":self.data.qvel[6:], - 'tau_est':self.actuator_tau} + return {"q": self.data.qpos[7:], + "dq": self.data.qvel[6:], + 'tau_est': self.actuator_tau} def getPose(self): return self.data.qpos[:3], self.data.qpos[3:7] def getIMU(self): return { - 'accel':np.array(self.data.sensordata[0:3]),\ - 'gyro': np.array(self.data.sensordata[3:6]) - } + 'accel': np.array(self.data.sensordata[0:3]), + 'gyro': np.array(self.data.sensordata[3:6]) + } def getFootContact(self): return self.data.sensordata[6:10] @@ -118,39 +120,39 @@ class Go2Sim: self.kv = kv self.tau_ff = tau_ff self.latest_command_stamp = time.time() - + def step(self): state = self.getJointStates() q, dq = state['q'], state['dq'] - tau = np.diag(self.kp)@(self.q_des-q).reshape(12,1)+ \ - np.diag(self.kv)@(self.dq_des-dq).reshape(12,1)+self.tau_ff.reshape(12,1) + tau = np.diag(self.kp) @ (self.q_des - q).reshape(12, 1) + \ + np.diag(self.kv) @ (self.dq_des - dq).reshape(12, 1) + self.tau_ff.reshape(12, 1) self.actuator_tau = tau self.data.ctrl[:] = tau.squeeze() self.step_counter += 1 mujoco.mj_step(self.model, self.data) # Render every render_ds_ratio steps (60Hz GUI update) - if self.render and (self.step_counter%self.render_ds_ratio)==0: + if self.render and (self.step_counter % self.render_ds_ratio) == 0: self.viewer.sync() def getSiteJacobian(self, site_name): - id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE,site_name) - assert id>0, 'The requested site could not be found' + id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE, site_name) + assert id > 0, 'The requested site could not be found' mujoco.mj_jacSite(self.model, self.data, self.jacp, self.jacr, id) return self.jacp, self.jacr def getDynamicsParams(self): mujoco.mj_fullM(self.model, self.M, self.data.qM) - nle = self.data.qfrc_bias.reshape(self.nv,1) + nle = self.data.qfrc_bias.reshape(self.nv, 1) return { - 'M':self.M, - 'nle':nle + 'M': self.M, + 'nle': nle } def getGravityInBody(self): _, q = self.getPose() R = Rotation.from_quat([q[1], q[2], q[3], q[0]]).as_matrix() - g_in_body = R.T@np.array([0.0, 0.0, -1.0]).reshape(3, 1) + g_in_body = R.T @ np.array([0.0, 0.0, -1.0]).reshape(3, 1) return g_in_body def overheat(self): diff --git a/Go2Py/sim/utils.py b/Go2Py/sim/utils.py index 993accb..4922f7d 100644 --- a/Go2Py/sim/utils.py +++ b/Go2Py/sim/utils.py @@ -10,6 +10,7 @@ def load_config(file_path): config = yaml.safe_load(file) return config + class NumpyMemMapDataPipe: def __init__(self, channel_name, force=False, dtype="uint8", shape=(640, 480, 3)): self.channel_name = channel_name diff --git a/Go2Py/simple_lowcmd_sender.py b/Go2Py/simple_lowcmd_sender.py index 2a621df..2932a1d 100644 --- a/Go2Py/simple_lowcmd_sender.py +++ b/Go2Py/simple_lowcmd_sender.py @@ -15,13 +15,13 @@ dp = DomainParticipant(0) tp = Topic(dp, "go2py/lowcmd", LowCmd) dw = DataWriter(dp, tp) cmd = LowCmd( - q = 12*[0.], - dq = 12*[0.], - tau_ff = 12*[0.], - kp = 12*[0.], - kv = 12*[0.], - e_stop = 0 + q=12 * [0.], + dq=12 * [0.], + tau_ff=12 * [0.], + kp=12 * [0.], + kv=12 * [0.], + e_stop=0 ) while True: dw.write(cmd) - time.sleep(0.01) \ No newline at end of file + time.sleep(0.01) diff --git a/deploy/launch_files/bridge.launch.py b/deploy/launch_files/bridge.launch.py index 6554524..1b36bbe 100644 --- a/deploy/launch_files/bridge.launch.py +++ b/deploy/launch_files/bridge.launch.py @@ -7,6 +7,7 @@ from launch.substitutions import Command, FindExecutable, LaunchConfiguration from ament_index_python.packages import get_package_share_directory import os + def generate_launch_description(): # go2_xacro_file = os.path.join( # get_package_share_directory("go2_description"), "xacro", "robot_virtual_arm.xacro" @@ -49,4 +50,4 @@ def generate_launch_description(): # ], # name="static_tf_pub_trunk_to_lidar", # ), - ]) \ No newline at end of file + ]) diff --git a/deploy/launch_files/hesai.launch.py b/deploy/launch_files/hesai.launch.py index 4900018..a5a7b78 100644 --- a/deploy/launch_files/hesai.launch.py +++ b/deploy/launch_files/hesai.launch.py @@ -33,4 +33,4 @@ def generate_launch_description(): # # 'tf_publish_rate': 0.0 # }] # ) - ]) \ No newline at end of file + ]) diff --git a/deploy/launch_files/robot_description.launch.py b/deploy/launch_files/robot_description.launch.py index 1f7b2ed..6817531 100644 --- a/deploy/launch_files/robot_description.launch.py +++ b/deploy/launch_files/robot_description.launch.py @@ -7,6 +7,7 @@ from launch.substitutions import Command, FindExecutable, LaunchConfiguration from ament_index_python.packages import get_package_share_directory import os + def generate_launch_description(): go2_xacro_file = os.path.join( get_package_share_directory("go2_description"), "xacro", "robot.xacro" @@ -15,22 +16,22 @@ def generate_launch_description(): [FindExecutable(name="xacro"), " ", go2_xacro_file, " DEBUG:=", 'false'] ) - return LaunchDescription([ + return LaunchDescription([ Node( - package="robot_state_publisher", - executable="robot_state_publisher", - name="robot_state_publisher", - output="screen", - parameters=[{"robot_description": robot_description}], - remappings=[ + package="robot_state_publisher", + executable="robot_state_publisher", + name="robot_state_publisher", + output="screen", + parameters=[{"robot_description": robot_description}], + remappings=[ ("/joint_states", "/go2/joint_states"), ("/robot_description", "/go2/robot_description"), - ], - ), + ], + ), Node( - package="tf2_ros", - executable="static_transform_publisher", - arguments=[ + package="tf2_ros", + executable="static_transform_publisher", + arguments=[ "0.15", "0", "0.15", @@ -40,7 +41,7 @@ def generate_launch_description(): "0.707107", "/trunk", "/go2/hesai_lidar", - ], - name="static_tf_pub_trunk_to_lidar", - ), - ]) \ No newline at end of file + ], + name="static_tf_pub_trunk_to_lidar", + ), + ]) diff --git a/deploy/ros2_nodes/lidar_node/launch/dashing_start.py b/deploy/ros2_nodes/lidar_node/launch/dashing_start.py index 84ba352..fb7ea6b 100755 --- a/deploy/ros2_nodes/lidar_node/launch/dashing_start.py +++ b/deploy/ros2_nodes/lidar_node/launch/dashing_start.py @@ -1,8 +1,10 @@ from launch import LaunchDescription from launch_ros.actions import Node from ament_index_python.packages import get_package_share_directory + + def generate_launch_description(): - rviz_config=get_package_share_directory('hesai_ros_driver')+'/rviz/rviz2.rviz' + rviz_config = get_package_share_directory('hesai_ros_driver') + '/rviz/rviz2.rviz' return LaunchDescription([ Node( package='hesai_ros_driver', @@ -16,6 +18,6 @@ def generate_launch_description(): node_namespace='rviz2', node_name='rviz2', node_executable='rviz2', - arguments=['-d',rviz_config] + arguments=['-d', rviz_config] ) ]) diff --git a/deploy/ros2_nodes/lidar_node/launch/start.py b/deploy/ros2_nodes/lidar_node/launch/start.py index b670762..5fcb9d2 100755 --- a/deploy/ros2_nodes/lidar_node/launch/start.py +++ b/deploy/ros2_nodes/lidar_node/launch/start.py @@ -2,11 +2,22 @@ from launch import LaunchDescription from launch_ros.actions import Node from ament_index_python.packages import get_package_share_directory + def generate_launch_description(): - rviz_config=get_package_share_directory('hesai_ros_driver')+'/rviz/rviz2.rviz' + rviz_config = get_package_share_directory('hesai_ros_driver') + '/rviz/rviz2.rviz' - return LaunchDescription([ - Node(namespace='hesai_ros_driver', package='hesai_ros_driver', executable='hesai_ros_driver_node', output='screen'), - Node(namespace='rviz2', package='rviz2', executable='rviz2', arguments=['-d',rviz_config]) - ]) + return LaunchDescription( + [ + Node( + namespace='hesai_ros_driver', + package='hesai_ros_driver', + executable='hesai_ros_driver_node', + output='screen'), + Node( + namespace='rviz2', + package='rviz2', + executable='rviz2', + arguments=[ + '-d', + rviz_config])]) diff --git a/deploy/ros2_nodes/sportmode_nav2/launch/cbf_navigation_launch.py b/deploy/ros2_nodes/sportmode_nav2/launch/cbf_navigation_launch.py index 8d04f98..1376065 100644 --- a/deploy/ros2_nodes/sportmode_nav2/launch/cbf_navigation_launch.py +++ b/deploy/ros2_nodes/sportmode_nav2/launch/cbf_navigation_launch.py @@ -55,7 +55,7 @@ def generate_launch_description(): # TODO(orduno) Substitute with `PushNodeRemapping` # https://github.com/ros2/launch_ros/issues/56 remappings = [('/tf', 'tf'), - ('/tf_static', 'tf_static'),('/cmd_vel', '/nav2/cmd_vel')] + ('/tf_static', 'tf_static'), ('/cmd_vel', '/nav2/cmd_vel')] # Create our own temporary YAML files that include substitutions param_substitutions = { @@ -239,7 +239,7 @@ def generate_launch_description(): name='velocity_smoother', parameters=[configured_params], remappings=remappings + - [('cmd_vel', 'cmd_vel_nav'), ('cmd_vel_smoothed', 'cmd_vel')]), + [('cmd_vel', 'cmd_vel_nav'), ('cmd_vel_smoothed', 'cmd_vel')]), ComposableNode( package='nav2_lifecycle_manager', plugin='nav2_lifecycle_manager::LifecycleManager', diff --git a/deploy/ros2_nodes/sportmode_nav2/launch/localization.launch.py b/deploy/ros2_nodes/sportmode_nav2/launch/localization.launch.py index 839f1ce..c7b62de 100644 --- a/deploy/ros2_nodes/sportmode_nav2/launch/localization.launch.py +++ b/deploy/ros2_nodes/sportmode_nav2/launch/localization.launch.py @@ -9,19 +9,19 @@ from launch_ros.actions import Node def generate_launch_description(): - + robot_localization = IncludeLaunchDescription( - PythonLaunchDescriptionSource([os.path.join( - get_package_share_directory('sportmode_nav2'), 'launch'), - '/ros_ekf.launch.py']) - ) + PythonLaunchDescriptionSource([os.path.join( + get_package_share_directory('sportmode_nav2'), 'launch'), + '/ros_ekf.launch.py']) + ) async_mapping = IncludeLaunchDescription( - PythonLaunchDescriptionSource([os.path.join( - get_package_share_directory('sportmode_nav2'), 'launch'), - '/localization_async.launch.py']) - ) + PythonLaunchDescriptionSource([os.path.join( + get_package_share_directory('sportmode_nav2'), 'launch'), + '/localization_async.launch.py']) + ) return LaunchDescription([ robot_localization, async_mapping, - ]) \ No newline at end of file + ]) diff --git a/deploy/ros2_nodes/sportmode_nav2/launch/localization_async.launch.py b/deploy/ros2_nodes/sportmode_nav2/launch/localization_async.launch.py index 70db31e..e955b2e 100644 --- a/deploy/ros2_nodes/sportmode_nav2/launch/localization_async.launch.py +++ b/deploy/ros2_nodes/sportmode_nav2/launch/localization_async.launch.py @@ -34,15 +34,15 @@ def generate_launch_description(): actual_params_file = PythonExpression(['"', params_file, '" if ', has_node_params, ' else "', default_params_file, '"']) - log_param_change = LogInfo(msg=['provided params_file ', params_file, + log_param_change = LogInfo(msg=['provided params_file ', params_file, ' does not contain slam_toolbox parameters. Using default: ', default_params_file], condition=UnlessCondition(has_node_params)) start_async_slam_toolbox_node = Node( parameters=[ - actual_params_file, - {'use_sim_time': use_sim_time} + actual_params_file, + {'use_sim_time': use_sim_time} ], package='slam_toolbox', executable='async_slam_toolbox_node', diff --git a/deploy/ros2_nodes/sportmode_nav2/launch/mapping.launch.py b/deploy/ros2_nodes/sportmode_nav2/launch/mapping.launch.py index 526fc21..cc9a2bf 100644 --- a/deploy/ros2_nodes/sportmode_nav2/launch/mapping.launch.py +++ b/deploy/ros2_nodes/sportmode_nav2/launch/mapping.launch.py @@ -9,19 +9,19 @@ from launch_ros.actions import Node def generate_launch_description(): - + robot_localization = IncludeLaunchDescription( - PythonLaunchDescriptionSource([os.path.join( - get_package_share_directory('sportmode_nav2'), 'launch'), - '/ros_ekf.launch.py']) - ) + PythonLaunchDescriptionSource([os.path.join( + get_package_share_directory('sportmode_nav2'), 'launch'), + '/ros_ekf.launch.py']) + ) async_mapping = IncludeLaunchDescription( - PythonLaunchDescriptionSource([os.path.join( - get_package_share_directory('sportmode_nav2'), 'launch'), - '/mapping_async.launch.py']) - ) + PythonLaunchDescriptionSource([os.path.join( + get_package_share_directory('sportmode_nav2'), 'launch'), + '/mapping_async.launch.py']) + ) return LaunchDescription([ robot_localization, async_mapping, - ]) \ No newline at end of file + ]) diff --git a/deploy/ros2_nodes/sportmode_nav2/launch/mapping_async.launch.py b/deploy/ros2_nodes/sportmode_nav2/launch/mapping_async.launch.py index 939c31a..b8c5cfb 100644 --- a/deploy/ros2_nodes/sportmode_nav2/launch/mapping_async.launch.py +++ b/deploy/ros2_nodes/sportmode_nav2/launch/mapping_async.launch.py @@ -34,15 +34,15 @@ def generate_launch_description(): actual_params_file = PythonExpression(['"', params_file, '" if ', has_node_params, ' else "', default_params_file, '"']) - log_param_change = LogInfo(msg=['provided params_file ', params_file, + log_param_change = LogInfo(msg=['provided params_file ', params_file, ' does not contain slam_toolbox parameters. Using default: ', default_params_file], condition=UnlessCondition(has_node_params)) start_async_slam_toolbox_node = Node( parameters=[ - actual_params_file, - {'use_sim_time': use_sim_time} + actual_params_file, + {'use_sim_time': use_sim_time} ], package='slam_toolbox', executable='async_slam_toolbox_node', diff --git a/deploy/ros2_nodes/sportmode_nav2/launch/navigation.launch.py b/deploy/ros2_nodes/sportmode_nav2/launch/navigation.launch.py index c48a741..0df875a 100644 --- a/deploy/ros2_nodes/sportmode_nav2/launch/navigation.launch.py +++ b/deploy/ros2_nodes/sportmode_nav2/launch/navigation.launch.py @@ -58,10 +58,10 @@ def generate_launch_description(): 'map_subscribe_transient_local': map_subscribe_transient_local} configured_params = RewrittenYaml( - source_file=params_file, - root_key=namespace, - param_rewrites=param_substitutions, - convert_types=True) + source_file=params_file, + root_key=namespace, + param_rewrites=param_substitutions, + convert_types=True) return LaunchDescription([ # Set env var to print messages to stdout immediately diff --git a/deploy/ros2_nodes/sportmode_nav2/launch/navigation_launch.py b/deploy/ros2_nodes/sportmode_nav2/launch/navigation_launch.py index 685b1c6..e1f9eb3 100644 --- a/deploy/ros2_nodes/sportmode_nav2/launch/navigation_launch.py +++ b/deploy/ros2_nodes/sportmode_nav2/launch/navigation_launch.py @@ -55,7 +55,7 @@ def generate_launch_description(): # TODO(orduno) Substitute with `PushNodeRemapping` # https://github.com/ros2/launch_ros/issues/56 remappings = [('/tf', 'tf'), - ('/tf_static', 'tf_static'),('/cmd_vel', '/go2/cmd_vel')] + ('/tf_static', 'tf_static'), ('/cmd_vel', '/go2/cmd_vel')] # Create our own temporary YAML files that include substitutions param_substitutions = { @@ -239,7 +239,7 @@ def generate_launch_description(): name='velocity_smoother', parameters=[configured_params], remappings=remappings + - [('cmd_vel', 'cmd_vel_nav'), ('cmd_vel_smoothed', 'cmd_vel')]), + [('cmd_vel', 'cmd_vel_nav'), ('cmd_vel_smoothed', 'cmd_vel')]), ComposableNode( package='nav2_lifecycle_manager', plugin='nav2_lifecycle_manager::LifecycleManager', diff --git a/deploy/ros2_nodes/sportmode_nav2/launch/ros_ekf.launch.py b/deploy/ros2_nodes/sportmode_nav2/launch/ros_ekf.launch.py index 1ee570a..7d80853 100644 --- a/deploy/ros2_nodes/sportmode_nav2/launch/ros_ekf.launch.py +++ b/deploy/ros2_nodes/sportmode_nav2/launch/ros_ekf.launch.py @@ -23,6 +23,7 @@ import pathlib import launch.actions from launch.actions import DeclareLaunchArgument + def generate_launch_description(): return LaunchDescription([ launch_ros.actions.Node( @@ -31,5 +32,5 @@ def generate_launch_description(): name='ekf_filter_node', output='screen', parameters=[os.path.join(get_package_share_directory("sportmode_nav2"), 'params', 'ros_ekf.yaml')], - ), -]) + ), + ]) diff --git a/examples/imu-noise-id-dataset-collector.py b/examples/imu-noise-id-dataset-collector.py index 0f6d9ea..294b0de 100644 --- a/examples/imu-noise-id-dataset-collector.py +++ b/examples/imu-noise-id-dataset-collector.py @@ -8,7 +8,7 @@ from rosbags.typesys.stores.ros2_foxy import ( builtin_interfaces__msg__Time as Time, sensor_msgs__msg__Imu as Imu, std_msgs__msg__Header as Header, - geometry_msgs__msg__Vector3 as Vector3, + geometry_msgs__msg__Vector3 as Vector3, geometry_msgs__msg__Quaternion as Quaternion, ) @@ -25,25 +25,25 @@ def main(): with Writer('imu_noise_identification_dataset') as writer: conn = writer.add_connection('/go2/imu', Imu.__msgtype__) - while(soc>10): + while (soc > 10): soc = robot.getBatteryState() imu = robot.getIMU() a = imu['accel'] g = imu['gyro'] accel = Vector3(a[0], a[1], a[2]) gyro = Vector3(g[0], g[1], g[2]) - q = Quaternion(0,0,0,1) - timestamp = time.time()*10**9 + q = Quaternion(0, 0, 0, 1) + timestamp = time.time() * 10**9 msg = Imu( - Header( + Header( stamp=Time(sec=int(timestamp // 10**9), nanosec=int(timestamp % 10**9)), frame_id='base_link'), - linear_acceleration=accel, - angular_velocity= gyro, - orientation=q, - orientation_covariance=np.zeros(9), - linear_acceleration_covariance=np.zeros(9), - angular_velocity_covariance=np.zeros(9) + linear_acceleration=accel, + angular_velocity=gyro, + orientation=q, + orientation_covariance=np.zeros(9), + linear_acceleration_covariance=np.zeros(9), + angular_velocity_covariance=np.zeros(9) ) writer.write(conn, timestamp, typestore.serialize_cdr(msg, msg.__msgtype__)) time.sleep(0.02) @@ -51,5 +51,6 @@ def main(): robot.close() exit() -if __name__=='__main__': - main() \ No newline at end of file + +if __name__ == '__main__': + main() diff --git a/go2py_messages/msg/dds_/_Go2pyHighCmd.py b/go2py_messages/msg/dds_/_Go2pyHighCmd.py index ddab137..905d567 100644 --- a/go2py_messages/msg/dds_/_Go2pyHighCmd.py +++ b/go2py_messages/msg/dds_/_Go2pyHighCmd.py @@ -25,5 +25,3 @@ class Go2pyHighCmd_(idl.IdlStruct, typename="go2py_messages.msg.dds_.Go2pyHighCm vx: types.float32 vy: types.float32 vz: types.float32 - - diff --git a/go2py_messages/msg/dds_/_Go2pyLowCmd.py b/go2py_messages/msg/dds_/_Go2pyLowCmd.py index 9e19dbc..af63386 100644 --- a/go2py_messages/msg/dds_/_Go2pyLowCmd.py +++ b/go2py_messages/msg/dds_/_Go2pyLowCmd.py @@ -27,5 +27,3 @@ class Go2pyLowCmd_(idl.IdlStruct, typename="go2py_messages.msg.dds_.Go2pyLowCmd_ kp: types.array[types.float32, 12] kd: types.array[types.float32, 12] tau: types.array[types.float32, 12] - - diff --git a/go2py_messages/msg/dds_/_Go2pyState.py b/go2py_messages/msg/dds_/_Go2pyState.py index d562faa..cdce3f9 100644 --- a/go2py_messages/msg/dds_/_Go2pyState.py +++ b/go2py_messages/msg/dds_/_Go2pyState.py @@ -35,5 +35,3 @@ class Go2pyState_(idl.IdlStruct, typename="go2py_messages.msg.dds_.Go2pyState_") motor_temp: types.array[types.float32, 12] wireless_remote: types.array[types.uint8, 40] soc: types.uint8 - - diff --git a/setup.py b/setup.py index 6ba2cfe..68855b7 100644 --- a/setup.py +++ b/setup.py @@ -14,5 +14,5 @@ setup( author_email='rk4342@nyu.edu', description='Toolkit for deployment using Unitree Go2.', install_requires=[ - ] -) \ No newline at end of file + ] +)