tracking

frontend/index.html

@@ -1065,7 +1065,13 @@
                 <div class="control-group">
                     <label style="display: flex; align-items: center; gap: 8px; cursor: pointer;">
                         <input type="checkbox" id="cam_follow" checked onchange="setCameraFollow()">
-                        Camera Follow Robot
+                        <span id="cam_follow_label">Camera Follow Robot</span>
+                    </label>
+                </div>
+                <div class="control-group" id="target_visibility_control" style="display: none;">
+                    <label style="display: flex; align-items: center; gap: 8px; cursor: pointer;">
+                        <input type="checkbox" id="target_visible" checked onchange="setTargetVisibility()">
+                        Show Target T
                     </label>
                 </div>
             </div>
@@ -1991,6 +1997,27 @@
             if (robotSelect.value !== expectedValue) {
                 robotSelect.value = expectedValue;
             }
+
+            // Update camera follow label based on robot type
+            const camFollowLabel = document.getElementById('cam_follow_label');
+            if (camFollowLabel) {
+                if (robot === 'ur5' || robot === 'ur5_t_push') {
+                    camFollowLabel.innerText = 'Camera Follow Object';
+                } else {
+                    camFollowLabel.innerText = 'Camera Follow Robot';
+                }
+            }
+
+            // Show/hide target visibility control (only for push-T scene)
+            const targetVisibilityControl = document.getElementById('target_visibility_control');
+            if (targetVisibilityControl) {
+                if (robot === 'ur5_t_push') {
+                    targetVisibilityControl.style.display = 'block';
+                } else {
+                    targetVisibilityControl.style.display = 'none';
+                }
+            }
+
             // Toggle controls based on robot type
             if (robot === 'ur5' || robot === 'ur5_t_push') {
                 locomotionControls.classList.add('hidden');
@@ -2120,6 +2147,11 @@
             send('camera_follow', { follow });
         }
 
+        function setTargetVisibility() {
+            const visible = document.getElementById('target_visible').checked;
+            send('toggle_target_visibility', { visible });
+        }
+
         // UR5 controls
         let currentControlMode = 'ik';
 

mujoco_server.py

@@ -98,6 +98,13 @@ episode_control_state = {
     "terminate": False,
     "truncate": False,
 }
+
+# Perception state - stores latest detected object poses
+perception_state = {
+    "poses": [],  # List of {object_id, position, orientation, confidence}
+    "last_update": 0,  # Timestamp of last update
+}
+perception_lock = threading.Lock()
 episode_control_lock = threading.Lock()
 
 # Latest teleoperation action (for trainer state)
@@ -154,6 +161,9 @@ cam.lookat = np.array([0, 0, 0.8])
 # Camera follow mode
 camera_follow = True
 
+# Target visibility for push-T scene
+t_target_visible = True
+
 def _reset_camera_for_current_robot() -> None:
     """Reset the orbit camera to the defaults for the active robot."""
     cam.azimuth = 135
