mirror of
https://github.com/jcreek/Sarpy.git
synced 2026-07-13 03:03:43 +00:00
Josh Creek
61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
import os
|
|
import rlgym
|
|
from stable_baselines3.ppo import PPO
|
|
from stable_baselines3.common.callbacks import EvalCallback
|
|
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 observationBuilders.CustomObsBuilderBluePerspective import (
|
|
# CustomObsBuilderBluePerspective,
|
|
# )
|
|
|
|
# set up the folders
|
|
model_name = "PPO-rl4"
|
|
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=VelocityBallToGoalReward(),
|
|
obs_builder=AdvancedObs(),
|
|
# obs_builder=CustomObsBuilderBluePerspective(),
|
|
)
|
|
|
|
# wrap the RLGym environment with the single instance wrapper
|
|
env = SB3SingleInstanceEnv(gym_env)
|
|
|
|
# create a PPO instance and start learning
|
|
learner = PPO(policy="MlpPolicy", env=env, verbose=1, tensorboard_log=logdir)
|
|
|
|
# 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
|
|
)
|
|
|
|
|
|
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}")
|