117 lines
4.3 KiB
Python
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 |