Files
IK_qp/kine_ctrl/rm75_mjc.py
LiuzhengSJ c23b5d0d6d try mujoco2
2026-05-29 14:21:03 +01:00

228 lines
6.8 KiB
Python

#!/usr/bin/env python3
"""
Thread-based MuJoCo controller for kinematic verification
No ROS dependency - pure Python threading
"""
import mujoco
import mujoco.viewer
import numpy as np
import threading
import time
from pathlib import Path
class MuJoCoRobotController:
"""
Thread-based robot controller for kinematic verification
Command and feedback via thread-safe queues
"""
def __init__(self, urdf_path, smoothness=0.1, enable_viewer=True):
"""
Args:
urdf_path: Path to URDF file
smoothness: Motion smoothness (0.05-0.2)
enable_viewer: Show MuJoCo viewer
"""
# Load model
self.model = mujoco.MjModel.from_xml_path(urdf_path)
self.data = mujoco.MjData(self.model)
# Robot info
self.n_joints = self.model.njnt
self.joint_names = [self.model.joint(i).name for i in range(self.n_joints)]
# Joint limits
self.joint_lower_limits = [self.model.jnt_range[i, 0] for i in range(self.n_joints)]
self.joint_upper_limits = [self.model.jnt_range[i, 1] for i in range(self.n_joints)]
print(f"Loaded robot: {self.n_joints} joints")
# Control parameters
self.smoothness = smoothness
self.j_cmd = self.data.qpos[:self.n_joints].copy()
self.j_fbk = self.data.qvel[:self.n_joints].copy()
# Control flags
self.running = False
self.simulation_thread = None
# Viewer
self.viewer = None
if enable_viewer:
try:
self.viewer = mujoco.viewer.launch_passive(self.model, self.data)
except Exception as e:
print(f"Viewer warning: {e}")
print("Robot controller ready")
def start(self):
"""Start the simulation thread"""
if self.running:
return
self.running = True
self.simulation_thread = threading.Thread(target=self._simulation_loop, daemon=True)
self.simulation_thread.start()
print("Simulation thread started")
def stop(self):
"""Stop the simulation thread"""
self.running = False
if self.simulation_thread:
self.simulation_thread.join(timeout=2.0)
if self.viewer:
self.viewer.close()
print("Simulation stopped")
def send_command(self, joint_positions):
"""
Send joint command to robot (non-blocking)
Args:
joint_positions: Array of target joint angles
"""
cmd = np.array(joint_positions[:self.n_joints])
# Apply joint limits
for i in range(self.n_joints):
cmd[i] = max(self.joint_lower_limits[i], min(self.joint_upper_limits[i], cmd[i]))
self.j_cmd = cmd
def get_feedback(self):
"""
Get current robot state
Returns:
Dictionary with positions, velocities, etc.
"""
return self.j_fbk
def _simulation_loop(self):
"""Main simulation loop (runs in separate thread)"""
last_time = time.time()
while self.running:
# Process commands
# Get current state
current_pos = self.data.qpos[:self.n_joints].copy()
self.j_fbk = self.data.qpos[:self.n_joints].copy()
current_vel = self.data.qvel[:self.n_joints].copy()
# Smooth interpolation
alpha = self.smoothness
next_positions = current_pos + alpha * (self.j_cmd - current_pos)
# Calculate velocities for smooth motion
dt = self.model.opt.timestep
target_velocities = (next_positions - current_pos) / dt
# Limit velocities
max_vel = 3.0
target_velocities = np.clip(target_velocities, -max_vel, max_vel)
# Apply control
self.data.qvel[:self.n_joints] = target_velocities
# Corrective forces
pos_error = self.j_cmd - current_pos
kp = 30.0
self.data.qfrc_applied[:self.n_joints] = kp * pos_error
# Step simulation
mujoco.mj_step(self.model, self.data)
# Sync viewer
if self.viewer:
self.viewer.sync()
# Maintain real-time speed
elapsed = time.time() - last_time
sleep_time = self.model.opt.timestep - elapsed
if sleep_time > 0:
time.sleep(sleep_time)
last_time = time.time()
print(f'j_cmd: {self.j_cmd}, and j_fbk: {self.j_fbk}')
def print_state(self):
"""Print current robot state"""
feedback = self.get_feedback()
if feedback:
print(f"Positions: {[f'{p:.3f}' for p in feedback['positions'][:4]]}...")
print(f"Velocities: {[f'{v:.3f}' for v in feedback['velocities'][:4]]}...")
# Example usage for kinematic verification
def verify_kinematics():
"""Test sequence to verify kinematic code"""
# Create controller
urdf_path = "/home/zl/Downloads/urdf_rm75/RM75-B.urdf"
robot = MuJoCoRobotController(urdf_path, smoothness=0.08, enable_viewer=True)
# Start simulation
robot.start()
time.sleep(1) # Wait for simulation to initialize
print("\n" + "=" * 60)
print("Kinematic Verification Test")
print("=" * 60)
# Test 1: Single joint movement
print("\n[Test 1] Moving joint 1 to 45 degrees...")
cmd = np.zeros(7)
cmd[0] = 0.785 # 45 degrees
robot.send_command(cmd)
time.sleep(1)
feedback = robot.get_feedback()
print(f" Result: joint_1 = {feedback} rad (expected 0.785)")
# Test 2: Multi-joint pose
print("\n[Test 2] Moving to complex pose...")
cmd = np.array([0.5, -0.4, 0.3, 0.2, 0, 0, 0])
robot.send_command(cmd)
feedback = robot.get_feedback()
print(f" Result: {feedback}")
# Test 3: Return to home
print("\n[Test 3] Returning to home...")
robot.send_command(np.zeros(7))
feedback = robot.get_feedback()
print(f" Result: home = {feedback}")
# Test 4: Continuous trajectory for kinematic verification
print("\n[Test 4] Testing trajectory following...")
trajectory = [
np.array([0.3, 0, 0, 0, 0, 0, 0]),
np.array([0.6, 0, 0, 0, 0, 0, 0]),
np.array([0.3, 0, 0, 0, 0, 0, 0]),
np.array([0, 0, 0, 0, 0, 0, 0]),
]
for i, cmd in enumerate(trajectory):
print(f" Step {i + 1}: sending {cmd[0]:.3f}")
robot.send_command(cmd)
feedback = robot.get_feedback()
print(f" Reached: {feedback}")
print("\n✓ Kinematic verification complete!")
print("\nInteractive mode - send commands using robot.send_command()")
# Keep running for interactive control
try:
while True:
time.sleep(0.1)
except KeyboardInterrupt:
print("\nShutting down...")
robot.stop()
if __name__ == "__main__":
verify_kinematics()