ensemble env fixes, load agent, sep main to experiments

rl_core/ensemble_experiment.py

@@ -0,0 +1,61 @@
+import os
+
+from meta_automl.utils import project_root
+from rl_core.generator import Generator
+
+if __name__ == '__main__':
+    task_type = 'classification'
+    pipeline_len = 5
+    n_episodes = 15000
+
+    data_folder_path = os.path.join(str(project_root()), 'MetaFEDOT/rl_core/data/')
+
+    train_datasets = {
+        'amazon': os.path.join(data_folder_path, 'amazon_train.csv'),
+        'australian': os.path.join(data_folder_path, 'australian_train.csv'),
+        'bank': os.path.join(data_folder_path, 'bank_train.csv'),
+        'blood': os.path.join(data_folder_path, 'blood_train.csv'),
+        'jasmine': os.path.join(data_folder_path, 'jasmine_train.csv'),
+        'kc1': os.path.join(data_folder_path, 'kc1_train.csv'),
+    }
+
+    primitives = ['scaling', 'simple_imputation', 'normalization', 'dt', 'logit', 'rf']
+
+    gen = Generator(task_type, state_dim=pipeline_len, n_episodes=n_episodes) \
+        .set_environment(env_name='ensemble', primitives=primitives) \
+        .set_dataloader(train_datasets) \
+        .set_agent(
+            eval_schedule=15,
+            critic_updates_per_actor=10,
+        ) \
+        .set_writer()
+
+    gen.fit()
+
+    gen.save_agent()
+
+    test_datasets = {
+        'amazon': os.path.join(data_folder_path, 'amazon_test.csv'),
+        'australian': os.path.join(data_folder_path, 'australian_test.csv'),
+        'bank': os.path.join(data_folder_path, 'bank_test.csv'),
+        'blood': os.path.join(data_folder_path, 'blood_test.csv'),
+        'jasmine': os.path.join(data_folder_path, 'jasmine_test.csv'),
+        'kc1': os.path.join(data_folder_path, 'kc1_test.csv'),
+    }
+
+    for name, dataset in test_datasets:
+        valid, not_valid = 0, 0
+
+        for _ in range(25):
+            pipeline, metric_value = gen.generate(path_to_dataset=dataset)
+
+            if pipeline:
+                pipeline.show()
+                print('Test metric:', metric_value)
+                valid += 1
+            else:
+                print('Pipeline generation failed!')
+                not_valid += 1
+
+        print(f'{valid=}, {not_valid=}')
+

rl_core/environments/ensemble.py

@@ -1,4 +1,5 @@
 from abc import ABC
+from copy import deepcopy
 
 import numpy as np
 import torch
@@ -10,6 +11,10 @@
 
 class EnsemblePipelineGenerationEnvironment(PipelineGenerationEnvironment, ABC):
     """ Linear Pipeline Generation Environment """
+    _meta_info = {
+        'name': 'ensemble',
+    }
+
     def __init__(self, state_dim: int, primitives: list):
         super().__init__(state_dim)
         self.primitives = ['pop'] + primitives + ['eop']
@@ -33,6 +38,7 @@ def update_state(self, action):
 
     def reset(self, **kwargs):
         self.pipeline = PipelineBuilder()
+        self.is_valid = False
         self.time_step = 0
         self.metric_value = 0
 
@@ -45,6 +51,7 @@ def reset(self, **kwargs):
         return self.state
 
     def _train_step(self, action):
+        terminated, truncated = False, False
         self.last_action = action
 
         if self.primitives[action] == 'eop' or self.position == self.state_dim:
@@ -61,12 +68,10 @@ def _train_step(self, action):
         elif self.primitives[action] == 'pop':
             self.time_step += 1
             reward = -0.005
-            done = False
 
         else:
             self.time_step += 1
             reward = -0.001
-            done = False
 
             self.update_state(action)
 
@@ -78,18 +83,19 @@ def _train_step(self, action):
                 self.pipeline.add_branch(primitive, branch_idx=self.branch_idx)
                 self.branch_idx += 1
 
-        reward, done, info = self._environment_response(reward, done)
+        reward, info = self._environment_response(reward)
 
-        return self.state, reward, done, info
+        return deepcopy(self.state), reward, terminated, truncated, info
 
     def _inference_step(self, action):
         raise NotImplementedError()
 
