44 lines
1.8 KiB
Python
44 lines
1.8 KiB
Python
"""
|
|
ml/callbacks.py - Stable-Baselines3 Custom Callbacks for Logging & Curriculum Advancement
|
|
"""
|
|
import os
|
|
import numpy as np
|
|
from stable_baselines3.common.callbacks import BaseCallback
|
|
|
|
|
|
class RewardLoggerCallback(BaseCallback):
|
|
"""Logs individual reward component averages to TensorBoard."""
|
|
|
|
def __init__(self, verbose: int = 0):
|
|
super().__init__(verbose)
|
|
|
|
def _on_step(self) -> bool:
|
|
# Pull component averages from the environment vector
|
|
for env_idx, env in enumerate(self.training_env.envs):
|
|
if hasattr(env, "get_reward_component_averages"):
|
|
averages = env.get_reward_component_averages()
|
|
for key, val in averages.items():
|
|
self.logger.record(f"reward_components/{key}", val)
|
|
return True
|
|
|
|
|
|
class CurriculumCallback(BaseCallback):
|
|
"""Monitors evaluation performance and automatically manages curriculum progression."""
|
|
|
|
def __init__(self, eval_freq: int = 10000, verbose: int = 1):
|
|
super().__init__(verbose)
|
|
self.eval_freq = eval_freq
|
|
|
|
def _on_step(self) -> bool:
|
|
if self.n_calls % self.eval_freq == 0:
|
|
for env in self.training_env.envs:
|
|
if hasattr(env, "get_current_robot_metrics"):
|
|
metrics = env.get_current_robot_metrics()
|
|
if metrics:
|
|
phase = metrics[0].get("phase_name", "UNKNOWN")
|
|
dist = metrics[0].get("distance_from_start", 0.0)
|
|
self.logger.record("curriculum/phase_idx", phase)
|
|
self.logger.record("curriculum/max_distance", dist)
|
|
if self.verbose > 0:
|
|
print(f"[CurriculumCallback] Step {self.num_timesteps}: Current Phase = {phase}, Max Dist = {dist:.2f}m")
|
|
return True |