feat(*): Log live goal rate to TensorBoard during training

training/train.py

@@ -16,7 +16,8 @@ import os
 import pathlib
 
 from stable_baselines3 import PPO
-from stable_baselines3.common.callbacks import CheckpointCallback
+from stable_baselines3.common.callbacks import BaseCallback, CheckpointCallback
+from stable_baselines3.common.utils import safe_mean
 from stable_baselines3.common.vec_env.vec_monitor import VecMonitor
 
 from cosmic_env import CosmicClashVecEnv
@@ -25,6 +26,27 @@ TRAINING_DIR = pathlib.Path(__file__).resolve().parent
 DEFAULT_GODOT_MACOS = "/Applications/Godot.app/Contents/MacOS/Godot"
 
 
+class GoalRateCallback(BaseCallback):
+    """Logs rollout/goal_rate: the fraction of completed episodes in the
+    current ep_info_buffer that ended in an actual goal, vs. timing out as a
+    draw. rollout/ep_rew_mean mixes dense reward-shaping (ball chasing/
+    touching) with the sparse terminal goal reward, so it can trend up from
+    better shaping alone without the policy finishing more episodes by
+    actually scoring — this isolates that. Requires VecMonitor(...,
+    info_keywords=("goal_scored",)), which copies ShipAIController.get_info()
+    into each completed episode's info["episode"] dict (see
+    training_mode.gd's _on_goal_scored / timeout branch)."""
+
+    def _on_step(self) -> bool:
+        return True
+
+    def _on_rollout_end(self) -> None:
+        if len(self.model.ep_info_buffer) == 0:
+            return
+        goal_rate = safe_mean([ep_info["goal_scored"] for ep_info in self.model.ep_info_buffer])
+        self.logger.record("rollout/goal_rate", goal_rate)
+
+
 def parse_args():
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument(
@@ -164,7 +186,7 @@ def main():
         speedup=args.speedup,
         **_curriculum_kwargs(args),
     )
-    env = VecMonitor(env)
+    env = VecMonitor(env, info_keywords=("goal_scored",))
 
     if args.resume:
         model = PPO.load(
@@ -204,11 +226,12 @@ def main():
         save_path=str(checkpoint_dir),
         name_prefix="ppo",
     )
+    goal_rate_callback = GoalRateCallback()
 
     try:
         model.learn(
             args.timesteps,
-            callback=checkpoint_callback,
+            callback=[checkpoint_callback, goal_rate_callback],
             tb_log_name=args.experiment,
             reset_num_timesteps=not args.resume,
         )