study e_1 corpus

environments/helpers.py

@@ -5,6 +5,8 @@ from typing import Tuple, Union
 import numpy as np
 from pathlib import Path
 
+from stable_baselines3 import PPO, DQN, A2C
+
 LEVELS_DIR = 'levels'
 
 TO_BE_AVERAGED = ['dirt_amount', 'dirty_tiles']
@@ -142,6 +144,8 @@ def asset_str(agent):
         return c.AGENT.value, 'idle'
 
 
+model_map = dict(PPO=PPO, DQN=DQN, A2C=A2C)
+
 if __name__ == '__main__':
     parsed_level = parse_level(Path(__file__).parent / 'factory' / 'levels' / 'simple.txt')
     y = one_hot_level(parsed_level)