File size: 1,982 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 70 71 72 73 | import os
from stable_baselines3 import PPO
import gym
class MegaLoader:
def __init__(self, env_name, algo, directory):
self.env_name = env_name
self.env = gym.make(self.env_name)
self.env.reset()
self.algorithm = algo
self.directory = directory
def model_selector(self, all_models_dir):
# print( os.listdir(all_models_dir)[-1])
return os.listdir(all_models_dir)[-1]
# return whole file name properly, with date and all
pass
def load(self, execution_number, no_of_episodes):
# Input :-
# 1. execution_number -
# 2. no_of_episodes -
# 3.
# Returns :-
# Dictionary with the following items
# 1. total_return [int]
# 2. return_per_episode [list]
# 3.
# 3.
# 3.
# 3.
total_return = 0
return_per_episode = []
info_per_episode = []
all_models_dir = f"model/{self.directory}/{execution_number}"
model_no = self.model_selector(all_models_dir)
model_dir = f"{all_models_dir}/{model_no}"
model = PPO.load(model_dir, env=self.env)
for ep in range(no_of_episodes):
obs = self.env.reset()
done = False
info = ""
episode_return = 0
while not done:
action, _states = model.predict(obs)
obs, rewards, done, info = self.env.step(action)
episode_return += rewards
# self.env.render()
return_per_episode.append(episode_return)
total_return += episode_return
info_per_episode.append(info)
model_test_info = {
"total_return": total_return,
"return_per_episode": return_per_episode,
"info_per_episode": info_per_episode,
}
return model_test_info |