Skip to content

Game Coordinator

Coordinator is the centerpiece of the game orchestration. It provides an interface between the agents and the worlds.

In detail it handles:

World initialiazation
Registration of new agents in the game
Agent-World communication (message verification and forwarding)
Recording (and storing) trajectories of agents (optional)
Detection of episode ends (either by reaching timout or agents reaching their respective goals)
Assigning rewards for each action and at the end of each episode
Removing agents from the game
Registering the GameReset requests and handelling the game resets.

To facilitate the communication the coordinator uses a TCP server to which agents connect. The communication is asynchronous and depends of the

`AIDojoCoordinator.coordinator.AgentServer(actions_queue, agent_response_queues, max_connections)`

Bases: Protocol

Class used for serving the agents when connecting to the game run by the GameCoordinator.

Attributes:

Name	Type	Description
`actions_queue`	`Queue`	Queue for actions from agents.
`answers_queues`	`dict`	Mapping of agent addresses to their response queues.
`max_connections`	`int`	Maximum allowed concurrent agent connections.
`current_connections`	`int`	Current number of connected agents.
`logger`	`Logger`	Logger for the AgentServer.

Initialize the AgentServer.

Parameters:

Name	Type	Description	Default
`actions_queue`	`Queue`	Queue for actions from agents.	required
`agent_response_queues`	`dict`	Mapping of agent addresses to their response queues.	required
`max_connections`	`int`	Maximum allowed concurrent agent connections.	required

Source code in AIDojoCoordinator/coordinator.py

def __init__(self, actions_queue, agent_response_queues, max_connections):
    """
    Initialize the AgentServer.

    Args:
        actions_queue (asyncio.Queue): Queue for actions from agents.
        agent_response_queues (dict): Mapping of agent addresses to their response queues.
        max_connections (int): Maximum allowed concurrent agent connections.
    """
    self.actions_queue = actions_queue
    self.answers_queues = agent_response_queues
    self.max_connections = max_connections
    self.current_connections = 0
    self.logger = logging.getLogger("AIDojo-AgentServer")

`call(reader, writer)` `async`

Allow the server instance to be called as a coroutine.

Parameters:

Name	Type	Description	Default
`reader`	`StreamReader`	Stream reader for the agent.	required
`writer`	`StreamWriter`	Stream writer for the agent.	required

Source code in AIDojoCoordinator/coordinator.py

async def __call__(self, reader, writer):
    """
    Allow the server instance to be called as a coroutine.

    Args:
        reader (asyncio.StreamReader): Stream reader for the agent.
        writer (asyncio.StreamWriter): Stream writer for the agent.
    """
    await self.handle_new_agent(reader, writer)

`handle_agent_quit(peername)` `async`

Helper function to handle agent disconnection.

Parameters:

Name	Type	Description	Default
`peername`	`tuple`	The address of the disconnecting agent.	required

Source code in AIDojoCoordinator/coordinator.py

async def handle_agent_quit(self, peername:tuple):
    """
    Helper function to handle agent disconnection.

    Args:
        peername (tuple): The address of the disconnecting agent.
    """
    # Send a quit message to the Coordinator
    self.logger.info(f"\tHandling agent quit for {peername}.")
    quit_message = Action(ActionType.QuitGame, parameters={}).to_json()
    await self.actions_queue.put((peername, quit_message))

`handle_new_agent(reader, writer)` `async`

Handle a new agent connection.

Parameters:

Name	Type	Description	Default
`reader`	`StreamReader`	Stream reader for the agent.	required
`writer`	`StreamWriter`	Stream writer for the agent.	required

Source code in AIDojoCoordinator/coordinator.py

