feat(*): Make training better & add better logging

rl-training.py

@@ -1,10 +1,19 @@
+import os
 import rlgym
 from stable_baselines3.ppo import PPO
+from stable_baselines3.common.callbacks import EvalCallback
 from rlgym_tools.sb3_utils import SB3SingleInstanceEnv
-import os
+from rlgym.utils.reward_functions.common_rewards import VelocityBallToGoalReward
+from rlgym.utils.reward_functions.common_rewards import BallYCoordinateReward
+from rlgym.utils.obs_builders import AdvancedObs
+
+# from observationBuilders.CustomObsBuilderBluePerspective import (
+#     CustomObsBuilderBluePerspective,
+# )
 
 # set up the folders
-models_dir = "models/PPO-rl"
+model_name = "PPO-rl4"
+models_dir = f"models/{model_name}"
 logdir = "logs"
 
 if not os.path.exists(models_dir):
@@ -14,7 +23,13 @@ if not os.path.exists(logdir):
     os.makedirs(logdir)
 
 # set up the RLGym environment
-gym_env = rlgym.make(use_injector=True, self_play=True)
+gym_env = rlgym.make(
+    use_injector=True,
+    self_play=True,
+    reward_fn=VelocityBallToGoalReward(),
+    obs_builder=AdvancedObs(),
+    # obs_builder=CustomObsBuilderBluePerspective(),
+)
 
 # wrap the RLGym environment with the single instance wrapper
 env = SB3SingleInstanceEnv(gym_env)
@@ -22,8 +37,24 @@ env = SB3SingleInstanceEnv(gym_env)
 # create a PPO instance and start learning
 learner = PPO(policy="MlpPolicy", env=env, verbose=1, tensorboard_log=logdir)
 
+# Define a callback to log rewards to TensorBoard
+eval_callback = EvalCallback(
+    env,
+    best_model_save_path=models_dir,
+    log_path=logdir,
+    eval_freq=10000,
+    deterministic=True,
+    render=False,  # Set to True to render the environment during evaluation
+)
+
+
 TIMESTEPS = 10000
 iters = 0
 for i in range(1, 30):
-    learner.learn(total_timesteps=TIMESTEPS, reset_num_timesteps=False, tb_log_name="PPO-rl")
+    learner.learn(
+        total_timesteps=TIMESTEPS,
+        reset_num_timesteps=False,
+        tb_log_name=model_name,
+        callback=eval_callback,
+    )
     learner.save(f"{models_dir}/{TIMESTEPS*i}")