pretrain logic and small fixes

ml/MetricsOverlay.py

@@ -1,7 +1,7 @@
 """
 ml/MetricsOverlay.py - Camera-Facing (Billboard) 3D Floating Text Overlay
 """
-from typing import List, Tuple, Optional
+from typing import List, Tuple, Optional, Dict
 import numpy as np
 import pybullet as p
 
@@ -21,20 +21,15 @@ class MetricsHUD:
             yaw = cam_info[8]
             pitch = cam_info[9]
             
-            # Orient the text normal toward the camera view direction
-            # PyBullet text default faces local +Z/-Y depending on roll,
-            # converting visualizer yaw/pitch to Euler angles (roll, pitch, yaw in radians)
-            roll_rad = 0.0
             pitch_rad = np.radians(pitch + 90.0)
             yaw_rad = np.radians(yaw)
             
             text_orientation = p.getQuaternionFromEuler(
-                [pitch_rad, roll_rad, yaw_rad], 
+                [pitch_rad, 0.0, yaw_rad], 
                 physicsClientId=self.client_id
             )
             return text_orientation
         except Exception:
-            # Fallback default orientation if camera info call fails
             return [0.0, 0.0, 0.0, 1.0]
 
     def update(
@@ -45,29 +40,48 @@ class MetricsHUD:
         cmd_vel: np.ndarray,
         fps: float = 0.0,
         height: float = 0.0,
-        roll_pitch: Tuple[float, float] = (0.0, 0.0)
+        roll_pitch: Tuple[float, float] = (0.0, 0.0),
+        mode: str = "direct",
+        phase: str = "STAND_ONLY",
+        distance: float = 0.0,
+        status: str = "ALIVE",
+        reward_components: Optional[Dict[str, float]] = None,
+        ep_step: int = 0,
     ) -> None:
-        """Updates floating black text block in 3D space with billboarding."""
+        """Updates floating text block in 3D space with expanded telemetry."""
         sorted_rewards = sorted(robot_rewards, reverse=True)
         top1 = f"{sorted_rewards[0]:+.2f}" if len(sorted_rewards) > 0 else "0.00"
-        top2 = f"{sorted_rewards[1]:+.2f}" if len(sorted_rewards) > 1 else "0.00"
-        top3 = f"{sorted_rewards[2]:+.2f}" if len(sorted_rewards) > 2 else "0.00"
 
         vx = cmd_vel[0] if len(cmd_vel) > 0 else 0.0
         vy = cmd_vel[1] if len(cmd_vel) > 1 else 0.0
-        omega = cmd_vel[3] if len(cmd_vel) > 3 else 0.0
+        omega = cmd_vel[2] if len(cmd_vel) > 2 else 0.0
 
-        hud_text = (
-            f"=== JACKBOT METRICS ===\n"
-            f"Episode:     {episode}\n"
-            f"Global Step: {step}\n"
-            f"FPS:         {fps:.1f}\n"
-            f"----------------------\n"
-            f"Top Rewards: [{top1}, {top2}, {top3}]\n"
-            f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]\n"
-            f"Height:      {height:.3f} m\n"
-            f"Roll/Pitch:  {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°"
-        )
+        lines = [
+            "=== JACKBOT TELEMETRY ===",
+            f"Mode:        {mode.upper()}",
+            f"Curriculum:  {phase}",
+            f"Status:      {status}",
+            f"Episode:     {episode} (Step {ep_step})",
+            f"Global Step: {step}",
+            f"FPS:         {fps:.1f}",
+            "-------------------------",
+            f"Episode Rew: {top1}",
+            f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]",
+            f"Height:      {height:.3f} m",
+            f"Roll/Pitch:  {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°",
+            f"Max Dist:    {distance:.2f} m",
+        ]
+
+        if reward_components:
+            lin_v = reward_components.get("lin_vel", 0.0)
+            stab = reward_components.get("stability", 0.0)
+            h_rew = reward_components.get("height", 0.0)
+            jit = reward_components.get("jitter_penalty", 0.0)
+            lines.append("--- Reward Components ---")
+            lines.append(f"LinVel: {lin_v:.2f} | Stab: {stab:.2f}")
+            lines.append(f"Height: {h_rew:.2f} | Jitter: {jit:+.3f}")
+
+        hud_text = "\n".join(lines)
 
         # Position above origin in simulation world
         text_position = [-0.8, -0.8, 1.2]
@@ -76,7 +90,7 @@ class MetricsHUD:
         # Calculate dynamic orientation to align text flat against camera plane
         text_orientation = self._get_camera_facing_orientation()
 
-        # Safely remove the old text to prevent PyBullet ghosting/overlapping
+        # Safely remove old text to prevent PyBullet ghosting/overlapping
         if self._text_id is not None:
             try:
                 p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)
@@ -88,7 +102,7 @@ class MetricsHUD:
             text=hud_text,
             textPosition=text_position,
             textColorRGB=text_color,
-            textSize=0.1,
+            textSize=0.085,
             textOrientation=text_orientation,
             physicsClientId=self.client_id
         )