Add homing functionality to mujoco_server.py and implement test script

- Introduced a state-based homing mechanism for UR5 robots, allowing for smoother joint movements towards home positions.
- Added a new `homing_state` dictionary to track the current state of the homing process, including joint errors and jogging status.
- Implemented logic to handle joint overshoot and stuck conditions during the homing sequence.
- Created a `test_home.py` script to simulate continuous home commands, testing the homing functionality and ensuring proper behavior during operation.
- Updated WebSocket message handling to support starting and stopping the homing process based on user input.

mujoco_server.py

@@ -122,6 +122,20 @@ last_teleop_command: dict[str, Any] = {
 }
 teleop_lock = threading.Lock()
 
+# Homing state
+homing_state = {
+    "active": False,
+    "current_joint": None,
+    "joint_errors": [],
+    "current_velocity": 0.2,
+    "previous_error": None,
+    "previous_abs_error": None,
+    "stuck_count": 0,
+    "is_jogging": False,
+    "current_direction": None,
+}
+homing_lock = threading.Lock()
+
 # WebSocket clients
 ws_clients = set()
 ws_clients_lock = threading.Lock()
@@ -990,10 +1004,9 @@ def handle_ws_message(ws, data):
         _schedule_robot_switch(robot, scene)
 
     elif msg_type == 'home':
-        # Stepwise homing for UR5: move robot toward home via joint jogging
-        # Each home message jogs the joint furthest from home
-        # If messages stop arriving, robot stops (jogging stops automatically)
-        with mujoco_lock:
+        # State-based homing: processes one step each time a 'home' message is received
+        # UI sends home messages continuously while button is pressed
+        with mujoco_lock, homing_lock:
             if env is not None and current_robot in ("ur5", "ur5_t_push"):
                 home_pose = getattr(env, "home_pose", None)
                 get_joints = getattr(env, "get_joint_positions", None)
@@ -1007,35 +1020,134 @@ def handle_ws_message(ws, data):
                         target_positions = list(home_pose)
 
                     current_positions = get_joints()
-
-                    # Stop any existing jog
-                    stop_jog()
-
-                    # Find the joint that's furthest from home and jog it
-                    # We jog one joint at a time for smoother motion
                     tolerance = 0.01  # rad (~0.57 degrees)
-                    max_error = 0.0
-                    max_error_joint = -1
 
