mirror of
https://github.com/linyiLYi/street-fighter-ai.git
synced 2025-04-05 07:30:42 +00:00
74 lines
2.1 KiB
Python
74 lines
2.1 KiB
Python
import os
|
|
|
|
import retro
|
|
import optuna
|
|
from stable_baselines3 import PPO
|
|
from stable_baselines3.common.evaluation import evaluate_policy
|
|
from stable_baselines3.common.monitor import Monitor
|
|
from stable_baselines3.common.vec_env import DummyVecEnv, VecFrameStack
|
|
|
|
from street_fighter_custom_wrapper import StreetFighterCustomWrapper
|
|
|
|
LOG_DIR = 'logs/'
|
|
OPT_DIR = 'optuna/'
|
|
os.makedirs(LOG_DIR, exist_ok=True)
|
|
os.makedirs(OPT_DIR, exist_ok=True)
|
|
|
|
def optimize_ppo(trial):
|
|
return {
|
|
'n_steps':trial.suggest_int('n_steps', 1024, 8192, log=True),
|
|
'gamma':trial.suggest_float('gamma', 0.9, 0.9999),
|
|
'learning_rate':trial.suggest_float('learning_rate', 5e-5, 1e-4, log=True),
|
|
'clip_range':trial.suggest_float('clip_range', 0.1, 0.4),
|
|
'gae_lambda':trial.suggest_float('gae_lambda', 0.8, 0.99)
|
|
}
|
|
|
|
def make_env(game, state, seed=0):
|
|
def _init():
|
|
env = retro.make(
|
|
game=game,
|
|
state=state,
|
|
use_restricted_actions=retro.Actions.FILTERED,
|
|
obs_type=retro.Observations.IMAGE
|
|
)
|
|
env = StreetFighterCustomWrapper(env)
|
|
env.seed(seed)
|
|
return env
|
|
return _init
|
|
|
|
def optimize_agent(trial):
|
|
game = "StreetFighterIISpecialChampionEdition-Genesis"
|
|
state = "ChampionX.Level1.ChunLiVsKen"
|
|
|
|
# try:
|
|
model_params = optimize_ppo(trial)
|
|
|
|
# Create environment
|
|
env = make_env(game, state)()
|
|
env = Monitor(env, LOG_DIR)
|
|
env = DummyVecEnv([lambda: env])
|
|
env = VecFrameStack(env, 4, channels_order='last')
|
|
|
|
# Create algo
|
|
model = PPO('CnnPolicy', env, tensorboard_log=LOG_DIR, verbose=0, **model_params)
|
|
model.learn(total_timesteps=100000)
|
|
|
|
# Evaluate model
|
|
mean_reward, _ = evaluate_policy(model, env, n_eval_episodes=5)
|
|
env.close()
|
|
|
|
SAVE_PATH = os.path.join(OPT_DIR, 'trial_{}_best_model'.format(trial.number))
|
|
model.save(SAVE_PATH)
|
|
|
|
return mean_reward
|
|
|
|
# except Exception as e:
|
|
# return -1
|
|
|
|
# Creating the experiment
|
|
study = optuna.create_study(direction='maximize')
|
|
study.optimize(optimize_agent, n_trials=10, n_jobs=1)
|
|
|
|
print(study.best_params)
|
|
print(study.best_trial)
|