feat(*): Use SB3 with multiple game instances

bot/src/agent.py

@@ -1,7 +1,7 @@
 import os
 from stable_baselines3 import PPO
 
-model_zip = "PPO-rl2/280000.zip"
+model_zip = "exit_save.zip"
 
 
 class Agent:

rl-training.py

@@ -1,102 +1,143 @@
 import os
-import rlgym
+import numpy as np
-from stable_baselines3.ppo import PPO
+import atexit
-from stable_baselines3.common.callbacks import EvalCallback
+
-from stable_baselines3.common.logger import TensorBoardOutputFormat, configure
+from stable_baselines3 import PPO
-from torch.utils.tensorboard import SummaryWriter
+from stable_baselines3.ppo import MlpPolicy
-from rlgym_tools.sb3_utils import SB3SingleInstanceEnv
+from stable_baselines3.common.callbacks import CheckpointCallback
-from rlgym.utils.reward_functions.common_rewards import VelocityBallToGoalReward
+from stable_baselines3.common.vec_env import VecMonitor, VecNormalize, VecCheckNan
-from rlgym.utils.reward_functions.common_rewards import BallYCoordinateReward
+
+from rlgym.envs import Match
+
+from rlgym.utils.action_parsers import DiscreteAction
+from rlgym.utils.reward_functions.common_rewards import (
+    VelocityBallToGoalReward,
+    BallYCoordinateReward,
+    EventReward,
+)
 from rlgym.utils.obs_builders import AdvancedObs
+from rlgym.utils.state_setters import DefaultState
 from rlgym.utils.terminal_conditions.common_conditions import (
     TimeoutCondition,
+    NoTouchTimeoutCondition,
     GoalScoredCondition,
 )
+from rlgym.utils.reward_functions.common_rewards.player_ball_rewards import (
+    VelocityPlayerToBallReward,
+)
+from rlgym.utils.reward_functions.common_rewards.ball_goal_rewards import (
+    VelocityBallToGoalReward,
+)
+from rlgym.utils.reward_functions import CombinedReward
+
+from rlgym_tools.sb3_utils import SB3MultipleInstanceEnv
+
 
 # from observationBuilders.CustomObsBuilderBluePerspective import (
 #     CustomObsBuilderBluePerspective,
 # )
 
-# Set up the folders
+if __name__ == "__main__":  # Required for multiprocessing
-model_name = "PPO-rl6"
+    # Set up the folders
-models_dir = f"models/{model_name}"
+    models_dir = "models"
-logdir = "logs"
+    log_dir = "logs"
+    if not os.path.exists(models_dir):
+        os.makedirs(models_dir)
+    if not os.path.exists(log_dir):
+        os.makedirs(log_dir)
 
-if not os.path.exists(models_dir):
+    frame_skip = 8  # Number of ticks to repeat an action
-    os.makedirs(models_dir)
+    half_life_seconds = (
+        5  # Easier to conceptualize, after this many seconds the reward discount is 0.5
+    )
 
-if not os.path.exists(logdir):
+    fps = 120 / frame_skip
-    os.makedirs(logdir)
+    gamma = np.exp(np.log(0.5) / (fps * half_life_seconds))  # Quick mafs
+    agents_per_match = 2
+    num_instances = 4
+    target_steps = 100_000
+    steps = target_steps // (num_instances * agents_per_match)
+    batch_size = steps
 
-# Set up the RLGym environment
+    def exit_save(model):
-gym_env = rlgym.make(
+        model.save("models/exit_save")
-    use_injector=True,
-    self_play=True,
-    reward_fn=BallYCoordinateReward(),
-    obs_builder=AdvancedObs(),
-    terminal_conditions=(TimeoutCondition(225), GoalScoredCondition()),
-)
 