-                    for joint_idx in range(min(len(target_positions), len(current_positions))):
-                        current = current_positions[joint_idx]
-                        target = target_positions[joint_idx]
-                        error = abs(target - current)
-
-                        if error > tolerance and error > max_error:
-                            max_error = error
-                            max_error_joint = joint_idx
-
-                    # Jog the joint with the largest error
-                    if max_error_joint >= 0:
-                        current = current_positions[max_error_joint]
-                        target = target_positions[max_error_joint]
-                        error = target - current
-
-                        direction = '+' if error > 0 else '-'
-                        # Use velocity proportional to error, but capped
-                        velocity = min(0.3, max(0.05, abs(error)))  # rad/s, 0.05 to 0.3 rad/s
-                        start_jog('joint', joint=max_error_joint + 1, direction=direction, velocity=velocity)
+                    # Initialize homing state if not active
+                    if not homing_state["active"]:
+                        # Compute differences for all joints
+                        joint_errors = []
+                        for joint_idx in range(min(len(target_positions), len(current_positions))):
+                            current = current_positions[joint_idx]
+                            target = target_positions[joint_idx]
+                            error = target - current
+                            # Wrap angle to [-pi, pi]
+                            error = (error + np.pi) % (2 * np.pi) - np.pi
+                            abs_error = abs(error)
+
+                            if abs_error > tolerance:
+                                joint_errors.append({
+                                    'index': joint_idx,
+                                    'error': error,
+                                    'abs_error': abs_error
+                                })
+
+                        if joint_errors:
+                            # Sort by absolute error (largest first)
+                            joint_errors.sort(key=lambda x: x['abs_error'], reverse=True)
+                            homing_state["active"] = True
+                            homing_state["joint_errors"] = joint_errors
+                            homing_state["current_joint"] = joint_errors[0]['index']
+                            homing_state["current_velocity"] = 0.15
+                            homing_state["previous_error"] = None
+                            homing_state["previous_abs_error"] = None
+                            homing_state["stuck_count"] = 0
+                            homing_state["is_jogging"] = False
+                            homing_state["current_direction"] = None
+                            print(f"🏠 Starting homing sequence for {len(joint_errors)} joints")
+                            print(f"→ Homing joint {homing_state['current_joint'] + 1}...")
+                        else:
+                            print("✓ Already at home position")
+                            return
+
+                    # Process current joint
+                    joint_idx = homing_state["current_joint"]
+                    current_pos = current_positions[joint_idx]
+                    target_pos = target_positions[joint_idx]
+
+                    # Compute shortest angular distance
+                    error = target_pos - current_pos
+                    error = (error + np.pi) % (2 * np.pi) - np.pi
+                    abs_error = abs(error)
+
+                    # Check if current joint reached target
+                    if abs_error <= tolerance:
+                        if homing_state["is_jogging"]:
+                            stop_jog()
+                            homing_state["is_jogging"] = False
+                        print(f"✓ Joint {joint_idx + 1} reached target (error: {abs_error:.4f} rad)")
+
+                        # Move to next joint
+                        homing_state["joint_errors"].pop(0)
+                        if homing_state["joint_errors"]:
+                            homing_state["current_joint"] = homing_state["joint_errors"][0]['index']
+                            homing_state["current_velocity"] = 0.15
+                            homing_state["previous_error"] = None
+                            homing_state["previous_abs_error"] = None
+                            homing_state["stuck_count"] = 0
+                            homing_state["is_jogging"] = False
+                            homing_state["current_direction"] = None
+                            print(f"→ Homing joint {homing_state['current_joint'] + 1}...")
+                        else:
+                            # All joints homed
+                            print("✓ Homing sequence complete!")
+                            homing_state["active"] = False
+                        return
+
+                    # Check if stuck
+                    if homing_state["previous_abs_error"] is not None:
+                        if abs(abs_error - homing_state["previous_abs_error"]) < 0.0001:
+                            homing_state["stuck_count"] += 1
+                            if homing_state["stuck_count"] > 20:  # Stuck for ~2 seconds
+                                if abs_error <= tolerance * 3:
+                                    print(f"⚠ Joint {joint_idx + 1} stuck at {abs_error:.4f}, close enough")
+                                    if homing_state["is_jogging"]:
+                                        stop_jog()
+                                        homing_state["is_jogging"] = False
+                                    # Move to next joint
+                                    homing_state["joint_errors"].pop(0)
+                                    if homing_state["joint_errors"]:
+                                        homing_state["current_joint"] = homing_state["joint_errors"][0]['index']
+                                        homing_state["current_velocity"] = 0.15
+                                        homing_state["previous_error"] = None
+                                        homing_state["previous_abs_error"] = None
+                                        homing_state["stuck_count"] = 0
+                                        homing_state["is_jogging"] = False
+                                        homing_state["current_direction"] = None
+                                        print(f"→ Homing joint {homing_state['current_joint'] + 1}...")
+                                    else:
+                                        homing_state["active"] = False
+                                    return
+                        else:
+                            homing_state["stuck_count"] = 0
+
+                    # Check for overshoot
+                    if homing_state["previous_error"] is not None:
+                        prev_err = homing_state["previous_error"]
+                        if (prev_err > 0 and error < 0) or (prev_err < 0 and error > 0):
+                            print(f"↩ Joint {joint_idx + 1} overshot! {prev_err:.4f} → {error:.4f}")
+                            homing_state["current_velocity"] = homing_state["current_velocity"] / 2.0
+                            homing_state["current_velocity"] = max(0.005, homing_state["current_velocity"])
+                            print(f"  Velocity halved to {homing_state['current_velocity']:.3f}")
+                            if homing_state["is_jogging"]:
+                                stop_jog()
+                                homing_state["is_jogging"] = False
+
+                    # Calculate velocity and direction
+                    direction = '+' if error > 0 else '-'
+                    desired_velocity = min(homing_state["current_velocity"], max(0.005, abs_error * 1.0))
+
+                    # Start/restart jogging if needed
+                    if homing_state["current_direction"] != direction or not homing_state["is_jogging"]:
+                        if homing_state["is_jogging"]:
+                            stop_jog()
+
+                        start_jog('joint', joint=joint_idx + 1, direction=direction, velocity=desired_velocity)
+                        homing_state["is_jogging"] = True
+                        homing_state["current_direction"] = direction
+                        print(f"  J{joint_idx + 1}: err={error:.4f}, dir={direction}, vel={desired_velocity:.3f}")
+
+                    homing_state["previous_error"] = error
+                    homing_state["previous_abs_error"] = abs_error
 
             elif env is not None:
                 # Instant homing for non-UR5 robots
@@ -1048,6 +1160,18 @@ def handle_ws_message(ws, data):
                         positions = list(home_pose)
                     set_joints(positions)
 
+    elif msg_type == 'stop_home':
+        # Stop homing when button is released
+        with mujoco_lock, homing_lock:
+            if homing_state["active"]:
+                if env is not None and current_robot in ("ur5", "ur5_t_push"):
+                    stop_jog = getattr(env, "stop_jog", None)
+                    if callable(stop_jog) and homing_state["is_jogging"]:
+                        stop_jog()
+                print("⏸ Homing stopped by user")
+                homing_state["active"] = False
+                homing_state["is_jogging"] = False
+
     elif msg_type == 'camera':
         payload = data.get('data', {})
         a = cam.azimuth * np.pi / 180.0
