""" visualizer.py Pusher-v5 시각화 및 대시보드 생성 유틸리티 모듈 """ import os import json import numpy as np from PIL import Image, ImageDraw, ImageFont import imageio import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt def draw_hud_on_frame( frame: np.ndarray, timestep: int, episode: int, step: int, step_reward: float, total_reward: float, reward_dist: float = None, reward_near: float = None, actions: np.ndarray = None, tag: str = "" ) -> np.ndarray: """ MuJoCo 렌더링 프레임 상단에 실시간 상태 정보 HUD(Heads-Up Display)를 오버레이합니다. """ img = Image.fromarray(frame).convert("RGBA") overlay = Image.new("RGBA", img.size, (0, 0, 0, 0)) draw = ImageDraw.Draw(overlay) # 기본 폰트 로드 try: font_large = ImageFont.truetype("arial.ttf", 16) font_small = ImageFont.truetype("arial.ttf", 12) font_bold = ImageFont.truetype("arialbd.ttf", 14) except Exception: font_large = ImageFont.load_default() font_small = ImageFont.load_default() font_bold = ImageFont.load_default() # 상단 HUD 패널 (반투명 블랙 글래스모피즘 박스) panel_w = 260 panel_h = 135 margin = 12 draw.rounded_rectangle( [margin, margin, margin + panel_w, margin + panel_h], radius=10, fill=(15, 23, 42, 210), # Slate-900 semi-transparent outline=(56, 189, 248, 180), # Cyan accent border width=2, ) # 타이틀 / 태그 title_text = f"Pusher-v5 PPO [{tag or f'Step {timestep:,}'}]" draw.text((margin + 12, margin + 8), title_text, fill=(56, 189, 248, 255), font=font_bold) # 정보 텍스트 라인 lines = [ (f"Global Step : {timestep:,}", (226, 232, 240, 255)), (f"Episode / Step : #{episode} (Step {step})", (203, 213, 225, 255)), (f"Cumulative Rwd : {total_reward:+.2f} (Step: {step_reward:+.3f})", (74, 222, 128, 255)), ] if reward_dist is not None: # Pusher-v5 distance reward is approx -norm(obj-goal)*1.25 dist_goal = abs(reward_dist) / 1.25 if reward_dist <= 0 else reward_dist lines.append((f"Dist (Obj -> Goal) : {dist_goal:.3f} m", (251, 191, 36, 255))) if reward_near is not None: dist_arm = abs(reward_near) if reward_near <= 0 else reward_near lines.append((f"Dist (Arm -> Obj) : {dist_arm:.3f} m", (168, 85, 247, 255))) y_offset = margin + 28 for line_text, color in lines: draw.text((margin + 12, y_offset), line_text, fill=color, font=font_small) y_offset += 18 # 하단 액션 바 게이지 오버레이 (옵션) if actions is not None and len(actions) > 0: act_panel_y = img.height - 32 draw.rounded_rectangle( [margin, act_panel_y, img.width - margin, img.height - 8], radius=6, fill=(15, 23, 42, 190), outline=(100, 116, 139, 120), width=1, ) act_text = "Act: " + " ".join([f"{a:+.1f}" for a in actions[:7]]) draw.text((margin + 10, act_panel_y + 4), act_text, fill=(148, 163, 184, 255), font=font_small) combined = Image.alpha_composite(img, overlay).convert("RGB") return np.array(combined) def save_video_and_gif(frames: list, output_base_path: str, fps: int = 30): """ 프레임 리스트를 MP4 비디오 및 GIF 애니메이션으로 저장합니다. """ if not frames: return os.makedirs(os.path.dirname(output_base_path), exist_ok=True) gif_path = f"{output_base_path}.gif" mp4_path = f"{output_base_path}.mp4" # Save GIF try: # 프레임이 너무 많을 경우 GIF 용량 최적화 (2프레임마다 1개 샘플링) sampled_frames = frames[::2] if len(frames) > 100 else frames imageio.mimsave(gif_path, sampled_frames, fps=max(15, fps // 2), loop=0) except Exception as e: print(f"[Warning] GIF 저장 실패: {e}") # Save MP4 try: imageio.mimsave(mp4_path, frames, fps=fps, quality=8) except Exception as e: # Fallback to default writer try: writer = imageio.get_writer(mp4_path, fps=fps) for f in frames: writer.append_data(f) writer.close() except Exception as e2: print(f"[Warning] MP4 저장 실패: {e2}") return gif_path, mp4_path def generate_training_plots(metrics: dict, output_path: str): """ 학습 과정의 지표들을 미려한 다크 테마 4분할 차트로 시각화하여 저장합니다. """ plt.style.use("dark_background") fig, axes = plt.subplots(2, 2, figsize=(15, 10), dpi=150) fig.patch.set_facecolor("#0f172a") for row in axes: for ax in row: ax.set_facecolor("#1e293b") ax.grid(True, linestyle="--", alpha=0.3, color="#94a3b8") ax.tick_params(colors="#cbd5e1") for spine in ax.spines.values(): spine.set_color("#475569") # 1. Episode Returns eval_timesteps = metrics.get("eval_timesteps", []) eval_rewards = metrics.get("eval_rewards", []) ax1 = axes[0, 0] ax1.set_title("Evaluation Episode Return over Timesteps", color="#38bdf8", fontsize=13, fontweight="bold", pad=10) ax1.set_xlabel("Timesteps", color="#cbd5e1") ax1.set_ylabel("Mean Return", color="#cbd5e1") if eval_timesteps and eval_rewards: ax1.plot(eval_timesteps, eval_rewards, color="#38bdf8", marker="o", linewidth=2.5, label="Mean Reward") # Moving average if len(eval_rewards) >= 3: window = min(5, len(eval_rewards)) smooth = np.convolve(eval_rewards, np.ones(window)/window, mode='valid') smooth_x = eval_timesteps[window-1:] ax1.plot(smooth_x, smooth, color="#f43f5e", linestyle="--", linewidth=2, label="Trend (MA)") ax1.legend(loc="lower right", facecolor="#1e293b", edgecolor="#475569") # 2. Distance to Goal (Object -> Target) eval_dist_goal = metrics.get("eval_dist_goal", []) ax2 = axes[0, 1] ax2.set_title("Final Object-to-Goal Distance (Lower is Better)", color="#34d399", fontsize=13, fontweight="bold", pad=10) ax2.set_xlabel("Timesteps", color="#cbd5e1") ax2.set_ylabel("Final Distance (meters)", color="#cbd5e1") if eval_timesteps and eval_dist_goal: ax2.plot(eval_timesteps, eval_dist_goal, color="#34d399", marker="s", linewidth=2.5, label="Distance to Goal") ax2.axhline(0.05, color="#f59e0b", linestyle=":", label="Success Threshold (~0.05m)") ax2.legend(loc="upper right", facecolor="#1e293b", edgecolor="#475569") # 3. Distance Arm to Object eval_dist_arm = metrics.get("eval_dist_arm", []) ax3 = axes[1, 0] ax3.set_title("Arm-to-Object Distance (Reaching Efficiency)", color="#a855f7", fontsize=13, fontweight="bold", pad=10) ax3.set_xlabel("Timesteps", color="#cbd5e1") ax3.set_ylabel("Distance (meters)", color="#cbd5e1") if eval_timesteps and eval_dist_arm: ax3.plot(eval_timesteps, eval_dist_arm, color="#a855f7", marker="^", linewidth=2.5, label="Arm-to-Object") ax3.legend(loc="upper right", facecolor="#1e293b", edgecolor="#475569") # 4. Episode Lengths or Loss eval_lens = metrics.get("eval_lengths", []) ax4 = axes[1, 1] ax4.set_title("Evaluation Episode Length", color="#fbbf24", fontsize=13, fontweight="bold", pad=10) ax4.set_xlabel("Timesteps", color="#cbd5e1") ax4.set_ylabel("Steps", color="#cbd5e1") if eval_timesteps and eval_lens: ax4.bar(eval_timesteps, eval_lens, width=max(500, (eval_timesteps[-1] - eval_timesteps[0])//(len(eval_timesteps)*2) if len(eval_timesteps)>1 else 500), color="#fbbf24", alpha=0.7, label="Ep Length") ax4.legend(loc="lower right", facecolor="#1e293b", edgecolor="#475569") plt.tight_layout(pad=2.5) os.makedirs(os.path.dirname(output_path), exist_ok=True) plt.savefig(output_path, dpi=150, facecolor=fig.get_facecolor(), edgecolor="none") plt.close() return output_path def generate_html_dashboard( metrics: dict, checkpoint_records: list, plot_rel_path: str, output_html_path: str ): """ 모든 비디오/GIF 및 메트릭 차트를 한눈에 탐색할 수 있는 인터랙티브 HTML 대시보드를 생성합니다. """ best_reward = max(metrics.get("eval_rewards", [0])) if metrics.get("eval_rewards") else 0 final_reward = metrics.get("eval_rewards", [0])[-1] if metrics.get("eval_rewards") else 0 best_dist = min(metrics.get("eval_dist_goal", [0])) if metrics.get("eval_dist_goal") else 0 total_timesteps = metrics.get("total_timesteps", 0) # HTML 템플릿 생성 cards_html = "" for rec in checkpoint_records: tag = rec.get("tag", "") step = rec.get("timestep", 0) reward = rec.get("reward", 0.0) dist = rec.get("dist_goal", 0.0) gif_rel = rec.get("gif_rel", "") mp4_rel = rec.get("mp4_rel", "") cards_html += f"""
Gymnasium MuJoCo Robotic Arm Manipulation • Stable-Baselines3 PPO