-# Wrap the RLGym environment with the single instance wrapper
+    def get_match():  # Need to use a function so that each instance can call it and produce their own objects
-env = SB3SingleInstanceEnv(gym_env)
+        return Match(
+            team_size=1,
+            tick_skip=frame_skip,
+            reward_function=CombinedReward(
+                (
+                    VelocityPlayerToBallReward(),
+                    VelocityBallToGoalReward(),
+                    EventReward(
+                        team_goal=100.0,
+                        concede=-100.0,
+                        shot=5.0,
+                        save=30.0,
+                        demo=10.0,
+                    ),
+                    BallYCoordinateReward(),
+                ),
+                (0.1, 1.0, 1.0, 0.1),
+            ),
+            # self_play=True,
+            terminal_conditions=[
+                TimeoutCondition(fps * 300),
+                NoTouchTimeoutCondition(fps * 20),
+                GoalScoredCondition(),
+            ],
+            obs_builder=AdvancedObs(),  # Not that advanced, good default
+            state_setter=DefaultState(),  # Resets to kickoff position
+            action_parser=DiscreteAction(),  # Discrete > Continuous don't @ me
+        )
 
-# Create a PPO instance
+    # Generate the environment (the Rocket League game used by RL Gym)
-learner = PPO(policy="MlpPolicy", env=env, verbose=1, tensorboard_log=logdir)
+    env = SB3MultipleInstanceEnv(
+        get_match, num_instances
+    )  # Start x instances, waiting 60 seconds between each
+    env = VecCheckNan(env)  # Optional
+    env = VecMonitor(env)  # Recommended, logs mean reward and ep_len to Tensorboard
+    env = VecNormalize(
+        env, norm_obs=False, gamma=gamma
+    )  # Highly recommended, normalizes rewards
 
-# Configure TensorBoard for custom logging
+    # Load the model that was last trained, or start a new one if the zip file doesn't exist
-tb_log_name = f"logs/{model_name}"
+    try:
-tensorboard_writer = SummaryWriter(log_dir=tb_log_name)
+        model = PPO.load(
+            f"{models_dir}/exit_save.zip",
+            env,
+            device="auto",  # Need to set device again (if using a specific one)
+        )
+    except:
+        model = PPO(
+            MlpPolicy,
+            env,
+            learning_rate=5e-5,  # Around this is fairly common for PPO
+            n_steps=steps,  # Number of steps to perform before optimizing network
+            batch_size=batch_size,  # Batch size as high as possible within reason
+            n_epochs=1,  # PPO calls for multiple epochs
+            gamma=gamma,  # Gamma as calculated using half-life
+            ent_coef=0.01,  # From PPO Atari
+            vf_coef=1.0,  # From PPO Atari
+            tensorboard_log=log_dir,  # `tensorboard --logdir out/logs` in terminal to see graphs
+            verbose=3,  # Print out all the info as we're going
+            device="auto",  # Uses GPU if available
+        )
 
-# Define a callback to log rewards to TensorBoard
+    # Save model every so often
-eval_callback = EvalCallback(
+    # Divide by num_envs (number of agents) because callback only increments every time all agents have taken a step
-    env,
+    # This saves to specified folder with a specified name
-    best_model_save_path=models_dir,
+    callback = CheckpointCallback(
-    log_path=logdir,
+        round(5_000_000 / env.num_envs), save_path=models_dir, name_prefix="rl_model"
-    eval_freq=10000,
+    )
-    deterministic=True,
-    render=False,  # Set to True to render the environment during evaluation
-)
 
+    atexit.register(exit_save, model)
 
-# Custom callback for saving the model every 100 episodes
+    try:
-def custom_callback(locals, globals):
+        while True:
-    if num_episodes % 100 == 0:
+            model.learn(25_000_000, callback=callback)
-        model_save_path = f"{models_dir}/episode_{num_episodes}"
+            model.save(f"{models_dir}/exit_save")
-        locals["self"].save(model_save_path)
+            model.save(f"mmr_models/{model.num_timesteps}")
-
+    except Exception as e:
-
+        print(e)
-# 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())