async def handle_new_agent(self, reader, writer):
    """
    Handle a new agent connection.

    Args:
        reader (asyncio.StreamReader): Stream reader for the agent.
        writer (asyncio.StreamWriter): Stream writer for the agent.
    """
    # get the peername of the writer
    peername = writer.get_extra_info("peername")
    queue_created = False
    try:
        self.logger.info(f"New connection from {peername}")
        # Check if the maximum number of connections has been reached
        if self.current_connections < self.max_connections:
            # increment the count of current connections
            self.current_connections += 1
            self.logger.info(f"New agent connected: {peername}. Current connections: {self.current_connections}")
            # Ensure a queue exists for this agent
            if peername not in self.answers_queues:
                self.answers_queues[peername] = asyncio.Queue(maxsize=2)
                queue_created = True
                self.logger.info(f"Created queue for agent {peername}")
                # Handle the new agent
                while True:
                    # Step 1: Read data from the agent
                    data = await reader.read(ProtocolConfig.BUFFER_SIZE)
                    if not data:
                        self.logger.info(f"Agent {peername} disconnected.")
                        await self.handle_agent_quit(peername)
                        break

                    raw_message = data.decode().strip()
                    self.logger.debug(f"Handler received from {peername}: {raw_message}")

                    # Step 2: Forward the message to the Coordinator
                    await self.actions_queue.put((peername, raw_message))

                    # Step 3: Get a matching response from the answers queue
                    response_queue = self.answers_queues[peername]
                    response = await response_queue.get()
                    self.logger.info(f"Sending response to agent {peername}: {response}")

                    # Step 4: Send the response to the agent
                    response = str(response).encode() + ProtocolConfig.END_OF_MESSAGE
                    writer.write(response)
                    await writer.drain()
            else:
                self.logger.warning(f"Queue for agent {peername} already exists. Closing connection.")
        else:
            self.logger.info(f"Max connections reached. Rejecting new connection from {writer.get_extra_info('peername')}")
    except ConnectionResetError:
        self.logger.warning(f"Connection reset by {peername}")
        await self.handle_agent_quit(peername)
    except asyncio.CancelledError:
        self.logger.debug("Connection handling cancelled.")
        raise  # Ensure the exception propagates
    except Exception as e:
        self.logger.error(f"Unexpected error with client {peername}: {e}")
        raise
    finally:
        try:
            if peername in self.answers_queues:
                # If the queue was created, remove it
                if queue_created:
                    self.answers_queues.pop(peername)
                    self.logger.info(f"Removed queue for agent {peername}")
                self.current_connections = max(0, self.current_connections - 1)
            writer.close()
            await writer.wait_closed()
        except Exception:
            # swallow exceptions on close to avoid crash on cleanup
            pass

`AIDojoCoordinator.coordinator.GameCoordinator(game_host, game_port, service_host, service_port, allowed_roles=['Attacker', 'Defender', 'Benign'], task_config_file=None)`

Class for creation, and management of agent interactions in AI Dojo.

Source code in AIDojoCoordinator/coordinator.py

def __init__(self, game_host: str, game_port: int, service_host:str, service_port:int, allowed_roles=["Attacker", "Defender", "Benign"], task_config_file:str=None) -> None:
    self.host = game_host
    self.port = game_port
    self.logger = logging.getLogger("AIDojo-GameCoordinator")

    self._tasks = set()
    self.shutdown_flag = asyncio.Event()
    self._reset_event = asyncio.Event()
    self._episode_end_event = asyncio.Event()
    self._episode_start_event = asyncio.Event()
    self._episode_rewards_condition = asyncio.Condition()
    self._reset_done_condition = asyncio.Condition()
    self._reset_lock = asyncio.Lock()
    self._agents_lock = asyncio.Lock()

    # for accessing configuration remotely
    self._service_host = service_host
    self._service_port = service_port
    # for reading configuration locally
    self._task_config_file = task_config_file
    self.logger = logging.getLogger("AIDojo-GameCoordinator")
    self.ALLOWED_ROLES = allowed_roles
    self._cyst_objects = None
    self._cyst_object_string = None

    # prepare agent communication
    self._agent_action_queue = asyncio.Queue()
    self._agent_response_queues = {}

    # agent information
    self.agents = {}
    # step counter per agent_addr (int)
    self._agent_steps = {}
    # reset request per agent_addr (bool)
    self._reset_requests = {}
    self._agent_status = {}
    self._episode_ends = {}
    self._agent_observations = {}
    # starting per agent_addr (dict)
    self._agent_starting_position = {}
    # current state per agent_addr (GameState)
    self._agent_states = {}
    # last action played by agent (Action)
    self._agent_last_action = {}
    # False positives per agent (due to added blocks)
    self._agent_false_positives = {}
    # agent status dict {agent_addr: int}
    self._agent_rewards = {}
    # trajectories per agent_addr
    self._agent_trajectories = {}
    # false_positives per agent_addr
    self._agent_false_positives = {}

