299 lines
12 KiB
Python
299 lines
12 KiB
Python
from legged_gym import LEGGED_GYM_ROOT_DIR
|
|
from typing import Union
|
|
import numpy as np
|
|
import time
|
|
import torch
|
|
|
|
import rclpy as rp
|
|
from unitree_hg.msg import LowCmd as LowCmdHG, LowState as LowStateHG
|
|
from unitree_go.msg import LowCmd as LowCmdGo, LowState as LowStateGo
|
|
from common.command_helper_ros import create_damping_cmd, create_zero_cmd, init_cmd_hg, init_cmd_go, MotorMode
|
|
from common.rotation_helper import get_gravity_orientation, transform_imu_data
|
|
from common.remote_controller import RemoteController, KeyMap
|
|
from config import Config
|
|
from common.crc import CRC
|
|
|
|
from enum import Enum
|
|
from ikctrl import IKCtrl
|
|
|
|
class Mode(Enum):
|
|
wait = 0
|
|
zero_torque = 1
|
|
default_pos = 2
|
|
damping = 3
|
|
policy = 4
|
|
null = 5
|
|
|
|
class Controller:
|
|
def __init__(self, config: Config) -> None:
|
|
self.config = config
|
|
self.remote_controller = RemoteController()
|
|
act_joint = ["left_shoulder_pitch_joint",
|
|
"left_shoulder_roll_joint",
|
|
"left_shoulder_yaw_joint",
|
|
"left_elbow_joint",
|
|
"left_wrist_roll_joint",
|
|
"left_wrist_pitch_joint",
|
|
"left_wrist_yaw_joint"]
|
|
self.ikctrl = IKCtrl('/input/unitree_ros/robots/g1_description/g1_29dof_with_hand_rev_1_0.urdf',
|
|
act_joint)
|
|
self.lim_lo_pin = self.ikctrl.robot.model.lowerPositionLimit
|
|
self.lim_hi_pin = self.ikctrl.robot.model.upperPositionLimit
|
|
|
|
self.pin_from_mot = np.zeros(29) # FIXME(ycho): hardcoded
|
|
self.mot_from_pin = np.zeros(43) # FIXME(ycho): hardcoded
|
|
self.mot_from_pin_act = np.zeros(7) # FIXME(ycho): hardcoded
|
|
for i_mot, j in enumerate( self.config.motor_joint ):
|
|
i_pin = (self.ikctrl.robot.index(j) - 1)
|
|
self.pin_from_mot[i_mot] = i_pin
|
|
self.mot_from_pin[i_pin] = i_mot
|
|
if j in act_joint:
|
|
i_act = act_joint.index(j)
|
|
self.mot_from_pin_act[i_act] = i_mot
|
|
|
|
|
|
# Initialize the policy network
|
|
self.policy = torch.jit.load(config.policy_path)
|
|
# Initializing process variables
|
|
self.qj = np.zeros(43, dtype=np.float32)
|
|
self.dqj = np.zeros(43, dtype=np.float32)
|
|
self.action = np.zeros(7, dtype=np.float32)
|
|
self.target_dof_pos = config.default_angles.copy()
|
|
self.obs = np.zeros(config.num_obs, dtype=np.float32)
|
|
self.cmd = np.array([0.0, 0, 0])
|
|
self.counter = 0
|
|
|
|
rp.init()
|
|
self._node = rp.create_node("low_level_cmd_sender")
|
|
|
|
if config.msg_type == "hg":
|
|
# g1 and h1_2 use the hg msg type
|
|
|
|
self.low_cmd = LowCmdHG()
|
|
self.low_state = LowStateHG()
|
|
|
|
self.lowcmd_publisher_ = self._node.create_publisher(LowCmdHG,
|
|
'lowcmd', 10)
|
|
self.lowstate_subscriber = self._node.create_subscription(LowStateHG,
|
|
'lowstate', self.LowStateHgHandler, 10)
|
|
self.mode_pr_ = MotorMode.PR
|
|
self.mode_machine_ = 0
|
|
|
|
# self.lowcmd_publisher_ = ChannelPublisher(config.lowcmd_topic, LowCmdHG)
|
|
# self.lowcmd_publisher_.Init()
|
|
|
|
# self.lowstate_subscriber = ChannelSubscriber(config.lowstate_topic, LowStateHG)
|
|
# self.lowstate_subscriber.Init(self.LowStateHgHandler, 10)
|
|
|
|
elif config.msg_type == "go":
|
|
raise ValueError(f"{config.msg_type} is not implemented yet.")
|
|
|
|
else:
|
|
raise ValueError("Invalid msg_type")
|
|
|
|
# wait for the subscriber to receive data
|
|
# self.wait_for_low_state()
|
|
|
|
# Initialize the command msg
|
|
if config.msg_type == "hg":
|
|
init_cmd_hg(self.low_cmd, self.mode_machine_, self.mode_pr_)
|
|
elif config.msg_type == "go":
|
|
init_cmd_go(self.low_cmd, weak_motor=self.config.weak_motor)
|
|
|
|
self.mode = Mode.wait
|
|
self._mode_change = True
|
|
self._timer = self._node.create_timer(self.config.control_dt, self.run_wrapper)
|
|
self._terminate = False
|
|
try:
|
|
rp.spin(self._node)
|
|
except KeyboardInterrupt:
|
|
print("KeyboardInterrupt")
|
|
finally:
|
|
self._node.destroy_timer(self._timer)
|
|
create_damping_cmd(self.low_cmd)
|
|
self.send_cmd(self.low_cmd)
|
|
self._node.destroy_node()
|
|
rp.shutdown()
|
|
print("Exit")
|
|
|
|
def LowStateHgHandler(self, msg: LowStateHG):
|
|
self.low_state = msg
|
|
self.mode_machine_ = self.low_state.mode_machine
|
|
self.remote_controller.set(self.low_state.wireless_remote)
|
|
|
|
def LowStateGoHandler(self, msg: LowStateGo):
|
|
self.low_state = msg
|
|
self.remote_controller.set(self.low_state.wireless_remote)
|
|
|
|
def send_cmd(self, cmd: Union[LowCmdGo, LowCmdHG]):
|
|
cmd.mode_machine = self.mode_machine_
|
|
cmd.crc = CRC().Crc(cmd)
|
|
size = len(cmd.motor_cmd)
|
|
# print(cmd.mode_machine)
|
|
# for i in range(size):
|
|
# print(i, cmd.motor_cmd[i].q,
|
|
# cmd.motor_cmd[i].dq,
|
|
# cmd.motor_cmd[i].kp,
|
|
# cmd.motor_cmd[i].kd,
|
|
# cmd.motor_cmd[i].tau)
|
|
self.lowcmd_publisher_.publish(cmd)
|
|
|
|
def wait_for_low_state(self):
|
|
while self.low_state.crc == 0:
|
|
print(self.low_state)
|
|
time.sleep(self.config.control_dt)
|
|
print("Successfully connected to the robot.")
|
|
|
|
def zero_torque_state(self):
|
|
if self.remote_controller.button[KeyMap.start] == 1:
|
|
self._mode_change = True
|
|
self.mode = Mode.default_pos
|
|
else:
|
|
create_zero_cmd(self.low_cmd)
|
|
self.send_cmd(self.low_cmd)
|
|
|
|
def prepare_default_pos(self):
|
|
# move time 2s
|
|
total_time = 2
|
|
self.counter = 0
|
|
self._num_step = int(total_time / self.config.control_dt)
|
|
|
|
dof_idx = self.config.leg_joint2motor_idx + self.config.arm_waist_joint2motor_idx
|
|
kps = self.config.kps + self.config.arm_waist_kps
|
|
kds = self.config.kds + self.config.arm_waist_kds
|
|
self._kps = [float(kp) for kp in kps]
|
|
self._kds = [float(kd) for kd in kds]
|
|
self._default_pos = np.concatenate((self.config.default_angles, self.config.arm_waist_target), axis=0)
|
|
self._dof_size = len(dof_idx)
|
|
self._dof_idx = dof_idx
|
|
|
|
# record the current pos
|
|
self._init_dof_pos = np.zeros(self._dof_size,
|
|
dtype=np.float32)
|
|
for i in range(self._dof_size):
|
|
self._init_dof_pos[i] = self.low_state.motor_state[dof_idx[i]].q
|
|
|
|
def move_to_default_pos(self):
|
|
# move to default pos
|
|
if self.counter < self._num_step:
|
|
alpha = self.counter / self._num_step
|
|
for j in range(self._dof_size):
|
|
motor_idx = self._dof_idx[j]
|
|
target_pos = self._default_pos[j]
|
|
self.low_cmd.motor_cmd[motor_idx].q = (self._init_dof_pos[j] *
|
|
(1 - alpha) + target_pos * alpha)
|
|
self.low_cmd.motor_cmd[motor_idx].dq = 0.0
|
|
self.low_cmd.motor_cmd[motor_idx].kp = self._kps[j]
|
|
self.low_cmd.motor_cmd[motor_idx].kd = self._kds[j]
|
|
self.low_cmd.motor_cmd[motor_idx].tau = 0.0
|
|
self.send_cmd(self.low_cmd)
|
|
self.counter += 1
|
|
else:
|
|
self._mode_change = True
|
|
self.mode = Mode.damping
|
|
|
|
def default_pos_state(self):
|
|
if self.remote_controller.button[KeyMap.A] != 1:
|
|
for i in range(len(self.config.leg_joint2motor_idx)):
|
|
motor_idx = self.config.leg_joint2motor_idx[i]
|
|
self.low_cmd.motor_cmd[motor_idx].q = float(self.config.default_angles[i])
|
|
self.low_cmd.motor_cmd[motor_idx].dq = 0.0
|
|
self.low_cmd.motor_cmd[motor_idx].kp = self._kps[i]
|
|
self.low_cmd.motor_cmd[motor_idx].kd = self._kds[i]
|
|
self.low_cmd.motor_cmd[motor_idx].tau = 0.0
|
|
for i in range(len(self.config.arm_waist_joint2motor_idx)):
|
|
motor_idx = self.config.arm_waist_joint2motor_idx[i]
|
|
self.low_cmd.motor_cmd[motor_idx].q = float(self.config.arm_waist_target[i])
|
|
self.low_cmd.motor_cmd[motor_idx].dq = 0.0
|
|
self.low_cmd.motor_cmd[motor_idx].kp = self._kps[i]
|
|
self.low_cmd.motor_cmd[motor_idx].kd = self._kds[i]
|
|
self.low_cmd.motor_cmd[motor_idx].tau = 0.0
|
|
self.send_cmd(self.low_cmd)
|
|
else:
|
|
self._mode_change = True
|
|
self.mode = Mode.policy
|
|
|
|
def run_policy(self):
|
|
if self.remote_controller.button[KeyMap.select] == 1:
|
|
self._mode_change = True
|
|
self.mode = Mode.null
|
|
return
|
|
self.counter += 1
|
|
|
|
# Get current joint positions
|
|
for i_mot in range(len(self.motor_joint)):
|
|
i_pin = self.pin_from_mot[i_mot]
|
|
self.qj[i_pin] = self.low_state.motor_state[i_mot].q
|
|
self.cmd[0] = self.remote_controller.ly
|
|
self.cmd[1] = self.remote_controller.lx * -1
|
|
self.cmd[2] = self.remote_controller.rx * -1
|
|
delta = np.concatenate([self.cmd,
|
|
[1,0,0,0]])
|
|
res_q = self.ikctrl(self.qj, delta, rel=True)
|
|
for i_act in range(len(res_q)):
|
|
i_mot = self.mot_from_pin_act[i_act]
|
|
i_pin = self.pin_from_mot[i_mot]
|
|
target_q = (
|
|
self.low_state.motor_state[i_mot].q + res_q[i_act]
|
|
)
|
|
target_q = np.clip(target_q,
|
|
self.lim_lo_pin[i_pin],
|
|
self.lim_hi_pin[i_pin])
|
|
self.low_cmd.motor_cmd[i_mot].q = target_q
|
|
self.low_cmd.motor_cmd[i_mot].dq = 0.0
|
|
self.low_cmd.motor_cmd[i_mot].kp = float(self.config.kps[i_mot])
|
|
self.low_cmd.motor_cmd[i_mot].kd = float(self.config.kps[i_mot])
|
|
self.low_cmd.motor_cmd[i_mot].tau = 0.0
|
|
# send the command
|
|
self.send_cmd(self.low_cmd)
|
|
|
|
def run_wrapper(self):
|
|
# print("hello", self.mode,
|
|
# self.mode == Mode.zero_torque)
|
|
if self.mode == Mode.wait:
|
|
if self.low_state.crc != 0:
|
|
self.mode = Mode.zero_torque
|
|
self.low_cmd.mode_machine = self.mode_machine_
|
|
print("Successfully connected to the robot.")
|
|
elif self.mode == Mode.zero_torque:
|
|
if self._mode_change:
|
|
print("Enter zero torque state.")
|
|
print("Waiting for the start signal...")
|
|
self._mode_change = False
|
|
self.zero_torque_state()
|
|
elif self.mode == Mode.default_pos:
|
|
if self._mode_change:
|
|
print("Moving to default pos.")
|
|
self._mode_change = False
|
|
self.prepare_default_pos()
|
|
self.move_to_default_pos()
|
|
elif self.mode == Mode.damping:
|
|
if self._mode_change:
|
|
print("Enter default pos state.")
|
|
print("Waiting for the Button A signal...")
|
|
self._mode_change = False
|
|
self.default_pos_state()
|
|
elif self.mode == Mode.policy:
|
|
if self._mode_change:
|
|
print("Run policy.")
|
|
self._mode_change = False
|
|
self.counter = 0
|
|
self.run_policy()
|
|
elif self.mode == Mode.null:
|
|
self._terminate = True
|
|
|
|
# time.sleep(self.config.control_dt)
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("config", type=str, help="config file name in the configs folder", default="g1.yaml")
|
|
args = parser.parse_args()
|
|
|
|
# Load config
|
|
config_path = f"{LEGGED_GYM_ROOT_DIR}/deploy/deploy_real/configs/{args.config}"
|
|
config = Config(config_path)
|
|
|
|
controller = Controller(config)
|