@@ -2126,7 +2250,7 @@ def index():
 
         <div class="rl-notifications" id="rl_notifications_list"></div>
 
-        <div class="hint" id="hint_box">
+        <div class="hint" id="hint_box" style="display: none;">
             Drag: Rotate Camera<br>
             Scroll: Zoom
         </div>
@@ -3042,12 +3166,11 @@ def index():
                     console.log('Starting Nova jogging homing...');
                     startNovaJoggingHome();
                 } else {
-                    // Local mode: send home message
-                    console.log('Starting local homing, sending home message...');
+                    // Local mode: send home messages continuously while button is pressed
+                    console.log('Starting local homing...');
                     send('home');
-                    console.log('Sent first home message');
 
-                    // Then send repeatedly while button is held (every 100ms)
+                    // Send home message every 100ms while button is pressed
                     if (homingInterval) clearInterval(homingInterval);
                     homingInterval = setInterval(() => {
                         send('home');
@@ -3067,10 +3190,14 @@ def index():
                         homingInterval = null;
                     }
                 } else {
+                    // Stop the home message interval
                     if (homingInterval) {
                         clearInterval(homingInterval);
                         homingInterval = null;
                     }
+                    // Send stop_home message to backend
+                    send('stop_home');
+                    console.log('Sent stop_home message');
                 }
             }
 

test_home.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python3
+"""
+Test the home function by sending continuous home commands.
+Simulates the UI button press behavior.
+"""
+
+import time
+import json
+from websocket import create_connection
+import threading
+
+WS_URL = "ws://localhost:3004/nova-sim/api/v1/ws"
+
+print("=" * 80)
+print("HOME FUNCTION TEST")
+print("=" * 80)
+
+# Connect to WebSocket
+print("\n1. Connecting to WebSocket...")
+ws = create_connection(WS_URL)
+print("   ✓ Connected")
+
+# Switch to UR5
+print("\n2. Switching to UR5...")
+ws.send(json.dumps({"type": "switch_robot", "data": {"robot": "ur5", "scene": None}}))
+time.sleep(2)
+print("   ✓ Switched to UR5")
+
+# Jog some joints away from home position
+print("\n3. Moving joints away from home...")
+for joint, direction, duration in [(1, "+", 2), (3, "-", 2), (6, "+", 3)]:
+    ws.send(json.dumps({
+        "type": "start_jog",
+        "data": {"jog_type": "joint", "params": {"joint": joint, "direction": direction, "velocity": 0.4}}
+    }))
+    time.sleep(duration)
+    ws.send(json.dumps({"type": "stop_jog"}))
+    time.sleep(0.5)
+print("   ✓ Joints moved away from home")
+
+# Send continuous home messages (simulating button press)
+print("\n4. Starting HOME sequence (sending messages every 100ms)...")
+print("   Press Ctrl+C to stop or wait 60 seconds")
+print("=" * 80)
+print("\n📺 WATCH SERVER CONSOLE FOR:")
+print("   🏠 Starting homing sequence for X joints")
+print("   → Homing joint X...")
+print("   J1: err=..., dir=..., vel=...")
+print("   ✓ Joint X reached target")
+print("   ✓ Homing sequence complete!")
+print("")
+
+stop_homing = False
+message_count = 0
+
+def send_home_messages():
+    """Send home messages every 100ms while running"""
+    global message_count
+    while not stop_homing:
+        ws.send(json.dumps({"type": "home"}))
+        message_count += 1
+        time.sleep(0.1)
+
+# Start sending home messages in background
+home_thread = threading.Thread(target=send_home_messages, daemon=True)
+home_thread.start()
+
+try:
+    # Wait for homing to complete (up to 60 seconds)
+    for i in range(12):
+        time.sleep(5)
+        print(f"   {(i+1)*5}s - sent {message_count} home messages")
+
+    print("\n   Timeout reached (60 seconds)")
+
+except KeyboardInterrupt:
+    print("\n\n   Interrupted by user")
+
+# Stop homing
+print("\n5. Stopping homing...")
+stop_homing = True
+time.sleep(0.2)
+ws.send(json.dumps({"type": "stop_home"}))
+print(f"   ✓ Sent stop_home message (total: {message_count} home messages)")
+
+# Close connection
+ws.close()
+
+print("\n" + "=" * 80)
+print("TEST COMPLETE")
+print("=" * 80)
+print("\nCheck the server console to verify:")
+print("  ✓ Sequential homing (one joint after another)")
+print("  ✓ Joints stopped when error <= 0.01 rad")
+print("  ✓ Overshoot detection and velocity halving (if applicable)")
+print("  ✓ 'Homing sequence complete!' was printed")
+print("=" * 80)