`add_false_positive(agent)`

Method for adding false positive to the agent.

Source code in AIDojoCoordinator/coordinator.py

def add_false_positive(self, agent:tuple)->None:
    """
    Method for adding false positive to the agent.
    """
    self.logger.debug(f"Adding false positive to {agent}")
    if agent in self._agent_false_positives:
        self._agent_false_positives[agent] += 1
    else:
        self._agent_false_positives[agent] = 1
    self.logger.debug(f"False positives for {agent}: {self._agent_false_positives[agent]}")

`convert_msg_dict_to_json(msg_dict)`

Helper function to create text-base messge from a dictionary. Used in the Agent-Game communication.

Source code in AIDojoCoordinator/coordinator.py

def convert_msg_dict_to_json(self, msg_dict:dict)->str:
    """
    Helper function to create text-base messge from a dictionary. Used in the Agent-Game communication.
    """
    try:
        # Convert message into string representation
        output_message = json.dumps(msg_dict)
    except Exception as e:
        self.logger.error(f"Error when converting msg to JSON:{e}")
        raise e
        # Send to anwer_queue
    return output_message

`create_agent_queue(agent_addr)` `async`

Creates a queue for the given agent address if it doesn't already exist.

Source code in AIDojoCoordinator/coordinator.py

async def create_agent_queue(self, agent_addr:tuple)->None:
    """
    Creates a queue for the given agent address if it doesn't already exist.
    """
    if agent_addr not in self._agent_response_queues:
        self._agent_response_queues[agent_addr] = asyncio.Queue()
        self.logger.info(f"Created queue for agent {agent_addr}. {len(self._agent_response_queues)} queues in total.")

`goal_check(agent_addr)`

Check if the goal conditons were satisfied in a given game state

Source code in AIDojoCoordinator/coordinator.py

def goal_check(self, agent_addr:tuple)->bool:
    """
    Check if the goal conditons were satisfied in a given game state
    """
    def goal_dict_satistfied(goal_dict:dict, known_dict: dict)-> bool:
        """
        Helper function for checking if a goal dictionary condition is satisfied
        """
        # check if we have all IPs that should have some values (are keys in goal_dict)
        if goal_dict.keys() <= known_dict.keys():
            try:
                # Check if values (sets) for EACH key (host) in goal_dict are subsets of known_dict, keep matching_keys
                matching_keys = [host for host in goal_dict.keys() if goal_dict[host]<= known_dict[host]]
                # Check we have the amount of mathing keys as in the goal_dict
                if len(matching_keys) == len(goal_dict.keys()):
                    return True
            except KeyError:
                # some keys are missing in the known_dict
                return False
        return False
    self.logger.debug(f"Checking goal for agent {agent_addr}.")
    goal_conditions = self._win_conditions_per_role[self.agents[agent_addr][1]]
    state = self._agent_states[agent_addr]
    # For each part of the state of the game, check if the conditions are met
    goal_reached = {}    
    goal_reached["networks"] = set(goal_conditions["known_networks"]) <= set(state.known_networks)
    goal_reached["known_hosts"] = set(goal_conditions["known_hosts"]) <= set(state.known_hosts)
    goal_reached["controlled_hosts"] = set(goal_conditions["controlled_hosts"]) <= set(state.controlled_hosts)
    goal_reached["services"] = goal_dict_satistfied(goal_conditions["known_services"], state.known_services)
    goal_reached["data"] = goal_dict_satistfied(goal_conditions["known_data"], state.known_data)
    goal_reached["known_blocks"] = goal_dict_satistfied(goal_conditions["known_blocks"], state.known_blocks)
    self.logger.debug(f"\t{goal_reached}")
    return all(goal_reached.values())

