import lcm

from time import sleep

import numpy as np
import threading
import signal
from leg_control_data_lcmt import leg_control_data_lcmt
from motor_ctrl_lcmt import motor_ctrl_lcmt
from motor_ctrl_state_lcmt import motor_ctrl_state_lcmt
from state_estimator_lcmt import state_estimator_lcmt
# 模型推理
import torch
import [Link] as nn
from rsl_rl.modules.actor_critic import *
from rsl_rl.[Link] import Estimator

class HardwareVisionNN([Link]):
def __init__(self, num_prop,
num_scan,
num_priv_latent,
num_priv_explicit,
num_hist,
num_actions,
tanh,
actor_hidden_dims=[512, 256, 128],
scan_encoder_dims=[128, 64, 32],
depth_encoder_hidden_dim=512,
activation='elu',
priv_encoder_dims=[64, 20]
):
super(HardwareVisionNN, self).__init__()

self.num_prop = num_prop
self.num_scan = num_scan
self.num_hist = num_hist
self.num_actions = num_actions
self.num_priv_latent = num_priv_latent
self.num_priv_explicit = num_priv_explicit
num_obs = num_prop + num_scan + num_hist*num_prop +
num_priv_latent + num_priv_explicit
self.num_obs = num_obs
activation = get_activation(activation)

[Link] = Actor(num_prop, num_scan, num_actions,

scan_encoder_dims, actor_hidden_dims, priv_encoder_dims,
num_priv_latent, num_priv_explicit, num_hist, activation,
tanh_encoder_output=tanh)
 [Link] = Estimator(input_dim=num_prop,
output_dim=num_priv_explicit, hidden_dims=[128, 64])

def forward(self, obs, depth_latent):

obs[:, self.num_prop+self.num_scan :
self.num_prop+self.num_scan+self.num_priv_explicit] =
[Link](obs[:, :self.num_prop])
return [Link](obs, hist_encoding=True, eval=False,
scandots_latent=depth_latent)
# return obs, depth_latent

def play():

"""
[1, 53+132+9+29+53*10]
53: obs
132:
9: estimator 3+3+3 imu/xyz 线速度
29:
53*10:history encoder, obs*10
"""

class RobotData:
def __init__(self):
self.q = [Link](12) # 12 个关节电机⻆度 弧度制
[Link] = [Link](12) # 电机⻆速度 弧度制
[Link] = [Link](12) # 电机扭矩 N.M
[Link] = [Link](4) # 机身姿态四元数 右⼿坐标系
simulator_state
[Link] = [Link](3) # 机身姿态横滚 俯仰 偏航⻆
弧度制 global_to_robot simulator_state
[Link] = [Link](3) # 加速度计值
[Link] = [Link](3) # ⻆速度计值
global_to_robot
self.ctrl_topic_interval = 0.0 # 控制 topic 通信延迟
self.err_flag = 0

class MotorCmd:
def __init__(self):
self.q_des = [Link](12) # 12 个关节电机期望⻆度 弧度制
self.qd_des = [Link](12) # 电机期望⻆速度 弧度制
self.kp_des = [Link](12) # 电机位置控制⽐例系数
 self.kd_des = [Link](12) # 电机速度控制⽐例系数
self.tau_des = [Link](12) # 电机期望前馈扭矩

class CustomInterface:
def __init__(self, loop_rate):
self.robot_data = RobotData()
self.motor_cmd = MotorCmd()

[Link] = 1.0 / loop_rate

[Link] = True
self.all_thread_done = False
self.mode_state = True

# self.lcm_instance = [Link]()

