""" LifeOS — Reward Curve Plotter (File 11 of 15) Visualises reward improvement across iterations from run_lifeos() results. Usage: from utils.plot_rewards import plot_reward_curve plot_reward_curve(results) # displays plot plot_reward_curve(results, "out.png") # saves to file IMPROVEMENT 2: Supports multi-line plotting by difficulty level. """ from __future__ import annotations from typing import Any, Dict, List, Optional def plot_reward_curve( results: List[Dict[str, Any]], save_path: Optional[str] = None, ) -> None: """ Plot reward improvement line chart from orchestrator results. Chart features: - Blue line with circle markers - Each point annotated with its reward value - Dashed reference lines at y=50 (acceptable) and y=80 (optimal) - Y-axis starts at 0 - X-axis shows integer iteration numbers Parameters ---------- results : list of result dicts from run_lifeos() — must have 'iteration' and 'reward' keys save_path : if given, saves PNG here; otherwise calls plt.show() """ try: import matplotlib matplotlib.use("Agg" if save_path else matplotlib.get_backend()) import matplotlib.pyplot as plt except ImportError: print("[Reward] matplotlib not installed. Run: pip install matplotlib") return if not results: print("[Reward] No results to plot.") return iterations = [r["iteration"] for r in results] rewards = [r["reward"] for r in results] fig, ax = plt.subplots(figsize=(10, 5)) # Main reward line ax.plot( iterations, rewards, color="royalblue", linewidth=2.5, marker="o", markersize=9, label="Reward", zorder=3, ) # Annotate each data point for x, y in zip(iterations, rewards): ax.annotate( f"{y:.0f}", xy=(x, y), xytext=(0, 12), textcoords="offset points", ha="center", fontsize=10, fontweight="bold", color="royalblue", ) # Reference lines ax.axhline( 50, color="darkorange", linestyle="--", linewidth=1.5, alpha=0.85, label="Acceptable (50)", zorder=2, ) ax.axhline( 80, color="green", linestyle="--", linewidth=1.5, alpha=0.85, label="Optimal (80)", zorder=2, ) # Shade the "below acceptable" region ax.axhspan(0, 50, alpha=0.04, color="red") ax.axhspan(50, 80, alpha=0.04, color="orange") ax.axhspan(80, max(max(rewards) + 20, 100), alpha=0.04, color="green") ax.set_xlabel("Iteration", fontsize=12) ax.set_ylabel("Reward Score", fontsize=12) ax.set_title( "LifeOS Reward Improvement Across Iterations", fontsize=14, fontweight="bold", ) ax.set_ylim(bottom=0, top=max(max(rewards) + 25, 105)) ax.set_xticks(iterations) ax.legend(loc="lower right", fontsize=10) ax.grid(axis="y", alpha=0.3, linestyle=":") ax.spines["top"].set_visible(False) ax.spines["right"].set_visible(False) plt.tight_layout() if save_path: plt.savefig(save_path, dpi=150, bbox_inches="tight") print(f"[Reward] Chart saved to {save_path}") else: plt.show() plt.close(fig) # Text summary if len(rewards) >= 2: delta = rewards[-1] - rewards[0] sign = "+" if delta >= 0 else "" print( f"[Reward] Score improved from {rewards[0]:.1f} to {rewards[-1]:.1f} " f"({sign}{delta:.1f} points across {len(rewards)} iterations)" ) elif rewards: print(f"[Reward] Single iteration score: {rewards[0]:.1f}") def plot_component_comparison( baseline_data: Dict[str, Any], trained_data: Dict[str, Any], save_path: Optional[str] = None, ) -> None: """ Plot a side-by-side bar chart comparing baseline vs trained per-component reward scores. Used by the Streamlit 'Before vs After' tab. Parameters ---------- baseline_data : loaded baseline_scores.json trained_data : loaded trained_scores.json save_path : if given, saves PNG; otherwise returns the figure """ try: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np except ImportError: print("[Reward] matplotlib/numpy not installed.") return None from agents.reward import COMPONENT_NAMES # Average final-iteration scores across all scenarios def _avg_components(data: Dict[str, Any]) -> Dict[str, float]: avgs = {c: 0.0 for c in COMPONENT_NAMES} avgs["total"] = 0.0 count = 0 # Handle both formats if "by_scenario" in data: scenarios = data["by_scenario"] else: scenarios = data for label, sdata in scenarios.items(): iters = sdata.get("iterations", []) if iters: last = iters[-1] for c in COMPONENT_NAMES: avgs[c] += float(last.get(c, 0.0)) avgs["total"] += float(last.get("total", 0.0)) count += 1 if count > 0: for k in avgs: avgs[k] /= count return avgs baseline_avgs = _avg_components(baseline_data) trained_avgs = _avg_components(trained_data) labels = COMPONENT_NAMES + ["total"] baseline_vals = [baseline_avgs.get(c, 0.0) for c in labels] trained_vals = [trained_avgs.get(c, 0.0) for c in labels] x = np.arange(len(labels)) width = 0.35 fig, ax = plt.subplots(figsize=(12, 5)) bars1 = ax.bar(x - width/2, baseline_vals, width, label="Baseline", color="#3b82f6", alpha=0.85) bars2 = ax.bar(x + width/2, trained_vals, width, label="Trained", color="#f97316", alpha=0.85) ax.set_ylabel("Score", fontsize=11) ax.set_title("LifeOS Reward — Baseline vs Trained (per component)", fontsize=13, fontweight="bold") ax.set_xticks(x) short_labels = [c.replace("_score", "") for c in labels] ax.set_xticklabels(short_labels, fontsize=9, rotation=30, ha="right") ax.legend(fontsize=10) ax.grid(axis="y", alpha=0.3, linestyle=":") ax.spines["top"].set_visible(False) ax.spines["right"].set_visible(False) # Annotate bars for bar in bars1: height = bar.get_height() if height != 0: ax.annotate(f"{height:.1f}", xy=(bar.get_x() + bar.get_width()/2, height), xytext=(0, 3), textcoords="offset points", ha="center", fontsize=8, color="#3b82f6") for bar in bars2: height = bar.get_height() if height != 0: ax.annotate(f"{height:.1f}", xy=(bar.get_x() + bar.get_width()/2, height), xytext=(0, 3), textcoords="offset points", ha="center", fontsize=8, color="#f97316") plt.tight_layout() if save_path: plt.savefig(save_path, dpi=150, bbox_inches="tight") plt.close(fig) print(f"[Reward] Comparison chart saved to {save_path}") return None else: return fig