aimclub · maypink · Sep 12, 2023 · Mar 17, 2023 · Mar 17, 2023 · Mar 21, 2023
diff --git a/experiments/mab/mab_synthetic_experiment_helper.py b/experiments/mab/mab_synthetic_experiment_helper.py
@@ -72,6 +72,8 @@ def log_action_values(next_pop: PopulationT, optimizer: EvoGraphOptimizer):
             if 0 not in stats_action_value_log.keys():
                 stats_action_value_log[0] = []
             stats_action_value_log[0].append(list(values))
+            # MAB agent can be saved here -- commented not to clog up the memory
+            # optimizer.mutation.agent.save()
 
         def log_action_values_with_clusters(next_pop: PopulationT, optimizer: EvoGraphOptimizer):
             obs_contexts = optimizer.mutation.agent.get_context(next_pop)

diff --git a/golem/core/optimisers/adaptive/mab_agents/mab_agent.py b/golem/core/optimisers/adaptive/mab_agents/mab_agent.py
@@ -1,4 +1,7 @@
+import os.path
+import _pickle as pickle
 import random
+import re
 from typing import Union, Sequence, Optional
 
 from mabwiser.mab import MAB, LearningPolicy
@@ -7,13 +10,15 @@
 from golem.core.dag.graph import Graph
 from golem.core.dag.graph_node import GraphNode
 from golem.core.optimisers.adaptive.operator_agent import OperatorAgent, ActType, ObsType, ExperienceBuffer
+from golem.core.paths import default_data_dir
 
 
 class MultiArmedBanditAgent(OperatorAgent):
     def __init__(self,
                  actions: Sequence[ActType],
                  n_jobs: int = 1,
-                 enable_logging: bool = True):
+                 enable_logging: bool = True,
+                 path_to_save: Optional[str] = None):
         super().__init__(enable_logging)
         self.actions = list(actions)
         self._indices = list(range(len(actions)))
@@ -22,6 +27,7 @@ def __init__(self,
                           learning_policy=LearningPolicy.UCB1(alpha=1.25),
                           n_jobs=n_jobs)
         self._initial_fit()
+        self._path_to_save = path_to_save
 
     def _initial_fit(self):
         n = len(self.actions)
@@ -51,3 +57,37 @@ def partial_fit(self, experience: ExperienceBuffer):
         self._dbg_log(obs, actions, rewards)
         arms = [self._arm_by_action[action] for action in actions]
         self._agent.partial_fit(decisions=arms, rewards=rewards)
+
+    def save(self, path_to_save: Optional[str] = None):
+        """ Saves bandit to specified file. """
+
+        if not path_to_save:
+            path_to_save = self._path_to_save
+
+        # if path was not specified at all
+        if not path_to_save:
+            path_to_save = os.path.join(default_data_dir(), 'MAB')
+
+        if not path_to_save.endswith('.pkl'):
+            os.makedirs(path_to_save, exist_ok=True)
+            mabs_num = [int(name.split('_')[0]) for name in os.listdir(path_to_save)
+                        if re.fullmatch(r'\d_mab.pkl', name)]
+            if not mabs_num:
+                max_saved_mab = 0
+            else:
+                max_saved_mab = max(mabs_num) + 1
+            path_to_file = os.path.join(path_to_save, f'{max_saved_mab}_mab.pkl')
+        else:
+            path_to_dir = os.path.dirname(path_to_save)
+            os.makedirs(path_to_dir, exist_ok=True)
+            path_to_file = path_to_save
+        with open(path_to_file, 'wb') as f:
+            pickle.dump(self, f)
+        self._log.info(f"MAB was saved to {path_to_file}")
+
+    @staticmethod
+    def load(path: str):
+        """ Loads bandit from the specified file. """
+        with open(path, 'rb') as f:
+            mab = pickle.load(f)
+        return mab
diff --git a/golem/core/optimisers/genetic/operators/mutation.py b/golem/core/optimisers/genetic/operators/mutation.py
@@ -50,7 +50,8 @@ def _init_operator_agent(graph_gen_params: GraphGenerationParams,
             agent = RandomAgent(actions=parameters.mutation_types)
         elif kind == MutationAgentTypeEnum.bandit:
             agent = MultiArmedBanditAgent(actions=parameters.mutation_types,
-                                          n_jobs=requirements.n_jobs)
+                                          n_jobs=requirements.n_jobs,
+                                          path_to_save=requirements.agent_dir)
         elif kind == MutationAgentTypeEnum.contextual_bandit:
             agent = ContextualMultiArmedBanditAgent(
                 actions=parameters.mutation_types,

diff --git a/golem/core/optimisers/optimization_parameters.py b/golem/core/optimisers/optimization_parameters.py
@@ -64,6 +64,7 @@ class OptimizationParameters:
 
     keep_history: bool = True
     history_dir: Optional[str] = field(default_factory=default_data_dir)
+    agent_dir: Optional[str] = field(default_factory=default_data_dir)
 
 
 @dataclass

diff --git a/test/unit/adaptive/test_mab_agent.py b/test/unit/adaptive/test_mab_agent.py
@@ -0,0 +1,42 @@
+import os.path
+from pathlib import Path
+
+import pytest
+from mabwiser.mab import MAB
+
+from golem.core.optimisers.adaptive.mab_agents.mab_agent import MultiArmedBanditAgent
+
+
+@pytest.mark.parametrize('path_to_save, file_name',
+                         [(os.path.join(Path(__file__).parent, 'test_mab.pkl'), 'test_mab.pkl'),
+                          (os.path.join(Path(__file__).parent), '0_mab.pkl')])
+def test_save_mab(path_to_save, file_name):
+    """ Tests if MAB is saved with specifying file_nam and without. """
+    mab = MultiArmedBanditAgent(actions=[0, 1, 2],
+                                n_jobs=1,
+                                path_to_save=path_to_save)
+    mab.save()
+    assert file_name in os.listdir(Path(__file__).parent)
+    os.remove(path_to_save if path_to_save.endswith('pkl') else os.path.join(path_to_save, file_name))
+
+
+def test_load_mab():
+    """ Tests if MAB is loaded. """
+    file_name = 'test_mab.pkl'
+    path_to_load = os.path.join(Path(__file__).parent, file_name)
+    # save mab to load it later
+    mab = MultiArmedBanditAgent(actions=[0, 1, 2],
+                                n_jobs=1,
+                                path_to_save=path_to_load)
+    mab.save()
+
+    loaded_mab = MultiArmedBanditAgent.load(path=path_to_load)
+    assert isinstance(loaded_mab, MultiArmedBanditAgent)
+
+    assert isinstance(loaded_mab._agent, MAB)
+    assert loaded_mab.__eq__(mab)
+    assert loaded_mab.actions == mab.actions
+    assert loaded_mab._enable_logging == mab._enable_logging
+    assert loaded_mab._path_to_save == mab._path_to_save
+
+    os.remove(path_to_load)