From 0fe8b7038bab27b71cc136ddb370d685d7c98c66 Mon Sep 17 00:00:00 2001 From: yangning wu Date: Thu, 23 May 2024 15:49:05 +0800 Subject: [PATCH] add joystick support of python simulation --- example/python/stand_go2.py | 4 +- simulate/src/joystick/readme.md | 5 + simulate_python/config.py | 4 + simulate_python/test/gamepad_test.py | 28 +++ simulate_python/unitree_mujoco.py | 12 +- simulate_python/unitree_sdk2py_bridge.py | 285 ++++++++++++++++++----- 6 files changed, 270 insertions(+), 68 deletions(-) create mode 100644 simulate/src/joystick/readme.md create mode 100644 simulate_python/test/gamepad_test.py diff --git a/example/python/stand_go2.py b/example/python/stand_go2.py index a84f0e8..154eea5 100644 --- a/example/python/stand_go2.py +++ b/example/python/stand_go2.py @@ -2,8 +2,8 @@ import time import sys import numpy as np -from unitree_sdk2py.core.channel import ChannelPublisher, ChannelFactortyInitialize -from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactortyInitialize +from unitree_sdk2py.core.channel import ChannelPublisher, ChannelSubscriber +from unitree_sdk2py.core.channel import ChannelFactoryInitialize from unitree_sdk2py.idl.default import unitree_go_msg_dds__LowCmd_ from unitree_sdk2py.idl.unitree_go.msg.dds_ import LowCmd_ from unitree_sdk2py.utils.crc import CRC diff --git a/simulate/src/joystick/readme.md b/simulate/src/joystick/readme.md new file mode 100644 index 0000000..241b1ec --- /dev/null +++ b/simulate/src/joystick/readme.md @@ -0,0 +1,5 @@ +This joystick tool is based on https://github.com/drewnoakes/joystick + +``` +sudo apt install joystick +``` \ No newline at end of file diff --git a/simulate_python/config.py b/simulate_python/config.py index c3aa867..b75b9ef 100644 --- a/simulate_python/config.py +++ b/simulate_python/config.py @@ -3,6 +3,10 @@ ROBOT_SCENE = "../unitree_robots/" + ROBOT + "/scene.xml" # Robot scene DOMAIN_ID = 1 # Domain id INTERFACE = "lo" # Interface +USE_JOYSTICK = 1 # Simulate Unitree WirelessController using a gamepad +JOYSTICK_TYPE = "xbox" # support "xbox" and "switch" gamepad layout +JOYSTICK_DEVICE = 0 # Joystick number + PRINT_SCENE_INFORMATION = True # Print link, joint and sensors information of robot ENABLE_ELASTIC_BAND = False # Virtual spring band, used for lifting h1 diff --git a/simulate_python/test/gamepad_test.py b/simulate_python/test/gamepad_test.py new file mode 100644 index 0000000..47ea77e --- /dev/null +++ b/simulate_python/test/gamepad_test.py @@ -0,0 +1,28 @@ +import pygame +import sys + +pygame.init() +pygame.joystick.init() + +joystick_count = pygame.joystick.get_count() +if joystick_count > 0: + joystick = pygame.joystick.Joystick(0) + joystick.init() +else: + sys.exit() + +while True: + # for event in pygame.event.get(): + # if event.type == pygame.JOYBUTTONDOWN: + # print(f"Button pressed： {event.button}") + # elif event.type == pygame.JOYBUTTONUP: + # print(f"Button released： {event.button}") + pygame.event.get() + # 读取轴位置 + axes = joystick.get_button(0) + print(joystick.get_numaxes(), joystick.get_numhats(),joystick.get_numbuttons(),joystick.get_numballs()) + print(f"axis： {axes}") + + + pygame.time.wait(100) + diff --git a/simulate_python/unitree_mujoco.py b/simulate_python/unitree_mujoco.py index 6a9f4d1..7649b7d 100755 --- a/simulate_python/unitree_mujoco.py +++ b/simulate_python/unitree_mujoco.py @@ -4,7 +4,7 @@ import mujoco.viewer from threading import Thread import threading -from unitree_sdk2py.core.channel import ChannelFactortyInitialize +from unitree_sdk2py.core.channel import ChannelFactoryInitialize from unitree_sdk2py_bridge import UnitreeSdk2Bridge, ElasticBand import config @@ -37,8 +37,14 @@ time.sleep(0.2) def SimulationThread(): global mj_data, mj_model - ChannelFactortyInitialize(config.DOMAIN_ID, config.INTERFACE) - unitree = UnitreeSdk2Bridge(mj_model, mj_data, config.PRINT_SCENE_INFORMATION) + ChannelFactoryInitialize(config.DOMAIN_ID, config.INTERFACE) + unitree = UnitreeSdk2Bridge(mj_model, mj_data) + + if config.USE_JOYSTICK: + unitree.SetupJoystick() + if config.PRINT_SCENE_INFORMATION: + unitree.PrintSceneInformation() + while viewer.is_running(): step_start = time.perf_counter() diff --git a/simulate_python/unitree_sdk2py_bridge.py b/simulate_python/unitree_sdk2py_bridge.py index 1b0b279..6e5164e 100644 --- a/simulate_python/unitree_sdk2py_bridge.py +++ b/simulate_python/unitree_sdk2py_bridge.py @@ -1,25 +1,30 @@ import mujoco import numpy as np +import pygame +import sys from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelPublisher from unitree_sdk2py.idl.unitree_go.msg.dds_ import LowCmd_ from unitree_sdk2py.idl.unitree_go.msg.dds_ import LowState_ from unitree_sdk2py.idl.unitree_go.msg.dds_ import SportModeState_ +from unitree_sdk2py.idl.unitree_go.msg.dds_ import WirelessController_ from unitree_sdk2py.idl.default import unitree_go_msg_dds__LowCmd_ from unitree_sdk2py.idl.default import unitree_go_msg_dds__LowState_ from unitree_sdk2py.idl.default import unitree_go_msg_dds__SportModeState_ +from unitree_sdk2py.idl.default import unitree_go_msg_dds__WirelessController_ from unitree_sdk2py.utils.thread import RecurrentThread TOPIC_LOWCMD = "rt/lowcmd" TOPIC_LOWSTATE = "rt/lowstate" TOPIC_HIGHSTATE = "rt/sportmodestate" +TOPIC_WIRELESS_CONTROLLER = "rt/wirelesscontroller" MOTOR_SENSOR_NUM = 3 class UnitreeSdk2Bridge: - def __init__(self, mj_model, mj_data, info): + def __init__(self, mj_model, mj_data): self.mj_model = mj_model self.mj_data = mj_data @@ -31,78 +36,129 @@ class UnitreeSdk2Bridge: # Check sensor for i in range(self.dim_motor_sensor, self.mj_model.nsensor): - name = mujoco.mj_id2name(self.mj_model, - mujoco._enums.mjtObj.mjOBJ_SENSOR, i) + name = mujoco.mj_id2name( + self.mj_model, mujoco._enums.mjtObj.mjOBJ_SENSOR, i + ) if name == "imu_quat": self.have_imu_ = True if name == "frame_pos": self.have_frame_sensor_ = True - - if info: - self.PrintSceneInformation() + # Unitree sdk2 message self.low_state = unitree_go_msg_dds__LowState_() self.low_state_puber = ChannelPublisher(TOPIC_LOWSTATE, LowState_) self.low_state_puber.Init() - self.lowStateThread = RecurrentThread(interval=0.002, - target=self.PublishLowState, - name="sim_pub_lowstate") + self.lowStateThread = RecurrentThread( + interval=self.dt, target=self.PublishLowState, name="sim_lowstate" + ) self.lowStateThread.Start() self.high_state = unitree_go_msg_dds__SportModeState_() - self.high_state_puber = ChannelPublisher(TOPIC_HIGHSTATE, - SportModeState_) + self.high_state_puber = ChannelPublisher(TOPIC_HIGHSTATE, SportModeState_) self.high_state_puber.Init() - self.HighStateThread = RecurrentThread(interval=self.dt, - target=self.PublishHighState, - name="highstate") + self.HighStateThread = RecurrentThread( + interval=self.dt, target=self.PublishHighState, name="sim_highstate" + ) self.HighStateThread.Start() + self.wireless_controller = unitree_go_msg_dds__WirelessController_() + self.wireless_controller_puber = ChannelPublisher( + TOPIC_WIRELESS_CONTROLLER, WirelessController_ + ) + self.wireless_controller_puber.Init() + self.WirelessControllerThread = RecurrentThread( + interval=0.01, + target=self.PublishWirelessController, + name="sim_wireless_controller", + ) + self.WirelessControllerThread.Start() + self.low_cmd_suber = ChannelSubscriber(TOPIC_LOWCMD, LowCmd_) self.low_cmd_suber.Init(self.LowCmdHandler, 10) + # joystick + self.key_map = { + "R1": 0, + "L1": 1, + "start": 2, + "select": 3, + "R2": 4, + "L2": 5, + "F1": 6, + "F2": 7, + "A": 8, + "B": 9, + "X": 10, + "Y": 11, + "up": 12, + "right": 13, + "down": 14, + "left": 15, + } + + self.joystick = None + + def LowCmdHandler(self, msg: LowCmd_): if self.mj_data != None: for i in range(self.num_motor): - self.mj_data.ctrl[i] = msg.motor_cmd[i].tau +\ - msg.motor_cmd[i].kp * (msg.motor_cmd[i].q - self.mj_data.sensordata[i]) +\ - msg.motor_cmd[i].kd * (msg.motor_cmd[i].dq - self.mj_data.sensordata[i + self.num_motor]) + self.mj_data.ctrl[i] = ( + msg.motor_cmd[i].tau + + msg.motor_cmd[i].kp + * (msg.motor_cmd[i].q - self.mj_data.sensordata[i]) + + msg.motor_cmd[i].kd + * ( + msg.motor_cmd[i].dq + - self.mj_data.sensordata[i + self.num_motor] + ) + ) def PublishLowState(self): if self.mj_data != None: for i in range(self.num_motor): self.low_state.motor_state[i].q = self.mj_data.sensordata[i] self.low_state.motor_state[i].dq = self.mj_data.sensordata[ - i + self.num_motor] - self.low_state.motor_state[ - i].tau_est = self.mj_data.sensordata[i + - 2 * self.num_motor] + i + self.num_motor + ] + self.low_state.motor_state[i].tau_est = self.mj_data.sensordata[ + i + 2 * self.num_motor + ] if self.have_frame_sensor_: - self.low_state.imu_state.quaternion[ - 0] = self.mj_data.sensordata[self.dim_motor_sensor + 0] - self.low_state.imu_state.quaternion[ - 1] = self.mj_data.sensordata[self.dim_motor_sensor + 1] - self.low_state.imu_state.quaternion[ - 2] = self.mj_data.sensordata[self.dim_motor_sensor + 2] - self.low_state.imu_state.quaternion[ - 3] = self.mj_data.sensordata[self.dim_motor_sensor + 3] + self.low_state.imu_state.quaternion[0] = self.mj_data.sensordata[ + self.dim_motor_sensor + 0 + ] + self.low_state.imu_state.quaternion[1] = self.mj_data.sensordata[ + self.dim_motor_sensor + 1 + ] + self.low_state.imu_state.quaternion[2] = self.mj_data.sensordata[ + self.dim_motor_sensor + 2 + ] + self.low_state.imu_state.quaternion[3] = self.mj_data.sensordata[ + self.dim_motor_sensor + 3 + ] - self.low_state.imu_state.gyroscope[ - 0] = self.mj_data.sensordata[self.dim_motor_sensor + 4] - self.low_state.imu_state.gyroscope[ - 1] = self.mj_data.sensordata[self.dim_motor_sensor + 5] - self.low_state.imu_state.gyroscope[ - 2] = self.mj_data.sensordata[self.dim_motor_sensor + 6] + self.low_state.imu_state.gyroscope[0] = self.mj_data.sensordata[ + self.dim_motor_sensor + 4 + ] + self.low_state.imu_state.gyroscope[1] = self.mj_data.sensordata[ + self.dim_motor_sensor + 5 + ] + self.low_state.imu_state.gyroscope[2] = self.mj_data.sensordata[ + self.dim_motor_sensor + 6 + ] - self.low_state.imu_state.accelerometer[ - 0] = self.mj_data.sensordata[self.dim_motor_sensor + 7] - self.low_state.imu_state.accelerometer[ - 1] = self.mj_data.sensordata[self.dim_motor_sensor + 8] - self.low_state.imu_state.accelerometer[ - 2] = self.mj_data.sensordata[self.dim_motor_sensor + 9] + self.low_state.imu_state.accelerometer[0] = self.mj_data.sensordata[ + self.dim_motor_sensor + 7 + ] + self.low_state.imu_state.accelerometer[1] = self.mj_data.sensordata[ + self.dim_motor_sensor + 8 + ] + self.low_state.imu_state.accelerometer[2] = self.mj_data.sensordata[ + self.dim_motor_sensor + 9 + ] self.low_state_puber.Write(self.low_state) @@ -110,44 +166,141 @@ class UnitreeSdk2Bridge: if self.mj_data != None: self.high_state.position[0] = self.mj_data.sensordata[ - self.dim_motor_sensor + 10] + self.dim_motor_sensor + 10 + ] self.high_state.position[1] = self.mj_data.sensordata[ - self.dim_motor_sensor + 11] + self.dim_motor_sensor + 11 + ] self.high_state.position[2] = self.mj_data.sensordata[ - self.dim_motor_sensor + 12] + self.dim_motor_sensor + 12 + ] self.high_state.velocity[0] = self.mj_data.sensordata[ - self.dim_motor_sensor + 13] + self.dim_motor_sensor + 13 + ] self.high_state.velocity[1] = self.mj_data.sensordata[ - self.dim_motor_sensor + 14] + self.dim_motor_sensor + 14 + ] self.high_state.velocity[2] = self.mj_data.sensordata[ - self.dim_motor_sensor + 15] + self.dim_motor_sensor + 15 + ] self.high_state_puber.Write(self.high_state) + def PublishWirelessController(self): + if self.joystick != None: + pygame.event.get() + key_state = [0]*16 + key_state[self.key_map["R1"]] = self.joystick.get_button(self.button_id["RB"]) + key_state[self.key_map["L1"]] = self.joystick.get_button(self.button_id["LB"]) + key_state[self.key_map["start"]] = self.joystick.get_button(self.button_id["START"]) + key_state[self.key_map["select"]] = self.joystick.get_button(self.button_id["SELECT"]) + key_state[self.key_map["R2"]] = (self.joystick.get_axis(self.axis_id["RT"])>0) + key_state[self.key_map["L2"]] = (self.joystick.get_axis(self.axis_id["LT"])>0) + key_state[self.key_map["F1"]] = 0 + key_state[self.key_map["F2"]] = 0 + key_state[self.key_map["A"]] = self.joystick.get_button(self.button_id["A"]) + key_state[self.key_map["B"]] = self.joystick.get_button(self.button_id["B"]) + key_state[self.key_map["X"]] = self.joystick.get_button(self.button_id["X"]) + key_state[self.key_map["Y"]] = self.joystick.get_button(self.button_id["Y"]) + key_state[self.key_map["up"]] = (self.joystick.get_hat(0)[1]>0) + key_state[self.key_map["right"]] = (self.joystick.get_hat(0)[0]>0) + key_state[self.key_map["down"]] = (self.joystick.get_hat(0)[1]<0) + key_state[self.key_map["left"]] = (self.joystick.get_hat(0)[0]<0) + + key_value = 0 + for i in range(16): + key_value += key_state[i] << i + + self.wireless_controller.keys = key_value + self.wireless_controller.lx = self.joystick.get_axis(self.axis_id["LX"]) + self.wireless_controller.ly = -self.joystick.get_axis(self.axis_id["LY"]) + self.wireless_controller.rx = self.joystick.get_axis(self.axis_id["RX"]) + self.wireless_controller.ry = -self.joystick.get_axis(self.axis_id["RY"]) + + self.wireless_controller_puber.Write(self.wireless_controller) + + def SetupJoystick(self, device_id=0, js_type="xbox"): + pygame.init() + pygame.joystick.init() + joystick_count = pygame.joystick.get_count() + if joystick_count > 0: + self.joystick = pygame.joystick.Joystick(device_id) + self.joystick.init() + else: + print("No gamepad detected.") + sys.exit() + + if js_type == "xbox": + self.axis_id = { + "LX": 0, # Left stick axis x + "LY": 1, # Left stick axis y + "RX": 3, # Right stick axis x + "RY": 4, # Right stick axis y + "LT": 2, # Left trigger + "RT": 5, # Right trigger + "DX": 6, # Directional pad x + "DY": 7, # Directional pad y + } + + self.button_id = { + "X": 2, + "Y": 3, + "B": 1, + "A": 0, + "LB": 4, + "RB": 5, + "SELECT": 6, + "START": 7, + } + + elif js_type == "stitch": + self.axis_id = { + "LX": 0, # Left stick axis x + "LY": 1, # Left stick axis y + "RX": 2, # Right stick axis x + "RY": 3, # Right stick axis y + "LT": 5, # Left trigger + "RT": 4, # Right trigger + "DX": 6, # Directional pad x + "DY": 7, # Directional pad y + } + + self.button_id = { + "X": 3, + "Y": 4, + "B": 1, + "A": 0, + "LB": 6, + "RB": 7, + "SELECT": 10, + "START": 11, + } + else: + print("Unsupported gamepad. ") + def PrintSceneInformation(self): print(" ") print("<<------------- Link ------------->> ") for i in range(self.mj_model.nbody): - name = mujoco.mj_id2name(self.mj_model, - mujoco._enums.mjtObj.mjOBJ_BODY, i) + name = mujoco.mj_id2name(self.mj_model, mujoco._enums.mjtObj.mjOBJ_BODY, i) if name: print("link_index:", i, ", name:", name) print(" ") print("<<------------- Joint ------------->> ") for i in range(self.mj_model.njnt): - name = mujoco.mj_id2name(self.mj_model, - mujoco._enums.mjtObj.mjOBJ_JOINT, i) + name = mujoco.mj_id2name(self.mj_model, mujoco._enums.mjtObj.mjOBJ_JOINT, i) if name: print("joint_index:", i, ", name:", name) print(" ") print("<<------------- Actuator ------------->>") for i in range(self.mj_model.nu): - name = mujoco.mj_id2name(self.mj_model, - mujoco._enums.mjtObj.mjOBJ_ACTUATOR, i) + name = mujoco.mj_id2name( + self.mj_model, mujoco._enums.mjtObj.mjOBJ_ACTUATOR, i + ) if name: print("actuator_index:", i, ", name:", name) print(" ") @@ -155,11 +308,18 @@ class UnitreeSdk2Bridge: print("<<------------- Sensor ------------->>") index = 0 for i in range(self.mj_model.nsensor): - name = mujoco.mj_id2name(self.mj_model, - mujoco._enums.mjtObj.mjOBJ_SENSOR, i) + name = mujoco.mj_id2name( + self.mj_model, mujoco._enums.mjtObj.mjOBJ_SENSOR, i + ) if name: - print("sensor_index:", index, ", name:", name, ", dim:", - self.mj_model.sensor_dim[i]) + print( + "sensor_index:", + index, + ", name:", + name, + ", dim:", + self.mj_model.sensor_dim[i], + ) index = index + self.mj_model.sensor_dim[i] print(" ") @@ -175,16 +335,15 @@ class ElasticBand: def Advance(self, x, dx): """ - Args: - δx: desired position - current position - dx: current velocity - """ + Args: + δx: desired position - current position + dx: current velocity + """ δx = self.point - x distance = np.linalg.norm(δx) direction = δx / distance v = np.dot(dx, direction) - f = (self.stiffness * - (distance - self.length) - self.damping * v) * direction + f = (self.stiffness * (distance - self.length) - self.damping * v) * direction return f def MujuocoKeyCallback(self, key):