2025-07-28 12:35:44 +08:00

60 lines
1.4 KiB
Python

from mini_bdx_runtime.hwi import HWI
from mini_bdx_runtime.rl_utils import make_action_dict, mujoco_joints_order
import time
import numpy as np
import mujoco
import mujoco_viewer
import pickle
hwi = HWI("/dev/ttyUSB0")
hwi.turn_on()
hwi.set_pid_all([1100, 0, 0])
# hwi.set_pid([500, 0, 0], "neck_pitch")
# hwi.set_pid([500, 0, 0], "head_pitch")
# hwi.set_pid([500, 0, 0], "head_yaw")
dt = 0.0001
model = mujoco.MjModel.from_xml_path("../../mini_bdx/robots/bdx/scene.xml")
model.opt.timestep = dt
data = mujoco.MjData(model)
mujoco.mj_step(model, data)
viewer = mujoco_viewer.MujocoViewer(model, data)
init_pos = list(hwi.init_pos.values())
# init_pos += [0, 0]
data.qpos[:13] = init_pos
data.ctrl[:13] = init_pos
dof = 7
a = 0.3
f = 3
recording = {}
recording["mujoco_vel"] = []
recording["robot_vel"] = []
try:
while True:
target = a * np.sin(2 * np.pi * f * time.time())
full_target = init_pos.copy()
full_target[dof] += target
#
data.ctrl[:13] = full_target
action_dict = make_action_dict(full_target, mujoco_joints_order)
hwi.set_position_all(action_dict)
mujoco.mj_step(model, data, 50)
mujoco_vel = data.qvel[dof]
robot_vel = hwi.get_present_velocities()[dof]
recording["mujoco_vel"].append(mujoco_vel)
recording["robot_vel"].append(robot_vel)
viewer.render()
except KeyboardInterrupt:
pickle.dump(recording, open("speeds.pkl", "wb"))