forked from ZhengLiu-cart/IK_qp
228 lines
6.8 KiB
Python
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() |