| 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): |
| |
| return os.listdir(all_models_dir)[-1] |
|
|
| |
| pass |
|
|
| def load(self, execution_number, no_of_episodes): |
| |
| |
| |
| |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| 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 |
|
|
| |
|
|
| 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 |
