mirror of
https://github.com/jcreek/Sarpy.git
synced 2026-07-13 03:03:43 +00:00
feat(*): Use SB3 with multiple game instances
This commit is contained in:
Josh Creek
+1
-1
@@ -1,7 +1,7 @@
|
||||
import os
|
||||
from stable_baselines3 import PPO
|
||||
|
||||
model_zip = "PPO-rl2/280000.zip"
|
||||
model_zip = "exit_save.zip"
|
||||
|
||||
|
||||
class Agent:
|
||||
|
||||
+123
-82
@@ -1,102 +1,143 @@
|
||||
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
|
||||
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,
|
||||
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
|
||||
model_name = "PPO-rl6"
|
||||
models_dir = f"models/{model_name}"
|
||||
logdir = "logs"
|
||||
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)
|
||||
|
||||
if not os.path.exists(models_dir):
|
||||
os.makedirs(models_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
|
||||
)
|
||||
|
||||
if not os.path.exists(logdir):
|
||||
os.makedirs(logdir)
|
||||
fps = 120 / frame_skip
|
||||
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
|
||||
gym_env = rlgym.make(
|
||||
use_injector=True,
|
||||
self_play=True,
|
||||
reward_fn=BallYCoordinateReward(),
|
||||
obs_builder=AdvancedObs(),
|
||||
terminal_conditions=(TimeoutCondition(225), GoalScoredCondition()),
|
||||
)
|
||||
def exit_save(model):
|
||||
model.save("models/exit_save")
|
||||
|
||||
# Wrap the RLGym environment with the single instance wrapper
|
||||
env = SB3SingleInstanceEnv(gym_env)
|
||||
def get_match(): # Need to use a function so that each instance can call it and produce their own objects
|
||||
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
|
||||
learner = PPO(policy="MlpPolicy", env=env, verbose=1, tensorboard_log=logdir)
|
||||
# Generate the environment (the Rocket League game used by RL Gym)
|
||||
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
|
||||
tb_log_name = f"logs/{model_name}"
|
||||
tensorboard_writer = SummaryWriter(log_dir=tb_log_name)
|
||||
# Load the model that was last trained, or start a new one if the zip file doesn't exist
|
||||
try:
|
||||
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
|
||||
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
|
||||
)
|
||||
# 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(5_000_000 / env.num_envs), save_path=models_dir, name_prefix="rl_model"
|
||||
)
|
||||
|
||||
atexit.register(exit_save, model)
|
||||
|
||||
# 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())
|
||||
try:
|
||||
while True:
|
||||
model.learn(25_000_000, callback=callback)
|
||||
model.save(f"{models_dir}/exit_save")
|
||||
model.save(f"mmr_models/{model.num_timesteps}")
|
||||
except Exception as e:
|
||||
print(e)
|
||||
|
||||
Reference in New Issue
Block a user