99 lines
3.8 KiB
Python
99 lines
3.8 KiB
Python
"""
|
|
ml/callbacks.py - Stable-Baselines3 Custom Callbacks for Logging & Curriculum Advancement
|
|
Fully compatible with SubprocVecEnv and DummyVecEnv.
|
|
"""
|
|
|
|
import numpy as np
|
|
from stable_baselines3.common.callbacks import BaseCallback
|
|
|
|
|
|
class RewardLoggerCallback(BaseCallback):
|
|
"""
|
|
Logs individual reward component averages to TensorBoard and prints
|
|
the best worker's performance breakdown to the console per iteration.
|
|
"""
|
|
|
|
def __init__(self, verbose: int = 1):
|
|
super().__init__(verbose)
|
|
self.iteration = 0
|
|
|
|
def _on_step(self) -> bool:
|
|
return True
|
|
|
|
def _on_rollout_end(self) -> None:
|
|
self.iteration += 1
|
|
if self.training_env is None:
|
|
return
|
|
|
|
try:
|
|
# Safely query method across all parallel worker processes
|
|
all_worker_averages = self.training_env.env_method("get_reward_component_averages")
|
|
except Exception:
|
|
return
|
|
|
|
if not all_worker_averages or len(all_worker_averages) == 0:
|
|
return
|
|
|
|
# 1. Log mean component values across ALL workers to TensorBoard
|
|
component_keys = all_worker_averages[0].keys()
|
|
for key in component_keys:
|
|
mean_val = float(np.mean([w.get(key, 0.0) for w in all_worker_averages]))
|
|
self.logger.record(f"reward_components/{key}", mean_val)
|
|
|
|
# 2. Identify the best performing worker of this iteration
|
|
worker_totals = [sum(w.values()) for w in all_worker_averages]
|
|
best_worker_idx = int(np.argmax(worker_totals))
|
|
best_averages = all_worker_averages[best_worker_idx]
|
|
best_total = worker_totals[best_worker_idx]
|
|
|
|
# 3. Print best worker breakdown to console
|
|
if self.verbose > 0:
|
|
print(f"\n" + "=" * 65)
|
|
print(f" ITERATION {self.iteration} | BEST WORKER (#{best_worker_idx}) REWARD BREAKDOWN")
|
|
print(f" Total Avg Reward / Step: {best_total:+.4f}")
|
|
print("-" * 65)
|
|
for key, val in best_averages.items():
|
|
print(f" • {key:<26}: {val:+.5f}")
|
|
print("=" * 65 + "\n")
|
|
|
|
|
|
class CurriculumCallback(BaseCallback):
|
|
"""
|
|
Monitors training metrics using SB3's native ep_info_buffer and
|
|
dynamically advances curriculum phases across worker processes.
|
|
"""
|
|
|
|
def __init__(self, reward_threshold: float = 100.0, verbose: int = 1):
|
|
super().__init__(verbose)
|
|
self.reward_threshold = reward_threshold
|
|
|
|
def _on_step(self) -> bool:
|
|
return True
|
|
|
|
def _on_rollout_end(self) -> None:
|
|
if self.training_env is None:
|
|
return
|
|
|
|
# SB3 natively records finished episode stats in self.model.ep_info_buffer
|
|
if hasattr(self.model, "ep_info_buffer") and len(self.model.ep_info_buffer) > 0:
|
|
recent_rewards = [ep_info["r"] for ep_info in self.model.ep_info_buffer]
|
|
mean_reward = float(np.mean(recent_rewards[-50:]))
|
|
|
|
try:
|
|
# Query current phase from worker 0
|
|
phases = self.training_env.get_attr("curriculum_phase")
|
|
current_phase = phases[0]
|
|
|
|
# Advance curriculum if mean reward exceeds threshold
|
|
if mean_reward >= self.reward_threshold:
|
|
if hasattr(current_phase, "next"):
|
|
next_phase = current_phase.next()
|
|
if next_phase != current_phase:
|
|
self.training_env.set_attr("curriculum_phase", next_phase)
|
|
if self.verbose > 0:
|
|
print(
|
|
f"\n[Curriculum] 🚀 Promoted workers to phase: {next_phase.name} "
|
|
f"(Mean Reward: {mean_reward:.2f})"
|
|
)
|
|
except Exception:
|
|
pass # Keep rollout loop running safely if phase check fails |