Callback restored

ml/run_train.py

@@ -16,7 +16,7 @@ from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback
 
 sys.path.append(str(Path(__file__).resolve().parent.parent))
 from ml.env import JackBotEnv
-from ml.callbacks import CurriculumCallback
+from ml.callbacks import CurriculumCallback, RewardLoggerCallback
 
 # Silence SB3's UserWarning about SubprocVecEnv vs DummyVecEnv
 warnings.filterwarnings("ignore", category=UserWarning, module="stable_baselines3")
@@ -100,6 +100,7 @@ def main():
         save_path=args.save_dir,
         name_prefix=f"jackbot_{ppo_name}",
     )
+    reward_logger_callback = RewardLoggerCallback(verbose=1)
     curriculum_callback = CurriculumCallback()
 
     eval_env = DummyVecEnv([lambda: JackBotEnv(use_gui=False, random_command=True)])
@@ -118,7 +119,7 @@ def main():
     try:
         model.learn(
             total_timesteps=args.total_timesteps,
-            callback=[checkpoint_callback, curriculum_callback, eval_callback],
+            callback=[checkpoint_callback, reward_logger_callback, curriculum_callback, eval_callback],
             progress_bar=True,
         )
         final_model_path = os.path.join(args.save_dir, f"jackbot_{ppo_name}_final.zip")