self.motor_data_lcm = [Link]("udpm://
[Link]:7667?ttl=255")
self.motor_ctrl_state_lcm = [Link]("udpm://
[Link]:7667?ttl=255")
self.robot_state_lcm = [Link]("udpm://
[Link]:7669?ttl=255")
self.motor_ctrl_lcm = [Link]("udpm://
[Link]:7667?ttl=255")

# self.motor_data_Lcm = [Link]("udpm://
[Link]:7667?ttl=255")
# self.motor_ctrl_state_Lcm =
self.get_lcm_url_port(7667, 255)
# self.robot_state_Lcm = self.get_lcm_url_port(7669,
255)
# self.motor_ctrl_Lcm = self.get_lcm_url_port(7667,
255)

self.motor_ctrl_state_lcm.subscribe("motor_ctrl_state",
self.handle_motor_ctrl_state_lcm)
self.motor_data_lcm.subscribe("leg_control_data",
self.handle_motor_data_lcm)
self.robot_state_lcm.subscribe("state_estimator",
self.handle_robot_state_lcm)

# Starting threads for LCM handling

self.motor_ctrl_state_lcm_thread =
[Link](target=self.motor_ctrl_state_lcm_thread)
self.motor_data_lcm_thread =
[Link](target=self.motor_data_lcm_thread)
self.robot_state_lcm_thread =
[Link](target=self.robot_state_lcm_thread)
self.user_code_control_thread =
[Link](target=[Link])

# self.lcm_instance =
[Link](self.get_lcm_url_port(7667, 255))

# self.lcm_instance.subscribe("motor_ctrl_state",
self.handle_motor_ctrl_state_lcm)
# self.lcm_instance.subscribe("leg_control_data",
self.handle_motor_data_lcm)
# self.lcm_instance.subscribe("state_estimator",
self.handle_robot_state_lcm)

# self.motor_ctrl_state_thread =
[Link](target=self.motor_ctrl_state_lcm_thread)
# self.motor_data_thread =
[Link](target=self.motor_data_lcm_thread)
# self.robot_state_thread =
[Link](target=self.robot_state_lcm_thread)
# self.user_code_control_thread =
[Link](target=[Link])

# def get_lcm_url_port(self, port, ttl):

# assert 0 <= ttl <= 255
# return f"udpm://[Link]:{port}?ttl={ttl}"

