import math import select import threading import time import numpy as np from go2_gym_deploy.lcm_types.leg_control_data_lcmt import leg_control_data_lcmt from go2_gym_deploy.lcm_types.rc_command_lcmt import rc_command_lcmt from go2_gym_deploy.lcm_types.state_estimator_lcmt import state_estimator_lcmt # 不调用相机 !!! # from go1_gym_deploy.lcm_types.camera_message_lcmt import camera_message_lcmt # from go1_gym_deploy.lcm_types.camera_message_rect_wide import camera_message_rect_wide def get_rpy_from_quaternion(q): w, x, y, z = q r = np.arctan2(2 * (w * x + y * z), 1 - 2 * (x ** 2 + y ** 2)) p = np.arcsin(2 * (w * y - z * x)) y = np.arctan2(2 * (w * z + x * y), 1 - 2 * (y ** 2 + z ** 2)) return np.array([r, p, y]) def get_rotation_matrix_from_rpy(rpy): """ Get rotation matrix from the given quaternion. Args: q (np.array[float[4]]): quaternion [w,x,y,z] Returns: np.array[float[3,3]]: rotation matrix. """ r, p, y = rpy R_x = np.array([[1, 0, 0], [0, math.cos(r), -math.sin(r)], [0, math.sin(r), math.cos(r)] ]) R_y = np.array([[math.cos(p), 0, math.sin(p)], [0, 1, 0], [-math.sin(p), 0, math.cos(p)] ]) R_z = np.array([[math.cos(y), -math.sin(y), 0], [math.sin(y), math.cos(y), 0], [0, 0, 1] ]) rot = np.dot(R_z, np.dot(R_y, R_x)) return rot class StateEstimator: def __init__(self, lc, use_cameras=False): # defaul use_cameras=True # 这里腿的顺序为什么要转换？ # reverse legs self.joint_idxs = [3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8] self.contact_idxs = [1, 0, 3, 2] # self.joint_idxs = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11] self.lc = lc self.joint_pos = np.zeros(12) self.joint_vel = np.zeros(12) self.tau_est = np.zeros(12) self.world_lin_vel = np.zeros(3) self.world_ang_vel = np.zeros(3) self.euler = np.zeros(3) self.R = np.eye(3) self.buf_idx = 0 self.smoothing_length = 12 self.deuler_history = np.zeros((self.smoothing_length, 3)) self.dt_history = np.zeros((self.smoothing_length, 1)) self.euler_prev = np.zeros(3) self.timuprev = time.time() self.body_lin_vel = np.zeros(3) self.body_ang_vel = np.zeros(3) self.smoothing_ratio = 0.2 self.contact_state = np.ones(4) self.mode = 0 self.ctrlmode_left = 0 self.ctrlmode_right = 0 self.left_stick = [0, 0] self.right_stick = [0, 0] self.left_upper_switch = 0 self.left_lower_left_switch = 0 self.left_lower_right_switch = 0 self.right_upper_switch = 0 self.right_lower_left_switch = 0 self.right_lower_right_switch = 0 self.left_upper_switch_pressed = 0 self.left_lower_left_switch_pressed = 0 self.left_lower_right_switch_pressed = 0 self.right_upper_switch_pressed = 0 self.right_lower_left_switch_pressed = 0 self.right_lower_right_switch_pressed = 0 # default trotting gait self.cmd_freq = 3.0 self.cmd_phase = 0.5 self.cmd_offset = 0.0 self.cmd_duration = 0.5 self.init_time = time.time() self.received_first_legdata = False self.imu_subscription = self.lc.subscribe("state_estimator_data", self._imu_cb) self.legdata_state_subscription = self.lc.subscribe("leg_control_data", self._legdata_cb) self.rc_command_subscription = self.lc.subscribe("rc_command", self._rc_command_cb) # -------------------------------------------------------------- # if use_cameras: # for cam_id in [1, 2, 3, 4, 5]: # self.camera_subscription = self.lc.subscribe(f"camera{cam_id}", self._camera_cb) # self.camera_names = ["front", "bottom", "left", "right", "rear"] # for cam_name in self.camera_names: # self.camera_subscription = self.lc.subscribe(f"rect_image_{cam_name}", self._rect_camera_cb) self.camera_image_left = None self.camera_image_right = None self.camera_image_front = None self.camera_image_bottom = None self.camera_image_rear = None self.body_loc = np.array([0, 0, 0]) self.body_quat = np.array([0, 0, 0, 1]) def get_body_linear_vel(self): self.body_lin_vel = np.dot(self.R.T, self.world_lin_vel) return self.body_lin_vel def get_body_angular_vel(self): self.body_ang_vel = self.smoothing_ratio * np.mean(self.deuler_history / self.dt_history, axis=0) + ( 1 - self.smoothing_ratio) * self.body_ang_vel return self.body_ang_vel def get_gravity_vector(self): grav = np.dot(self.R.T, np.array([0, 0, -1])) return grav def get_contact_state(self): return self.contact_state[self.contact_idxs] def get_rpy(self): return self.euler def get_command(self): MODES_LEFT = ["body_height", "lat_vel", "stance_width"] MODES_RIGHT = ["step_frequency", "footswing_height", "body_pitch"] if self.left_upper_switch_pressed: self.ctrlmode_left = (self.ctrlmode_left + 1) % 3 self.left_upper_switch_pressed = False if self.right_upper_switch_pressed: self.ctrlmode_right = (self.ctrlmode_right + 1) % 3 self.right_upper_switch_pressed = False MODE_LEFT = MODES_LEFT[self.ctrlmode_left] MODE_RIGHT = MODES_RIGHT[self.ctrlmode_right] # always in use cmd_x = 1 * self.left_stick[1] cmd_yaw = -1 * self.right_stick[0] # default values cmd_y = 0. # -1 * self.left_stick[0] cmd_height = 0. cmd_footswing = 0.08 cmd_stance_width = 0.33 cmd_stance_length = 0.40 cmd_ori_pitch = 0. cmd_ori_roll = 0. cmd_freq = 3.0 # joystick commands if MODE_LEFT == "body_height": cmd_height = 0.3 * self.left_stick[0] elif MODE_LEFT == "lat_vel": cmd_y = 0.6 * self.left_stick[0] elif MODE_LEFT == "stance_width": cmd_stance_width = 0.275 + 0.175 * self.left_stick[0] if MODE_RIGHT == "step_frequency": min_freq = 2.0 max_freq = 4.0 cmd_freq = (1 + self.right_stick[1]) / 2 * (max_freq - min_freq) + min_freq elif MODE_RIGHT == "footswing_height": cmd_footswing = max(0, self.right_stick[1]) * 0.32 + 0.03 elif MODE_RIGHT == "body_pitch": cmd_ori_pitch = -0.4 * self.right_stick[1] # gait buttons if self.mode == 0: # Press Button 'A' -> 'Bound' self.cmd_phase = 0.5 self.cmd_offset = 0.0 self.cmd_bound = 0.0 self.cmd_duration = 0.5 elif self.mode == 1: # Press Button 'B' -> 'Trot' self.cmd_phase = 0.0 self.cmd_offset = 0.0 self.cmd_bound = 0.0 self.cmd_duration = 0.5 elif self.mode == 2: # Press Button 'X' -> 'Pace' self.cmd_phase = 0.0 self.cmd_offset = 0.5 self.cmd_bound = 0.0 self.cmd_duration = 0.5 elif self.mode == 3: # Press Button 'Y' -> 'Pronk' self.cmd_phase = 0.0 self.cmd_offset = 0.0 self.cmd_bound = 0.5 self.cmd_duration = 0.5 else: # Default Gait -> 'Trot' self.cmd_phase = 0.0 self.cmd_offset = 0.0 self.cmd_bound = 0.0 self.cmd_duration = 0.5 return np.array([cmd_x, cmd_y, cmd_yaw, cmd_height, cmd_freq, self.cmd_phase, self.cmd_offset, self.cmd_bound, self.cmd_duration, cmd_footswing, cmd_ori_pitch, cmd_ori_roll, cmd_stance_width, cmd_stance_length, 0, 0, 0, 0, 0]) def get_buttons(self): return np.array([self.left_lower_left_switch, self.left_upper_switch, self.right_lower_right_switch, self.right_upper_switch]) def get_dof_pos(self): # print("dofposquery", self.joint_pos[self.joint_idxs]) return self.joint_pos[self.joint_idxs] def get_dof_vel(self): return self.joint_vel[self.joint_idxs] def get_tau_est(self): return self.tau_est[self.joint_idxs] def get_yaw(self): return self.euler[2] def get_body_loc(self): return np.array(self.body_loc) def get_body_quat(self): return np.array(self.body_quat) def get_camera_front(self): return self.camera_image_front def get_camera_bottom(self): return self.camera_image_bottom def get_camera_rear(self): return self.camera_image_rear def get_camera_left(self): return self.camera_image_left def get_camera_right(self): return self.camera_image_right def _legdata_cb(self, channel, data): # print("update legdata") if not self.received_first_legdata: self.received_first_legdata = True print(f"First legdata: {time.time() - self.init_time}") msg = leg_control_data_lcmt.decode(data) # print(msg.q) self.joint_pos = np.array(msg.q) self.joint_vel = np.array(msg.qd) self.tau_est = np.array(msg.tau_est) # print(f"update legdata {msg.id}") def _imu_cb(self, channel, data): # print("update imu") msg = state_estimator_lcmt.decode(data) self.euler = np.array(msg.rpy) self.R = get_rotation_matrix_from_rpy(self.euler) self.contact_state = 1.0 * (np.array(msg.contact_estimate) > 200) self.deuler_history[self.buf_idx % self.smoothing_length, :] = msg.rpy - self.euler_prev self.dt_history[self.buf_idx % self.smoothing_length] = time.time() - self.timuprev self.timuprev = time.time() self.buf_idx += 1 self.euler_prev = np.array(msg.rpy) def _sensor_cb(self, channel, data): pass def _rc_command_cb(self, channel, data): msg = rc_command_lcmt.decode(data) self.left_upper_switch_pressed = ((msg.left_upper_switch and not self.left_upper_switch) or self.left_upper_switch_pressed) self.left_lower_left_switch_pressed = ((msg.left_lower_left_switch and not self.left_lower_left_switch) or self.left_lower_left_switch_pressed) self.left_lower_right_switch_pressed = ((msg.left_lower_right_switch and not self.left_lower_right_switch) or self.left_lower_right_switch_pressed) self.right_upper_switch_pressed = ((msg.right_upper_switch and not self.right_upper_switch) or self.right_upper_switch_pressed) self.right_lower_left_switch_pressed = ((msg.right_lower_left_switch and not self.right_lower_left_switch) or self.right_lower_left_switch_pressed) self.right_lower_right_switch_pressed = ((msg.right_lower_right_switch and not self.right_lower_right_switch) or self.right_lower_right_switch_pressed) self.mode = msg.mode self.right_stick = msg.right_stick self.left_stick = msg.left_stick self.left_upper_switch = msg.left_upper_switch self.left_lower_left_switch = msg.left_lower_left_switch self.left_lower_right_switch = msg.left_lower_right_switch self.right_upper_switch = msg.right_upper_switch self.right_lower_left_switch = msg.right_lower_left_switch self.right_lower_right_switch = msg.right_lower_right_switch # print(self.right_stick, self.left_stick) # 是否要删除下面的camera相关函数？ # -------------------------------------------------- # def _camera_cb(self, channel, data): # msg = camera_message_lcmt.decode(data) # img = np.fromstring(msg.data, dtype=np.uint8) # img = img.reshape((3, 200, 464)).transpose(1, 2, 0) # cam_id = int(channel[-1]) # if cam_id == 1: # self.camera_image_front = img # elif cam_id == 2: # self.camera_image_bottom = img # elif cam_id == 3: # self.camera_image_left = img # elif cam_id == 4: # self.camera_image_right = img # elif cam_id == 5: # self.camera_image_rear = img # else: # print("Image received from camera with unknown ID#!") # #im = Image.fromarray(img).convert('RGB') # #im.save("test_image_" + channel + ".jpg") # #print(channel) # def _rect_camera_cb(self, channel, data): # message_types = [camera_message_rect_wide, camera_message_rect_wide, camera_message_rect_wide, # camera_message_rect_wide, camera_message_rect_wide] # image_shapes = [(116, 100, 3), (116, 100, 3), (116, 100, 3), (116, 100, 3), (116, 100, 3)] # cam_name = channel.split("_")[-1] # # print(f"received py from {cam_name}") # cam_id = self.camera_names.index(cam_name) + 1 # msg = message_types[cam_id - 1].decode(data) # img = np.fromstring(msg.data, dtype=np.uint8) # img = np.flip(np.flip( # img.reshape((image_shapes[cam_id - 1][2], image_shapes[cam_id - 1][1], image_shapes[cam_id - 1][0])), # axis=0), axis=1).transpose(1, 2, 0) # # print(img.shape) # # img = np.flip(np.flip(img.reshape(image_shapes[cam_id - 1]), axis=0), axis=1)[:, :, # # [2, 1, 0]] # .transpose(1, 2, 0) # if cam_id == 1: # self.camera_image_front = img # elif cam_id == 2: # self.camera_image_bottom = img # elif cam_id == 3: # self.camera_image_left = img # elif cam_id == 4: # self.camera_image_right = img # elif cam_id == 5: # self.camera_image_rear = img # else: # print("Image received from camera with unknown ID#!") # -------------------------------------------------- def poll(self, cb=None): t = time.time() try: while True: timeout = 0.01 rfds, wfds, efds = select.select([self.lc.fileno()], [], [], timeout) if rfds: # print("message received!") self.lc.handle() # print(f'Freq {1. / (time.time() - t)} Hz'); t = time.time() else: continue # print(f'waiting for message... Freq {1. / (time.time() - t)} Hz'); t = time.time() # if cb is not None: # cb() except KeyboardInterrupt: pass def spin(self): self.run_thread = threading.Thread(target=self.poll, daemon=False) self.run_thread.start() def close(self): self.lc.unsubscribe(self.legdata_state_subscription) if __name__ == "__main__": import lcm lc = lcm.LCM("udpm://239.255.76.67:7667?ttl=255") se = StateEstimator(lc) se.poll()