Files
JackBot/ml/MetricsOverlay.py
T

117 lines
4.3 KiB
Python

"""
ml/MetricsOverlay.py - Camera-Facing (Billboard) 3D Floating Text Overlay
"""
from typing import List, Tuple, Optional, Dict
import numpy as np
import pybullet as p
class MetricsHUD:
"""Renders real-time telemetry as black text floating in 3D, always facing the active camera."""
def __init__(self, physics_client_id: int = 0):
self.client_id = physics_client_id
self._text_id: Optional[int] = None
def _get_camera_facing_orientation(self) -> List[float]:
"""Calculates a quaternion that points the text towards the current GUI camera."""
try:
cam_info = p.getDebugVisualizerCamera(physicsClientId=self.client_id)
# cam_info index 8: yaw, index 9: pitch
yaw = cam_info[8]
pitch = cam_info[9]
pitch_rad = np.radians(pitch + 90.0)
yaw_rad = np.radians(yaw)
text_orientation = p.getQuaternionFromEuler(
[pitch_rad, 0.0, yaw_rad],
physicsClientId=self.client_id
)
return text_orientation
except Exception:
return [0.0, 0.0, 0.0, 1.0]
def update(
self,
episode: int,
step: int,
robot_rewards: List[float],
cmd_vel: np.ndarray,
fps: float = 0.0,
height: float = 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 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"
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[2] if len(cmd_vel) > 2 else 0.0
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]
text_color = [0, 0, 0] # Pure black
# Calculate dynamic orientation to align text flat against camera plane
text_orientation = self._get_camera_facing_orientation()
# 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)
except Exception:
pass
# Draw fresh text
self._text_id = p.addUserDebugText(
text=hud_text,
textPosition=text_position,
textColorRGB=text_color,
textSize=0.085,
textOrientation=text_orientation,
physicsClientId=self.client_id
)
def reset(self) -> None:
"""Removes the active debug text item so a new episode starts with a clean overlay."""
if self._text_id is not None:
try:
p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)
except Exception:
pass
self._text_id = None