@@ -223,6 +233,7 @@ AVAILABLE_ACTIONS = [
     "switch_robot",
     "camera",
     "camera_follow",
+    "toggle_target_visibility",
     "teleop_action",
     "home",
     "stop_home",
@@ -568,7 +579,8 @@ def _get_scene_objects():
         t_object_id = getattr(env, 't_object_body_id', -1)
         if t_object_id >= 0:
             pos = env.data.xpos[t_object_id]
-            quat = env._get_body_quat(t_object_id) if hasattr(env, '_get_body_quat') else np.array([1.0, 0.0, 0.0, 0.0])
+            # MuJoCo stores quaternions in data.xquat as [w, x, y, z]
+            quat = env.data.xquat[t_object_id]
             scene_objects.append({
                 'name': 't_object',
                 'position': {'x': float(pos[0]), 'y': float(pos[1]), 'z': float(pos[2])},
@@ -579,7 +591,7 @@ def _get_scene_objects():
         t_target_id = getattr(env, 't_target_body_id', -1)
         if t_target_id >= 0:
             pos = env.data.xpos[t_target_id]
-            quat = env._get_body_quat(t_target_id) if hasattr(env, '_get_body_quat') else np.array([1.0, 0.0, 0.0, 0.0])
+            quat = env.data.xquat[t_target_id]
             scene_objects.append({
                 'name': 't_target',
                 'position': {'x': float(pos[0]), 'y': float(pos[1]), 'z': float(pos[2])},
@@ -591,7 +603,7 @@ def _get_scene_objects():
             box_id = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_BODY, "box")
             if box_id >= 0:
                 pos = env.data.xpos[box_id]
-                quat = env._get_body_quat(box_id) if hasattr(env, '_get_body_quat') else np.array([1.0, 0.0, 0.0, 0.0])
+                quat = env.data.xquat[box_id]
                 scene_objects.append({
                     'name': 'box',
                     'position': {'x': float(pos[0]), 'y': float(pos[1]), 'z': float(pos[2])},
@@ -887,8 +899,112 @@ def _consume_episode_control_flags():
     return terminate, truncate
 
 
+def _quat_to_mat(quat):
+    """Convert quaternion (w, x, y, z) to 3x3 rotation matrix.
+    
+    Args:
+        quat: dict with keys 'w', 'x', 'y', 'z' or tuple/list [w, x, y, z]
+    
+    Returns:
+        3x3 rotation matrix as numpy array
+    """
+    if isinstance(quat, dict):
+        w, x, y, z = quat['w'], quat['x'], quat['y'], quat['z']
+    else:
+        w, x, y, z = quat
+    
+    # Normalize quaternion
+    norm = math.sqrt(w*w + x*x + y*y + z*z)
+    if norm < 1e-10:
+        return np.eye(3)
+    w, x, y, z = w/norm, x/norm, y/norm, z/norm
+    
+    # Convert to rotation matrix
+    mat = np.array([
+        [1 - 2*(y*y + z*z),     2*(x*y - w*z),     2*(x*z + w*y)],
+        [    2*(x*y + w*z), 1 - 2*(x*x + z*z),     2*(y*z - w*x)],
+        [    2*(x*z - w*y),     2*(y*z + w*x), 1 - 2*(x*x + y*y)]
+    ])
+    return mat
+
+
+def _render_perception_frames(renderer_obj, model, data):
+    """Render bounding boxes for detected object poses.
+
+    Draws wireframe boxes at detected poses, color-coded by confidence.
+    Green (high confidence >0.8), Yellow (medium 0.5-0.8), Red (low <0.5).
+    """
+    with perception_lock:
+        poses = perception_state.get("poses", [])
+        last_update = perception_state.get("last_update", 0)
+
+    # Don't render if data is too old (> 1 second)
+    if time.time() - last_update > 1.0:
+        return
+
+    # Render bounding box for each detected pose
+    for pose in poses:
+        try:
+            pos = pose.get("position", {})
+            orientation = pose.get("orientation", {})
+            confidence = pose.get("confidence", 0.0)
+            dimensions = pose.get("dimensions")
+
+            # Skip if no dimensions provided
+            if not dimensions:
+                continue
+
+            # Extract position
+            px = float(pos.get("x", 0.0))
+            py = float(pos.get("y", 0.0))
+            pz = float(pos.get("z", 0.0))
+
+            # Extract dimensions
+            if isinstance(dimensions, dict):
+                width = float(dimensions.get("width", 0.05))
+                height = float(dimensions.get("height", 0.05))
+                depth = float(dimensions.get("depth", 0.05))
+            elif isinstance(dimensions, (list, tuple)):
+                width, height, depth = float(dimensions[0]), float(dimensions[1]), float(dimensions[2])
+            else:
+                continue
+
+            # Color based on confidence: high=green, medium=yellow, low=red
+            if confidence > 0.8:
+                rgba = [0.0, 1.0, 0.0, 0.5]  # Green
+            elif confidence > 0.5:
+                rgba = [1.0, 1.0, 0.0, 0.5]  # Yellow
+            else:
+                rgba = [1.0, 0.0, 0.0, 0.5]  # Red
+
+            # Convert quaternion to rotation matrix
+            rot_mat = _quat_to_mat(orientation)
+
+            # Render as box geometry
+            # MuJoCo box size is half-extents, so divide by 2
+            mujoco.mjv_initGeom(
+                renderer_obj.scene.geoms[renderer_obj.scene.ngeom],
+                type=mujoco.mjtGeom.mjGEOM_BOX,
+                size=[width/2, height/2, depth/2],
+                pos=[px, py, pz],
+                mat=rot_mat.flatten(),
+                rgba=rgba
+            )
+            renderer_obj.scene.ngeom += 1
+
+            # Limit to max geoms to avoid overflow
+            if renderer_obj.scene.ngeom >= renderer_obj.scene.maxgeom - 10:
+                break
+
+        except Exception as e:
+            # Silently skip invalid poses
+            pass
+
+
+
+
 def simulation_loop():
-    global running, latest_frame, camera_follow, renderer, needs_robot_switch, env
+    global running, latest_frame, camera_follow, t_target_visible, renderer, needs_robot_switch, env
 
     print("Starting simulation loop...")
 
@@ -933,8 +1049,24 @@ def simulation_loop():
             for _ in range(physics_steps_per_frame):
                 env.step_with_controller(dt=sim_dt)
 
-            # Update camera to follow robot (not for UR5 which is stationary)
-            if camera_follow and current_robot not in ("ur5", "ur5_t_push"):
+            # Update camera to follow robot or object
+            if current_robot in ("ur5", "ur5_t_push"):
+                # For UR5: follow t_object when enabled, table center when disabled
+                if camera_follow:
+                    # Follow t_object if it exists
+                    t_object_id = getattr(env, 't_object_body_id', -1)
+                    if t_object_id >= 0:
+                        object_pos = env.data.xpos[t_object_id]
+                        cam.lookat[0] = object_pos[0]
+                        cam.lookat[1] = object_pos[1]
+                        cam.lookat[2] = object_pos[2]
+                else:
+                    # Center on table
+                    cam.lookat[0] = 0.5
+                    cam.lookat[1] = 0.0
+                    cam.lookat[2] = 0.42
+            elif camera_follow:
+                # For locomotion robots: follow robot base
                 robot_pos = env.data.qpos[:3]
                 cam.lookat[0] = robot_pos[0]
                 cam.lookat[1] = robot_pos[1]
@@ -943,6 +1075,24 @@ def simulation_loop():
             # Render
             if renderer is not None:
                 renderer.update_scene(env.data, camera=cam)
+
+                # Hide t_target if visibility is disabled (push-T scene only)
+                if current_robot == "ur5_t_push" and not t_target_visible:
+                    t_target_id = getattr(env, 't_target_body_id', -1)
+                    if t_target_id >= 0:
+                        # Set alpha to 0 for all geoms belonging to t_target body
+                        for i in range(renderer.scene.ngeom):
+                            geom = renderer.scene.geoms[i]
+                            # Check if this geom belongs to t_target body by comparing position
+                            target_pos = env.data.xpos[t_target_id]
+                            geom_pos = geom.pos
+                            # If geom is very close to target body, it likely belongs to it
+                            dist = np.linalg.norm(np.array(geom_pos) - target_pos)
+                            if dist < 0.15:  # Within 15cm - likely part of the T-shape
+                                geom.rgba[3] = 0.0  # Set alpha to 0 (invisible)
+
+                # Add perception visualization markers
+                _render_perception_frames(renderer, env.model, env.data)
                 frame = renderer.render()
             else:
                 return
@@ -988,6 +1138,22 @@ def simulation_loop():
                         cam_obj.azimuth = -yaw + 90.0
                         cam_obj.elevation = -pitch
             renderer_obj.update_scene(env.data, camera=cam_obj)
+
+            # Hide t_target if visibility is disabled (push-T scene only)
+            if current_robot == "ur5_t_push" and not t_target_visible:
+                t_target_id = getattr(env, 't_target_body_id', -1)
+                if t_target_id >= 0:
+                    # Set alpha to 0 for all geoms belonging to t_target body
+                    for i in range(renderer_obj.scene.ngeom):
+                        geom = renderer_obj.scene.geoms[i]
+                        # Check if this geom belongs to t_target body by comparing position
+                        target_pos = env.data.xpos[t_target_id]
+                        geom_pos = geom.pos
+                        # If geom is very close to target body, it likely belongs to it
+                        dist = np.linalg.norm(np.array(geom_pos) - target_pos)
+                        if dist < 0.15:  # Within 15cm - likely part of the T-shape
+                            geom.rgba[3] = 0.0  # Set alpha to 0 (invisible)
+
             overlay_frame = renderer_obj.render()
             if overlay_frame is None:
                 continue
@@ -1048,7 +1214,7 @@ def generate_overlay_frames(name: str):
 
 def handle_ws_message(ws, data):
     """Handle incoming WebSocket message."""
-    global needs_robot_switch, camera_follow, last_teleop_action
+    global needs_robot_switch, camera_follow, t_target_visible, last_teleop_action
 
     msg_type = data.get('type')
 
@@ -1342,6 +1508,10 @@ def handle_ws_message(ws, data):
         payload = data.get('data', {})
         camera_follow = payload.get('follow', True)
 
+    elif msg_type == 'toggle_target_visibility':
+        payload = data.get('data', {})
+        t_target_visible = payload.get('visible', True)
+
     elif msg_type == 'teleop_action':
         payload = data.get('data', {})
         # Accept both old format (dx/dy/dz) and new format (vx/vy/vz)
@@ -1457,6 +1627,16 @@ def handle_ws_message(ws, data):
             if env is not None and current_robot in ("ur5", "ur5_t_push"):
                 env.set_use_orientation(use)
 
+    elif msg_type == 'perception_update':
+        # Handle perception updates from training system
+        payload = data.get('data', {})
+        poses = payload.get('poses', [])
+        timestamp = payload.get('timestamp', time.time())
+        with perception_lock:
+            perception_state["poses"] = poses
+            perception_state["last_update"] = timestamp
+        # print(f"[Perception] Received {len(poses)} pose(s) at {timestamp:.2f}")
+
     elif msg_type == 'set_nova_mode':
         payload = data.get('data', {})
         enabled = bool(payload.get('enabled', False))