Sarpy/rl-training.py

import os
import numpy as np
import atexit

from stable_baselines3 import PPO
from stable_baselines3.ppo import MlpPolicy
from stable_baselines3.common.callbacks import CheckpointCallback
from stable_baselines3.common.vec_env import VecMonitor, VecNormalize, VecCheckNan

from rlgym.envs import Match

from rlgym.utils.action_parsers import DiscreteAction
from rlgym.utils.reward_functions.common_rewards import (
    VelocityBallToGoalReward,
    LiuDistanceBallToGoalReward,
    BallYCoordinateReward,
    EventReward,
    RewardIfClosestToBall,
    LiuDistancePlayerToBallReward,
    FaceBallReward,
    TouchBallReward,
    AlignBallGoal,
    SaveBoostReward,
)
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 rlgym_tools.extra_rewards.kickoff_reward import KickoffReward
from rlgym_tools.extra_obs.advanced_padder import AdvancedObsPadder

if __name__ == "__main__":  # Required for multiprocessing
    # Set up the folders
    models_dir = "models"
    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)

    frame_skip = 8  # Number of ticks to repeat an action
    half_life_seconds = (
        5  # Easier to conceptualize, after this many seconds the reward discount is 0.5
    )

    fps = 120 / frame_skip
    gamma = np.exp(np.log(0.5) / (fps * half_life_seconds))  # Quick mafs
    print(f"fps={fps}, gamma={gamma})")

    def exit_save(model):
        model.save("models/exit_save")

    team_size = 3

    def get_match():  # Need to use a function so that each instance can call it and produce their own objects
        return Match(
            team_size=team_size,
            tick_skip=frame_skip,
            reward_function=CombinedReward(
                (
                    VelocityPlayerToBallReward(),
                    VelocityBallToGoalReward(),
                    EventReward(
                        team_goal=1000.0,
                        concede=-100.0,
                        shot=10.0,
                        save=60.0,
                        demo=20.0,
                    ),
                    KickoffReward(),
                    SaveBoostReward(),
                ),
                (1.0, 1.0, 1.0, 1.0, 1.0),
            ),
            # reward_function=CombinedReward(
            #     (
            #         VelocityPlayerToBallReward(),
            #         VelocityBallToGoalReward(),
            #         LiuDistanceBallToGoalReward(),
            #         EventReward(
            #             team_goal=10000.0,
            #             concede=-10000.0,
            #             shot=10.0,
            #             save=60.0,
            #             demo=20.0,
            #         ),
            #         BallYCoordinateReward(),
            #         RewardIfClosestToBall(LiuDistancePlayerToBallReward()),
            #         LiuDistancePlayerToBallReward(),
            #         FaceBallReward(),
            #         TouchBallReward(),
            #         AlignBallGoal(),
            #         KickoffReward(),
            #     ),
            #     (50.0, 10.0, 20.0, 100.0, 0.1, 10.0, 1.0, 0.2, 10.0, 20.0, 10.0),
            # ),
            spawn_opponents=True,
            terminal_conditions=[
                TimeoutCondition(fps * 300),
                NoTouchTimeoutCondition(fps * 10),
                GoalScoredCondition(),
            ],
            obs_builder=AdvancedObsPadder(
                team_size=team_size
            ),  # Not that advanced, good default
            state_setter=DefaultState(),  # Resets to kickoff position
            action_parser=DiscreteAction(),  # Discrete > Continuous don't @ me
        )

    # Generate the environment (the Rocket League game used by RL Gym)
    env = SB3MultipleInstanceEnv(
        get_match, 14
    )  # Start 14 instances, waiting 30 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

    # Save model every so often
    # Divide by num_envs (number of agents) because callback only increments every time all agents have taken a step
    # This saves to specified folder with a specified name
    callback = CheckpointCallback(
        round(1_000_000 / env.num_envs), save_path=models_dir, name_prefix="sarpy_model"
    )

    try:
        while True:
            # Load the model that was last trained, or start a new one if the zip file doesn't exist
            try:
                # Now, if one wants to load a trained model from a checkpoint, use this function
                # This will contain all the attributes of the original model
                # Any attribute can be overwritten by using the custom_objects parameter,
                # which includes n_envs (number of agents), which has to be overwritten to use a different amount
                model = PPO.load(
                    f"{models_dir}/exit_save.zip",
                    env,
                    custom_objects=dict(
                        n_envs=env.num_envs, _last_obs=None
                    ),  # Need this to change number of agents
                    device="auto",  # Need to set device again (if using a specific one)
                    force_reset=True,  # Make SB3 reset the env so it doesn't think we're continuing from last state
                )
                print("Loaded exit_save.zip model")
            except:
                model = PPO(
                    MlpPolicy,
                    env,
                    n_epochs=32,  # PPO calls for multiple epochs
                    learning_rate=1e-5,  # Around this is fairly common for PPO
                    ent_coef=0.01,  # From PPO Atari
                    vf_coef=1.0,  # From PPO Atari
                    gamma=gamma,  # Gamma as calculated using half-life
                    verbose=3,  # Print out all the info as we're going
                    batch_size=4096,  # Batch size as high as possible within reason
                    n_steps=4096,  # Number of steps to perform before optimizing network
                    tensorboard_log="out/logs",  # `tensorboard --logdir out/logs` in terminal to see graphs
                    device="auto",  # Uses GPU if available
                )
                print("Created new model")

            atexit.register(exit_save, model)

            print("Learning...")
            model.learn(100_000_000, callback=callback, reset_num_timesteps=False)
            model.save(f"{models_dir}/exit_save")
            print("Saved exit_save model")

    except Exception as e:
        print(e)