"""Module for handling configurable observation spaces in PrimAITE.""" import logging from abc import ABC, abstractmethod from enum import Enum from typing import List, Tuple import numpy as np from gym import spaces from primaite.common.enums import FileSystemState, HardwareState, SoftwareState from primaite.environment.primaite_env import Primaite from primaite.nodes.active_node import ActiveNode from primaite.nodes.service_node import ServiceNode _LOGGER = logging.getLogger(__name__) class AbstractObservationComponent(ABC): """Represents a part of the PrimAITE observation space.""" @abstractmethod def __init__(self, env: Primaite): _LOGGER.info(f"Initialising {self} observation component") self.env: Primaite = env self.space: spaces.Space self.current_observation: np.ndarray # type might be too restrictive? return NotImplemented @abstractmethod def update(self): """Look at the environment and update the current observation value.""" self.current_observation = NotImplemented class NodeLinkTable(AbstractObservationComponent): """Table with nodes/links as rows and hardware/software status as cols. #todo: write full description """ _FIXED_PARAMETERS = 4 _MAX_VAL = 1_000_000 _DATA_TYPE = np.int64 def __init__(self, env: Primaite): super().__init__() # 1. Define the shape of your observation space component num_items = self.env.num_links + self.env.num_nodes num_columns = self.env.num_services + self._FIXED_PARAMETERS observation_shape = (num_items, num_columns) # 2. Create Observation space self.space = spaces.Box( low=0, high=self._MAX_VAL, shape=observation_shape, dtype=self._DATA_TYPE, ) # 3. Initialise Observation with zeroes self.current_observation = np.zeroes(observation_shape, dtype=self._DATA_TYPE) def update_obs(self): """Update the observation. todo: complete description.. """ item_index = 0 nodes = self.env.nodes links = self.env.links # Do nodes first for _, node in nodes.items(): self.current_observation[item_index][0] = int(node.node_id) self.current_observation[item_index][1] = node.hardware_state.value if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): self.current_observation[item_index][2] = node.software_state.value self.current_observation[item_index][ 3 ] = node.file_system_state_observed.value else: self.current_observation[item_index][2] = 0 self.current_observation[item_index][3] = 0 service_index = 4 if isinstance(node, ServiceNode): for service in self.env.services_list: if node.has_service(service): self.current_observation[item_index][ service_index ] = node.get_service_state(service).value else: self.current_observation[item_index][service_index] = 0 service_index += 1 else: # Not a service node for service in self.env.services_list: self.current_observation[item_index][service_index] = 0 service_index += 1 item_index += 1 # Now do links for _, link in links.items(): self.current_observation[item_index][0] = int(link.get_id()) self.current_observation[item_index][1] = 0 self.current_observation[item_index][2] = 0 self.current_observation[item_index][3] = 0 protocol_list = link.get_protocol_list() protocol_index = 0 for protocol in protocol_list: self.current_observation[item_index][ protocol_index + 4 ] = protocol.get_load() protocol_index += 1 item_index += 1 class NodeStatuses(AbstractObservationComponent): """todo: complete description.""" _DATA_TYPE = np.int64 def __init__(self): super().__init__() # 1. Define the shape of your observation space component shape = [ len(HardwareState) + 1, len(SoftwareState) + 1, len(FileSystemState) + 1, ] services_shape = [len(SoftwareState) + 1] * self.env.num_services shape = shape + services_shape # 2. Create Observation space self.space = spaces.MultiDiscrete(shape) # 3. Initialise observation with zeroes self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE) def update_obs(self): """todo: complete description.""" obs = [] for _, node in self.env.nodes.items(): hardware_state = node.hardware_state.value software_state = 0 file_system_state = 0 service_states = [0] * self.env.num_services if isinstance(node, ActiveNode): software_state = node.software_state.value file_system_state = node.file_system_state_observed.value if isinstance(node, ServiceNode): for i, service in enumerate(self.env.services_list): if node.has_service(service): service_states[i] = node.get_service_state(service).value obs.extend([hardware_state, software_state, file_system_state, *service_states]) self.current_observation[:] = obs class LinkTrafficLevels(AbstractObservationComponent): """todo: complete description.""" _DATA_TYPE = np.int64 def __init__( self, combine_service_traffic: bool = False, quantisation_levels: int = 5 ): super().__init__() self._combine_service_traffic: bool = combine_service_traffic self._quantisation_levels: int = quantisation_levels self._entries_per_link: int = 1 if not self._combine_service_traffic: self._entries_per_link = self.env.num_services # 1. Define the shape of your observation space component shape = ( [self._quantisation_levels] * self.env.num_links * self._entries_per_link ) # 2. Create Observation space self.space = spaces.MultiDiscrete(shape) # 3. Initialise observation with zeroes self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE) def update_obs(self): """todo: complete description.""" obs = [] for _, link in self.env.links.items(): bandwidth = link.bandwidth if self._combine_service_traffic: loads = [link.get_current_load()] else: loads = [protocol.get_load() for protocol in link.protocol_list] for load in loads: if load <= 0: traffic_level = 0 elif load >= bandwidth: traffic_level = self._quantisation_levels - 1 else: traffic_level = (load / bandwidth) // ( 1 / (self._quantisation_levels - 1) ) + 1 obs.append(int(traffic_level)) self.current_observation[:] = obs class ObservationsHandler: """todo: complete description.""" class registry(Enum): """todo: complete description.""" NODE_LINK_TABLE: NodeLinkTable NODE_STATUSES: NodeStatuses LINK_TRAFFIC_LEVELS: LinkTrafficLevels def __init__(self): """todo: complete description.""" """Initialise the handler without any components yet. They""" self.registered_obs_components: List[AbstractObservationComponent] = [] self.space: spaces.Space self.current_observation: Tuple[np.ndarray] # i can access the registry items like this: # self.registry.LINK_TRAFFIC_LEVELS def update_obs(self): """todo: complete description.""" current_obs = [] for obs in self.registered_obs_components: obs.update_obs() current_obs.append(obs.current_observation) self.current_observation = tuple(current_obs) def register(self, obs_component: AbstractObservationComponent): """todo: complete description.""" self.registered_obs_components.append(obs_component) self.update_space() def deregister(self, obs_component: AbstractObservationComponent): """todo: complete description.""" self.registered_obs_components.remove(obs_component) self.update_space() def update_space(self): """todo: complete description.""" component_spaces = [] for obs_comp in self.registered_obs_components: component_spaces.append(obs_comp.space) self.space = spaces.Tuple(component_spaces)