from collections import deque from typing import Optional import numpy as np from omni.isaac.core.articulations import Articulation from omni.isaac.core.utils.prims import define_prim from omni.isaac.sensor import ContactSensor, IMUSensor import Go2Py from Go2Py.sim.isaacsim.lcm_types.unitree_lowlevel import UnitreeLowState class UnitreeGo2(Articulation): def __init__( self, prim_path: str, name: str = "go2_quadruped", physics_dt: Optional[float] = 1 / 400.0, position: Optional[np.ndarray] = None, orientation: Optional[np.ndarray] = None, usd_path: Optional[str] = None, ) -> None: """ [Summary] initialize robot, set up sensors and interfaces Args: prim_path {str} -- prim path of the robot on the stage name {str} -- name of the quadruped physics_dt {float} -- physics downtime of the controller position {np.ndarray} -- position of the robot orientation {np.ndarray} -- orientation of the robot """ self._prim_path = prim_path prim = define_prim(self._prim_path, "Xform") self.usd_path = usd_path # breakpoint() if self.usd_path is None: prim.GetReferences().AddReference(Go2Py.GO2_USD_PATH) else: prim.GetReferences().AddReference(self.usd_path) breakpoint() super().__init__( prim_path=self._prim_path+'/Go2/base', name=name, position=position, orientation=orientation, ) # contact sensor setup self.feet_order = ["FL", "RL", "FR", "RR"] self.feet_path = [ self._prim_path + "/Go2/base/FL_calf", self._prim_path + "/Go2/base/FR_calf", self._prim_path + "/Go2/base/RL_calf", self._prim_path + "/Go2/base/RR_calf", ] self.color = [(1, 0, 0, 1), (0, 1, 0, 1), (0, 0, 1, 1), (1, 1, 0, 1)] self._contact_sensors = [None] * 4 for i in range(4): self._contact_sensors[i] = ContactSensor( prim_path=self.feet_path[i] + "/sensor", min_threshold=0, max_threshold=1000000, radius=0.043, dt=physics_dt, ) self.foot_force = np.zeros(4) self.enable_foot_filter = True self._FILTER_WINDOW_SIZE = 20 self._foot_filters = [deque(), deque(), deque(), deque()] # imu sensor setup self.imu_path = self._prim_path + "/Go2/base/base" self._imu_sensor = IMUSensor( prim_path=self.imu_path + "/imu_sensor", name="imu", dt=physics_dt, translation=np.array([0, 0, 0]), orientation=np.array([1, 0, 0, 0]), ) self.accel = np.zeros((3,)) self.gyro = np.zeros((3,)) # Translation maps between Isaac and Bullet # self.bullet_joint_order = ['FL_hip_joint', 'FL_thigh_joint', 'FL_calf_joint', # 'RL_hip_joint', 'RL_thigh_joint', 'RL_calf_joint', # 'FR_hip_joint', 'FR_thigh_joint', 'FR_calf_joint', # 'RR_hip_joint', 'RR_thigh_joint', 'RR_calf_joint'] self.bullet_joint_order = [ "FL_hip_joint", "FL_thigh_joint", "FL_calf_joint", "FR_hip_joint", "FR_thigh_joint", "FR_calf_joint", "RL_hip_joint", "RL_thigh_joint", "RL_calf_joint", "RR_hip_joint", "RR_thigh_joint", "RR_calf_joint", ] self.isaac_joint_order = [ "FL_hip_joint", "FR_hip_joint", "RL_hip_joint", "RR_hip_joint", "FL_thigh_joint", "FR_thigh_joint", "RL_thigh_joint", "RR_thigh_joint", "FL_calf_joint", "FR_calf_joint", "RL_calf_joint", "RR_calf_joint", ] self.isaac_name_2_index = {s: i for i, s in enumerate(self.isaac_joint_order)} self.bullet_name_2_index = {s: i for i, s in enumerate(self.bullet_joint_order)} self.to_bullet_index = np.array( [self.isaac_name_2_index[id] for id in self.bullet_joint_order] ) self.to_isaac_index = np.array( [self.bullet_name_2_index[id] for id in self.isaac_joint_order] ) self.tau_est = np.zeros((12,)) self.state = UnitreeLowState() self.init_pos = np.array([0.0, 0.0, 0.6]) self.init_quat = np.array([0.0, 0.0, 0.0, 1.0]) self.init_joint_pos = np.array( [0.2, 0.8, -1.5, -0.2, 0.8, -1.5, 0.2, 1.0, -1.6, -0.2, 1.0, -1.6] ) return def toIsaacOrder(self, x): return x[self.to_isaac_index, ...] def toBulletOrder(self, x): return x[self.to_bullet_index, ...] def setState(self, pos, quat, q) -> None: """[Summary] Set the kinematic state of the robot. Args: pos {ndarray} -- The position of the robot (x, y, z) quat {ndarray} -- The orientation of the robot (qx, qy, qz, qw) q {ndarray} -- Joint angles of the robot in standard Pinocchio order Raises: RuntimeError: When the DC Toolbox interface has not been configured. """ self.set_world_pose(position=pos, orientation=quat[[3, 0, 1, 2]]) self.set_linear_velocity(np.zeros((3,))) self.set_angular_velocity(np.zeros((3,))) self.set_joint_positions( positions=np.asarray(self.toIsaacOrder(q), dtype=np.float32) ) self.set_joint_velocities( velocities=np.asarray(self.toIsaacOrder(np.zeros((12,))), dtype=np.float32) ) self.set_joint_efforts(np.zeros_like(q)) return def initialize(self, physics_sim_view=None) -> None: """[summary] initialize dc interface, set up drive mode and initial robot state """ super().initialize(physics_sim_view=physics_sim_view) self.get_articulation_controller().set_effort_modes("force") self.get_articulation_controller().switch_control_mode("effort") self.setState(self.init_pos, self.init_quat, self.init_joint_pos) for i in range(4): self._contact_sensors[i].initialize() return def readContactSensor(self) -> None: """[summary] Updates processed contact sensor data from the robot feet, store them in member variable foot_force """ # Order: FL, RL, FR, RR for i in range(len(self.feet_path)): frame = self._contact_sensors[i].get_current_frame() if "force" in frame: if self.enable_foot_filter: self._foot_filters[i].append(frame["force"]) if len(self._foot_filters[i]) > self._FILTER_WINDOW_SIZE: self._foot_filters[i].popleft() self.foot_force[i] = np.mean(self._foot_filters[i]) else: self.foot_force[i] = frame["force"] return self.foot_force def readIMUSensor(self) -> None: """[summary] Updates processed imu sensor data from the robot body, store them in member variable base_lin and ang_vel """ frame = self._imu_sensor.get_current_frame() self.accel = frame["lin_acc"] self.gyro = frame["ang_vel"] return self.accel, self.gyro def readStates(self): contact_forces = self.readContactSensor() accel, gyro = self.readIMUSensor() # joint pos and vel from the DC interface joint_state = super().get_joints_state() joint_pos = self.toBulletOrder(joint_state.positions) joint_vel = self.toBulletOrder(joint_state.velocities) # base frame base_pose = self.get_world_pose() base_pos = base_pose[0] base_quat = base_pose[1][[1, 2, 3, 0]] base_lin_vel = self.get_linear_velocity() base_ang_vel = self.get_angular_velocity() # assign to state objects self.state.accel = accel.tolist() self.state.gyro = gyro.tolist() self.state.q = joint_pos.tolist() self.state.dq = joint_vel.tolist() self.state.tau_est = self.tau_est self.state.quaternion = base_quat.tolist() self.state.gt_pos = base_pos.tolist() self.state.gt_quat = base_quat.tolist() self.state.gt_lin_vel = base_lin_vel.tolist() self.state.gt_ang_vel = base_ang_vel.tolist() self.state.contact_state = contact_forces.tolist() # projected_gravity = pp.SO3(base_quat).matrix().T @ np.array( # [0.0, 0.0, -1.0] # ).reshape(3, 1) # self.state.gravity = ( # projected_gravity.squeeze().numpy().astype(np.float32).tolist() # ) return self.state def setCommands(self, cmd): """[summary] sets the joint torques Argument: action {np.ndarray} -- Joint torque command """ q = np.array(self.state.q) dq = np.array(self.state.dq) kp = np.array(cmd.kp) kd = np.array(cmd.kd) tau_ff = np.array(cmd.tau_ff) q_des = np.array(cmd.q_des) dq_des = np.array(cmd.dq_des) tau = (q_des - q) * kp + (dq_des - dq) * kd + tau_ff self.set_joint_efforts(np.asarray(self.toIsaacOrder(tau), dtype=np.float32)) self.tau_est = tau return def step(self, cmd): self.readStates() self.setCommands(cmd) return self.state