`is_agent_benign(agent_addr)`

Check if the agent is benign (defender, normal)

Source code in AIDojoCoordinator/coordinator.py

def is_agent_benign(self, agent_addr:tuple)->bool:
    """
    Check if the agent is benign (defender, normal)
    """
    if agent_addr not in self.agents:
        return False
    return self.agents[agent_addr][1].lower() in ["defender", "benign"]

`register_agent(agent_id, agent_role, agent_initial_view)` `async`

Domain specific method of the environment. Creates the initial state of the agent.

Source code in AIDojoCoordinator/coordinator.py

async def register_agent(self, agent_id:tuple, agent_role:str, agent_initial_view:dict)->GameState:
    """
    Domain specific method of the environment. Creates the initial state of the agent.
    """
    raise NotImplementedError

`remove_agent(agent_id, agent_state)` `async`

Domain specific method of the environment. Creates the initial state of the agent.

Source code in AIDojoCoordinator/coordinator.py

async def remove_agent(self, agent_id:tuple, agent_state:GameState)->bool:
    """
    Domain specific method of the environment. Creates the initial state of the agent.
    """
    raise NotImplementedError

`run()`

Wrapper for ayncio run function. Starts all tasks in AIDojo

Source code in AIDojoCoordinator/coordinator.py

def run(self)->None:
    """
    Wrapper for ayncio run function. Starts all tasks in AIDojo
    """
    try:
        asyncio.run(self.start_tasks())
    except Exception as e:
        self.logger.error(f"Unexpected error: {e}")
    finally:
        self.logger.info(f"{__class__.__name__} has exited.")

`run_game()` `async`

Task responsible for reading messages from the agent queue and processing them based on the ActionType.

Source code in AIDojoCoordinator/coordinator.py

async def run_game(self):
    """
    Task responsible for reading messages from the agent queue and processing them based on the ActionType.
    """
    while not self.shutdown_flag.is_set():
        # Read message from the queue
        agent_addr, message = await self._agent_action_queue.get()
        if message is not None:
            self.logger.info(f"Coordinator received from agent {agent_addr}: {message}.")
            try:  # Convert message to Action
                action = Action.from_json(message)
                self.logger.debug(f"\tConverted to: {action}.")
            except Exception as e:
                self.logger.error(
                    f"Error when converting msg to Action using Action.from_json():{e}, {message}"
                )
            match action.type:  # process action based on its type
                case ActionType.JoinGame:
                    self.logger.debug(f"About agent {agent_addr}. Start processing of ActionType.JoinGame by {agent_addr}")
                    self.logger.debug(f"{action.type}, {action.type.value}, {action.type == ActionType.JoinGame}")
                    self._spawn_task(self._process_join_game_action, agent_addr, action)
                case ActionType.QuitGame:
                    self.logger.debug(f"About agent {agent_addr}. Start processing of ActionType.QuitGame by {agent_addr}")
                    self._spawn_task(self._process_quit_game_action, agent_addr)
                case ActionType.ResetGame:
                    self.logger.debug(f"About agent {agent_addr}. Start processing of ActionType.ResetGame by {agent_addr}")
                    self._spawn_task(self._process_reset_game_action, agent_addr, action)
                case ActionType.ExfiltrateData | ActionType.FindData | ActionType.ScanNetwork | ActionType.FindServices | ActionType.ExploitService:
                    self.logger.debug(f"About agent {agent_addr}. Start processing of {action.type} by {agent_addr}")
                    self._spawn_task(self._process_game_action, agent_addr, action)
                case ActionType.BlockIP:
                    self.logger.debug(f"About agent {agent_addr}. Start processing of {action.type} by {agent_addr}")
                    self._spawn_task(self._process_game_action, agent_addr, action)
                case _:
                    self.logger.warning(f"About agent {agent_addr}. Unsupported action type: {action}!")
    self.logger.info("\tAction processing task stopped.")

`shutdown_signal_handler()` `async`

Handle shutdown signals.

Source code in AIDojoCoordinator/coordinator.py