-    def _environment_response(self, reward: float, done: bool) -> (int, bool, dict):
+    def _environment_response(self, reward: float) -> (int, bool, dict):
         info = {
             'pipeline': self.pipeline.build(),
             'time_step': self.time_step,
             'metric_value': self.metric_value,
+            'is_valid': self.is_valid
         }
 
-        return reward, done, info
+        return reward, info

rl_core/environments/linear.py

@@ -15,6 +15,7 @@ class LinearPipelineGenerationEnvironment(PipelineGenerationEnvironment, ABC):
     _meta_info = {
         'name': 'linear',
     }
+
     def __init__(self, state_dim: int, primitives: list):
         super().__init__(state_dim)
         self.primitives = ['pop'] + primitives + ['eop']

rl_core/generator.py

@@ -57,7 +57,6 @@ def __init__(self, task_type: str, state_dim: int, n_episodes: int, batch_size:
     def _get_default_environment(task_type, state_dim):
         # TODO: Adding tags "reinforce" into Fedot
         primitives = OperationTypesRepository('all').suitable_operation(task_type=task_type)
-        # primitives = ['rf', 'dt', 'scaling', 'knn']
 
         return LinearPipelineGenerationEnvironment(state_dim=state_dim, primitives=primitives)
 
@@ -76,7 +75,9 @@ def set_environment(self, env_name: str, primitives: list[str] = None):
 
         if not primitives:
             primitives = OperationTypesRepository('all').suitable_operation(task_type=self.task_type)
-            primitives.remove('lgbm')
+
+            for d_primitves in ['lgbm', 'knn']:
+                primitives.remove(d_primitves)
 
         self.env = env(state_dim=self.state_dim, primitives=primitives)
         return self
@@ -86,7 +87,7 @@ def set_dataloader(self, datasets: dict = None, split_ratio: float = 0.7):
         return self
 
     def set_agent(self, hidden_dims=(128, 128), lr=1e-3, gamma=0.99, replay_buffer_size=500, train_schedule=32,
-                  eval_schedule=100, critic_updates_per_actor=4, path_to_wieghts=None):
+                  eval_schedule=100, critic_updates_per_actor=4, path_to_weights=None):
         self.hidden_dim = hidden_dims[-1]
         self.train_schedule = train_schedule
         self.eval_schedule = eval_schedule
@@ -101,8 +102,8 @@ def set_agent(self, hidden_dims=(128, 128), lr=1e-3, gamma=0.99, replay_buffer_s
             critic_rb_size=replay_buffer_size
         )
 
-        if path_to_wieghts:
-            self.load_agent(path_to_wieghts)
+        if path_to_weights:
+            self.load_agent(path_to_weights)
 
         return self
 
@@ -196,7 +197,7 @@ def generate(self, path_to_dataset):
             return None, 0
 
     def save_agent(self):
-        name = f'{self.env._meta_info["name"]}_{self.state_dim}_a2c_{self.hidden_dim}'
+        name = f'{self.env._meta_info["name"]}_{self.state_dim}_a2c_{self.hidden_dim}_{self.n_episodes}'
         path = f'{project_root()}/MetaFEDOT/rl_core/agent/pretrained/{name}'
         self.agent.save(path)
 

rl_core/main.py → rl_core/linear_experiment.py

@@ -3,10 +3,12 @@
 from meta_automl.utils import project_root
 from rl_core.generator import Generator
 
+os.environ['OPENBLAS_NUM_THREADS'] = '1'
+
 if __name__ == '__main__':
     task_type = 'classification'
     pipeline_len = 5
-    n_episodes = 250
+    n_episodes = 500
 
     scoring_train_path = os.path.join(str(project_root()), 'MetaFEDOT/rl_core/data/scoring_train.csv')
     airlines_train_path = os.path.join(str(project_root()), 'MetaFEDOT/rl_core/data/airlines.csv')
@@ -16,10 +18,15 @@
         'airlines': airlines_train_path
     }
 
+    primitives = ['scaling', 'simple_imputation', 'normalization', 'dt', 'logit', 'rf', 'knn']
+
     gen = Generator(task_type, state_dim=pipeline_len, n_episodes=n_episodes) \
-        .set_environment(env_name='linear') \
+        .set_environment(env_name='linear', primitives=primitives) \
         .set_dataloader(train_datasets) \
-        .set_agent(eval_schedule=15) \
+        .set_agent(
+            eval_schedule=10,
+            critic_updates_per_actor=10,
+        ) \
         .set_writer()
 
     gen.fit()