Runtime controller that bridges policy outputs to Unitree low-level commands.
| 42 | |
| 43 | |
| 44 | class Controller: |
| 45 | """Runtime controller that bridges policy outputs to Unitree low-level commands.""" |
| 46 | |
| 47 | def __init__(self, config: G1Config) -> None: |
| 48 | self.config = config |
| 49 | self.remote_controller = RemoteController() |
| 50 | |
| 51 | # Initialize the policy network |
| 52 | |
| 53 | self.policy = torch.jit.load(config.policy_path) |
| 54 | # Initializing process variables |
| 55 | self.qj = np.zeros(config.num_actions, dtype=np.float32) |
| 56 | self.dqj = np.zeros(config.num_actions, dtype=np.float32) |
| 57 | self.action = np.zeros(config.num_actions, dtype=np.float32) |
| 58 | self.target_body_dof_pos = config.default_body_angles.copy() |
| 59 | self.obs = np.zeros(config.num_obs, dtype=np.float32) |
| 60 | self.obs_queue = deque( |
| 61 | [deepcopy(self.obs) for _ in range(self.config.obs_len_history)], |
| 62 | maxlen=self.config.obs_len_history, |
| 63 | ) |
| 64 | self.cmd = np.array([0.0, 0, 0]) |
| 65 | self.counter = 0 |
| 66 | self.last_action = np.zeros((config.num_actions,), dtype=np.float32) |
| 67 | |
| 68 | if config.msg_type == "hg": |
| 69 | # g1 and h1_2 use the hg msg type |
| 70 | self.low_cmd = unitree_hg_msg_dds__LowCmd_() |
| 71 | self.low_state = unitree_hg_msg_dds__LowState_() |
| 72 | self.mode_pr_ = MotorMode.PR |
| 73 | self.mode_machine_ = 0 |
| 74 | |
| 75 | self.lowcmd_publisher_ = ChannelPublisher(config.lowcmd_topic, LowCmdHG) |
| 76 | self.lowcmd_publisher_.Init() |
| 77 | |
| 78 | self.lowstate_subscriber = ChannelSubscriber(config.lowstate_topic, LowStateHG) |
| 79 | self.lowstate_subscriber.Init(self.LowStateHgHandler, 10) |
| 80 | |
| 81 | elif config.msg_type == "go": |
| 82 | # h1 uses the go msg type |
| 83 | self.low_cmd = unitree_go_msg_dds__LowCmd_() |
| 84 | self.low_state = unitree_go_msg_dds__LowState_() |
| 85 | |
| 86 | self.lowcmd_publisher_ = ChannelPublisher(config.lowcmd_topic, LowCmdGo) |
| 87 | self.lowcmd_publisher_.Init() |
| 88 | |
| 89 | self.lowstate_subscriber = ChannelSubscriber(config.lowstate_topic, LowStateGo) |
| 90 | self.lowstate_subscriber.Init(self.LowStateGoHandler, 10) |
| 91 | |
| 92 | else: |
| 93 | raise ValueError("Invalid msg_type") |
| 94 | |
| 95 | # wait for the subscriber to receive data |
| 96 | self.wait_for_low_state() |
| 97 | |
| 98 | # Initialize the command msg |
| 99 | if config.msg_type == "hg": |
| 100 | init_cmd_hg(self.low_cmd, self.mode_machine_, self.mode_pr_) |
| 101 | elif config.msg_type == "go": |