async def shutdown_signal_handler(self):
    """Handle shutdown signals."""
    self.logger.info("Shutdown signal received. Setting shutdown flag.")
    self.shutdown_flag.set()

`start_tasks()` `async`

High level funciton to start all the other asynchronous tasks. - Reads the conf of the coordinator - Creates queues - Start the main part of the coordinator - Start a server that listens for agents

Source code in AIDojoCoordinator/coordinator.py

async def start_tasks(self):
    """
    High level funciton to start all the other asynchronous tasks.
    - Reads the conf of the coordinator
    - Creates queues
    - Start the main part of the coordinator
    - Start a server that listens for agents
    """
    loop = asyncio.get_running_loop()

    # Set up signal handlers for graceful shutdown
    loop.add_signal_handler(
        signal.SIGINT, lambda: asyncio.create_task(self.shutdown_signal_handler())
    )
    loop.add_signal_handler(
        signal.SIGTERM, lambda: asyncio.create_task(self.shutdown_signal_handler())
    )


    # initialize the game objects
    if self._service_host: #get the task config using REST API
        self.logger.info(f"Fetching task configuration from {self._service_host}:{self._service_port}")
        await self._fetch_initialization_objects()
    elif self._task_config_file: # load task config locally from a file
        self.logger.info(f"Loading task configuration from file: {self._task_config_file}")
        self._load_initialization_objects()
    else:
        raise ValueError("Task configuration not specified")


    # Read configuration
    self._starting_positions_per_role = self._get_starting_position_per_role()
    self._win_conditions_per_role = self._get_win_condition_per_role()
    self._goal_description_per_role = self._get_goal_description_per_role()
    self._steps_limit_per_role = self._get_max_steps_per_role()
    self.logger.debug(f"Timeouts set to:{self._steps_limit_per_role}")
    if self.task_config.get_use_global_defender():
        self._global_defender = GlobalDefender()
    else:
        self._global_defender = None
    self._use_dynamic_ips = self.task_config.get_use_dynamic_addresses()
    self.logger.info(f"Change IP every episode set to: {self._use_dynamic_ips}")
    self._rewards = self.task_config.get_rewards(["step", "success", "fail", "false_positive"])
    self.logger.info(f"Rewards set to:{self._rewards}")
    self._min_required_players = self.task_config.get_required_num_players()
    self.logger.info(f"Min player requirement set to:{self._min_required_players}")
    # start server for agent communication
    self._spawn_task(self.start_tcp_server)

    # start episode rewards task
    self._spawn_task(self._assign_rewards_episode_end)

    # start episode rewards task
    self._spawn_task(self._reset_game)

    # start action processing task
    self._spawn_task(self.run_game)

    while not self.shutdown_flag.is_set():
        # just wait until user terminates
        await asyncio.sleep(1)
    self.logger.debug("Final cleanup started")
    # make sure there are no running tasks left
    for task in self._tasks:
        task.cancel()  # Cancel each active task
    await asyncio.gather(*self._tasks, return_exceptions=True)  # Wait for all tasks to finish
    self.logger.info("All tasks shut down.")

`start_tcp_server()` `async`

Starts TPC sever for the agent communication.

Source code in AIDojoCoordinator/coordinator.py

async def start_tcp_server(self):
    """
    Starts TPC sever for the agent communication.
    """
    try:
        self.logger.info("Starting the server listening for agents")
        server = await asyncio.start_server(
            AgentServer(
                self._agent_action_queue,
                self._agent_response_queues,
                max_connections=self._min_required_players
            ),
            self.host,
            self.port
        )
        addrs = ", ".join(str(sock.getsockname()) for sock in server.sockets)
        self.logger.info(f"\tServing on {addrs}")
        while not self.shutdown_flag.is_set():
            await asyncio.sleep(1)
    except asyncio.CancelledError:
        self.logger.debug("\tStopping TCP server task.")
    except Exception as e:
        self.logger.error(f"TCP server failed: {e}")
    finally:
        server.close()
        await server.wait_closed()
        self.logger.info("\tTCP server task stopped")