learn-to-follow

1
from pogema_toolbox.create_env import create_env_base, Environment
2
from pogema_toolbox.evaluator import evaluation
3
from pogema import BatchAStarAgent
4

5
from pogema_toolbox.eval_utils import initialize_wandb, save_evaluation_results
6

7
from pathlib import Path
8
import wandb
9

10
import yaml
11

12
from pogema_toolbox.registry import ToolboxRegistry
13

14
from follower.inference import FollowerInference, FollowerInferenceConfig
15
from follower.preprocessing import follower_preprocessor
16
from follower_cpp.inference import FollowerConfigCPP, FollowerInferenceCPP
17
from follower_cpp.preprocessing import follower_cpp_preprocessor
18

19
PROJECT_NAME = 'pogema-toolbox'
20
BASE_PATH = Path('experiments')
21

22

23
def main(disable_wandb=False):
24
    ToolboxRegistry.register_env('Pogema-v0', create_env_base, Environment)
25
    ToolboxRegistry.register_algorithm('A*', BatchAStarAgent)
26
    ToolboxRegistry.register_algorithm('Follower', FollowerInference, FollowerInferenceConfig,
27
                                       follower_preprocessor)
28
    ToolboxRegistry.register_algorithm('FollowerLite', FollowerInferenceCPP, FollowerConfigCPP,
29
                                       follower_cpp_preprocessor)
30

31
    with open("env/test-maps.yaml", 'r') as f:
32
        maps_to_register = yaml.safe_load(f)
33
    ToolboxRegistry.register_maps(maps_to_register)
34

35
    folder_names = [
36
        '01-random-20x20',
37
        '02-mazes',
38
        '03-den520d',
39
        '04-Paris_1',
40
        '05-warehouse',
41
    ]
42

43
    for folder in folder_names:
44
        config_path = BASE_PATH / folder / f"{Path(folder).name}.yaml"
45
        eval_dir = BASE_PATH / folder
46

47
        with open(config_path) as f:
48
            evaluation_config = yaml.safe_load(f)
49
        if folder == 'eval-fast':
50
            disable_wandb = True
51

52
        initialize_wandb(evaluation_config, eval_dir, disable_wandb, PROJECT_NAME)
53
        evaluation(evaluation_config, eval_dir=eval_dir)
54
        save_evaluation_results(eval_dir)
55
        wandb.finish()
56

57

58
if __name__ == '__main__':
59
    main()
60