From 4122853d880bbcf2b0c0a4712db8c2974e961836 Mon Sep 17 00:00:00 2001
From: Josh Creek <8179928+jcreek@users.noreply.github.com>
Date: Wed, 27 Sep 2023 18:15:52 +0100
Subject: [PATCH] feat(*): Update training to use episodes

---
 rl-training.py | 76 +++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 59 insertions(+), 17 deletions(-)

diff --git a/rl-training.py b/rl-training.py
index 91298bc..bae852f 100644
--- a/rl-training.py
+++ b/rl-training.py
@@ -2,17 +2,23 @@ import os
 import rlgym
 from stable_baselines3.ppo import PPO
 from stable_baselines3.common.callbacks import EvalCallback
+from stable_baselines3.common.logger import TensorBoardOutputFormat, configure
+from torch.utils.tensorboard import SummaryWriter
 from rlgym_tools.sb3_utils import SB3SingleInstanceEnv
 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 rlgym.utils.terminal_conditions.common_conditions import (
+    TimeoutCondition,
+    GoalScoredCondition,
+)
 
 # from observationBuilders.CustomObsBuilderBluePerspective import (
 #     CustomObsBuilderBluePerspective,
 # )
 
-# set up the folders
-model_name = "PPO-rl4"
+# Set up the folders
+model_name = "PPO-rl6"
 models_dir = f"models/{model_name}"
 logdir = "logs"
 
@@ -22,21 +28,25 @@ if not os.path.exists(models_dir):
 if not os.path.exists(logdir):
     os.makedirs(logdir)
 
-# set up the RLGym environment
+# Set up the RLGym environment
 gym_env = rlgym.make(
     use_injector=True,
     self_play=True,
-    reward_fn=VelocityBallToGoalReward(),
+    reward_fn=BallYCoordinateReward(),
     obs_builder=AdvancedObs(),
-    # obs_builder=CustomObsBuilderBluePerspective(),
+    terminal_conditions=(TimeoutCondition(225), GoalScoredCondition()),
 )
 
-# wrap the RLGym environment with the single instance wrapper
+# Wrap the RLGym environment with the single instance wrapper
 env = SB3SingleInstanceEnv(gym_env)
 
-# create a PPO instance and start learning
+# Create a PPO instance
 learner = PPO(policy="MlpPolicy", env=env, verbose=1, tensorboard_log=logdir)
 
+# Configure TensorBoard for custom logging
+tb_log_name = f"logs/{model_name}"
+tensorboard_writer = SummaryWriter(log_dir=tb_log_name)
+
 # Define a callback to log rewards to TensorBoard
 eval_callback = EvalCallback(
     env,
@@ -48,13 +58,45 @@ eval_callback = EvalCallback(
 )
 
 
-TIMESTEPS = 10000
-iters = 0
-for i in range(1, 30):
-    learner.learn(
-        total_timesteps=TIMESTEPS,
-        reset_num_timesteps=False,
-        tb_log_name=model_name,
-        callback=eval_callback,
-    )
-    learner.save(f"{models_dir}/{TIMESTEPS*i}")
+# Custom callback for saving the model every 100 episodes
+def custom_callback(locals, globals):
+    if num_episodes % 100 == 0:
+        model_save_path = f"{models_dir}/episode_{num_episodes}"
+        locals["self"].save(model_save_path)
+
+
+# Training loop
+EPISODES = 1000  # Set the total number of episodes
+num_episodes = 0
+
+while num_episodes < EPISODES:
+    obs = env.reset()
+    episode_reward = 0
+    episode_length = 0
+
+    while True:
+        action, _ = learner.predict(obs)
+        obs, reward, done, _ = env.step(action)
+        episode_reward += sum(reward)
+        episode_length += 1
+
+        print(done)
+
+        if done.any():  # Check if any element in 'done' is True
+            num_episodes += 1  # Increment the episode count
+            break
+
+    # Calculate the average reward per episode
+    avg_episode_reward = episode_reward / episode_length
+
+    # Log custom episode-related information to TensorBoard
+    print(episode_reward)
+    tensorboard_writer.add_scalar("Episode/Reward", float(episode_reward), num_episodes)
+    tensorboard_writer.add_scalar("Episode/Length", episode_length, num_episodes)
+
+    # Run the evaluation callback
+    if num_episodes % 50 == 0:
+        eval_callback.on_epoch_end()
+
+    # Use the custom callback for saving the model
+    custom_callback(locals(), globals())