def control(self):
while [Link]:
if not self.mode_state:
sleep(1)
print("Motor control mode has not been
activated successfully")
continue
self.user_code()
if self.robot_data.err_flag & 0x02:
print("Error: Communication lost over 500ms!")
else:
self.motor_cmd_send()

def spin(self):
# self.motor_ctrl_state_thread.start()
# self.motor_data_thread.start()
# self.robot_state_thread.start()
# self.user_code_control_thread.start()
self.motor_ctrl_state_lcm_thread.start()
 self.motor_data_lcm_thread.start()
self.robot_state_lcm_thread.start()
self.user_code_control_thread.start()

while not self.all_thread_done:

sleep(1)
print("~ Exit ~")

def stop(self):
[Link] = False
self.motor_ctrl_state_lcm_thread.join()
self.motor_data_lcm_thread.join()
self.robot_state_lcm_thread.join()
self.user_code_control_thread.join()
self.all_thread_done = True

def user_code(self):
raise NotImplementedError("This method should be
overridden by subclasses")

def motor_cmd_send(self):
sig = [1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1]
motor_ctrl = motor_ctrl_lcmt()
for i in range(12):
motor_ctrl.q_des[i] = self.motor_cmd.q_des[i] *
sig[i]
motor_ctrl.qd_des[i] = self.motor_cmd.qd_des[i] *
sig[i]
motor_ctrl.kp_des[i] = self.motor_cmd.kp_des[i]
motor_ctrl.kd_des[i] = self.motor_cmd.kd_des[i]
motor_ctrl.tau_des[i] = self.motor_cmd.tau_des[i] *
sig[i]
self.motor_ctrl_lcm.publish("motor_ctrl",
motor_ctrl.encode())
# self.lcm_instance.publish("motor_ctrl",
motor_ctrl.encode())

def motor_ctrl_state_lcm_thread(self):
while [Link]:
# self.lcm_instance.handle_timeout(1000)
self.motor_ctrl_state_lcm.handle_timeout(1000)

def motor_data_lcm_thread(self):
while [Link]:
# self.lcm_instance.handle_timeout(1000)
self.motor_data_lcm.handle_timeout(1000)
 def robot_state_lcm_thread(self):
while [Link]:
# self.lcm_instance.handle_timeout(1000)
self.robot_state_lcm.handle_timeout(1000)

def handle_motor_ctrl_state_lcm(self, channel, data):

msg = motor_ctrl_state_lcmt.decode(data)
self.robot_data.err_flag = msg.err_flag
self.robot_data.ctrl_topic_interval =
msg.ctrl_topic_interval
self.mode_state = True

def handle_motor_data_lcm(self, channel, data):

msg = leg_control_data_lcmt.decode(data)
sig = [1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1]
for i in range(12):
self.robot_data.q[i] = msg.q[i] * sig[i]
self.robot_data.qd[i] = [Link][i] * sig[i]
self.robot_data.tau[i] = msg.tau_est[i] * sig[i]

def handle_robot_state_lcm(self, channel, data):

msg = state_estimator_lcmt.decode(data)
for i in range(3):
self.robot_data.omega[i] = [Link][i]
self.robot_data.rpy[i] = [Link][i]
self.robot_data.acc[i] = [Link][i]
for i in range(4):
self.robot_data.quat[i] = [Link][i]

# def handle_robot_state_lcm(self, channel, data):

# print(f"Received data on channel: {channel}")
# try:
# msg = state_estimator_lcmt.decode(data)
# print(f"Decoded data:
omegaWorld={[Link]}, rpy={[Link]},
aWorld={[Link]}, quat={[Link]}")
# self.robot_data.omega = [Link]
# self.robot_data.rpy = [Link]
# self.robot_data.acc = [Link]
# self.robot_data.quat = [Link]
# print(f"Updated robot state:
rpy={self.robot_data.rpy}, acc={self.robot_data.acc},
quat={self.robot_data.quat}, omega={self.robot_data.omega}")
# except Exception as e:
# print(f"Error decoding data: {e}")
class ExampleMotorCtrl(CustomInterface):
def __init__(self, loop_rate):
super().__init__(loop_rate)
[Link] = 0

def user_code(self):
# 例如实现⼀些特定的电机控制逻辑
t = min([Link] / 1500.0, 2)
if [Link] % 100000 == 0:
print(f"interval:
{self.robot_data.ctrl_topic_interval:.4f}")
obs = [Link](1, n_proprio + num_scan +
n_priv_explicit + n_priv_latent + history_len*n_proprio,
device=device)
depth_latent = [Link](1, 32, device=device)

with torch.no_grad():
output = model(obs,depth_latent)
print(output)

q = self.robot_data.q # 假设这是当前的 q 值
qd = self.robot_data.qd

self.motor_cmd.q_des = output[0, :12].cpu().numpy()

self.motor_cmd.qd_des = [Link](12)
self.motor_cmd.kp_des = [Link](12) * 20
self.motor_cmd.kd_des = [Link](12) * 0.5

tau_des = self.motor_cmd.tau_des -
(self.motor_cmd.q_des - q) * self.motor_cmd.kp_des -
(self.motor_cmd.qd_des - qd) * self.motor_cmd.kd_des
self.motor_cmd.tau_des = tau_des

print(f"rpy [3]: {self.robot_data.rpy}")

print(f"acc [3]: {self.robot_data.acc}")
print(f"quat[4]: {self.robot_data.quat}")
print(f"omeg[3]: {self.robot_data.omega}")
print(f"q [12]: {self.robot_data.q}")
print(f"qd [12]: {self.robot_data.qd}")
print(f"tau[12]: {self.robot_data.tau}")
print(f"ctrl[12]: {self.motor_cmd.q_des}")
[Link] += 1

def signal_callback_handler(signum, frame):

global io
 [Link]()
[Link]([Link], signal_callback_handler)

if __name__ == "__main__":
device = "cuda:0" if [Link].is_available() else "cpu"
print(device)

n_priv_explicit = 3 + 3 + 3
n_priv_latent = 4 + 1 + 12 +12
num_scan = 132
num_actions = 12
n_proprio = 3 + 2 + 3 + 4 + 36 + 4 +1
history_len = 10

model = [Link]("001-01-14000-
base_jit.pt").to(device)

io = ExampleMotorCtrl(500)
[Link]([Link], signal_callback_handler)
[Link]()