sunxianghui
340 lines
10 KiB
Python
Executable File
340 lines
10 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""ROS 控制指令 → 本地 MuJoCo → 记录并绘制 角度/速度/力矩。
|
||
|
||
订 /cb_<hand>_hand_control_cmd（0～255），用与仿真节点相同的 O6 映射写 ctrl，
|
||
从 MuJoCo 读 q / v / actuator_force，Ctrl+C 存 CSV + 三曲线图。
|
||
|
||
用法：
|
||
# 终端1（可先不开 run_sim.sh，本脚本自带 MuJoCo）
|
||
python3 tools/ros_mujoco_record_qvt.py --hand left --finger middle --viewer
|
||
|
||
# 终端2 发 ROS 指令，或加 --self-sweep 自动发中指慢扫
|
||
python3 tools/ros_mujoco_record_qvt.py --hand left --finger middle --viewer --self-sweep
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import sys
|
||
import threading
|
||
import time
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
|
||
ROOT = Path(__file__).resolve().parents[1]
|
||
PKG = (
|
||
ROOT
|
||
/ "src/linkerhand-sim/linker_hand_mujoco_ros2/linker_hand_mujoco_ros2"
|
||
)
|
||
sys.path.insert(0, str(PKG.parent))
|
||
|
||
from linker_hand_mujoco_ros2.utils.mapping import ( # noqa: E402
|
||
O6_JOINT_MAP,
|
||
O6_MIMIC,
|
||
apply_mimic,
|
||
range_to_arc_left,
|
||
range_to_arc_right,
|
||
)
|
||
|
||
FINGER_CFG = {
|
||
"index": {
|
||
"mcp": "index_mcp_pitch",
|
||
"mcp_act": "index_mcp_pitch_pos",
|
||
"ros_idx": 2,
|
||
},
|
||
"middle": {
|
||
"mcp": "middle_mcp_pitch",
|
||
"mcp_act": "middle_mcp_pitch_pos",
|
||
"ros_idx": 3,
|
||
},
|
||
"ring": {
|
||
"mcp": "ring_mcp_pitch",
|
||
"mcp_act": "ring_mcp_pitch_pos",
|
||
"ros_idx": 4,
|
||
},
|
||
"pinky": {
|
||
"mcp": "pinky_mcp_pitch",
|
||
"mcp_act": "pinky_mcp_pitch_pos",
|
||
"ros_idx": 5,
|
||
},
|
||
}
|
||
|
||
|
||
def map_position_array(position, joint_map):
|
||
mapped = [0.0] * len(joint_map)
|
||
for target_idx, source_idx in joint_map.items():
|
||
if source_idx < len(position):
|
||
mapped[target_idx] = position[source_idx]
|
||
return mapped
|
||
|
||
|
||
def u8_to_ctrl(position_u8, hand_type: str):
|
||
if hand_type == "left":
|
||
tmp = range_to_arc_left(position_u8, "O6")
|
||
else:
|
||
tmp = range_to_arc_right(position_u8, "O6")
|
||
res = map_position_array(tmp, O6_JOINT_MAP)
|
||
from linker_hand_mujoco_ros2.utils.mapping import O6_JOINT_ARC
|
||
|
||
res = apply_mimic(res, O6_MIMIC, O6_JOINT_ARC)
|
||
return res
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--hand", choices=["left", "right"], default="left")
|
||
parser.add_argument("--finger", choices=list(FINGER_CFG.keys()), default="middle")
|
||
parser.add_argument("--cmd-topic", default="")
|
||
parser.add_argument("--out", type=Path, default=None)
|
||
parser.add_argument("--viewer", action="store_true")
|
||
parser.add_argument(
|
||
"--self-sweep",
|
||
action="store_true",
|
||
help="本进程自动往控制话题发 255→0→255 慢扫（方便自测）",
|
||
)
|
||
parser.add_argument(
|
||
"--exit-after-sweep",
|
||
action="store_true",
|
||
help="配合 --self-sweep：扫完自动结束并画图",
|
||
)
|
||
parser.add_argument("--duration", type=float, default=8.0, help="self-sweep 单向时长")
|
||
parser.add_argument(
|
||
"--record-duration",
|
||
type=float,
|
||
default=0.0,
|
||
help=">0 时录满该秒数自动结束（配合 middle_step_cmd_pub）",
|
||
)
|
||
parser.add_argument("--label", default="ros_sim")
|
||
args = parser.parse_args()
|
||
if args.self_sweep and not args.exit_after_sweep:
|
||
# 自测默认扫完就退出，避免一直挂着
|
||
args.exit_after_sweep = True
|
||
|
||
cfg = FINGER_CFG[args.finger]
|
||
cmd_topic = args.cmd_topic or f"/cb_{args.hand}_hand_control_cmd"
|
||
out_dir = args.out or (ROOT / f"reports/O6_{args.finger}_ros")
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
xml = (
|
||
PKG
|
||
/ f"urdf/O6/linker_hand_o6_{args.hand}/linker_hand_o6_{args.hand}.xml"
|
||
)
|
||
|
||
import mujoco
|
||
import mujoco.viewer
|
||
|
||
model = mujoco.MjModel.from_xml_path(str(xml))
|
||
data = mujoco.MjData(model)
|
||
model.opt.disableflags = 0
|
||
|
||
mcp_jnt = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, cfg["mcp"])
|
||
mcp_act = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, cfg["mcp_act"])
|
||
if mcp_jnt < 0 or mcp_act < 0:
|
||
raise RuntimeError(f"joint/actuator not found for {args.finger}")
|
||
qadr = int(model.jnt_qposadr[mcp_jnt])
|
||
dadr = int(model.jnt_dofadr[mcp_jnt])
|
||
mcp_lo, mcp_hi = map(float, model.jnt_range[mcp_jnt])
|
||
|
||
ctrl_lock = threading.Lock()
|
||
ctrl_values = np.zeros(model.nu)
|
||
last_cmd_u8 = [255.0] * 6
|
||
running = True
|
||
sweep_done = threading.Event()
|
||
logs = {"t": [], "cmd_rad": [], "q": [], "v": [], "tau": [], "cmd_u8": []}
|
||
t0 = time.perf_counter()
|
||
t_stop = (
|
||
t0 + args.record_duration if args.record_duration > 0 else None
|
||
)
|
||
|
||
import signal
|
||
|
||
def _stop(*_):
|
||
nonlocal running
|
||
running = False
|
||
|
||
signal.signal(signal.SIGINT, _stop)
|
||
signal.signal(signal.SIGTERM, _stop)
|
||
|
||
import rclpy
|
||
from rclpy.node import Node
|
||
from sensor_msgs.msg import JointState
|
||
|
||
class CmdNode(Node):
|
||
def __init__(self):
|
||
super().__init__("ros_mujoco_record_qvt")
|
||
self.create_subscription(JointState, cmd_topic, self.on_cmd, 50)
|
||
self.get_logger().info(
|
||
f"MuJoCo+ROS record | topic={cmd_topic} | finger={args.finger} | "
|
||
f"joint={cfg['mcp']} range=[{mcp_lo:.2f},{mcp_hi:.2f}]"
|
||
)
|
||
self.get_logger().info("Ctrl+C to stop and plot")
|
||
|
||
def on_cmd(self, msg: JointState):
|
||
nonlocal last_cmd_u8
|
||
pos = list(msg.position)
|
||
if len(pos) < 6:
|
||
pos = pos + [255.0] * (6 - len(pos))
|
||
last_cmd_u8 = pos[:6]
|
||
mapped = u8_to_ctrl(last_cmd_u8, args.hand)
|
||
with ctrl_lock:
|
||
n = min(len(mapped), model.nu)
|
||
ctrl_values[:n] = mapped[:n]
|
||
|
||
rclpy.init()
|
||
node = CmdNode()
|
||
|
||
def spin_ros():
|
||
while running and rclpy.ok():
|
||
rclpy.spin_once(node, timeout_sec=0.01)
|
||
|
||
ros_thread = threading.Thread(target=spin_ros, daemon=True)
|
||
ros_thread.start()
|
||
|
||
pub = None
|
||
if args.self_sweep:
|
||
pub = node.create_publisher(JointState, cmd_topic, 10)
|
||
time.sleep(0.3)
|
||
|
||
def sweep_pub():
|
||
nonlocal running
|
||
half = args.duration
|
||
ros_i = cfg["ros_idx"]
|
||
# 255 -> 0
|
||
t_start = time.time()
|
||
while running and time.time() - t_start < half:
|
||
u = (time.time() - t_start) / half
|
||
mid = 255.0 * (1.0 - u)
|
||
m = JointState()
|
||
m.position = [255.0, 255.0, 255.0, 255.0, 255.0, 255.0]
|
||
m.position[ros_i] = mid
|
||
pub.publish(m)
|
||
time.sleep(0.01)
|
||
# 0 -> 255
|
||
t_start = time.time()
|
||
while running and time.time() - t_start < half:
|
||
u = (time.time() - t_start) / half
|
||
mid = 255.0 * u
|
||
m = JointState()
|
||
m.position = [255.0, 255.0, 255.0, 255.0, 255.0, 255.0]
|
||
m.position[ros_i] = mid
|
||
pub.publish(m)
|
||
time.sleep(0.01)
|
||
print("self-sweep finished", flush=True)
|
||
sweep_done.set()
|
||
if args.exit_after_sweep:
|
||
running = False
|
||
|
||
threading.Thread(target=sweep_pub, daemon=True).start()
|
||
|
||
viewer = None
|
||
if args.viewer:
|
||
viewer = mujoco.viewer.launch_passive(model, data)
|
||
|
||
try:
|
||
while running:
|
||
if viewer is not None and not viewer.is_running():
|
||
break
|
||
if t_stop is not None and time.perf_counter() >= t_stop:
|
||
break
|
||
with ctrl_lock:
|
||
data.ctrl[:] = ctrl_values
|
||
mujoco.mj_step(model, data)
|
||
if viewer is not None:
|
||
viewer.sync()
|
||
|
||
t = time.perf_counter() - t0
|
||
q = float(data.qpos[qadr])
|
||
v = float(data.qvel[dadr])
|
||
tau = float(data.actuator_force[mcp_act])
|
||
cmd_rad = float(ctrl_values[mcp_act]) if mcp_act < len(ctrl_values) else 0.0
|
||
logs["t"].append(t)
|
||
logs["cmd_rad"].append(cmd_rad)
|
||
logs["q"].append(q)
|
||
logs["v"].append(v)
|
||
logs["tau"].append(tau)
|
||
logs["cmd_u8"].append(float(last_cmd_u8[cfg["ros_idx"]]))
|
||
time.sleep(0.001)
|
||
except KeyboardInterrupt:
|
||
pass
|
||
finally:
|
||
running = False
|
||
if viewer is not None:
|
||
viewer.close()
|
||
node.destroy_node()
|
||
if rclpy.ok():
|
||
rclpy.shutdown()
|
||
|
||
if len(logs["t"]) < 2:
|
||
print("no samples")
|
||
return
|
||
|
||
for k in logs:
|
||
logs[k] = np.asarray(logs[k], dtype=float)
|
||
|
||
stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
stem = f"{args.label}_{args.finger}_{args.hand}_{stamp}"
|
||
csv_path = out_dir / f"{stem}.csv"
|
||
png_path = out_dir / f"{stem}_qvt.png"
|
||
|
||
np.savetxt(
|
||
csv_path,
|
||
np.column_stack(
|
||
[
|
||
logs["t"],
|
||
logs["cmd_u8"],
|
||
logs["cmd_rad"],
|
||
logs["q"],
|
||
logs["v"],
|
||
logs["tau"],
|
||
]
|
||
),
|
||
delimiter=",",
|
||
header="t_s,cmd_u8,cmd_rad,q_rad,v_rad_s,tau_Nm",
|
||
comments="",
|
||
)
|
||
|
||
import os
|
||
|
||
os.environ.setdefault("MPLCONFIGDIR", str(ROOT / ".mplconfig"))
|
||
Path(os.environ["MPLCONFIGDIR"]).mkdir(parents=True, exist_ok=True)
|
||
import matplotlib
|
||
|
||
matplotlib.use("Agg")
|
||
import matplotlib.pyplot as plt
|
||
|
||
fig, axes = plt.subplots(3, 1, figsize=(10, 8), sharex=True)
|
||
fig.suptitle(
|
||
f"ROS->MuJoCo {args.hand} {cfg['mcp']} "
|
||
f"range=[{mcp_lo:.2f}, {mcp_hi:.2f}] rad",
|
||
fontsize=12,
|
||
)
|
||
axes[0].plot(logs["t"], np.rad2deg(logs["cmd_rad"]), "k--", lw=1.2, label="cmd")
|
||
axes[0].plot(logs["t"], np.rad2deg(logs["q"]), "C0", lw=1.6, label="q")
|
||
axes[0].set_ylabel("angle (deg)")
|
||
axes[0].legend(loc="best")
|
||
axes[0].grid(True, alpha=0.3)
|
||
|
||
axes[1].plot(logs["t"], np.rad2deg(logs["v"]), "C1", lw=1.6, label="v")
|
||
axes[1].set_ylabel("velocity (deg/s)")
|
||
axes[1].legend(loc="best")
|
||
axes[1].grid(True, alpha=0.3)
|
||
|
||
axes[2].plot(logs["t"], logs["tau"], "C3", lw=1.6, label="tau (actuator_force)")
|
||
axes[2].set_ylabel("torque (N·m)")
|
||
axes[2].set_xlabel("time (s)")
|
||
axes[2].legend(loc="best")
|
||
axes[2].grid(True, alpha=0.3)
|
||
fig.tight_layout()
|
||
fig.savefig(png_path, dpi=140)
|
||
plt.close(fig)
|
||
|
||
print(f"CSV -> {csv_path}")
|
||
print(f"plot -> {png_path}")
|
||
print(f"samples={len(logs['t'])}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|