File size: 1,833 Bytes
22b374f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | import os
import datetime
import gym
import json
from stable_baselines3 import PPO
import optuna
def check_json_exists(filename):
try:
open(filename, 'r')
except:
main = {
"env" : "RoadRunner-v0",
"algo" : "PPO",
"policy" : "MlpPolicy",
"counter" : 0,
}
with open(filename, "w") as f:
json.dump(main, f)
check_json_exists("previous_request.json")
class MegaTrainer:
def __init__(self, env_name, algo):
with open("previous_request.json", "r") as f:
self.data = json.loads(f.read())
self.env_name = env_name
self.algorithm = algo
self.counter = self.data["counter"] + 1
self.env = gym.make(self.env_name)
self.env.reset()
self.data["counter"] += 1
with open("previous_request.json", "w") as f:
json.dump(self.data, f)
def make_directories(self, models_dir, logdir):
if(not os.path.exists(models_dir)):
os.makedirs(models_dir)
if(not os.path.exists(logdir)):
os.makedirs(logdir)
def learn(self, timestep, iterations, hyps, directory, trial):
model_dir = f"model/{directory}/{self.counter}"
logdir = f"logs/{directory}"
self.make_directories(model_dir, logdir)
model = PPO(**hyps)
for i in range(1, iterations+1):
if(trial.should_prune()):
raise optuna.TrialPruned()
dt = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S-%f")
model.learn(total_timesteps=timestep, reset_num_timesteps=False, tb_log_name=str(self.counter))
model.save(model_dir+f"/{dt}_{timestep*i}")
return self.counter
if(__name__=="__main__"):
load_from_json("previous_request.json")
|