mirror of
https://github.com/jcreek/Sarpy.git
synced 2026-07-13 03:03:43 +00:00
Josh Creek
103 lines
3.1 KiB
Python
103 lines
3.1 KiB
Python
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-rl6"
|
|
models_dir = f"models/{model_name}"
|
|
logdir = "logs"
|
|
|
|
if not os.path.exists(models_dir):
|
|
os.makedirs(models_dir)
|
|
|
|
if not os.path.exists(logdir):
|
|
os.makedirs(logdir)
|
|
|
|
# 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()),
|
|
)
|
|
|
|
# Wrap the RLGym environment with the single instance wrapper
|
|
env = SB3SingleInstanceEnv(gym_env)
|
|
|
|
# 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,
|
|
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
|
|
)
|
|
|
|
|
|
# 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())
|