From cc39cca0ab4c93d046157fdb6b60c4a1b6aee437 Mon Sep 17 00:00:00 2001
From: sp <stefan.pranger@iaik.tugraz.at>
Date: Sat, 30 Dec 2023 11:43:33 +0100
Subject: [PATCH] fixed typo, removed ws

---
 examples/shields/rl/15_train_eval_tune.py | 59 +++++++++++------------
 1 file changed, 29 insertions(+), 30 deletions(-)

diff --git a/examples/shields/rl/15_train_eval_tune.py b/examples/shields/rl/15_train_eval_tune.py
index 315f7b1..f3a17db 100644
--- a/examples/shields/rl/15_train_eval_tune.py
+++ b/examples/shields/rl/15_train_eval_tune.py
@@ -19,23 +19,23 @@ from shieldhandlers import MiniGridShieldHandler, create_shield_query
 
 from torch.utils.tensorboard import SummaryWriter
 from callbacks import MyCallbacks
- 
+
 
 def shielding_env_creater(config):
     name = config.get("name", "MiniGrid-LavaCrossingS9N3-v0")
     framestack = config.get("framestack", 4)
     args = config.get("args", None)
     args.grid_path = F"{args.expname}_{args.grid_path}_{config.worker_index}.txt"
-    args.prism_path = F"{args.expname}_{args.prism_path}_{config.worker_index}.prism"   
-    shielding = config.get("shielding", False)   
-    shield_creator = MiniGridShieldHandler(grid_file=args.grid_path, 
+    args.prism_path = F"{args.expname}_{args.prism_path}_{config.worker_index}.prism"
+    shielding = config.get("shielding", False)
+    shield_creator = MiniGridShieldHandler(grid_file=args.grid_path,
                                            grid_to_prism_path=args.grid_to_prism_binary_path,
                                            prism_path=args.prism_path,
                                            formula=args.formula,
                                            shield_value=args.shield_value,
                                            prism_config=args.prism_config,
                                            shield_comparision=args.shield_comparision)
-    
+
     prob_forward = args.prob_forward
     prob_direct = args.prob_direct
     prob_next = args.prob_next
@@ -47,8 +47,8 @@ def shielding_env_creater(config):
                         config.vector_index if hasattr(config, "vector_index") else 0,
                         framestack=framestack
                         )
-    
-    
+
+
     return env
 
 
@@ -57,10 +57,10 @@ def register_minigrid_shielding_env(args):
     register_env(env_name, shielding_env_creater)
 
     ModelCatalog.register_custom_model(
-        "shielding_model", 
+        "shielding_model",
         TorchActionMaskModel
     )
-    
+
 def trial_name_creator(trial : Trial):
     return "trial"
 
@@ -78,7 +78,7 @@ def ppo(args):
                                   "shielding": args.shielding is ShieldingConfig.Full or args.shielding is ShieldingConfig.Training,
                                   },)
         .framework("torch")
-        .callbacks(MyCallbacks, ShieldInfoCallback(logdir, [1,12])
+        .callbacks(MyCallbacks, ShieldInfoCallback(logdir, [1,12]))
         .evaluation(evaluation_config={
                                        "evaluation_interval": 1,
                                         "evaluation_duration": 10,
@@ -106,25 +106,24 @@ def ppo(args):
                        ),
                         run_config=air.RunConfig(
                                 stop = {"episode_reward_mean": 94,
-                                        "timesteps_total": args.steps,}, 
+                                        "timesteps_total": args.steps,},
                                 checkpoint_config=air.CheckpointConfig(checkpoint_at_end=True,
-                                                                       num_to_keep=1, 
+                                                                       num_to_keep=1,
                                                                        checkpoint_score_attribute="episode_reward_mean",
                                                                        ),
-                                
+
                                storage_path=F"{logdir}",
                                name=test_name(args),
-                               
-                               
-    )
-                        ,
+
+
+                        ),
     param_space=config,)
-    
+
     results = tuner.fit()
     best_result = results.get_best_result()
-    
+
     import pprint
-    
+
     metrics_to_print = [
     "episode_reward_mean",
     "episode_reward_max",
@@ -134,14 +133,14 @@ def ppo(args):
     pprint.pprint({k: v for k, v in best_result.metrics.items() if k in metrics_to_print})
 
    # algo = Algorithm.from_checkpoint(best_result.checkpoint)
-    
-       
+
+
     # eval_log_dir = F"{logdir}-eval"
-        
+
     # writer = SummaryWriter(log_dir=eval_log_dir)
     # csv_logger = CSVLogger(config=config, logdir=eval_log_dir)
-    
-    
+
+
     # for i in range(args.evaluations):
     #     eval_result = algo.evaluate()
     #     print(pretty_print(eval_result))
@@ -149,23 +148,23 @@ def ppo(args):
     #     # logger.on_result(eval_result)
 
     #     csv_logger.on_result(eval_result)
-        
+
     #     evaluation = eval_result['evaluation']
     #     epsiode_reward_mean = evaluation['episode_reward_mean']
     #     episode_len_mean = evaluation['episode_len_mean']
     #     print(epsiode_reward_mean)
     #     writer.add_scalar("evaluation/episode_reward_mean", epsiode_reward_mean, i)
     #     writer.add_scalar("evaluation/episode_len_mean", episode_len_mean, i)
-        
-    
+
+
 def main():
     ray.init(num_cpus=3)
     import argparse
     args = parse_arguments(argparse)
 
     ppo(args)
-   
+
     ray.shutdown()
 
 if __name__ == '__main__':
-    main()
\ No newline at end of file
+    main()