Merge pull request #103 from stratosphereips/sebas-fix-conceptual-agent

Sebas fix conceptual agent

Author: eldraco

.gitignore

@@ -153,4 +153,4 @@ agents/mlruns*
 agents/*/*/mlruns/
 agents/*/*/logs
 aim/*
-wandb/
+wandb

README.md

@@ -2,25 +2,40 @@
 Agents located in this repository should be used in the [Network Security Game](https://github.com/stratosphereips/NetSecGame) environment. They are intended for navigation and problem solving in the adversarial network-security based environment where they play the role of attackers or defenders.
 
 ## Installation
-We recommend to use virtual environment when installing the agents:
+Agents need their own set of libraries which are installed separatedly from the AiDojo environment.
+
+To run an agent you need to install
+- The library of the AIDojoCoordinator
+- The libraries needed by your agent
+
+We recommend to use virtual environment when installing.
+
 ```bash
 python -m venv aidojo-agents
 ```
+
 To activat the venv, run:
 ```
 source aidojo-agents/bin/activate
 ```
-This project requires components of the [Network Security Game](https://github.com/stratosphereips/NetSecGame) to run properly so make sure it is installed first.
 
-To install the all agents, run 
-```
-pip install .
-```
-It is possible to install only subset of agents with following command:
+Be sure you are in the directory of this _NetSecGameAgents_ repository.
+
+### Install the libraries of the AiDojoCoordinator
+Agents requires components of the [NeSecGame](https://github.com/stratosphereips/NetSecGame) to run properly so make sure it is installed first.
+The code for NetSecGame is assumed to be in the previous directory
+
+- `python -m pip install -e ..`
+
+To install the required packages for each agent, you can run 
 ```
-pip install -e .[<name-of-the-agent>] 
+python -m pip install -e .[<name-of-the-agent>] 
 ```
-For example `pip install -e .[tui,llm]`
+
+For example `python -m pip install -e ".[tui,llm]"`
+
+For a complete list of agents to install the dependencies see the pyproject.toml file.
+
 
 ## Runing the agent
 To run the agents, use
@@ -79,6 +94,7 @@ Agents of each type are stored in the corresponding directory within this reposi
     ├── benign
         ├── benign_random
 ```
+
 ### Agent utils
 Utility functions in [`agent_utils.py`](./agents/agent_utils.py) can be used by any agent to evaluate a `GameState`, and generate a set of valid `Actions` in a `GameState`, etc. 
 Additionally, there are several files with utils functions that can be used by any agents:
@@ -117,12 +133,5 @@ If you want to export the local mlflow to a remote mlflow you can use our util
 python utils/export_import_mlflow_exp.py --experiment_id 783457873620024898 --run_id 5f2e4a205b7745259a4ddedc12d71a74 --remote_mlflow_url http://127.0.0.1:8000 --mlruns_dir ./mlruns
 ```
 
-## Install
-
-- create new env
-- install numpy
-- install coor `pip install -e ..`
-- optionally install mlflow
-
 ## About us
 This code was developed at the [Stratosphere Laboratory at the Czech Technical University in Prague](https://www.stratosphereips.org/) as part of the [AIDojo Project](https://www.stratosphereips.org/ai-dojo).

agents/agent_utils.py

@@ -0,0 +1,18 @@
+# Conceptual Attacker Agent
+
+The conceptual attacker agent is a modification to the Q-learning attacker to avoid depending on IP addresses to play the game, and instead convert each IP address into a concept, just as humans do when they attack a network.
+
+# Install
+Install the dependencies of this agent with 
+
+```python -m venv venv
+source venv/bin/activate
+python -m pip install -e ".[conceptual_q_learning]"
+```
+
+# Run the Agent
+If the NetSecGame server is running in localhost, port 9000/TCP, then:
+
+```
+python -m agents.attackers.conceptual_q_learning.q_agent --host localhost --port 9000 --episodes 1 --experiment_id test-1 --env_conf ../AIDojoCoordinator/netsecenv_conf.yaml
+```

agents/attackers/conceptual_q_learning/README.md

@@ -0,0 +1,68 @@
+import argparse
+import pickle
+from colorama import Fore, init
+
+#q_values = {}
+#states = {}
+
+def load_q_table():
+    global q_values
+    global states
+    print(f'Loading file {args.file}')
+    with open(args.file, "rb") as f:
+        data = pickle.load(f)
+        q_values = data["q_table"]
+        states = data["state_mapping"]
+    print(f'Len of qtable: {len(q_values)}')
+
+def show_q_table():
+    """
+    Show details about a state in the qtable
+    """
+    # Get max valid state id
+    max_state = len(states) - 1
+    
+    # Validate state range
+    if args.state_id > max_state:
+        print(f"Error: state_id {args.state_id} is out of range. Max state is {max_state}")
+        return
+    
+    last_state = min(args.last_state_id, max_state) if args.last_state_id > 0 else args.state_id
+    
+    print(f"Showing states from {args.state_id} to {last_state} (max available: {max_state})")
+
+    for state in range(args.state_id, last_state + 1):
+        try:
+            print(f'\n-------------------------------------')
+            print(f'State {state}: {list(states.items())[state]}')
+            filtered_items = {key: value for key, value in q_values.items() if key[0] == state}
+
+            sorted_items = dict(sorted(filtered_items.items(), key=lambda item: item[1], reverse=True))
+
+            # Identify the maximum value 
+            max_value = next(iter(sorted_items.values()), None)
+
+            for index, (key, value) in enumerate(sorted_items.items()):
+                if value == max_value:
+                    print(Fore.RED + f'\t{key} -> {value}' + Fore.RESET)
+                else:
+                    if not args.only_top:
+                        print(Fore.GREEN + f'\t{key} -> {value}' + Fore.RESET)
+        except IndexError:
+            print(f"Error: Could not access state {state}")
+            continue
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('You can train the agent, or test it. \n Test is also to use the agent. \n During training and testing the performance is logged.')
+    parser.add_argument("--file", help="Q-table file to load", default="q_agent_marl.pickle", required=False, type=str)
+    parser.add_argument("--state_id", help="ID of the state to print", default=0, required=False, type=int)
+    parser.add_argument("--last_state_id", help="Last ID of the state to print", default=0, required=False, type=int)
+    parser.add_argument("--only_top", help="Print only the top action, the one to be taken if greedy", default=False, required=False, type=bool)
+    args = parser.parse_args()
+
+    # For the colorama
+    init(strip=False)  # Changed from autoreset=True
+
+    load_q_table()
+
+    show_q_table()