From 2fe46a7b9b6587117ec949373d137c947bcd28a8 Mon Sep 17 00:00:00 2001 From: Ilya Orson Sandoval Date: Tue, 5 Mar 2024 15:57:13 +0000 Subject: [PATCH 1/4] ignore evaluation files --- .gitignore | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 8a6d3516..c01d503a 100644 --- a/.gitignore +++ b/.gitignore @@ -31,4 +31,6 @@ CybORG/Emulator/AWS/Config/awsconfig.ini *.swp *.pstats -.ipynb_checkpoints/ \ No newline at end of file +.ipynb_checkpoints/ + +CybORG/Evaluation/ From ac3e1915c60d93f4e44e4ff183d46864d8c91177 Mon Sep 17 00:00:00 2001 From: Ilya Orson Sandoval Date: Tue, 12 Mar 2024 11:58:55 +0000 Subject: [PATCH 2/4] Update to gymnasium --- CybORG/Agents/Wrappers/ChallengeWrapper.py | 8 ++++---- CybORG/Agents/Wrappers/OpenAIGymWrapper.py | 10 +++++++--- CybORG/Agents/training_example.py | 3 ++- CybORG/Tests/test_sim/test_sim_Cyborg.py | 3 ++- .../test_wrappers/test_ChallengeWrapper.py | 3 ++- .../test_wrappers/test_OpenAIGymWrapper.py | 17 +++++++++++------ Requirements.txt | 2 +- 7 files changed, 29 insertions(+), 17 deletions(-) diff --git a/CybORG/Agents/Wrappers/ChallengeWrapper.py b/CybORG/Agents/Wrappers/ChallengeWrapper.py index 9e043a7f..58af7c23 100644 --- a/CybORG/Agents/Wrappers/ChallengeWrapper.py +++ b/CybORG/Agents/Wrappers/ChallengeWrapper.py @@ -1,4 +1,4 @@ -from gym import Env +from gymnasium import Env from CybORG.Agents.Wrappers import BaseWrapper, OpenAIGymWrapper, BlueTableWrapper,RedTableWrapper,EnumActionWrapper @@ -26,13 +26,13 @@ def __init__(self, agent_name: str, env, agent=None, self.step_counter = None def step(self,action=None): - obs, reward, done, info = self.env.step(action=action) + obs, reward, terminated, truncated, info = self.env.step(action=action) self.step_counter += 1 if self.max_steps is not None and self.step_counter >= self.max_steps: - done = True + truncated = True - return obs, reward, done, info + return obs, reward, terminated, truncated, info def reset(self): self.step_counter = 0 diff --git a/CybORG/Agents/Wrappers/OpenAIGymWrapper.py b/CybORG/Agents/Wrappers/OpenAIGymWrapper.py index 7eac1dfb..56e347bc 100644 --- a/CybORG/Agents/Wrappers/OpenAIGymWrapper.py +++ b/CybORG/Agents/Wrappers/OpenAIGymWrapper.py @@ -1,5 +1,5 @@ import numpy as np -from gym import spaces, Env +from gymnasium import spaces, Env from typing import Union, List from prettytable import PrettyTable @@ -22,13 +22,17 @@ def __init__(self, agent_name: str, env: BaseWrapper = None, agent: BaseAgent = self.metadata = {} self.action = None - def step(self, action: Union[int, List[int]] = None) -> (object, float, bool, dict): + # 'gym' done = 'gymnasium' terminated or truncated + # https://gymnasium.farama.org/content/migration-guide/ + def step(self, action: Union[int, List[int]] = None) -> tuple[object, float, bool, dict]: self.action = action result = self.env.step(self.agent_name, action) result.observation = self.observation_change(result.observation) result.action_space = self.action_space_change(result.action_space) + terminated = result.done + truncated = False info = vars(result) - return np.array(result.observation), result.reward, result.done, info + return np.array(result.observation), result.reward, terminated, truncated, info def reset(self, agent=None): result = self.env.reset(self.agent_name) diff --git a/CybORG/Agents/training_example.py b/CybORG/Agents/training_example.py index 720c7644..278957c6 100644 --- a/CybORG/Agents/training_example.py +++ b/CybORG/Agents/training_example.py @@ -28,7 +28,8 @@ def run_training_example(scenario): reward = 0 for j in range(MAX_STEPS_PER_GAME): # step in 1 game action = agent.get_action(observation, action_space) - next_observation, r, done, info = cyborg.step(action=action) + next_observation, r, terminated, truncated, info = cyborg.step(action=action) + done = terminated or truncated action_space = info.get('action_space') reward += r diff --git a/CybORG/Tests/test_sim/test_sim_Cyborg.py b/CybORG/Tests/test_sim/test_sim_Cyborg.py index 90c9cb58..1d02c1c2 100644 --- a/CybORG/Tests/test_sim/test_sim_Cyborg.py +++ b/CybORG/Tests/test_sim/test_sim_Cyborg.py @@ -267,7 +267,8 @@ def test_agent_train(create_cyborg_sim): reward = 0 for j in range(20): action = agent.get_action(observation, action_space) - next_observation, r, done, info = cyborg.step(action) + next_observation, r, terminated, truncated, info = cyborg.step(action) + done = terminated or truncated action_space = info['action_space'] reward += r diff --git a/CybORG/Tests/test_sim/test_wrappers/test_ChallengeWrapper.py b/CybORG/Tests/test_sim/test_wrappers/test_ChallengeWrapper.py index be2bfc9c..1ddc2b93 100644 --- a/CybORG/Tests/test_sim/test_wrappers/test_ChallengeWrapper.py +++ b/CybORG/Tests/test_sim/test_wrappers/test_ChallengeWrapper.py @@ -26,7 +26,8 @@ def test_ChallengeWrapper_reset(cyborg): def test_ChallengeWrapper_step(cyborg): cyborg.reset() - obs,reward,done,info = cyborg.step(action=0) + obs,reward,terminated,truncated,info = cyborg.step(action=0) + done = terminated or truncated expected_observation = np.array([0 for x in range(52)]) assert all(obs == expected_observation) assert reward == 0 diff --git a/CybORG/Tests/test_sim/test_wrappers/test_OpenAIGymWrapper.py b/CybORG/Tests/test_sim/test_wrappers/test_OpenAIGymWrapper.py index 8eecec59..a61479c9 100644 --- a/CybORG/Tests/test_sim/test_wrappers/test_OpenAIGymWrapper.py +++ b/CybORG/Tests/test_sim/test_wrappers/test_OpenAIGymWrapper.py @@ -1,6 +1,6 @@ import pytest import inspect -from gym import spaces +from gymnasium import spaces from CybORG import CybORG from CybORG.Agents.Wrappers.OpenAIGymWrapper import OpenAIGymWrapper from CybORG.Agents.Wrappers.FixedFlatWrapper import FixedFlatWrapper @@ -17,7 +17,8 @@ def test_steps(): env=FixedFlatWrapper(EnumActionWrapper(CybORG(path, 'sim')))) cyborg.reset() action = cyborg.action_space.sample() - obs, reward, done, info = cyborg.step(action) + obs, reward, terminated, truncated, info = cyborg.step(action) + done = terminated or truncated # assert isinstance(obs, object) # Redundant because everything in python is an object assert obs is not None @@ -44,7 +45,8 @@ def test_steps_multi_discrete(): env=FixedFlatWrapper(EnumActionWrapper(CybORG(path, 'sim')))) cyborg.reset() action = cyborg.action_space.sample() - obs, reward, done, info = cyborg.step(action) + obs, reward, terminated, truncated, info = cyborg.step(action) + done = terminated or truncated # assert isinstance(obs, object) # Redundant because everything in python is an object assert obs is not None @@ -80,7 +82,8 @@ def test_steps_random(): for i in range(MAX_EPS): for j in range(MAX_STEPS_PER_GAME): action = cyborg.action_space.sample() - obs, rew, done, info = cyborg.step(action) + obs, rew, terminated, truncated, info = cyborg.step(action) + done = terminated or truncated if done or j == MAX_STEPS_PER_GAME-1: break cyborg.reset() @@ -104,7 +107,8 @@ def test_steps_random_multi_discrete(): for i in range(MAX_EPS): for j in range(MAX_STEPS_PER_GAME): action = cyborg.action_space.sample() - obs, rew, done, info = cyborg.step(action) + obs, rew, terminated, truncated, info = cyborg.step(action) + done = terminated or truncated if done or j == MAX_STEPS_PER_GAME-1: break cyborg.reset() @@ -128,7 +132,8 @@ def test_get_observation(cyborg): method_obs = cyborg.get_observation(cyborg.agent_name) assert all(step_obs == method_obs) - step_obs, reward, done, info = cyborg.step() + step_obs, reward, terminated, truncated, info = cyborg.step() + done = terminated or truncated method_obs = cyborg.get_observation(cyborg.agent_name) assert all(step_obs == method_obs) diff --git a/Requirements.txt b/Requirements.txt index f8a8712c..4abf2c93 100644 --- a/Requirements.txt +++ b/Requirements.txt @@ -2,6 +2,6 @@ pytest paramiko pyyaml numpy -gym +gymnasium prettytable docutils From ebc173947fa6afe529097db245695119ac45510d Mon Sep 17 00:00:00 2001 From: Ilya Orson Sandoval Date: Tue, 12 Mar 2024 11:59:33 +0000 Subject: [PATCH 3/4] Add BlueTable wrapper fix from John Cardiff --- CybORG/Agents/Wrappers/BlueTableWrapper.py | 107 ++++++++++----------- 1 file changed, 53 insertions(+), 54 deletions(-) diff --git a/CybORG/Agents/Wrappers/BlueTableWrapper.py b/CybORG/Agents/Wrappers/BlueTableWrapper.py index ff963d11..21d8d565 100644 --- a/CybORG/Agents/Wrappers/BlueTableWrapper.py +++ b/CybORG/Agents/Wrappers/BlueTableWrapper.py @@ -6,9 +6,10 @@ from CybORG.Agents.Wrappers.BaseWrapper import BaseWrapper from CybORG.Agents.Wrappers.TrueTableWrapper import TrueTableWrapper + class BlueTableWrapper(BaseWrapper): - def __init__(self,env=None,agent=None,output_mode='table'): - super().__init__(env,agent) + def __init__(self, env=None, agent=None, output_mode='table'): + super().__init__(env, agent) self.env = TrueTableWrapper(env=env, agent=agent) self.agent = agent @@ -16,41 +17,40 @@ def __init__(self,env=None,agent=None,output_mode='table'): self.output_mode = output_mode self.blue_info = {} - def reset(self, agent='Blue'): + def reset(self, agent='Blue'): result = self.env.reset(agent) obs = result.observation if agent == 'Blue': self._process_initial_obs(obs) - obs = self.observation_change(obs,baseline=True) + obs = self.observation_change(obs, baseline=True) result.observation = obs return result def step(self, agent=None, action=None) -> Results: result = self.env.step(agent, action) - obs = result.observation + obs = result.observation # NOTE The original Blue agent observation. Should be equivalent to self.get_observation("Blue") if agent == 'Blue': obs = self.observation_change(obs) result.observation = obs result.action_space = self.action_space_change(result.action_space) return result - def get_table(self,output_mode='blue_table'): + def get_table(self, output_mode='blue_table'): if output_mode == 'blue_table': return self._create_blue_table(success=None) elif output_mode == 'true_table': return self.env.get_table() - def observation_change(self,observation,baseline=False): + def observation_change(self, observation, baseline=False): obs = observation if type(observation) == dict else observation.data obs = deepcopy(observation) success = obs['success'] - self._process_last_action() anomaly_obs = self._detect_anomalies(obs) if not baseline else obs del obs['success'] # TODO check what info is for baseline info = self._process_anomalies(anomaly_obs) - if baseline: + if baseline: # NOTE: baseline refers to the first observation by the blue agent for host in info: info[host][-2] = 'None' info[host][-1] = 'No' @@ -82,14 +82,14 @@ def _process_initial_obs(self, obs): subnet = interface['Subnet'] ip = str(interface['IP Address']) hostname = host['System info']['Hostname'] - self.blue_info[hostname] = [str(subnet),str(ip),hostname, 'None','No'] + self.blue_info[hostname] = [str(subnet), str(ip), hostname, 'None', 'No'] return self.blue_info def _process_last_action(self): action = self.get_last_action(agent='Blue') if action is not None: name = action.__class__.__name__ - hostname = action.get_params()['hostname'] if name in ('Restore','Remove') else None + hostname = action.get_params()['hostname'] if name in ('Restore', 'Remove') else None if name == 'Restore': self.blue_info[hostname][-1] = 'No' @@ -98,13 +98,14 @@ def _process_last_action(self): if compromised != 'No': self.blue_info[hostname][-1] = 'Unknown' - def _detect_anomalies(self,obs): + def _detect_anomalies(self, obs): if self.baseline is None: - raise TypeError('BlueTableWrapper was unable to establish baseline. This usually means the environment was not reset before calling the step method.') + raise TypeError( + 'BlueTableWrapper was unable to establish baseline. This usually means the environment was not reset before calling the step method.') anomaly_dict = {} - for hostid,host in obs.items(): + for hostid, host in obs.items(): if hostid == 'success': continue @@ -114,7 +115,7 @@ def _detect_anomalies(self,obs): host_anomalies = {} if 'Files' in host: - baseline_files = host_baseline.get('Files',[]) + baseline_files = host_baseline.get('Files', []) anomalous_files = [] for f in host['Files']: if f not in baseline_files: @@ -123,7 +124,7 @@ def _detect_anomalies(self,obs): host_anomalies['Files'] = anomalous_files if 'Processes' in host: - baseline_processes = host_baseline.get('Processes',[]) + baseline_processes = host_baseline.get('Processes', []) anomalous_processes = [] for p in host['Processes']: if p not in baseline_processes: @@ -136,16 +137,18 @@ def _detect_anomalies(self,obs): return anomaly_dict - def _process_anomalies(self,anomaly_dict): + def _process_anomalies(self, anomaly_dict): info = deepcopy(self.blue_info) for hostid, host_anomalies in anomaly_dict.items(): assert len(host_anomalies) > 0 if 'Processes' in host_anomalies: - connection_type = self._interpret_connections(host_anomalies['Processes']) - info[hostid][-2] = connection_type - if connection_type == 'Exploit': - info[hostid][-1] = 'User' - self.blue_info[hostid][-1] = 'User' + # NOTE John Cardiff fix + if "Connections" in host_anomalies['Processes'][-1]: + connection_type = self._interpret_connections(host_anomalies['Processes']) + info[hostid][-2] = connection_type + if connection_type == 'Exploit': + info[hostid][-1] = 'User' + self.blue_info[hostid][-1] = 'User' if 'Files' in host_anomalies: malware = [f['Density'] >= 0.9 for f in host_anomalies['Files']] if any(malware): @@ -154,47 +157,43 @@ def _process_anomalies(self,anomaly_dict): return info - def _interpret_connections(self,activity:list): + def _interpret_connections(self, activity: list): num_connections = len(activity) - ports = set([item['Connections'][0]['local_port'] \ - for item in activity if 'Connections' in item]) + for item in activity if 'Connections' in item]) port_focus = len(ports) remote_ports = set([item['Connections'][0].get('remote_port') \ - for item in activity if 'Connections' in item]) + for item in activity if 'Connections' in item]) if None in remote_ports: remote_ports.remove(None) - if num_connections >= 3 and port_focus >=3: + if num_connections >= 3 and port_focus >= 3: anomaly = 'Scan' elif 4444 in remote_ports: anomaly = 'Exploit' elif num_connections >= 3 and port_focus == 1: anomaly = 'Exploit' - elif 'Service Name' in activity[0]: - anomaly = 'None' else: anomaly = 'Scan' return anomaly # def _malware_analysis(self,obs,hostname): - # anomaly_dict = {hostname: {'Files': []}} - # if hostname in obs: - # if 'Files' in obs[hostname]: - # files = obs[hostname]['Files'] - # else: - # return anomaly_dict - # else: - # return anomaly_dict + # anomaly_dict = {hostname: {'Files': []}} + # if hostname in obs: + # if 'Files' in obs[hostname]: + # files = obs[hostname]['Files'] + # else: + # return anomaly_dict + # else: + # return anomaly_dict - # for f in files: - # if f['Density'] >= 0.9: - # anomaly_dict[hostname]['Files'].append(f) - - # return anomaly_dict + # for f in files: + # if f['Density'] >= 0.9: + # anomaly_dict[hostname]['Files'].append(f) + # return anomaly_dict def _create_blue_table(self, success): table = PrettyTable([ @@ -203,10 +202,10 @@ def _create_blue_table(self, success): 'Hostname', 'Activity', 'Compromised' - ]) - for hostid in self.info: + ]) + for hostid in self.info: # NOTE extracted from the anomalies detection method table.add_row(self.info[hostid]) - + table.sortby = 'Hostname' table.success = success return table @@ -219,11 +218,11 @@ def _create_vector(self, success): # Activity activity = row[3] if activity == 'None': - value = [0,0] + value = [0, 0] elif activity == 'Scan': - value = [1,0] + value = [1, 0] elif activity == 'Exploit': - value = [1,1] + value = [1, 1] else: raise ValueError('Table had invalid Access Level') proto_vector.extend(value) @@ -235,16 +234,16 @@ def _create_vector(self, success): elif compromised == 'Unknown': value = [1, 0] elif compromised == 'User': - value = [0,1] + value = [0, 1] elif compromised == 'Privileged': - value = [1,1] + value = [1, 1] else: raise ValueError('Table had invalid Access Level') proto_vector.extend(value) return np.array(proto_vector) - def get_attr(self,attribute:str): + def get_attr(self, attribute: str): return self.env.get_attr(attribute) def get_observation(self, agent: str): @@ -255,13 +254,13 @@ def get_observation(self, agent: str): return output - def get_agent_state(self,agent:str): + def get_agent_state(self, agent: str): return self.get_attr('get_agent_state')(agent) - def get_action_space(self,agent): + def get_action_space(self, agent): return self.env.get_action_space(agent) - def get_last_action(self,agent): + def get_last_action(self, agent): return self.get_attr('get_last_action')(agent) def get_ip_map(self): From d6464bac738551c9f70cf8ba5b454614ae13f5b3 Mon Sep 17 00:00:00 2001 From: Ilya Orson Sandoval Date: Mon, 18 Mar 2024 14:01:40 +0000 Subject: [PATCH 4/4] Update reset() to return info as well --- CybORG/Agents/Wrappers/BaseWrapper.py | 4 ++-- CybORG/Agents/Wrappers/BlueTableWrapper.py | 4 ++-- CybORG/Agents/Wrappers/ChallengeWrapper.py | 4 ++-- CybORG/Agents/Wrappers/IntListToAction.py | 4 ++-- CybORG/Agents/Wrappers/OpenAIGymWrapper.py | 7 ++++--- CybORG/Agents/Wrappers/RedTableWrapper.py | 4 ++-- CybORG/Agents/Wrappers/ReduceActionSpaceWrapper.py | 4 ++-- CybORG/Agents/Wrappers/RewardShape.py | 4 ++-- CybORG/Agents/Wrappers/TrueTableWrapper.py | 4 ++-- CybORG/CybORG.py | 6 ++++-- .../Tests/test_sim/test_wrappers/test_ChallengeWrapper.py | 2 +- .../Tests/test_sim/test_wrappers/test_OpenAIGymWrapper.py | 2 +- 12 files changed, 26 insertions(+), 23 deletions(-) diff --git a/CybORG/Agents/Wrappers/BaseWrapper.py b/CybORG/Agents/Wrappers/BaseWrapper.py index 8c22531e..1e528202 100644 --- a/CybORG/Agents/Wrappers/BaseWrapper.py +++ b/CybORG/Agents/Wrappers/BaseWrapper.py @@ -18,8 +18,8 @@ def step(self, agent=None, action=None) -> Results: result.action_space = self.action_space_change(result.action_space) return result - def reset(self, agent=None): - result = self.env.reset(agent) + def reset(self, agent=None, **kwargs): + result = self.env.reset(agent, **kwargs) result.action_space = self.action_space_change(result.action_space) result.observation = self.observation_change(result.observation) return result diff --git a/CybORG/Agents/Wrappers/BlueTableWrapper.py b/CybORG/Agents/Wrappers/BlueTableWrapper.py index 21d8d565..6394b18b 100644 --- a/CybORG/Agents/Wrappers/BlueTableWrapper.py +++ b/CybORG/Agents/Wrappers/BlueTableWrapper.py @@ -17,8 +17,8 @@ def __init__(self, env=None, agent=None, output_mode='table'): self.output_mode = output_mode self.blue_info = {} - def reset(self, agent='Blue'): - result = self.env.reset(agent) + def reset(self, agent='Blue', **kwargs): + result = self.env.reset(agent, **kwargs) obs = result.observation if agent == 'Blue': self._process_initial_obs(obs) diff --git a/CybORG/Agents/Wrappers/ChallengeWrapper.py b/CybORG/Agents/Wrappers/ChallengeWrapper.py index 58af7c23..b5b03168 100644 --- a/CybORG/Agents/Wrappers/ChallengeWrapper.py +++ b/CybORG/Agents/Wrappers/ChallengeWrapper.py @@ -34,9 +34,9 @@ def step(self,action=None): return obs, reward, terminated, truncated, info - def reset(self): + def reset(self, **kwargs): self.step_counter = 0 - return self.env.reset() + return self.env.reset(**kwargs) def get_attr(self,attribute:str): return self.env.get_attr(attribute) diff --git a/CybORG/Agents/Wrappers/IntListToAction.py b/CybORG/Agents/Wrappers/IntListToAction.py index ad538b62..a936ec3d 100644 --- a/CybORG/Agents/Wrappers/IntListToAction.py +++ b/CybORG/Agents/Wrappers/IntListToAction.py @@ -30,8 +30,8 @@ def step(self, agent=None, action: list = None) -> Results: result.action_name = str(action_obj) return result - def reset(self, agent=None): - result = self.env.reset(agent) + def reset(self, agent=None, **kwargs): + result = self.env.reset(agent, **kwargs) result.action_space, result.selection_masks = self.action_space_change(result.action_space) self.selection_mask = result.selection_masks result.observation = self.observation_change(result.observation) diff --git a/CybORG/Agents/Wrappers/OpenAIGymWrapper.py b/CybORG/Agents/Wrappers/OpenAIGymWrapper.py index 56e347bc..f66d511c 100644 --- a/CybORG/Agents/Wrappers/OpenAIGymWrapper.py +++ b/CybORG/Agents/Wrappers/OpenAIGymWrapper.py @@ -34,11 +34,12 @@ def step(self, action: Union[int, List[int]] = None) -> tuple[object, float, boo info = vars(result) return np.array(result.observation), result.reward, terminated, truncated, info - def reset(self, agent=None): - result = self.env.reset(self.agent_name) + def reset(self, agent=None, **kwargs): + result = self.env.reset(self.agent_name, **kwargs) result.action_space = self.action_space_change(result.action_space) result.observation = self.observation_change(result.observation) - return np.array(result.observation) + # gymnasium API: observation, info = env.reset(seed=None, options={}) + return np.array(result.observation), vars(result) def render(self): # TODO: If FixedFlatWrapper it will error out! diff --git a/CybORG/Agents/Wrappers/RedTableWrapper.py b/CybORG/Agents/Wrappers/RedTableWrapper.py index 0054848b..de3cdac9 100644 --- a/CybORG/Agents/Wrappers/RedTableWrapper.py +++ b/CybORG/Agents/Wrappers/RedTableWrapper.py @@ -20,13 +20,13 @@ def __init__(self, env=None, agent=None, output_mode='table'): self.output_mode = output_mode self.success = None - def reset(self, agent=None): + def reset(self, agent=None, **kwargs): self.red_info = {} self.known_subnets = set() self.step_counter = -1 self.id_tracker = -1 self.success = None - result = self.env.reset(agent) + result = self.env.reset(agent, **kwargs) if agent =='Red': obs = self.observation_change(result.observation) result.observation = obs diff --git a/CybORG/Agents/Wrappers/ReduceActionSpaceWrapper.py b/CybORG/Agents/Wrappers/ReduceActionSpaceWrapper.py index a7589c52..384a2f85 100644 --- a/CybORG/Agents/Wrappers/ReduceActionSpaceWrapper.py +++ b/CybORG/Agents/Wrappers/ReduceActionSpaceWrapper.py @@ -50,6 +50,6 @@ def action_space_change(self, action_space: dict) -> dict: def get_attr(self,attribute:str): return self.env.get_attr(attribute) - def reset(self, agent=None): + def reset(self, agent=None, **kwargs): self.fixed_size = {} - return self.env.reset(agent) + return self.env.reset(agent, **kwargs) diff --git a/CybORG/Agents/Wrappers/RewardShape.py b/CybORG/Agents/Wrappers/RewardShape.py index 35a172f2..e35ad2a6 100644 --- a/CybORG/Agents/Wrappers/RewardShape.py +++ b/CybORG/Agents/Wrappers/RewardShape.py @@ -44,8 +44,8 @@ def step(self, agent=None, action=None) -> Results: return result - def reset(self, agent=None): - result = self.env.reset(agent) + def reset(self, agent=None, **kwargs): + result = self.env.reset(agent, **kwargs) self.action_buffer = deque(maxlen=2) self.observation_buffer = deque(maxlen=2) if torch.is_tensor(result.observation): diff --git a/CybORG/Agents/Wrappers/TrueTableWrapper.py b/CybORG/Agents/Wrappers/TrueTableWrapper.py index fb492f20..ebc5ebb4 100644 --- a/CybORG/Agents/Wrappers/TrueTableWrapper.py +++ b/CybORG/Agents/Wrappers/TrueTableWrapper.py @@ -12,10 +12,10 @@ def __init__(self,env=None,agent=None, observer_mode=True): self.step_counter = -1 self.observer_mode = observer_mode - def reset(self, agent=None): + def reset(self, agent=None, **kwargs): self.scanned_ips = set() self.step_counter = -1 - result = self.env.reset(agent) + result = self.env.reset(agent, **kwargs) result.observation = self.observation_change(result.observation) return result diff --git a/CybORG/CybORG.py b/CybORG/CybORG.py index 2d950ce4..fc057b1e 100644 --- a/CybORG/CybORG.py +++ b/CybORG/CybORG.py @@ -153,7 +153,7 @@ def get_agent_state(self, agent_name) -> dict: """ return self.environment_controller.get_agent_state(agent_name).data - def reset(self, agent: str = None) -> Results: + def reset(self, agent: str = None, seed=None, **kwargs) -> Results: """ Resets CybORG and gets initial observation and action-space for the specified agent. @@ -172,7 +172,9 @@ def reset(self, agent: str = None) -> Results: Results The initial observation and actions of an agent. """ - return self.environment_controller.reset(agent=agent) + if seed: + self.set_seed(seed) + return self.environment_controller.reset(agent=agent, **kwargs) def shutdown(self, **kwargs) -> bool: """ diff --git a/CybORG/Tests/test_sim/test_wrappers/test_ChallengeWrapper.py b/CybORG/Tests/test_sim/test_wrappers/test_ChallengeWrapper.py index 1ddc2b93..fffc7660 100644 --- a/CybORG/Tests/test_sim/test_wrappers/test_ChallengeWrapper.py +++ b/CybORG/Tests/test_sim/test_wrappers/test_ChallengeWrapper.py @@ -20,7 +20,7 @@ def cyborg(request,agents = {'Blue':BlueMonitorAgent,'Red':B_lineAgent},seed = 1 return cyborg def test_ChallengeWrapper_reset(cyborg): - obs = cyborg.reset() + obs, info = cyborg.reset() expected_observation = np.array([0 for x in range(52)]) assert all(obs == expected_observation) diff --git a/CybORG/Tests/test_sim/test_wrappers/test_OpenAIGymWrapper.py b/CybORG/Tests/test_sim/test_wrappers/test_OpenAIGymWrapper.py index a61479c9..9a9c08df 100644 --- a/CybORG/Tests/test_sim/test_wrappers/test_OpenAIGymWrapper.py +++ b/CybORG/Tests/test_sim/test_wrappers/test_OpenAIGymWrapper.py @@ -128,7 +128,7 @@ def test_get_attr(cyborg): assert cyborg.get_attr(attribute) == cyborg.env.get_attr(attribute) def test_get_observation(cyborg): - step_obs = cyborg.reset() + step_obs, info = cyborg.reset() method_obs = cyborg.get_observation(cyborg.agent_name) assert all(step_obs == method_obs)