Merge pull request #66 from F-Secure/sorin/refactor_explorer_agent_to…

…_migrate_to_ludwig05

feat: add initial building blocks for reinforcement learning

change_analyzer/envs/app_env.py

@@ -55,7 +55,7 @@ def close(self) -> None:
             current_open_windows = self._get_current_open_windows()
 
             for w in current_open_windows:
-                if "F-Secure" in w or "Software updates" in w:
+                if "F-Secure" in w or "Software updates" in w or "WithSecure" in w:
                     self._kill_app_based_on_window_title(window_title=w)
 
             # while not current_open_windows == initial_windows_without_sut:

change_analyzer/main.py

@@ -10,6 +10,7 @@
 
 from change_analyzer.agents.random_agent import RandomAgent
 from change_analyzer.agents.replay_agent import ReplayAgent
+from change_analyzer.agents.explorer_agent import ExplorerAgent
 from change_analyzer.wrappers.enhanced_monitor import EnhancedMonitor
 from change_analyzer.wrappers.sequence_recorder import SequenceRecorder
 
@@ -48,7 +49,7 @@ def reset() -> WebDriver:
     return driver
 
 
-def run(config: str, steps: int = 0, csv_folder:str = ""):
+def run(config: str, steps: int = 0, csv_folder:str = "", strategy:str = ""):
     CONFIG.read(config)
 
     env = gym.make(
@@ -66,6 +67,8 @@ def run(config: str, steps: int = 0, csv_folder:str = ""):
         env.reset()
         if csv_folder:
             ReplayAgent(env, csv_folder).run()
+        elif strategy == 'rl':
+            ExplorerAgent(env, int(steps)).run()
         else:
             RandomAgent(env, int(steps)).run()
     finally:
@@ -89,8 +92,13 @@ def main():
         help="path to the folder within recordings which has the targeted csv file",
         required=False,
     )
+    parser.add_argument(
+        "--strategy",
+        help="define the agent strategy, either random or rl (reinforcement learning)",
+        required=False,
+    )
     args = parser.parse_args()
-    run(args.config, args.steps, args.csv_folder)
+    run(args.config, args.steps, args.csv_folder, args.strategy)
 
 
 if __name__ == "__main__":

change_analyzer/spaces/discrete_app_action_space.py

@@ -83,7 +83,18 @@ def update_current_actions(self):
             #     )
         self._logger.info(f"Found {len(self.actions)} actions on the screen")
 
+    def get_action_based_on_string(self, action_string: str) -> AppAction:
+        """
+        Get action based on it's associated string
+        :param action_string: the string of the action (element to use)
+        :return: the actual AppAction
+        """
+        for action in list(self.actions):
+            if action_string in str(action):
+                return action
+
     def sample(self) -> AppAction:
+        # return self.get_action_based_on_string('Manual scanning')
         return random.choice(list(self.actions))
 
     def contains(self, x: str) -> bool:

change_analyzer/wrappers/sequence_recorder.py

@@ -26,9 +26,13 @@ class SequenceRecorder(Wrapper, TransparentWrapperMixin):
 
     def __init__(self, env: gym.Env, directory: str, sequence_id: str) -> None:
         super(SequenceRecorder, self).__init__(env)
-        self._sequence_id = sequence_id
-        os.makedirs(directory, exist_ok=True)
-        self._csv_file = f"{directory}/{sequence_id}.csv"
+        # The ifs regarding sequence_id and directory are in place, to allow independent access to sequence_recorder
+        # functions, for example to get enriched page_source within Explorer agent - therefore we can have empty strings
+        if sequence_id:
+            self._sequence_id = sequence_id
+        if directory:
+            os.makedirs(directory, exist_ok=True)
+            self._csv_file = f"{directory}/{sequence_id}.csv"
 
     def step(self, action: AppAction) -> Tuple[Dict, float, bool, WebDriver]:
         current_action = str(action)

config_ludwig_model.json

@@ -0,0 +1,48 @@
+{    
+	"input_features": [
+		{
+			"name": "ActionImage", 
+			"type": "image",
+			"preprocessing": {
+				"height": 150,
+				"width": 150,
+				"resize_method": "interpolate",
+				"scaling": "pixel_normalization"
+			}
+		},
+		{
+			"name": "ActionText",
+			"type": "text"
+		},
+		{
+			"name": "DistanceFromCenter",
+			"type": "numerical"
+		},
+		{
+			"name": "DistanceFromTopLeftCorner",
+			"type": "numerical"
+		},
+		{
+			"name": "PreviousSteps",
+			"type": "text"
+		},
+		{
+			"name": "ActionImageMainColorR",
+			"type": "numerical"
+		},
+		{
+			"name": "ActionImageMainColorG",
+			"type": "numerical"
+		},
+		{
+			"name": "ActionImageMainColorB",
+			"type": "numerical"
+		}
+	],
+	"output_features": [
+		{
+			"name": "Reward",
+			"type": "numerical"
+		}
+	]
+}

setup.py

@@ -24,7 +24,7 @@ def _read_long_description():
     "xmldiff",
     "beautifulsoup4",
     "lxml",
-    "ludwig==0.5.2",
+    "ludwig==0.5.5",
     "torch==1.11.0",
     "torchaudio==0.11.0",
     "torchinfo==1.7.0",