153 lines
5.2 KiB
Python
153 lines
5.2 KiB
Python
import time
|
||
from google.protobuf import timestamp_pb2
|
||
import matplotlib.pyplot as plt # 新增
|
||
|
||
from clients._path_setup import ensure_paths
|
||
|
||
ensure_paths()
|
||
|
||
from clients.base_client import RobotClientBase
|
||
from cmvr.api import humanoid_robot_command_pb2 as pb
|
||
from cmvr.api import humanoid_robot_service_pb2_grpc as rpc
|
||
|
||
# 固定关节顺序
|
||
JOINT_ORDER = [
|
||
"L_SHOULDER_P", "L_SHOULDER_R", "L_SHOULDER_Y",
|
||
"L_ELBOW_R", "L_WRIST_P", "L_WRIST_Y", "L_WRIST_R",
|
||
"R_SHOULDER_P", "R_SHOULDER_R", "R_SHOULDER_Y",
|
||
"R_ELBOW_R", "R_WRIST_P", "R_WRIST_Y", "R_WRIST_R",
|
||
"HEAD_P", "HEAD_Y", "HEAD_R","WAIST_Y","WAIST_R"
|
||
]
|
||
|
||
|
||
class GetJointStateClient(RobotClientBase):
|
||
"""Continuous client to fetch robot joint states"""
|
||
|
||
def __init__(self, *args, **kwargs):
|
||
super().__init__(*args, **kwargs)
|
||
# 时间轴
|
||
self.time_history = []
|
||
self._start_time = None
|
||
# 记录每个关节的角度、速度
|
||
self.pos_history = {name: [] for name in JOINT_ORDER}
|
||
self.vel_history = {name: [] for name in JOINT_ORDER}
|
||
|
||
def send(self, interval=0.5, device_id="hc01"):
|
||
"""Continuously fetch joint states and print in fixed order"""
|
||
try:
|
||
if self._start_time is None:
|
||
self._start_time = time.time()
|
||
|
||
while True:
|
||
# Construct request
|
||
req = pb.JointRequest()
|
||
req.header.device_id = device_id
|
||
timestamp = timestamp_pb2.Timestamp()
|
||
timestamp.GetCurrentTime()
|
||
req.header.timestamp.CopyFrom(timestamp)
|
||
|
||
# Call RPC with timeout
|
||
try:
|
||
resp = self.stub.getJointState(req, timeout=10)
|
||
except Exception as e:
|
||
print(f"RPC call failed: {e}")
|
||
time.sleep(interval)
|
||
continue
|
||
|
||
# ====== 1. 同时取 position 和 velocity ======
|
||
joint_pos_dict = {}
|
||
joint_vel_dict = {}
|
||
for state in resp.state:
|
||
# 这里就用你说的三元 zip
|
||
for name, pos, vel in zip(state.name, state.position, state.velocity):
|
||
joint_pos_dict[name] = round(pos, 12)
|
||
joint_vel_dict[name] = round(vel, 12)
|
||
|
||
# ====== 2. 记录到历史数据里，用于之后画图 ======
|
||
t = time.time() - self._start_time
|
||
self.time_history.append(t)
|
||
for name in JOINT_ORDER:
|
||
self.pos_history[name].append(joint_pos_dict.get(name, 0.0))
|
||
self.vel_history[name].append(joint_vel_dict.get(name, 0.0))
|
||
|
||
# 保持原来的打印逻辑（打印角度）
|
||
joint_list = [{"joint_name": name, "rad": joint_pos_dict.get(name, 0.0)}
|
||
for name in JOINT_ORDER]
|
||
|
||
# Print timestamp
|
||
print(f"[{time.strftime('%H:%M:%S')}]")
|
||
# Print each joint, last one without comma
|
||
for i, j in enumerate(joint_list):
|
||
if i < len(joint_list) - 1:
|
||
print(f"{j},")
|
||
else:
|
||
print(f"{j}")
|
||
|
||
time.sleep(interval)
|
||
|
||
except KeyboardInterrupt:
|
||
print("\nStopped fetching joint states.")
|
||
# 例子：退出时画几个关节的角度 & 速度曲线
|
||
try:
|
||
# 你可以根据需要改关节列表
|
||
self.plot_pos_and_vel([ "L_SHOULDER_R"])
|
||
except Exception as e:
|
||
print(f"绘图失败: {e}")
|
||
|
||
# ====== 3. 绘制速度和角度（同一张图，两行子图） ======
|
||
def plot_pos_and_vel(self, joint_names):
|
||
"""
|
||
在一张图中绘制多个关节的角度和速度：
|
||
- 上 subplot：角度（rad）
|
||
- 下 subplot：速度（rad/s）
|
||
|
||
joint_names: str 或 [str, ...]
|
||
"""
|
||
if isinstance(joint_names, str):
|
||
joint_names = [joint_names]
|
||
|
||
# 检查关节名合法性
|
||
invalid = [n for n in joint_names if n not in self.pos_history]
|
||
if invalid:
|
||
print(f"未知关节名: {invalid}")
|
||
print(f"可选关节: {list(self.pos_history.keys())}")
|
||
return
|
||
|
||
if not self.time_history:
|
||
print("还没有采集到任何数据，无法绘图。")
|
||
return
|
||
|
||
t = self.time_history
|
||
|
||
plt.figure(figsize=(10, 6))
|
||
|
||
# 上：角度
|
||
plt.subplot(2, 1, 1)
|
||
for name in joint_names:
|
||
plt.plot(t, self.pos_history[name], label=name)
|
||
plt.ylabel("Position (rad)")
|
||
plt.title("Joint Position")
|
||
plt.grid(True)
|
||
plt.legend()
|
||
|
||
# 下：速度
|
||
plt.subplot(2, 1, 2)
|
||
for name in joint_names:
|
||
plt.plot(t, self.vel_history[name], label=name)
|
||
plt.xlabel("Time (s)")
|
||
plt.ylabel("Velocity (rad/s)")
|
||
plt.title("Joint Velocity")
|
||
plt.grid(True)
|
||
plt.legend()
|
||
|
||
plt.tight_layout()
|
||
plt.show()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
client = GetJointStateClient()
|
||
try:
|
||
client.send(interval=0.005)
|
||
finally:
|
||
client.close()
|