Repository: automl/DACBench (language: python)

Issue:
Problem in diff_training reward functions
`get_log_diff_training_reward` and `get_diff_training_reward` functions fail in `reset` `get_reward` because `prev_training_loss` is not set.
Add strict benchmark-specific functional tests
Current tests test the functionality of the sgd benchmark are very general/superficial level.

We should add some 'easy to fail' tests that allows us to detect (accidental) low-level functional changes to the benchmark.

For a specific configuration and policy (sequence of learning rates) log/store various metrics, e.g., the training/validation losses, values of state features, etc. (initial and after every step), and use these to test that this functional behavior remains invariant.
Add strict benchmark-specific functional tests
Current tests test the functionality of the sgd benchmark are very general/superficial level.

We should add some 'easy to fail' tests that allows us to detect (accidental) low-level functional changes to the benchmark.

For a specific configuration and policy (sequence of learning rates) log/store various metrics, e.g., the training/validation losses, values of state features, etc. (initial and after every step), and use these to test that this functional behavior remains invariant.
Add strict benchmark-specific functional tests
Current tests test the functionality of the sgd benchmark are very general/superficial level.

We should add some 'easy to fail' tests that allows us to detect (accidental) low-level functional changes to the benchmark.

For a specific configuration and policy (sequence of learning rates) log/store various metrics, e.g., the training/validation losses, values of state features, etc. (initial and after every step), and use these to test that this functional behavior remains invariant.
Add strict benchmark-specific functional tests
Current tests test the functionality of the sgd benchmark are very general/superficial level.

We should add some 'easy to fail' tests that allows us to detect (accidental) low-level functional changes to the benchmark.

For a specific configuration and policy (sequence of learning rates) log/store various metrics, e.g., the training/validation losses, values of state features, etc. (initial and after every step), and use these to test that this functional behavior remains invariant.


Discussion / hints:

Fix patch:
diff --git a/dacbench/abstract_benchmark.py b/dacbench/abstract_benchmark.py
index bf023cc89..10df527e5 100644
--- a/dacbench/abstract_benchmark.py
+++ b/dacbench/abstract_benchmark.py
@@ -1,4 +1,7 @@
 import json
+from types import FunctionType
+from typing import Callable
+
 import numpy as np
 from gym import spaces
 from functools import partial
@@ -7,6 +10,7 @@
 # from dacbench import ModuleLogger
 
 
+
 class AbstractBenchmark:
     """
     Abstract template for benchmark classes
@@ -61,17 +65,7 @@ def save_config(self, path):
         if "action_space" in self.config:
             conf["action_space"] = self.space_to_list(conf["action_space"])
 
-        if "reward_function" in self.config:
-            conf["reward_function"] = [
-                conf["reward_function"].__module__,
-                conf["reward_function"].__name__,
-            ]
-
-        if "state_method" in self.config:
-            conf["state_method"] = [
-                conf["state_method"].__module__,
-                conf["state_method"].__name__,
-            ]
+        conf = AbstractBenchmark.__stringify_functions(conf)
 
         for k in self.config.keys():
             if isinstance(self.config[k], np.ndarray) or isinstance(
@@ -141,6 +135,61 @@ def dejson_wrappers(self, wrapper_list):
 
             self.wrap_funcs.append(partial(func, *args))
 
+    @staticmethod
+    def __import_from(module : str, name : str):
+        """
+        Imports the class / function / ... with name from module
+        Parameters
+        ----------
+        module
+        name
+
+        Returns
+        -------
+        the imported object
+        """
+        module = __import__(module, fromlist=[name])
+        return getattr(module, name)
+
+    @staticmethod
+    def __decorate_config_with_functions(conf : dict):
+        """
+        Replaced the stringified functions with the callable objects
+        Parameters
+        ----------
+        config
+
+        Returns
+        -------
+
+        """
+        for key, value in {k: v for k, v in conf.items() if isinstance(v, list) and len(v) == 3 and v[0] == 'function'}.items():
+            _, module_name, function_name = value
+            conf[key] = AbstractBenchmark.__import_from(module_name, function_name)
+        return conf
+
+    @staticmethod
+    def __stringify_functions(conf : dict) -> dict:
+        """
+        Replaced all callables in the config with a
+        triple ('function', module_name, function_name)
+
+        Parameters
+        ----------
+        config
+
+        Returns
+        -------
+        modified dict
+        """
+        for key, value in {k: v for k, v in conf.items() if isinstance(v, FunctionType)}.items():
+            conf[key] = [
+                'function',
+                conf[key].__module__,
+                conf[key].__name__
+            ]
+        return conf
+
     def space_to_list(self, space):
         res = []
         if isinstance(space, spaces.Box):
@@ -180,19 +229,31 @@ def jsonify_dict_space(self, dict_space):
         for k in dict_space.keys():
             keys.append(k)
             value = dict_space[k]
-            if not isinstance(value, spaces.Box):
-                raise ValueError("Only Dict spaces made up of Boxes are supported")
+            if not isinstance(value, (spaces.Box, spaces.Discrete)):
+                raise ValueError(f"Only Dict spaces made up of Box spaces or discrete spaces are supported but got {type(value)}")
+
+            if isinstance(value, spaces.Box):
+                low = value.low.tolist()
+                high = value.high.tolist()
+                arguments.append([low, high])
 
-            low = value.low.tolist()
-            high = value.high.tolist()
-            arguments.append([low, high])
+            if isinstance(value, spaces.Discrete):
+                n = value.n
+                arguments.append([n])
         return [keys, arguments]
 
     def dictify_json(self, dict_list):
         dict_space = {}
-        for i in range(len(dict_list[0])):
-            args = [np.array(arg) for arg in dict_list[1][i]]
-            dict_space[dict_list[0][i]] = spaces.Box(*args, dtype=np.float32)
+        types, args = dict_list
+        for type, args_ in zip(types, args):
+            prepared_args = map(np.array, args_)
+            if type == "box":
+                dict_space[type] = spaces.Box(*prepared_args, dtype=np.float32)
+            elif type == "discrete":
+                dict_space[type] = spaces.Discrete(*prepared_args)
+            else:
+                raise TypeError(f"Currently only Discrete and Box spaces are allowed in Dict spaces got {type}")
+
         return dict_space
 
     def read_config_file(self, path):
@@ -235,19 +296,7 @@ def read_config_file(self, path):
             self.dejson_wrappers(self.config["wrappers"])
             del self.config["wrappers"]
 
-        import importlib
-
-        if "reward_function" in self.config:
-            self.config["reward_function"] = getattr(
-                importlib.import_module(self.config["reward_function"][0]),
-                self.config["reward_function"][1],
-            )
-
-        if "state_method" in self.config:
-            self.config["state_method"] = getattr(
-                importlib.import_module(self.config["state_method"][0]),
-                self.config["state_method"][1],
-            )
+        self.config = AbstractBenchmark.__decorate_config_with_functions(self.config)
 
         for k in self.config.keys():
             if type(self.config[k]) == list:
diff --git a/dacbench/benchmarks/sgd_benchmark.py b/dacbench/benchmarks/sgd_benchmark.py
index 99db0d886..a60cfc3a4 100755
--- a/dacbench/benchmarks/sgd_benchmark.py
+++ b/dacbench/benchmarks/sgd_benchmark.py
@@ -3,9 +3,16 @@
 
 import numpy as np
 from gym import spaces
+from torch.nn import NLLLoss
 
 from dacbench.abstract_benchmark import AbstractBenchmark, objdict
 from dacbench.envs import SGDEnv
+from dacbench.envs.sgd import Reward
+
+
+def __default_loss_function(**kwargs):
+    return NLLLoss(reduction = 'none', **kwargs)
+
 
 HISTORY_LENGTH = 40
 INPUT_DIM = 10
@@ -13,7 +20,7 @@
 INFO = {
     "identifier": "LR",
     "name": "Learning Rate Adaption for Neural Networks",
-    "reward": "Validation Loss",
+    "reward": "Negative Log Differential Validation Loss",
     "state_description": [
         "Predictive Change Variance (Discounted Average)",
         "Predictive Change Variance (Uncertainty)",
@@ -22,9 +29,13 @@
         "Current Learning Rate",
         "Training Loss",
         "Validation Loss",
+        "Step",
+        "Alignment",
+        "Crashed"
     ],
 }
 
+
 SGD_DEFAULTS = objdict(
     {
         "action_space_class": "Box",
@@ -46,29 +57,58 @@
                 "currentLR": spaces.Box(low=0, high=1, shape=(1,)),
                 "trainingLoss": spaces.Box(low=0, high=np.inf, shape=(1,)),
                 "validationLoss": spaces.Box(low=0, high=np.inf, shape=(1,)),
+                "step": spaces.Box(low=0, high=np.inf, shape=(1,)),
+                "alignment": spaces.Box(low=0, high=1, shape=(1,)),
+                "crashed": spaces.Discrete(2),
             }
         ],
-        "reward_range": (-(10 ** 9), 0),
-        "cutoff": 5e1,
+        "reward_type": Reward.LogDiffTraining,
+        "cutoff": 1e3,
         "lr": 1e-3,
+        "discount_factor": 0.9,
+        "optimizer": "rmsprop",
+        "loss_function": __default_loss_function,
+        "loss_function_kwargs": {},
+        "val_loss_function": __default_loss_function,
+        "val_loss_function_kwargs": {},
         "training_batch_size": 64,
         "validation_batch_size": 64,
-        "no_cuda": True,
+        "train_validation_ratio": 0.8,
+        "no_cuda": False,
         "beta1": 0.9,
         "beta2": 0.999,
         "seed": 0,
-        "instance_set_path": "../instance_sets/sgd/sgd_train.csv",
+        "cd_paper_reconstruction": False,
+        "cd_bias_correction": True,
+        "terminate_on_crash": False,
+        "crash_penalty": 0.0,
+        "instance_set_path": "../instance_sets/sgd/sgd_train_100instances.csv",
         "benchmark_info": INFO,
+        "features": [
+            "predictiveChangeVarDiscountedAverage",
+            "predictiveChangeVarUncertainty",
+            "lossVarDiscountedAverage",
+            "lossVarUncertainty",
+            "currentLR",
+            "trainingLoss",
+            "validationLoss",
+            "step",
+            "alignment",
+            "crashed"
+        ],
     }
 )
 
+# Set reward range based on the chosen reward type
+SGD_DEFAULTS.reward_range = SGD_DEFAULTS['reward_type'].func.frange
+
 
 class SGDBenchmark(AbstractBenchmark):
     """
     Benchmark with default configuration & relevant functions for SGD
     """
 
-    def __init__(self, config_path=None):
+    def __init__(self, config_path=None, **kwargs):
         """
         Initialize SGD Benchmark
 
@@ -85,6 +125,9 @@ def __init__(self, config_path=None):
             if key not in self.config:
                 self.config[key] = SGD_DEFAULTS[key]
 
+        for k in kwargs:
+            self.config[k] = kwargs[k]
+
     def get_environment(self):
         """
         Return SGDEnv env with current configuration
@@ -117,14 +160,23 @@ def read_instance_set(self):
         with open(path, "r") as fh:
             reader = csv.DictReader(fh, delimiter=";")
             for row in reader:
+                if "_" in row["dataset"]:
+                    dataset_info = row["dataset"].split("_")
+                    dataset_name = dataset_info[0]
+                    dataset_size = int(dataset_info[1])
+                else:
+                    dataset_name = row["dataset"]
+                    dataset_size = None
                 instance = [
-                    row["dataset"],
+                    dataset_name,
                     int(row["seed"]),
                     row["architecture"],
+                    int(row["steps"]),
+                    dataset_size
                 ]
                 self.config["instance_set"][int(row["ID"])] = instance
 
-    def get_benchmark(self, seed=0):
+    def get_benchmark(self, instance_set_path=None, seed=0):
         """
         Get benchmark from the LTO paper
 
@@ -139,6 +191,8 @@ def get_benchmark(self, seed=0):
             SGD environment
         """
         self.config = objdict(SGD_DEFAULTS.copy())
+        if instance_set_path is not None:
+            self.config["instance_set_path"] = instance_set_path
         self.config.seed = seed
         self.read_instance_set()
         return SGDEnv(self.config)
diff --git a/dacbench/envs/sgd.py b/dacbench/envs/sgd.py
old mode 100755
new mode 100644
index c4e723249..5f15abcc7
--- a/dacbench/envs/sgd.py
+++ b/dacbench/envs/sgd.py
@@ -3,16 +3,45 @@
 import warnings
 import json
 from functools import reduce
+from enum import IntEnum, auto
 
+import numpy as np
 import torch
 from backpack import backpack, extend
 from backpack.extensions import BatchGrad
+from numpy import float32
 from torchvision import datasets, transforms
 from dacbench import AbstractEnv
+import random
 
 warnings.filterwarnings("ignore")
 
 
+def reward_range(frange):
+    def wrapper(f):
+        f.frange = frange
+        return f
+    return wrapper
+
+
+class Reward(IntEnum):
+    TrainingLoss = auto()
+    ValidationLoss = auto()
+    LogTrainingLoss = auto()
+    LogValidationLoss = auto()
+    DiffTraining = auto()
+    DiffValidation = auto()
+    LogDiffTraining = auto()
+    LogDiffValidation = auto()
+    FullTraining = auto()
+
+    def __call__(self, f):
+        if hasattr(self, 'func'):
+            raise ValueError('Can not assign the same reward to a different function!')
+        self.func = f
+        return f
+
+
 class SGDEnv(AbstractEnv):
     """
     Environment to control the learning rate of adam
@@ -33,11 +62,27 @@ def __init__(self, config):
         self.validation_batch_size = config.validation_batch_size
         self.no_cuda = config.no_cuda
         self.current_batch_size = config.training_batch_size
+        self.on_features = config.features
+        self.cd_paper_reconstruction = config.cd_paper_reconstruction
+        self.cd_bias_correction = config.cd_bias_correction
+        self.crashed = False
+        self.terminate_on_crash = config.terminate_on_crash
+        self.crash_penalty = config.crash_penalty
+
+        if isinstance(config.reward_type, Reward):
+            self.reward_type = config.reward_type
+        elif isinstance(config.reward_type, str):
+            try:
+                self.reward_type = getattr(Reward, config.reward_type)
+            except AttributeError:
+                raise ValueError(f'{config.reward_type} is not a valid reward type!')
+        else:
+            raise ValueError(f'Type {type(config.reward_type)} is not valid!')
 
         self.use_cuda = not self.no_cuda and torch.cuda.is_available()
         self.device = torch.device("cuda" if self.use_cuda else "cpu")
 
-        self.training_validation_ratio = 0.8
+        self.training_validation_ratio = config.train_validation_ratio
         # self.test_dataset = None
         self.train_dataset = None
         self.validation_dataset = None
@@ -52,14 +97,34 @@ def __init__(self, config):
 
         self.current_training_loss = None
         self.loss_batch = None
+        self.prev_training_loss = None
+        self._current_validation_loss = torch.zeros(
+            1, device=self.device, requires_grad=False
+        )
+        self._current_validation_loss.calculated = False
+        self.prev_validation_loss = torch.zeros(
+            1, device=self.device, requires_grad=False
+        )
 
         self.model = None
-
+        self.val_model = None
+        # TODO:
+        """
+        TODO: Samuel Mueller (PhD student in our group) also uses backpack and has ran into a similar memory leak.
+        He solved it calling this custom made RECURSIVE memory_cleanup function:
+        # from backpack import memory_cleanup
+        # def recursive_backpack_memory_cleanup(module: torch.nn.Module):
+        #   memory_cleanup(module)
+        #   for m in module.modules():
+        #      memory_cleanup(m)
+        (calling this after computing the training loss/gradients and after validation loss should suffice)
+        """
         self.parameter_count = 0
         self.layer_sizes = []
 
-        self.loss_function = torch.nn.NLLLoss(reduction="none")
+        self.loss_function = config.loss_function(**config.loss_function_kwargs)
         self.loss_function = extend(self.loss_function)
+        self.val_loss_function = config.loss_function(**config.val_loss_function_kwargs)
 
         self.initial_lr = config.lr * torch.ones(
             1, device=self.device, requires_grad=False
@@ -68,18 +133,25 @@ def __init__(self, config):
             1, device=self.device, requires_grad=False
         )
 
-        # Adam parameters
+        self.optimizer_name = config.optimizer
+
         self.beta1 = config.beta1
         self.beta2 = config.beta2
+        self.epsilon = 1.0e-06
+        # RMSprop parameters
+        self.beta2 = config.beta2
         self.m = 0
         self.v = 0
-        self.epsilon = 1.0e-08
+        # Momentum parameters
+        self.sgd_momentum_v = 0
+        self.sgd_rho = 0.9
+
         self.t = 0
         self.step_count = torch.zeros(1, device=self.device, requires_grad=False)
 
-        self.prev_descent = None
+        self.prev_direction = None
+        self.current_direction = None
 
-        self.learning_rate = 0.001
         self.predictiveChangeVarDiscountedAverage = torch.zeros(
             1, device=self.device, requires_grad=False
         )
@@ -92,7 +164,7 @@ def __init__(self, config):
         self.lossVarUncertainty = torch.zeros(
             1, device=self.device, requires_grad=False
         )
-        self.discount_factor = 0.9
+        self.discount_factor = config.discount_factor
         self.firstOrderMomentum = torch.zeros(
             1, device=self.device, requires_grad=False
         )
@@ -100,18 +172,80 @@ def __init__(self, config):
             1, device=self.device, requires_grad=False
         )
 
-        self.writer = None
+        if self.optimizer_name=="adam":
+            self.get_optimizer_direction = self.get_adam_direction
+        elif self.optimizer_name=="rmsprop":
+            self.get_optimizer_direction = self.get_rmsprop_direction
+        elif self.optimizer_name=="momentum":
+            self.get_optimizer_direction = self.get_momentum_direction
+        else:
+            raise NotImplementedError
 
         if "reward_function" in config.keys():
-            self.get_reward = config["reward_function"]
+            self._get_reward = config["reward_function"]
         else:
-            self.get_reward = self.get_default_reward
+            self._get_reward = self.reward_type.func
 
         if "state_method" in config.keys():
             self.get_state = config["state_method"]
         else:
             self.get_state = self.get_default_state
 
+        self.reward_range = self.reward_type.func.frange
+
+    def get_reward(self):
+        return self._get_reward(self)
+
+    @reward_range([-(10**9), 0])
+    @Reward.TrainingLoss
+    def get_training_reward(self):
+        return -self.current_training_loss.item()
+
+    @reward_range([-(10**9), 0])
+    @Reward.ValidationLoss
+    def get_validation_reward(self):
+        return -self.current_validation_loss.item()
+
+    @reward_range([-(10**9), (10**9)])
+    @Reward.LogTrainingLoss
+    def get_log_training_reward(self):
+        return -torch.log(self.current_training_loss).item()
+
+    @reward_range([-(10**9), (10**9)])
+    @Reward.LogValidationLoss
+    def get_log_validation_reward(self):
+        return -torch.log(self.current_validation_loss).item()
+
+    @reward_range([-(10**9), (10**9)])
+    @Reward.LogDiffTraining
+    def get_log_diff_training_reward(self):
+        return -(torch.log(self.current_training_loss) - torch.log(self.prev_training_loss)).item()
+
+    @reward_range([-(10**9), (10**9)])
+    @Reward.LogDiffValidation
+    def get_log_diff_validation_reward(self):
+        return -(torch.log(self.current_validation_loss) - torch.log(self.prev_validation_loss)).item()
+
+    @reward_range([-(10**9), (10**9)])
+    @Reward.DiffTraining
+    def get_diff_training_reward(self):
+        return (self.current_training_loss - self.prev_training_loss).item()
+
+    @reward_range([-(10**9), (10**9)])
+    @Reward.DiffValidation
+    def get_diff_validation_reward(self):
+        return (self.current_validation_loss - self.prev_validation_loss).item()
+
+    @reward_range([-(10**9), 0])
+    @Reward.FullTraining
+    def get_full_training_reward(self):
+        return -self._get_full_training_loss(loader=self.train_loader).item()
+
+    @property
+    def crash(self):
+        self.crashed = True
+        return self.get_state(self), self.crash_penalty, self.terminate_on_crash, {}
+
     def seed(self, seed=None, seed_action_space=False):
         """
         Set rng seed
@@ -126,6 +260,10 @@ def seed(self, seed=None, seed_action_space=False):
         (seed,) = super().seed(seed, seed_action_space)
         if seed is not None:
             torch.manual_seed(seed)
+            torch.cuda.manual_seed(seed)
+            torch.cuda.manual_seed_all(seed)
+            torch.backends.cudnn.benchmark = False
+            torch.backends.cudnn.deterministic = True
         return [seed]
 
     def step(self, action):
@@ -146,19 +284,26 @@ def step(self, action):
 
         self.step_count += 1
         index = 0
-        if not isinstance(action, float) and not isinstance(action, int):
+
+        if not isinstance(action, int) and not isinstance(action, float):
             action = action.item()
         if not isinstance(action, numbers.Number):
             action = action[0]
 
-        action = torch.Tensor([action]).to(self.device)
-        new_lr = 10 ** (-action)
+        if np.isnan(action):
+            return self.crash
+
+        new_lr = torch.Tensor([action]).to(self.device)
         self.current_lr = new_lr
-        delta_w = torch.mul(
-            new_lr,
-            self.firstOrderMomentum
-            / (torch.sqrt(self.secondOrderMomentum) + self.epsilon),
-        )
+
+        direction = self.get_optimizer_direction()
+        if any(np.isnan(direction)):
+            return self.crash
+
+        self.current_direction = direction
+
+        delta_w = torch.mul(new_lr, direction)
+
         for i, p in enumerate(self.model.parameters()):
             layer_size = self.layer_sizes[i]
             p.data = p.data - delta_w[index : index + layer_size].reshape(
@@ -166,9 +311,22 @@ def step(self, action):
             )
             index += layer_size
 
-        self._set_zero_grad()
-        reward = self.get_reward(self)
-        return self.get_state(self), reward, done, {}
+        self.model.zero_grad()
+
+        self.prev_training_loss = self.current_training_loss
+        if self._current_validation_loss.calculated:
+            self.prev_validation_loss = self.current_validation_loss
+
+        self.train_network()
+        reward = self.get_reward()
+        if np.isnan(reward):
+            return self.crash
+
+        state = self.get_state(self)
+        for value in state.values():
+            if np.isnan(value):
+                return self.crash
+        return state, reward, done, {}
 
     def _architecture_constructor(self, arch_str):
         layer_specs = []
@@ -204,32 +362,59 @@ def reset(self):
         dataset = self.instance[0]
         instance_seed = self.instance[1]
         construct_model = self._architecture_constructor(self.instance[2])
+        self.n_steps = self.instance[3]
+        dataset_size = self.instance[4]
+
+        self.crashed = False
 
         self.seed(instance_seed)
 
         self.model = construct_model().to(self.device)
+        self.val_model = construct_model().to(self.device)
+
+        def init_weights(m):
+            if type(m) == torch.nn.Linear or type(m) == torch.nn.Conv2d:
+                torch.nn.init.xavier_normal(m.weight)
+                m.bias.data.fill_(0.0)
 
-        self.training_validation_ratio = 0.8
+        if self.cd_paper_reconstruction:
+            self.model.apply(init_weights)
 
-        train_dataloader_args = {"batch_size": self.batch_size}
-        validation_dataloader_args = {"batch_size": self.validation_batch_size}
+        train_dataloader_args = {"batch_size": self.batch_size, "drop_last": True}
+        validation_dataloader_args = {"batch_size": self.validation_batch_size, "drop_last": True}
         if self.use_cuda:
             param = {"num_workers": 1, "pin_memory": True, "shuffle": True}
             train_dataloader_args.update(param)
             validation_dataloader_args.update(param)
 
         if dataset == "MNIST":
-            transform = transforms.Compose(
-                [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
-            )
+            transform = transforms.Compose([
+                transforms.ToTensor(),
+                transforms.Normalize((0.1307,), (0.3081,))
+            ])
 
             train_dataset = datasets.MNIST(
                 "../data", train=True, download=True, transform=transform
             )
             # self.test_dataset = datasets.MNIST('../data', train=False, transform=transform)
+        elif dataset == "CIFAR":
+            transform = transforms.Compose([
+                transforms.ToTensor(),
+                transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+            ])
+
+            train_dataset = datasets.CIFAR10(
+                "../data", train=True, download=True, transform=transform
+            )
+            # self.test_dataset = datasets.MNIST('../data', train=False, transform=transform)
         else:
             raise NotImplementedError
 
+        if dataset_size is not None:
+            train_dataset = torch.utils.data.Subset(
+                train_dataset, range(0, dataset_size)
+            )
+
         training_dataset_limit = math.floor(
             len(train_dataset) * self.training_validation_ratio
         )
@@ -264,23 +449,57 @@ def reset(self):
 
         self.model = extend(self.model)
 
-        self._set_zero_grad()
+        self.model.zero_grad()
         self.model.train()
+        self.val_model.eval()
 
         self.current_training_loss = None
         self.loss_batch = None
 
-        # Adam parameters
+        # Momentum parameters
         self.m = 0
         self.v = 0
+        self.sgd_momentum_v = 0
+
         self.t = 0
+
         self.step_count = torch.zeros(1, device=self.device, requires_grad=False)
 
         self.current_lr = self.initial_lr
-        self.prev_descent = torch.zeros(
+        self.prev_direction = torch.zeros(
+            (self.parameter_count,), device=self.device, requires_grad=False
+        )
+        self.current_direction = torch.zeros(
             (self.parameter_count,), device=self.device, requires_grad=False
         )
-        self.get_default_reward(self)
+
+        self.predictiveChangeVarDiscountedAverage = torch.zeros(
+            1, device=self.device, requires_grad=False
+        )
+        self.predictiveChangeVarUncertainty = torch.zeros(
+            1, device=self.device, requires_grad=False
+        )
+        self.lossVarDiscountedAverage = torch.zeros(
+            1, device=self.device, requires_grad=False
+        )
+        self.lossVarUncertainty = torch.zeros(
+            1, device=self.device, requires_grad=False
+        )
+        self.firstOrderMomentum = torch.zeros(
+            1, device=self.device, requires_grad=False
+        )
+        self.secondOrderMomentum = torch.zeros(
+            1, device=self.device, requires_grad=False
+        )
+
+        self._current_validation_loss = torch.zeros(
+            1, device=self.device, requires_grad=False
+        )
+        self._current_validation_loss.calculated = False
+        self.prev_validation_loss = torch.zeros(
+            1, device=self.device, requires_grad=False
+        )
+        self.train_network()
 
         return self.get_state(self)
 
@@ -322,39 +541,50 @@ def get_default_state(self, _):
             Environment state
 
         """
-        gradients = self._get_gradients()
-        self.firstOrderMomentum, self.secondOrderMomentum = self._get_momentum(
-            gradients
-        )
-        (
-            predictiveChangeVarDiscountedAverage,
-            predictiveChangeVarUncertainty,
-        ) = self._get_predictive_change_features(
-            self.current_lr, self.firstOrderMomentum, self.secondOrderMomentum
-        )
-        lossVarDiscountedAverage, lossVarUncertainty = self._get_loss_features()
-
-        state = {
-            "predictiveChangeVarDiscountedAverage": predictiveChangeVarDiscountedAverage,
-            "predictiveChangeVarUncertainty": predictiveChangeVarUncertainty,
-            "lossVarDiscountedAverage": lossVarDiscountedAverage,
-            "lossVarUncertainty": lossVarUncertainty,
-            "currentLR": self.current_lr,
-            "trainingLoss": self.current_training_loss,
-            "validationLoss": self.current_validation_loss,
-        }
+        self.gradients = self._get_gradients()
+        self.firstOrderMomentum, self.secondOrderMomentum, self.sgdMomentum = self._get_momentum(self.gradients)
+
+        if 'predictiveChangeVarDiscountedAverage' in self.on_features or 'predictiveChangeVarUncertainty' in self.on_features:
+            predictiveChangeVarDiscountedAverage, predictiveChangeVarUncertainty = \
+                self._get_predictive_change_features(self.current_lr)
+
+        if 'lossVarDiscountedAverage' in self.on_features or 'lossVarUncertainty' in self.on_features:
+            lossVarDiscountedAverage, lossVarUncertainty = self._get_loss_features()
+
+        if 'alignment' in self.on_features:
+            alignment = self._get_alignment()
+
+        state = {}
+
+        if 'predictiveChangeVarDiscountedAverage' in self.on_features:
+            state["predictiveChangeVarDiscountedAverage"] = predictiveChangeVarDiscountedAverage.item()
+        if 'predictiveChangeVarUncertainty' in self.on_features:
+            state["predictiveChangeVarUncertainty"] = predictiveChangeVarUncertainty.item()
+        if 'lossVarDiscountedAverage' in self.on_features:
+            state["lossVarDiscountedAverage"] = lossVarDiscountedAverage.item()
+        if 'lossVarUncertainty' in self.on_features:
+            state["lossVarUncertainty"] = lossVarUncertainty.item()
+        if 'currentLR' in self.on_features:
+            state["currentLR"] = self.current_lr.item()
+        if 'trainingLoss' in self.on_features:
+            if self.crashed:
+                state["trainingLoss"] = 0.0
+            else:
+                state["trainingLoss"] = self.current_training_loss.item()
+        if 'validationLoss' in self.on_features:
+            if self.crashed:
+                state["validationLoss"] = 0.0
+            else:
+                state["validationLoss"] = self.current_validation_loss.item()
+        if 'step' in self.on_features:
+            state["step"] = self.step_count.item()
+        if 'alignment' in self.on_features:
+            state["alignment"] = alignment.item()
+        if 'crashed' in self.on_features:
+            state["crashed"] = self.crashed
 
         return state
 
-    def _set_zero_grad(self):
-        index = 0
-        for i, p in enumerate(self.model.parameters()):
-            if p.grad is None:
-                continue
-            layer_size = self.layer_sizes[i]
-            p.grad.zero_()
-            index += layer_size
-
     def _train_batch_(self):
         (data, target) = self.train_loader_it.next()
         data, target = data.to(self.device), target.to(self.device)
@@ -366,50 +596,54 @@ def _train_batch_(self):
             loss.mean().backward()
 
         loss_value = loss.mean()
-        reward = self._get_validation_loss()
+
         self.loss_batch = loss
         self.current_training_loss = torch.unsqueeze(loss_value.detach(), dim=0)
         self.train_batch_index += 1
+        self._current_validation_loss.calculated = False
 
-        return reward
-
-    def get_default_reward(self, _):
+    def train_network(self):
         try:
-            reward = self._train_batch_()
+            self._train_batch_()
         except StopIteration:
             self.train_batch_index = 0
             self.epoch_index += 1
             self.train_loader_it = iter(self.train_loader)
-            reward = self._train_batch_()
-
-        return reward
+            self._train_batch_()
 
-    def _get_val_loss(self):
-        self.model.eval()
-        validation_loss = torch.zeros(1, device=self.device, requires_grad=False)
+    def _get_full_training_loss(self, loader):
+        for target_param, param in zip(self.val_model.parameters(), self.model.parameters()):
+            target_param.data.copy_(param.data)
+        loss = torch.zeros(1, device=self.device, requires_grad=False)
         with torch.no_grad():
-            for data, target in self.validation_loader:
+            for data, target in loader:
                 data, target = data.to(self.device), target.to(self.device)
-                output = self.model(data)
-                validation_loss += self.loss_function(output, target).mean()
+                output = self.val_model(data)
+                loss += self.val_loss_function(output, target).sum().detach().detach()
 
-        validation_loss /= len(self.validation_loader.dataset)
-        self.model.train()
-        return validation_loss
+        loss /= len(loader.dataset)
+        return loss
+
+    @property
+    def current_validation_loss(self):
+        if not self._current_validation_loss.calculated:
+            self._current_validation_loss = self._get_validation_loss()
+            self._current_validation_loss.calculated = True
+        return self._current_validation_loss
 
     def _get_validation_loss_(self):
-        self.model.eval()
-        (data, target) = self.validation_loader_it.next()
-        data, target = data.to(self.device), target.to(self.device)
-        output = self.model(data)
-        validation_loss = self.loss_function(output, target).mean()
-        validation_loss = torch.unsqueeze(validation_loss.detach(), dim=0)
-        self.current_validation_loss = validation_loss
-        self.model.train()
+        with torch.no_grad():
+            (data, target) = self.validation_loader_it.next()
+            data, target = data.to(self.device), target.to(self.device)
+            output = self.val_model(data)
+            validation_loss = self.val_loss_function(output, target).mean()
+            validation_loss = torch.unsqueeze(validation_loss.detach(), dim=0)
 
-        return -validation_loss.item()  # negative because it is the reward
+        return validation_loss
 
     def _get_validation_loss(self):
+        for target_param, param in zip(self.val_model.parameters(), self.model.parameters()):
+            target_param.data.copy_(param.data)
         try:
             validation_loss = self._get_validation_loss_()
         except StopIteration:
@@ -436,16 +670,25 @@ def _get_momentum(self, gradients):
         bias_corrected_m = self.m / (1 - self.beta1 ** self.t)
         bias_corrected_v = self.v / (1 - self.beta2 ** self.t)
 
-        return bias_corrected_m, bias_corrected_v
+        self.sgd_momentum_v = self.sgd_rho * self.sgd_momentum_v + gradients
+
+        return bias_corrected_m, bias_corrected_v, self.sgd_momentum_v
+
+    def get_adam_direction(self):
+        return self.firstOrderMomentum / (torch.sqrt(self.secondOrderMomentum) + self.epsilon)
 
-    def _get_adam_feature(self, learning_rate, m, v):
-        epsilon = 1.0e-8
-        return torch.mul(learning_rate, m / (torch.sqrt(v) + epsilon))
+    def get_rmsprop_direction(self):
+        return self.gradients / (torch.sqrt(self.secondOrderMomentum) + self.epsilon)
+
+    def get_momentum_direction(self):
+        return self.sgd_momentum_v
 
     def _get_loss_features(self):
+        if self.crashed:
+            return torch.tensor(0.0), torch.tensor(0.0)
+        bias_correction = (1 - self.discount_factor ** (self.c_step+1)) if self.cd_bias_correction else 1
         with torch.no_grad():
             loss_var = torch.log(torch.var(self.loss_batch))
-
             self.lossVarDiscountedAverage = (
                 self.discount_factor * self.lossVarDiscountedAverage
                 + (1 - self.discount_factor) * loss_var
@@ -453,12 +696,15 @@ def _get_loss_features(self):
             self.lossVarUncertainty = (
                 self.discount_factor * self.lossVarUncertainty
                 + (1 - self.discount_factor)
-                * (loss_var - self.lossVarDiscountedAverage) ** 2
+                * (loss_var - self.lossVarDiscountedAverage/bias_correction) ** 2
             )
 
-        return self.lossVarDiscountedAverage, self.lossVarUncertainty
+        return self.lossVarDiscountedAverage/bias_correction, self.lossVarUncertainty/bias_correction
 
-    def _get_predictive_change_features(self, lr, m, v):
+    def _get_predictive_change_features(self, lr):
+        if self.crashed:
+            return torch.tensor(0.0), torch.tensor(0.0)
+        bias_correction = (1 - self.discount_factor ** (self.c_step+1)) if self.cd_bias_correction else 1
         batch_gradients = []
         for i, (name, param) in enumerate(self.model.named_parameters()):
             grad_batch = param.grad_batch.reshape(
@@ -468,7 +714,8 @@ def _get_predictive_change_features(self, lr, m, v):
 
         batch_gradients = torch.cat(batch_gradients, dim=1)
 
-        update_value = self._get_adam_feature(lr, m, v)
+        update_value = torch.mul(lr, self.get_optimizer_direction())
+
         predictive_change = torch.log(
             torch.var(-1 * torch.matmul(batch_gradients, update_value))
         )
@@ -480,10 +727,150 @@ def _get_predictive_change_features(self, lr, m, v):
         self.predictiveChangeVarUncertainty = (
             self.discount_factor * self.predictiveChangeVarUncertainty
             + (1 - self.discount_factor)
-            * (predictive_change - self.predictiveChangeVarDiscountedAverage) ** 2
+            * (predictive_change - self.predictiveChangeVarDiscountedAverage/bias_correction) ** 2
         )
 
         return (
-            self.predictiveChangeVarDiscountedAverage,
-            self.predictiveChangeVarUncertainty,
+            self.predictiveChangeVarDiscountedAverage/bias_correction,
+            self.predictiveChangeVarUncertainty/bias_correction,
         )
+
+    def _get_alignment(self):
+        if self.crashed:
+            return torch.tensor(0.0)
+        alignment = torch.mean(torch.sign(torch.mul(self.prev_direction, self.current_direction)))
+        alignment = torch.unsqueeze(alignment, dim=0)
+        self.prev_direction = self.current_direction
+        return alignment
+
+    def generate_instance_file(self, file_name, mode='test', n=100):
+        header = ['ID', 'dataset', 'architecture', 'seed', 'steps']
+
+        # dataset name, architecture, dataset size, sample dimension, number of max pool layers, hidden layers, test architecture convolutional layers
+        architectures = [
+            ('MNIST',
+             'Conv2d(1, {0}, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d({0}, {1}, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d({1}, {2}, 3, 1, 1)-ReLU-Flatten-Linear({3}, 10)-LogSoftmax(1)',
+             60000,
+             28,
+             2,
+             3,
+             [20, 50, 500]
+            ),
+            ('CIFAR',
+             'Conv2d(3, {0}, 3, 1, 1)-MaxPool2d(2, 2)-ReLU-Conv2d({0}, {1}, 3, 1, 1)-ReLU-MaxPool2d(2, 2)-Conv2d({1}, {2}, 3, 1, 1)-ReLU-MaxPool2d(2, 2)-Conv2d({2}, {3}, 3, 1, 1)-ReLU-Flatten-Linear({4}, 10)-LogSoftmax(1)',
+             60000,
+             32,
+             3,
+             4,
+             [32, 32, 64, 64]
+            )
+        ]
+        if mode is 'test':
+
+            seed_list = [random.randrange(start=0, stop=1e9) for _ in range(n)]
+
+            for i in range(len(architectures)):
+
+                fname = file_name + "_" + architectures[i][0].lower() + ".csv"
+
+                steps = int(1e8)
+
+                conv = architectures[i][6]
+                hidden_layers = architectures[i][5]
+
+                sample_size = architectures[i][3]
+                pool_layer_count = architectures[i][4]
+                linear_layer_size = conv[-1] * pow(sample_size / pow(2, pool_layer_count), 2)
+                linear_layer_size = int(round(linear_layer_size))
+
+                dataset = architectures[i][0]
+
+                if hidden_layers == 3:
+                    architecture = architectures[i][1].format(conv[0], conv[1], conv[2], linear_layer_size)
+                else:
+                    architecture = architectures[i][1].format(conv[0], conv[1], conv[2], conv[3], linear_layer_size)
+
+                # args = conv
+                # args.append(linear_layer_size)
+                # # architecture = architectures[i][1].format(**conv)
+                # args = {0: conv[0], 1: conv[1], 2: conv[2], 3: linear_layer_size}
+                # architecture = architectures[i][1].format(**args)
+
+                with open(fname, 'w', encoding='UTF8') as f:
+                    for h in header:
+                        f.write(h + ";")
+
+                    f.write("\n")
+
+                    for id in range(0, n):
+                        f.write(str(id) + ";")
+
+                        f.write(dataset + ";")
+                        f.write(architecture + ";")
+
+                        seed = seed_list[id]
+                        f.write(str(seed) + ";")
+
+                        f.write(str(steps) + ";")
+
+                        f.write("\n")
+                    f.close()
+
+        else:
+            dataset_index = 0
+
+            dataset_size_start = 0.1
+            dataset_size_stop = 0.5
+
+            steps_start = 300
+            steps_stop = 1000
+
+            conv1_start = 2
+            conv1_stop = 10
+            conv2_start = 5
+            conv2_stop = 25
+            conv3_start = 50
+            conv3_stop = 250
+
+            dataset_list = [dataset_index for _ in range(n)]
+
+            dataset_size_list = [random.uniform(dataset_size_start, dataset_size_stop) for _ in range(n)]
+
+            seed_list = [random.randrange(start=0, stop=1e9) for _ in range(n)]
+
+            steps_list = [random.randrange(start=steps_start, stop=steps_stop) for _ in range(n)]
+
+            conv1_list = [random.randrange(start=conv1_start, stop=conv1_stop) for _ in range(n)]
+            conv2_list = [random.randrange(start=conv2_start, stop=conv2_stop) for _ in range(n)]
+            conv3_list = [random.randrange(start=conv3_start, stop=conv3_stop) for _ in range(n)]
+
+            fname = file_name + ".csv"
+            with open(fname, 'w', encoding='UTF8') as f:
+                for h in header:
+                    f.write(h + ";")
+
+                f.write("\n")
+
+                for id in range(0, n):
+                    f.write(str(id) + ";")
+
+                    sample_size = architectures[dataset_list[id]][3]
+                    pool_layer_count = architectures[dataset_list[id]][4]
+                    linear_layer_size = conv3_list[id] * pow(sample_size / pow(2, pool_layer_count), 2)
+                    linear_layer_size = int(round(linear_layer_size))
+
+                    dataset_size = int(dataset_size_list[id] * architectures[dataset_list[id]][2])
+                    dataset = architectures[dataset_list[id]][0] + "_" + str(dataset_size)
+                    architecture = architectures[dataset_list[id]][1].format(conv1_list[id], conv2_list[id], conv3_list[id], linear_layer_size)
+
+                    f.write(dataset + ";")
+                    f.write(architecture + ";")
+
+                    seed = seed_list[id]
+                    f.write(str(seed) + ";")
+
+                    steps = steps_list[id]
+                    f.write(str(steps) + ";")
+
+                    f.write("\n")
+                f.close()
diff --git a/dacbench/instance_sets/sgd/sgd_cnn.csv b/dacbench/instance_sets/sgd/sgd_cnn.csv
deleted file mode 100644
index e4a3a5487..000000000
--- a/dacbench/instance_sets/sgd/sgd_cnn.csv
+++ /dev/null
@@ -1,2 +0,0 @@
-ID;dataset;architecture;seed
-0;MNIST;Conv2d(1,8,3,1,1)-ReLU-Conv2d(8,1,3,1,1)-ReLU-Flatten-Linear(784,10)-LogSoftmax(1);790283118
\ No newline at end of file
diff --git a/dacbench/instance_sets/sgd/sgd_train.csv b/dacbench/instance_sets/sgd/sgd_train.csv
deleted file mode 100755
index 8254d2d74..000000000
--- a/dacbench/instance_sets/sgd/sgd_train.csv
+++ /dev/null
@@ -1,101 +0,0 @@
-ID;dataset;architecture;seed
-0;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);790283118
-1;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);600984959
-2;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);880665794
-3;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);36686381
-4;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);869863829
-5;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);773652426
-6;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);769733396
-7;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);684065927
-8;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);825057925
-9;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);400897614
-10;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);493280011
-11;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);58560565
-12;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);997460616
-13;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);658214425
-14;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);139665789
-15;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);128611529
-16;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);948145792
-17;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);43514665
-18;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);450105660
-19;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);625597733
-20;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);483847251
-21;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);185871259
-22;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);246674446
-23;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);228791258
-24;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);94192574
-25;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);392753888
-26;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);585785321
-27;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);910845204
-28;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);874003896
-29;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);652159358
-30;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);737313245
-31;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);836659660
-32;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);832500222
-33;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);837524123
-34;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);702126849
-35;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);138518203
-36;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);27314884
-37;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);492511750
-38;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);306323598
-39;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);304960628
-40;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);295419467
-41;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);485049497
-42;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);152581202
-43;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);267120130
-44;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);684800058
-45;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);920499534
-46;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);191011437
-47;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);114306679
-48;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);322168682
-49;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);85197090
-50;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);694363206
-51;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);280861838
-52;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);412483556
-53;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);758610160
-54;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);564883708
-55;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);134583239
-56;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);268500986
-57;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);49895890
-58;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);534793682
-59;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);807339335
-60;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);357938588
-61;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);978988301
-62;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);244695222
-63;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);723745050
-64;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);401660921
-65;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);193691143
-66;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);817808161
-67;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);420726681
-68;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);818325297
-69;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);762846618
-70;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);247819560
-71;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);65975392
-72;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);365323617
-73;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);220974336
-74;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);545152297
-75;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);540515443
-76;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);182877407
-77;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);772993267
-78;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);506541183
-79;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);846729895
-80;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);500669472
-81;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);354316502
-82;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);522512370
-83;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);650663702
-84;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);51975067
-85;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);614076064
-86;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);743019217
-87;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);850850823
-88;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);538961714
-89;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);685542560
-90;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);914126583
-91;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);359536450
-92;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);882758375
-93;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);194860564
-94;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);653268158
-95;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);272680258
-96;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);43573507
-97;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);767368191
-98;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);195798585
-99;MNIST;Flatten-Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);206649358
\ No newline at end of file
diff --git a/dacbench/instance_sets/sgd/sgd_train_100instances.csv b/dacbench/instance_sets/sgd/sgd_train_100instances.csv
new file mode 100755
index 000000000..0f1005c7e
--- /dev/null
+++ b/dacbench/instance_sets/sgd/sgd_train_100instances.csv
@@ -0,0 +1,101 @@
+ID;dataset;architecture;seed;steps;
+0;MNIST_14240;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 67, 3, 1, 1)-ReLU-Flatten-Linear(3283, 10)-LogSoftmax(1);85370663;698;
+1;MNIST_10048;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 81, 3, 1, 1)-ReLU-Flatten-Linear(3969, 10)-LogSoftmax(1);597236641;696;
+2;MNIST_21440;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 23, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(23, 152, 3, 1, 1)-ReLU-Flatten-Linear(7448, 10)-LogSoftmax(1);455413943;572;
+3;MNIST_16000;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 82, 3, 1, 1)-ReLU-Flatten-Linear(4018, 10)-LogSoftmax(1);666364192;585;
+4;MNIST_20480;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 13, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(13, 135, 3, 1, 1)-ReLU-Flatten-Linear(6615, 10)-LogSoftmax(1);738646438;498;
+5;MNIST_25152;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 20, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(20, 198, 3, 1, 1)-ReLU-Flatten-Linear(9702, 10)-LogSoftmax(1);671251256;549;
+6;MNIST_12160;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 126, 3, 1, 1)-ReLU-Flatten-Linear(6174, 10)-LogSoftmax(1);583709876;987;
+7;MNIST_14048;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 17, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(17, 191, 3, 1, 1)-ReLU-Flatten-Linear(9359, 10)-LogSoftmax(1);58254785;946;
+8;MNIST_28160;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 14, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(14, 205, 3, 1, 1)-ReLU-Flatten-Linear(10045, 10)-LogSoftmax(1);200533670;753;
+9;MNIST_26592;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 198, 3, 1, 1)-ReLU-Flatten-Linear(9702, 10)-LogSoftmax(1);975255077;994;
+10;MNIST_17056;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 198, 3, 1, 1)-ReLU-Flatten-Linear(9702, 10)-LogSoftmax(1);736032395;703;
+11;MNIST_6880;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 72, 3, 1, 1)-ReLU-Flatten-Linear(3528, 10)-LogSoftmax(1);897524533;811;
+12;MNIST_13568;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 20, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(20, 59, 3, 1, 1)-ReLU-Flatten-Linear(2891, 10)-LogSoftmax(1);73411480;622;
+13;MNIST_20128;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 231, 3, 1, 1)-ReLU-Flatten-Linear(11319, 10)-LogSoftmax(1);747370677;745;
+14;MNIST_16320;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 128, 3, 1, 1)-ReLU-Flatten-Linear(6272, 10)-LogSoftmax(1);590418655;805;
+15;MNIST_10432;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 23, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(23, 169, 3, 1, 1)-ReLU-Flatten-Linear(8281, 10)-LogSoftmax(1);759601910;891;
+16;MNIST_21440;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 203, 3, 1, 1)-ReLU-Flatten-Linear(9947, 10)-LogSoftmax(1);279035288;328;
+17;MNIST_25376;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 11, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(11, 200, 3, 1, 1)-ReLU-Flatten-Linear(9800, 10)-LogSoftmax(1);390813115;866;
+18;MNIST_10368;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 76, 3, 1, 1)-ReLU-Flatten-Linear(3724, 10)-LogSoftmax(1);417406150;885;
+19;MNIST_18624;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 163, 3, 1, 1)-ReLU-Flatten-Linear(7987, 10)-LogSoftmax(1);676716673;630;
+20;MNIST_28384;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 61, 3, 1, 1)-ReLU-Flatten-Linear(2989, 10)-LogSoftmax(1);483634587;495;
+21;MNIST_24480;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 145, 3, 1, 1)-ReLU-Flatten-Linear(7105, 10)-LogSoftmax(1);587269322;614;
+22;MNIST_16736;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 137, 3, 1, 1)-ReLU-Flatten-Linear(6713, 10)-LogSoftmax(1);896937394;463;
+23;MNIST_13568;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 220, 3, 1, 1)-ReLU-Flatten-Linear(10780, 10)-LogSoftmax(1);482762874;986;
+24;MNIST_11552;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 168, 3, 1, 1)-ReLU-Flatten-Linear(8232, 10)-LogSoftmax(1);62829004;561;
+25;MNIST_10560;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 170, 3, 1, 1)-ReLU-Flatten-Linear(8330, 10)-LogSoftmax(1);432423572;882;
+26;MNIST_9728;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 14, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(14, 110, 3, 1, 1)-ReLU-Flatten-Linear(5390, 10)-LogSoftmax(1);90250538;933;
+27;MNIST_13536;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 13, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(13, 209, 3, 1, 1)-ReLU-Flatten-Linear(10241, 10)-LogSoftmax(1);700199045;773;
+28;MNIST_7424;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 13, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(13, 141, 3, 1, 1)-ReLU-Flatten-Linear(6909, 10)-LogSoftmax(1);922181636;891;
+29;MNIST_23136;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 23, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(23, 148, 3, 1, 1)-ReLU-Flatten-Linear(7252, 10)-LogSoftmax(1);88633046;927;
+30;MNIST_7008;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 61, 3, 1, 1)-ReLU-Flatten-Linear(2989, 10)-LogSoftmax(1);216976841;869;
+31;MNIST_14848;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 130, 3, 1, 1)-ReLU-Flatten-Linear(6370, 10)-LogSoftmax(1);650853430;605;
+32;MNIST_6368;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 145, 3, 1, 1)-ReLU-Flatten-Linear(7105, 10)-LogSoftmax(1);502921667;960;
+33;MNIST_9376;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 10, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(10, 77, 3, 1, 1)-ReLU-Flatten-Linear(3773, 10)-LogSoftmax(1);850545979;309;
+34;MNIST_15328;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 235, 3, 1, 1)-ReLU-Flatten-Linear(11515, 10)-LogSoftmax(1);283063259;447;
+35;MNIST_20960;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 10, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(10, 50, 3, 1, 1)-ReLU-Flatten-Linear(2450, 10)-LogSoftmax(1);632368576;862;
+36;MNIST_19488;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 17, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(17, 184, 3, 1, 1)-ReLU-Flatten-Linear(9016, 10)-LogSoftmax(1);27639030;643;
+37;MNIST_12128;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 14, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(14, 79, 3, 1, 1)-ReLU-Flatten-Linear(3871, 10)-LogSoftmax(1);923343467;374;
+38;MNIST_24960;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 81, 3, 1, 1)-ReLU-Flatten-Linear(3969, 10)-LogSoftmax(1);665581259;493;
+39;MNIST_12992;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 114, 3, 1, 1)-ReLU-Flatten-Linear(5586, 10)-LogSoftmax(1);56307477;761;
+40;MNIST_10592;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 11, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(11, 208, 3, 1, 1)-ReLU-Flatten-Linear(10192, 10)-LogSoftmax(1);942298166;582;
+41;MNIST_8576;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 166, 3, 1, 1)-ReLU-Flatten-Linear(8134, 10)-LogSoftmax(1);313386172;719;
+42;MNIST_25984;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 18, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(18, 140, 3, 1, 1)-ReLU-Flatten-Linear(6860, 10)-LogSoftmax(1);840627960;350;
+43;MNIST_9856;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 132, 3, 1, 1)-ReLU-Flatten-Linear(6468, 10)-LogSoftmax(1);222275973;486;
+44;MNIST_26560;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 99, 3, 1, 1)-ReLU-Flatten-Linear(4851, 10)-LogSoftmax(1);972703876;803;
+45;MNIST_16576;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 61, 3, 1, 1)-ReLU-Flatten-Linear(2989, 10)-LogSoftmax(1);109829453;990;
+46;MNIST_21344;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 22, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(22, 224, 3, 1, 1)-ReLU-Flatten-Linear(10976, 10)-LogSoftmax(1);383971483;436;
+47;MNIST_13888;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 105, 3, 1, 1)-ReLU-Flatten-Linear(5145, 10)-LogSoftmax(1);643171914;832;
+48;MNIST_9792;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 132, 3, 1, 1)-ReLU-Flatten-Linear(6468, 10)-LogSoftmax(1);670675758;437;
+49;MNIST_11712;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 231, 3, 1, 1)-ReLU-Flatten-Linear(11319, 10)-LogSoftmax(1);718247435;662;
+50;MNIST_12768;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 94, 3, 1, 1)-ReLU-Flatten-Linear(4606, 10)-LogSoftmax(1);870508044;842;
+51;MNIST_6048;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 24, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(24, 162, 3, 1, 1)-ReLU-Flatten-Linear(7938, 10)-LogSoftmax(1);384966093;506;
+52;MNIST_17088;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 204, 3, 1, 1)-ReLU-Flatten-Linear(9996, 10)-LogSoftmax(1);263374537;815;
+53;MNIST_27040;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 228, 3, 1, 1)-ReLU-Flatten-Linear(11172, 10)-LogSoftmax(1);845919665;592;
+54;MNIST_26048;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 142, 3, 1, 1)-ReLU-Flatten-Linear(6958, 10)-LogSoftmax(1);892572882;691;
+55;MNIST_26944;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 14, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(14, 104, 3, 1, 1)-ReLU-Flatten-Linear(5096, 10)-LogSoftmax(1);493715397;434;
+56;MNIST_16384;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 161, 3, 1, 1)-ReLU-Flatten-Linear(7889, 10)-LogSoftmax(1);559166445;318;
+57;MNIST_18400;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 20, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(20, 184, 3, 1, 1)-ReLU-Flatten-Linear(9016, 10)-LogSoftmax(1);681400868;783;
+58;MNIST_29440;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 99, 3, 1, 1)-ReLU-Flatten-Linear(4851, 10)-LogSoftmax(1);132784964;608;
+59;MNIST_12416;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 13, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(13, 246, 3, 1, 1)-ReLU-Flatten-Linear(12054, 10)-LogSoftmax(1);179219961;333;
+60;MNIST_15360;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 14, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(14, 69, 3, 1, 1)-ReLU-Flatten-Linear(3381, 10)-LogSoftmax(1);342942351;870;
+61;MNIST_12864;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 183, 3, 1, 1)-ReLU-Flatten-Linear(8967, 10)-LogSoftmax(1);138470507;414;
+62;MNIST_6848;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 13, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(13, 195, 3, 1, 1)-ReLU-Flatten-Linear(9555, 10)-LogSoftmax(1);395086689;447;
+63;MNIST_21248;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 207, 3, 1, 1)-ReLU-Flatten-Linear(10143, 10)-LogSoftmax(1);551353010;631;
+64;MNIST_6112;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 17, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(17, 76, 3, 1, 1)-ReLU-Flatten-Linear(3724, 10)-LogSoftmax(1);526563120;832;
+65;MNIST_17024;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 155, 3, 1, 1)-ReLU-Flatten-Linear(7595, 10)-LogSoftmax(1);547892100;520;
+66;MNIST_14272;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 10, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(10, 143, 3, 1, 1)-ReLU-Flatten-Linear(7007, 10)-LogSoftmax(1);958733851;360;
+67;MNIST_17376;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 178, 3, 1, 1)-ReLU-Flatten-Linear(8722, 10)-LogSoftmax(1);46920774;971;
+68;MNIST_20320;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 17, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(17, 99, 3, 1, 1)-ReLU-Flatten-Linear(4851, 10)-LogSoftmax(1);48914259;892;
+69;MNIST_27744;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 203, 3, 1, 1)-ReLU-Flatten-Linear(9947, 10)-LogSoftmax(1);606982068;873;
+70;MNIST_19584;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 187, 3, 1, 1)-ReLU-Flatten-Linear(9163, 10)-LogSoftmax(1);909692949;481;
+71;MNIST_17280;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 159, 3, 1, 1)-ReLU-Flatten-Linear(7791, 10)-LogSoftmax(1);709429373;665;
+72;MNIST_10912;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 205, 3, 1, 1)-ReLU-Flatten-Linear(10045, 10)-LogSoftmax(1);826615641;927;
+73;MNIST_6496;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 201, 3, 1, 1)-ReLU-Flatten-Linear(9849, 10)-LogSoftmax(1);453979089;326;
+74;MNIST_17248;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 23, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(23, 174, 3, 1, 1)-ReLU-Flatten-Linear(8526, 10)-LogSoftmax(1);330914630;371;
+75;MNIST_14432;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 224, 3, 1, 1)-ReLU-Flatten-Linear(10976, 10)-LogSoftmax(1);656456878;447;
+76;MNIST_17568;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 242, 3, 1, 1)-ReLU-Flatten-Linear(11858, 10)-LogSoftmax(1);692559468;453;
+77;MNIST_8256;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 78, 3, 1, 1)-ReLU-Flatten-Linear(3822, 10)-LogSoftmax(1);94848173;340;
+78;MNIST_13984;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 17, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(17, 121, 3, 1, 1)-ReLU-Flatten-Linear(5929, 10)-LogSoftmax(1);767999222;894;
+79;MNIST_12544;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 18, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(18, 91, 3, 1, 1)-ReLU-Flatten-Linear(4459, 10)-LogSoftmax(1);732740509;365;
+80;MNIST_29216;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 148, 3, 1, 1)-ReLU-Flatten-Linear(7252, 10)-LogSoftmax(1);955411453;552;
+81;MNIST_20480;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 22, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(22, 113, 3, 1, 1)-ReLU-Flatten-Linear(5537, 10)-LogSoftmax(1);108292162;318;
+82;MNIST_17728;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 89, 3, 1, 1)-ReLU-Flatten-Linear(4361, 10)-LogSoftmax(1);628500738;844;
+83;MNIST_28288;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 148, 3, 1, 1)-ReLU-Flatten-Linear(7252, 10)-LogSoftmax(1);784353714;714;
+84;MNIST_6208;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 13, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(13, 189, 3, 1, 1)-ReLU-Flatten-Linear(9261, 10)-LogSoftmax(1);555927508;330;
+85;MNIST_7776;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 104, 3, 1, 1)-ReLU-Flatten-Linear(5096, 10)-LogSoftmax(1);121374477;716;
+86;MNIST_9952;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 227, 3, 1, 1)-ReLU-Flatten-Linear(11123, 10)-LogSoftmax(1);614230241;346;
+87;MNIST_10400;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 102, 3, 1, 1)-ReLU-Flatten-Linear(4998, 10)-LogSoftmax(1);803575145;367;
+88;MNIST_14848;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 18, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(18, 67, 3, 1, 1)-ReLU-Flatten-Linear(3283, 10)-LogSoftmax(1);100099230;622;
+89;MNIST_21632;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 169, 3, 1, 1)-ReLU-Flatten-Linear(8281, 10)-LogSoftmax(1);341604808;511;
+90;MNIST_23840;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 83, 3, 1, 1)-ReLU-Flatten-Linear(4067, 10)-LogSoftmax(1);896839395;444;
+91;MNIST_7680;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 18, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(18, 157, 3, 1, 1)-ReLU-Flatten-Linear(7693, 10)-LogSoftmax(1);300655330;705;
+92;MNIST_11008;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 152, 3, 1, 1)-ReLU-Flatten-Linear(7448, 10)-LogSoftmax(1);800030713;467;
+93;MNIST_22368;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 61, 3, 1, 1)-ReLU-Flatten-Linear(2989, 10)-LogSoftmax(1);854754649;537;
+94;MNIST_13600;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 22, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(22, 63, 3, 1, 1)-ReLU-Flatten-Linear(3087, 10)-LogSoftmax(1);316005183;863;
+95;MNIST_11104;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 146, 3, 1, 1)-ReLU-Flatten-Linear(7154, 10)-LogSoftmax(1);662691777;891;
+96;MNIST_13856;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 11, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(11, 238, 3, 1, 1)-ReLU-Flatten-Linear(11662, 10)-LogSoftmax(1);888322309;710;
+97;MNIST_18944;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 161, 3, 1, 1)-ReLU-Flatten-Linear(7889, 10)-LogSoftmax(1);574142479;792;
+98;MNIST_8192;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 231, 3, 1, 1)-ReLU-Flatten-Linear(11319, 10)-LogSoftmax(1);211625200;390;
+99;MNIST_14048;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 24, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(24, 54, 3, 1, 1)-ReLU-Flatten-Linear(2646, 10)-LogSoftmax(1);498679091;643;
\ No newline at end of file
diff --git a/dacbench/instance_sets/sgd/sgd_train_100instances_deprecated.csv b/dacbench/instance_sets/sgd/sgd_train_100instances_deprecated.csv
new file mode 100644
index 000000000..0f1234b7f
--- /dev/null
+++ b/dacbench/instance_sets/sgd/sgd_train_100instances_deprecated.csv
@@ -0,0 +1,101 @@
+ID;dataset;architecture;seed;steps;
+0;MNIST_14256;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 67, 3, 1, 1)-ReLU-Flatten-Linear(3283, 10)-LogSoftmax(1);85370663;698;
+1;MNIST_10059;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 81, 3, 1, 1)-ReLU-Flatten-Linear(3969, 10)-LogSoftmax(1);597236641;696;
+2;MNIST_21445;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 23, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(23, 152, 3, 1, 1)-ReLU-Flatten-Linear(7448, 10)-LogSoftmax(1);455413943;572;
+3;MNIST_16015;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 82, 3, 1, 1)-ReLU-Flatten-Linear(4018, 10)-LogSoftmax(1);666364192;585;
+4;MNIST_20497;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 13, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(13, 135, 3, 1, 1)-ReLU-Flatten-Linear(6615, 10)-LogSoftmax(1);738646438;498;
+5;MNIST_25179;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 20, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(20, 198, 3, 1, 1)-ReLU-Flatten-Linear(9702, 10)-LogSoftmax(1);671251256;549;
+6;MNIST_12174;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 126, 3, 1, 1)-ReLU-Flatten-Linear(6174, 10)-LogSoftmax(1);583709876;987;
+7;MNIST_14071;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 17, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(17, 191, 3, 1, 1)-ReLU-Flatten-Linear(9359, 10)-LogSoftmax(1);58254785;946;
+8;MNIST_28162;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 14, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(14, 205, 3, 1, 1)-ReLU-Flatten-Linear(10045, 10)-LogSoftmax(1);200533670;753;
+9;MNIST_26592;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 198, 3, 1, 1)-ReLU-Flatten-Linear(9702, 10)-LogSoftmax(1);975255077;994;
+10;MNIST_17081;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 198, 3, 1, 1)-ReLU-Flatten-Linear(9702, 10)-LogSoftmax(1);736032395;703;
+11;MNIST_6883;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 72, 3, 1, 1)-ReLU-Flatten-Linear(3528, 10)-LogSoftmax(1);897524533;811;
+12;MNIST_13582;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 20, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(20, 59, 3, 1, 1)-ReLU-Flatten-Linear(2891, 10)-LogSoftmax(1);73411480;622;
+13;MNIST_20133;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 231, 3, 1, 1)-ReLU-Flatten-Linear(11319, 10)-LogSoftmax(1);747370677;745;
+14;MNIST_16350;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 128, 3, 1, 1)-ReLU-Flatten-Linear(6272, 10)-LogSoftmax(1);590418655;805;
+15;MNIST_10441;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 23, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(23, 169, 3, 1, 1)-ReLU-Flatten-Linear(8281, 10)-LogSoftmax(1);759601910;891;
+16;MNIST_21442;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 203, 3, 1, 1)-ReLU-Flatten-Linear(9947, 10)-LogSoftmax(1);279035288;328;
+17;MNIST_25407;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 11, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(11, 200, 3, 1, 1)-ReLU-Flatten-Linear(9800, 10)-LogSoftmax(1);390813115;866;
+18;MNIST_10375;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 76, 3, 1, 1)-ReLU-Flatten-Linear(3724, 10)-LogSoftmax(1);417406150;885;
+19;MNIST_18651;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 163, 3, 1, 1)-ReLU-Flatten-Linear(7987, 10)-LogSoftmax(1);676716673;630;
+20;MNIST_28387;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 61, 3, 1, 1)-ReLU-Flatten-Linear(2989, 10)-LogSoftmax(1);483634587;495;
+21;MNIST_24489;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 145, 3, 1, 1)-ReLU-Flatten-Linear(7105, 10)-LogSoftmax(1);587269322;614;
+22;MNIST_16764;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 137, 3, 1, 1)-ReLU-Flatten-Linear(6713, 10)-LogSoftmax(1);896937394;463;
+23;MNIST_13599;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 220, 3, 1, 1)-ReLU-Flatten-Linear(10780, 10)-LogSoftmax(1);482762874;986;
+24;MNIST_11563;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 168, 3, 1, 1)-ReLU-Flatten-Linear(8232, 10)-LogSoftmax(1);62829004;561;
+25;MNIST_10560;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 170, 3, 1, 1)-ReLU-Flatten-Linear(8330, 10)-LogSoftmax(1);432423572;882;
+26;MNIST_9747;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 14, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(14, 110, 3, 1, 1)-ReLU-Flatten-Linear(5390, 10)-LogSoftmax(1);90250538;933;
+27;MNIST_13566;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 13, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(13, 209, 3, 1, 1)-ReLU-Flatten-Linear(10241, 10)-LogSoftmax(1);700199045;773;
+28;MNIST_7436;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 13, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(13, 141, 3, 1, 1)-ReLU-Flatten-Linear(6909, 10)-LogSoftmax(1);922181636;891;
+29;MNIST_23145;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 23, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(23, 148, 3, 1, 1)-ReLU-Flatten-Linear(7252, 10)-LogSoftmax(1);88633046;927;
+30;MNIST_7037;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 61, 3, 1, 1)-ReLU-Flatten-Linear(2989, 10)-LogSoftmax(1);216976841;869;
+31;MNIST_14864;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 130, 3, 1, 1)-ReLU-Flatten-Linear(6370, 10)-LogSoftmax(1);650853430;605;
+32;MNIST_6389;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 145, 3, 1, 1)-ReLU-Flatten-Linear(7105, 10)-LogSoftmax(1);502921667;960;
+33;MNIST_9382;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 10, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(10, 77, 3, 1, 1)-ReLU-Flatten-Linear(3773, 10)-LogSoftmax(1);850545979;309;
+34;MNIST_15342;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 235, 3, 1, 1)-ReLU-Flatten-Linear(11515, 10)-LogSoftmax(1);283063259;447;
+35;MNIST_20967;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 10, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(10, 50, 3, 1, 1)-ReLU-Flatten-Linear(2450, 10)-LogSoftmax(1);632368576;862;
+36;MNIST_19502;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 17, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(17, 184, 3, 1, 1)-ReLU-Flatten-Linear(9016, 10)-LogSoftmax(1);27639030;643;
+37;MNIST_12150;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 14, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(14, 79, 3, 1, 1)-ReLU-Flatten-Linear(3871, 10)-LogSoftmax(1);923343467;374;
+38;MNIST_24988;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 81, 3, 1, 1)-ReLU-Flatten-Linear(3969, 10)-LogSoftmax(1);665581259;493;
+39;MNIST_13019;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 114, 3, 1, 1)-ReLU-Flatten-Linear(5586, 10)-LogSoftmax(1);56307477;761;
+40;MNIST_10594;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 11, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(11, 208, 3, 1, 1)-ReLU-Flatten-Linear(10192, 10)-LogSoftmax(1);942298166;582;
+41;MNIST_8597;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 166, 3, 1, 1)-ReLU-Flatten-Linear(8134, 10)-LogSoftmax(1);313386172;719;
+42;MNIST_25993;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 18, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(18, 140, 3, 1, 1)-ReLU-Flatten-Linear(6860, 10)-LogSoftmax(1);840627960;350;
+43;MNIST_9860;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 132, 3, 1, 1)-ReLU-Flatten-Linear(6468, 10)-LogSoftmax(1);222275973;486;
+44;MNIST_26577;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 99, 3, 1, 1)-ReLU-Flatten-Linear(4851, 10)-LogSoftmax(1);972703876;803;
+45;MNIST_16576;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 61, 3, 1, 1)-ReLU-Flatten-Linear(2989, 10)-LogSoftmax(1);109829453;990;
+46;MNIST_21365;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 22, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(22, 224, 3, 1, 1)-ReLU-Flatten-Linear(10976, 10)-LogSoftmax(1);383971483;436;
+47;MNIST_13892;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 105, 3, 1, 1)-ReLU-Flatten-Linear(5145, 10)-LogSoftmax(1);643171914;832;
+48;MNIST_9794;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 132, 3, 1, 1)-ReLU-Flatten-Linear(6468, 10)-LogSoftmax(1);670675758;437;
+49;MNIST_11715;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 231, 3, 1, 1)-ReLU-Flatten-Linear(11319, 10)-LogSoftmax(1);718247435;662;
+50;MNIST_12776;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 94, 3, 1, 1)-ReLU-Flatten-Linear(4606, 10)-LogSoftmax(1);870508044;842;
+51;MNIST_6065;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 24, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(24, 162, 3, 1, 1)-ReLU-Flatten-Linear(7938, 10)-LogSoftmax(1);384966093;506;
+52;MNIST_17105;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 204, 3, 1, 1)-ReLU-Flatten-Linear(9996, 10)-LogSoftmax(1);263374537;815;
+53;MNIST_27050;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 228, 3, 1, 1)-ReLU-Flatten-Linear(11172, 10)-LogSoftmax(1);845919665;592;
+54;MNIST_26067;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 142, 3, 1, 1)-ReLU-Flatten-Linear(6958, 10)-LogSoftmax(1);892572882;691;
+55;MNIST_26954;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 14, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(14, 104, 3, 1, 1)-ReLU-Flatten-Linear(5096, 10)-LogSoftmax(1);493715397;434;
+56;MNIST_16397;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 161, 3, 1, 1)-ReLU-Flatten-Linear(7889, 10)-LogSoftmax(1);559166445;318;
+57;MNIST_18417;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 20, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(20, 184, 3, 1, 1)-ReLU-Flatten-Linear(9016, 10)-LogSoftmax(1);681400868;783;
+58;MNIST_29469;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 99, 3, 1, 1)-ReLU-Flatten-Linear(4851, 10)-LogSoftmax(1);132784964;608;
+59;MNIST_12426;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 13, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(13, 246, 3, 1, 1)-ReLU-Flatten-Linear(12054, 10)-LogSoftmax(1);179219961;333;
+60;MNIST_15386;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 14, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(14, 69, 3, 1, 1)-ReLU-Flatten-Linear(3381, 10)-LogSoftmax(1);342942351;870;
+61;MNIST_12876;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 183, 3, 1, 1)-ReLU-Flatten-Linear(8967, 10)-LogSoftmax(1);138470507;414;
+62;MNIST_6861;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 13, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(13, 195, 3, 1, 1)-ReLU-Flatten-Linear(9555, 10)-LogSoftmax(1);395086689;447;
+63;MNIST_21264;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 207, 3, 1, 1)-ReLU-Flatten-Linear(10143, 10)-LogSoftmax(1);551353010;631;
+64;MNIST_6136;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 17, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(17, 76, 3, 1, 1)-ReLU-Flatten-Linear(3724, 10)-LogSoftmax(1);526563120;832;
+65;MNIST_17048;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 155, 3, 1, 1)-ReLU-Flatten-Linear(7595, 10)-LogSoftmax(1);547892100;520;
+66;MNIST_14280;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 10, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(10, 143, 3, 1, 1)-ReLU-Flatten-Linear(7007, 10)-LogSoftmax(1);958733851;360;
+67;MNIST_17402;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 178, 3, 1, 1)-ReLU-Flatten-Linear(8722, 10)-LogSoftmax(1);46920774;971;
+68;MNIST_20330;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 17, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(17, 99, 3, 1, 1)-ReLU-Flatten-Linear(4851, 10)-LogSoftmax(1);48914259;892;
+69;MNIST_27753;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 203, 3, 1, 1)-ReLU-Flatten-Linear(9947, 10)-LogSoftmax(1);606982068;873;
+70;MNIST_19586;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 187, 3, 1, 1)-ReLU-Flatten-Linear(9163, 10)-LogSoftmax(1);909692949;481;
+71;MNIST_17283;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 159, 3, 1, 1)-ReLU-Flatten-Linear(7791, 10)-LogSoftmax(1);709429373;665;
+72;MNIST_10940;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 205, 3, 1, 1)-ReLU-Flatten-Linear(10045, 10)-LogSoftmax(1);826615641;927;
+73;MNIST_6516;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 201, 3, 1, 1)-ReLU-Flatten-Linear(9849, 10)-LogSoftmax(1);453979089;326;
+74;MNIST_17272;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 23, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(23, 174, 3, 1, 1)-ReLU-Flatten-Linear(8526, 10)-LogSoftmax(1);330914630;371;
+75;MNIST_14462;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 224, 3, 1, 1)-ReLU-Flatten-Linear(10976, 10)-LogSoftmax(1);656456878;447;
+76;MNIST_17598;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 242, 3, 1, 1)-ReLU-Flatten-Linear(11858, 10)-LogSoftmax(1);692559468;453;
+77;MNIST_8266;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 78, 3, 1, 1)-ReLU-Flatten-Linear(3822, 10)-LogSoftmax(1);94848173;340;
+78;MNIST_14006;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 17, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(17, 121, 3, 1, 1)-ReLU-Flatten-Linear(5929, 10)-LogSoftmax(1);767999222;894;
+79;MNIST_12575;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 18, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(18, 91, 3, 1, 1)-ReLU-Flatten-Linear(4459, 10)-LogSoftmax(1);732740509;365;
+80;MNIST_29246;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 148, 3, 1, 1)-ReLU-Flatten-Linear(7252, 10)-LogSoftmax(1);955411453;552;
+81;MNIST_20505;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 22, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(22, 113, 3, 1, 1)-ReLU-Flatten-Linear(5537, 10)-LogSoftmax(1);108292162;318;
+82;MNIST_17754;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 89, 3, 1, 1)-ReLU-Flatten-Linear(4361, 10)-LogSoftmax(1);628500738;844;
+83;MNIST_28308;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 148, 3, 1, 1)-ReLU-Flatten-Linear(7252, 10)-LogSoftmax(1);784353714;714;
+84;MNIST_6223;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 13, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(13, 189, 3, 1, 1)-ReLU-Flatten-Linear(9261, 10)-LogSoftmax(1);555927508;330;
+85;MNIST_7796;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 104, 3, 1, 1)-ReLU-Flatten-Linear(5096, 10)-LogSoftmax(1);121374477;716;
+86;MNIST_9961;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 227, 3, 1, 1)-ReLU-Flatten-Linear(11123, 10)-LogSoftmax(1);614230241;346;
+87;MNIST_10411;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 102, 3, 1, 1)-ReLU-Flatten-Linear(4998, 10)-LogSoftmax(1);803575145;367;
+88;MNIST_14862;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 18, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(18, 67, 3, 1, 1)-ReLU-Flatten-Linear(3283, 10)-LogSoftmax(1);100099230;622;
+89;MNIST_21663;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 169, 3, 1, 1)-ReLU-Flatten-Linear(8281, 10)-LogSoftmax(1);341604808;511;
+90;MNIST_23871;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 83, 3, 1, 1)-ReLU-Flatten-Linear(4067, 10)-LogSoftmax(1);896839395;444;
+91;MNIST_7708;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 18, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(18, 157, 3, 1, 1)-ReLU-Flatten-Linear(7693, 10)-LogSoftmax(1);300655330;705;
+92;MNIST_11033;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 152, 3, 1, 1)-ReLU-Flatten-Linear(7448, 10)-LogSoftmax(1);800030713;467;
+93;MNIST_22369;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 61, 3, 1, 1)-ReLU-Flatten-Linear(2989, 10)-LogSoftmax(1);854754649;537;
+94;MNIST_13620;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 22, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(22, 63, 3, 1, 1)-ReLU-Flatten-Linear(3087, 10)-LogSoftmax(1);316005183;863;
+95;MNIST_11106;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 146, 3, 1, 1)-ReLU-Flatten-Linear(7154, 10)-LogSoftmax(1);662691777;891;
+96;MNIST_13878;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 11, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(11, 238, 3, 1, 1)-ReLU-Flatten-Linear(11662, 10)-LogSoftmax(1);888322309;710;
+97;MNIST_18961;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 161, 3, 1, 1)-ReLU-Flatten-Linear(7889, 10)-LogSoftmax(1);574142479;792;
+98;MNIST_8197;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 231, 3, 1, 1)-ReLU-Flatten-Linear(11319, 10)-LogSoftmax(1);211625200;390;
+99;MNIST_14075;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 24, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(24, 54, 3, 1, 1)-ReLU-Flatten-Linear(2646, 10)-LogSoftmax(1);498679091;643;
diff --git a/dacbench/instance_sets/sgd/sgd_val_100instances.csv b/dacbench/instance_sets/sgd/sgd_val_100instances.csv
new file mode 100755
index 000000000..70ae44210
--- /dev/null
+++ b/dacbench/instance_sets/sgd/sgd_val_100instances.csv
@@ -0,0 +1,101 @@
+ID;dataset;architecture;seed;steps;
+0;MNIST_20144;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 10, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(10, 129, 3, 1, 1)-ReLU-Flatten-Linear(6321, 10)-LogSoftmax(1);481702131;508;
+1;MNIST_23398;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 123, 3, 1, 1)-ReLU-Flatten-Linear(6027, 10)-LogSoftmax(1);36639344;789;
+2;MNIST_8327;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 24, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(24, 212, 3, 1, 1)-ReLU-Flatten-Linear(10388, 10)-LogSoftmax(1);944517857;686;
+3;MNIST_22333;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 234, 3, 1, 1)-ReLU-Flatten-Linear(11466, 10)-LogSoftmax(1);3996140;311;
+4;MNIST_22678;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 18, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(18, 105, 3, 1, 1)-ReLU-Flatten-Linear(5145, 10)-LogSoftmax(1);709754768;632;
+5;MNIST_13380;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 248, 3, 1, 1)-ReLU-Flatten-Linear(12152, 10)-LogSoftmax(1);3357919;867;
+6;MNIST_25715;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 24, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(24, 181, 3, 1, 1)-ReLU-Flatten-Linear(8869, 10)-LogSoftmax(1);129335604;372;
+7;MNIST_14186;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 24, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(24, 232, 3, 1, 1)-ReLU-Flatten-Linear(11368, 10)-LogSoftmax(1);124261310;441;
+8;MNIST_24155;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 152, 3, 1, 1)-ReLU-Flatten-Linear(7448, 10)-LogSoftmax(1);947558146;702;
+9;MNIST_10898;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 13, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(13, 244, 3, 1, 1)-ReLU-Flatten-Linear(11956, 10)-LogSoftmax(1);410482975;677;
+10;MNIST_11599;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 242, 3, 1, 1)-ReLU-Flatten-Linear(11858, 10)-LogSoftmax(1);920568712;804;
+11;MNIST_7570;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 87, 3, 1, 1)-ReLU-Flatten-Linear(4263, 10)-LogSoftmax(1);904913090;559;
+12;MNIST_25019;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 202, 3, 1, 1)-ReLU-Flatten-Linear(9898, 10)-LogSoftmax(1);686667384;727;
+13;MNIST_22864;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 17, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(17, 113, 3, 1, 1)-ReLU-Flatten-Linear(5537, 10)-LogSoftmax(1);467546562;708;
+14;MNIST_6792;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 23, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(23, 136, 3, 1, 1)-ReLU-Flatten-Linear(6664, 10)-LogSoftmax(1);437889825;586;
+15;MNIST_8278;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 134, 3, 1, 1)-ReLU-Flatten-Linear(6566, 10)-LogSoftmax(1);31864339;967;
+16;MNIST_6840;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 197, 3, 1, 1)-ReLU-Flatten-Linear(9653, 10)-LogSoftmax(1);556715875;798;
+17;MNIST_26588;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 130, 3, 1, 1)-ReLU-Flatten-Linear(6370, 10)-LogSoftmax(1);607394867;680;
+18;MNIST_19170;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 20, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(20, 207, 3, 1, 1)-ReLU-Flatten-Linear(10143, 10)-LogSoftmax(1);94497579;858;
+19;MNIST_16935;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 10, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(10, 74, 3, 1, 1)-ReLU-Flatten-Linear(3626, 10)-LogSoftmax(1);329442673;651;
+20;MNIST_19957;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 121, 3, 1, 1)-ReLU-Flatten-Linear(5929, 10)-LogSoftmax(1);148097105;705;
+21;MNIST_9105;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 23, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(23, 214, 3, 1, 1)-ReLU-Flatten-Linear(10486, 10)-LogSoftmax(1);461232490;786;
+22;MNIST_27093;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 22, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(22, 249, 3, 1, 1)-ReLU-Flatten-Linear(12201, 10)-LogSoftmax(1);83613096;712;
+23;MNIST_27441;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 198, 3, 1, 1)-ReLU-Flatten-Linear(9702, 10)-LogSoftmax(1);676186936;570;
+24;MNIST_24649;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 23, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(23, 236, 3, 1, 1)-ReLU-Flatten-Linear(11564, 10)-LogSoftmax(1);99750884;769;
+25;MNIST_19788;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 10, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(10, 78, 3, 1, 1)-ReLU-Flatten-Linear(3822, 10)-LogSoftmax(1);144954009;351;
+26;MNIST_29794;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 23, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(23, 83, 3, 1, 1)-ReLU-Flatten-Linear(4067, 10)-LogSoftmax(1);808335337;900;
+27;MNIST_18885;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 246, 3, 1, 1)-ReLU-Flatten-Linear(12054, 10)-LogSoftmax(1);52875917;560;
+28;MNIST_27723;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 14, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(14, 165, 3, 1, 1)-ReLU-Flatten-Linear(8085, 10)-LogSoftmax(1);350127024;932;
+29;MNIST_13086;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 90, 3, 1, 1)-ReLU-Flatten-Linear(4410, 10)-LogSoftmax(1);416188290;687;
+30;MNIST_14995;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 10, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(10, 57, 3, 1, 1)-ReLU-Flatten-Linear(2793, 10)-LogSoftmax(1);231668807;715;
+31;MNIST_27656;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 178, 3, 1, 1)-ReLU-Flatten-Linear(8722, 10)-LogSoftmax(1);82002826;778;
+32;MNIST_23815;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 98, 3, 1, 1)-ReLU-Flatten-Linear(4802, 10)-LogSoftmax(1);418004922;841;
+33;MNIST_20910;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 22, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(22, 137, 3, 1, 1)-ReLU-Flatten-Linear(6713, 10)-LogSoftmax(1);921680356;519;
+34;MNIST_22852;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 179, 3, 1, 1)-ReLU-Flatten-Linear(8771, 10)-LogSoftmax(1);710784163;354;
+35;MNIST_20001;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 219, 3, 1, 1)-ReLU-Flatten-Linear(10731, 10)-LogSoftmax(1);295223054;941;
+36;MNIST_26044;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 18, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(18, 85, 3, 1, 1)-ReLU-Flatten-Linear(4165, 10)-LogSoftmax(1);290585899;899;
+37;MNIST_16858;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 206, 3, 1, 1)-ReLU-Flatten-Linear(10094, 10)-LogSoftmax(1);774610992;837;
+38;MNIST_25781;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 227, 3, 1, 1)-ReLU-Flatten-Linear(11123, 10)-LogSoftmax(1);93781380;961;
+39;MNIST_6871;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 234, 3, 1, 1)-ReLU-Flatten-Linear(11466, 10)-LogSoftmax(1);709064253;499;
+40;MNIST_24111;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 22, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(22, 100, 3, 1, 1)-ReLU-Flatten-Linear(4900, 10)-LogSoftmax(1);657550928;888;
+41;MNIST_12988;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 23, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(23, 157, 3, 1, 1)-ReLU-Flatten-Linear(7693, 10)-LogSoftmax(1);461349277;563;
+42;MNIST_27149;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 90, 3, 1, 1)-ReLU-Flatten-Linear(4410, 10)-LogSoftmax(1);282398476;347;
+43;MNIST_21970;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 130, 3, 1, 1)-ReLU-Flatten-Linear(6370, 10)-LogSoftmax(1);443582722;524;
+44;MNIST_19190;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 24, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(24, 111, 3, 1, 1)-ReLU-Flatten-Linear(5439, 10)-LogSoftmax(1);249869496;367;
+45;MNIST_7517;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 14, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(14, 163, 3, 1, 1)-ReLU-Flatten-Linear(7987, 10)-LogSoftmax(1);383320070;720;
+46;MNIST_20480;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 185, 3, 1, 1)-ReLU-Flatten-Linear(9065, 10)-LogSoftmax(1);67420305;418;
+47;MNIST_13436;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 23, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(23, 121, 3, 1, 1)-ReLU-Flatten-Linear(5929, 10)-LogSoftmax(1);99435385;537;
+48;MNIST_9924;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 18, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(18, 144, 3, 1, 1)-ReLU-Flatten-Linear(7056, 10)-LogSoftmax(1);428831551;743;
+49;MNIST_7033;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 10, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(10, 190, 3, 1, 1)-ReLU-Flatten-Linear(9310, 10)-LogSoftmax(1);520133527;825;
+50;MNIST_19767;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 206, 3, 1, 1)-ReLU-Flatten-Linear(10094, 10)-LogSoftmax(1);605821281;366;
+51;MNIST_16686;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 14, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(14, 207, 3, 1, 1)-ReLU-Flatten-Linear(10143, 10)-LogSoftmax(1);979695168;489;
+52;MNIST_19190;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 228, 3, 1, 1)-ReLU-Flatten-Linear(11172, 10)-LogSoftmax(1);492709606;642;
+53;MNIST_18459;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 23, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(23, 77, 3, 1, 1)-ReLU-Flatten-Linear(3773, 10)-LogSoftmax(1);542175131;943;
+54;MNIST_20051;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 11, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(11, 113, 3, 1, 1)-ReLU-Flatten-Linear(5537, 10)-LogSoftmax(1);246484771;786;
+55;MNIST_7393;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 143, 3, 1, 1)-ReLU-Flatten-Linear(7007, 10)-LogSoftmax(1);308781085;536;
+56;MNIST_9561;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 217, 3, 1, 1)-ReLU-Flatten-Linear(10633, 10)-LogSoftmax(1);8602731;813;
+57;MNIST_7989;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 120, 3, 1, 1)-ReLU-Flatten-Linear(5880, 10)-LogSoftmax(1);913510679;930;
+58;MNIST_24002;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 11, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(11, 146, 3, 1, 1)-ReLU-Flatten-Linear(7154, 10)-LogSoftmax(1);346708151;610;
+59;MNIST_22117;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 161, 3, 1, 1)-ReLU-Flatten-Linear(7889, 10)-LogSoftmax(1);270916049;986;
+60;MNIST_9423;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 84, 3, 1, 1)-ReLU-Flatten-Linear(4116, 10)-LogSoftmax(1);574711805;549;
+61;MNIST_26100;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 13, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(13, 122, 3, 1, 1)-ReLU-Flatten-Linear(5978, 10)-LogSoftmax(1);110935159;423;
+62;MNIST_7479;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 11, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(11, 81, 3, 1, 1)-ReLU-Flatten-Linear(3969, 10)-LogSoftmax(1);154821165;803;
+63;MNIST_16574;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 70, 3, 1, 1)-ReLU-Flatten-Linear(3430, 10)-LogSoftmax(1);423832342;963;
+64;MNIST_27590;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 141, 3, 1, 1)-ReLU-Flatten-Linear(6909, 10)-LogSoftmax(1);983337432;929;
+65;MNIST_7103;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 238, 3, 1, 1)-ReLU-Flatten-Linear(11662, 10)-LogSoftmax(1);945051840;651;
+66;MNIST_7771;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 242, 3, 1, 1)-ReLU-Flatten-Linear(11858, 10)-LogSoftmax(1);921571008;649;
+67;MNIST_9667;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 22, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(22, 93, 3, 1, 1)-ReLU-Flatten-Linear(4557, 10)-LogSoftmax(1);304836284;793;
+68;MNIST_8543;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 246, 3, 1, 1)-ReLU-Flatten-Linear(12054, 10)-LogSoftmax(1);959569119;687;
+69;MNIST_22037;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 53, 3, 1, 1)-ReLU-Flatten-Linear(2597, 10)-LogSoftmax(1);351111802;882;
+70;MNIST_28916;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 178, 3, 1, 1)-ReLU-Flatten-Linear(8722, 10)-LogSoftmax(1);786039844;847;
+71;MNIST_22329;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 124, 3, 1, 1)-ReLU-Flatten-Linear(6076, 10)-LogSoftmax(1);394068764;771;
+72;MNIST_6855;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 87, 3, 1, 1)-ReLU-Flatten-Linear(4263, 10)-LogSoftmax(1);240023450;752;
+73;MNIST_21105;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 16, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(16, 171, 3, 1, 1)-ReLU-Flatten-Linear(8379, 10)-LogSoftmax(1);376055264;426;
+74;MNIST_13951;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 229, 3, 1, 1)-ReLU-Flatten-Linear(11221, 10)-LogSoftmax(1);971977512;356;
+75;MNIST_18672;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 75, 3, 1, 1)-ReLU-Flatten-Linear(3675, 10)-LogSoftmax(1);361570205;406;
+76;MNIST_21145;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 208, 3, 1, 1)-ReLU-Flatten-Linear(10192, 10)-LogSoftmax(1);197549324;811;
+77;MNIST_29286;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 22, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(22, 69, 3, 1, 1)-ReLU-Flatten-Linear(3381, 10)-LogSoftmax(1);260209140;389;
+78;MNIST_29460;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 20, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(20, 81, 3, 1, 1)-ReLU-Flatten-Linear(3969, 10)-LogSoftmax(1);356942507;988;
+79;MNIST_6334;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 222, 3, 1, 1)-ReLU-Flatten-Linear(10878, 10)-LogSoftmax(1);40903392;559;
+80;MNIST_25406;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 10, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(10, 135, 3, 1, 1)-ReLU-Flatten-Linear(6615, 10)-LogSoftmax(1);61305260;368;
+81;MNIST_9749;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 12, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(12, 99, 3, 1, 1)-ReLU-Flatten-Linear(4851, 10)-LogSoftmax(1);891164401;723;
+82;MNIST_9736;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 18, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(18, 107, 3, 1, 1)-ReLU-Flatten-Linear(5243, 10)-LogSoftmax(1);801129243;696;
+83;MNIST_17115;Conv2d(1, 4, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(4, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 174, 3, 1, 1)-ReLU-Flatten-Linear(8526, 10)-LogSoftmax(1);46885188;681;
+84;MNIST_22467;Conv2d(1, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 138, 3, 1, 1)-ReLU-Flatten-Linear(6762, 10)-LogSoftmax(1);852371428;824;
+85;MNIST_14879;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 96, 3, 1, 1)-ReLU-Flatten-Linear(4704, 10)-LogSoftmax(1);965525701;410;
+86;MNIST_7227;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 21, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(21, 151, 3, 1, 1)-ReLU-Flatten-Linear(7399, 10)-LogSoftmax(1);758954152;413;
+87;MNIST_11278;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 77, 3, 1, 1)-ReLU-Flatten-Linear(3773, 10)-LogSoftmax(1);385890773;800;
+88;MNIST_23751;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 22, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(22, 158, 3, 1, 1)-ReLU-Flatten-Linear(7742, 10)-LogSoftmax(1);671450452;787;
+89;MNIST_25538;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 212, 3, 1, 1)-ReLU-Flatten-Linear(10388, 10)-LogSoftmax(1);416173568;661;
+90;MNIST_6066;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 201, 3, 1, 1)-ReLU-Flatten-Linear(9849, 10)-LogSoftmax(1);501417673;317;
+91;MNIST_24147;Conv2d(1, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 7, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(7, 94, 3, 1, 1)-ReLU-Flatten-Linear(4606, 10)-LogSoftmax(1);249501706;590;
+92;MNIST_27508;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 6, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(6, 126, 3, 1, 1)-ReLU-Flatten-Linear(6174, 10)-LogSoftmax(1);750051938;787;
+93;MNIST_11600;Conv2d(1, 8, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(8, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 214, 3, 1, 1)-ReLU-Flatten-Linear(10486, 10)-LogSoftmax(1);519068173;357;
+94;MNIST_24135;Conv2d(1, 2, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(2, 22, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(22, 196, 3, 1, 1)-ReLU-Flatten-Linear(9604, 10)-LogSoftmax(1);604336783;841;
+95;MNIST_7090;Conv2d(1, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 9, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(9, 146, 3, 1, 1)-ReLU-Flatten-Linear(7154, 10)-LogSoftmax(1);807875335;607;
+96;MNIST_6363;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 122, 3, 1, 1)-ReLU-Flatten-Linear(5978, 10)-LogSoftmax(1);4800486;872;
+97;MNIST_18093;Conv2d(1, 3, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(3, 19, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(19, 123, 3, 1, 1)-ReLU-Flatten-Linear(6027, 10)-LogSoftmax(1);840441126;853;
+98;MNIST_29418;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 15, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(15, 94, 3, 1, 1)-ReLU-Flatten-Linear(4606, 10)-LogSoftmax(1);500493748;576;
+99;MNIST_23511;Conv2d(1, 5, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(5, 22, 3, 1, 1)-MaxPool2d(2, 2)-Conv2d(22, 229, 3, 1, 1)-ReLU-Flatten-Linear(11221, 10)-LogSoftmax(1);239701752;486;
diff --git a/dacbench/wrappers/observation_wrapper.py b/dacbench/wrappers/observation_wrapper.py
index 381d23cfd..ba56ee39e 100644
--- a/dacbench/wrappers/observation_wrapper.py
+++ b/dacbench/wrappers/observation_wrapper.py
@@ -22,7 +22,6 @@ def __init__(self, env):
         super(ObservationWrapper, self).__init__(env)
         obs_sample = self.flatten(self.env.observation_space.sample())
         size = len(obs_sample)
-        print(obs_sample)
         self.observation_space = spaces.Box(
             low=-np.inf * np.ones(size), high=np.inf * np.ones(size)
         )


Test patch:
diff --git a/dacbench/instance_sets/sgd/sgd_test.csv b/dacbench/instance_sets/sgd/sgd_test.csv
deleted file mode 100755
index 03c49fd8a..000000000
--- a/dacbench/instance_sets/sgd/sgd_test.csv
+++ /dev/null
@@ -1,101 +0,0 @@
-ID;dataset;architecture;seed
-0;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);991677751
-1;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);89253990
-2;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);14715380
-3;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);882182227
-4;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);348248224
-5;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);925327980
-6;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);288846000
-7;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);73028194
-8;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);387407830
-9;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);845740298
-10;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);125592004
-11;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);693454908
-12;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);33977867
-13;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);593005528
-14;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);429370976
-15;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);240533823
-16;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);216869516
-17;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);17078890
-18;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);226536992
-19;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);544815759
-20;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);647899944
-21;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);978624377
-22;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);581178319
-23;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);213867989
-24;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);413049050
-25;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);692878661
-26;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);163104440
-27;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);802040094
-28;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);188974351
-29;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);744117712
-30;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);864886481
-31;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);982107761
-32;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);596495567
-33;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);307918545
-34;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);327318216
-35;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);857321738
-36;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);611512479
-37;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);568932993
-38;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);517624937
-39;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);41136094
-40;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);369727269
-41;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);699494276
-42;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);846801016
-43;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);140372469
-44;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);710584291
-45;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);991716301
-46;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);707047413
-47;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);534572021
-48;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);522820494
-49;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);837714476
-50;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);748133095
-51;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);494418427
-52;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);67704707
-53;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);416225126
-54;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);183291606
-55;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);438195403
-56;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);430349608
-57;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);725210557
-58;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);766824199
-59;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);383724069
-60;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);60862474
-61;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);624959056
-62;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);31232852
-63;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);564516808
-64;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);660885855
-65;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);201506331
-66;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);594836560
-67;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);567790985
-68;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);790011941
-69;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);184735941
-70;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);133686464
-71;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);405943423
-72;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);335056357
-73;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);256148150
-74;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);943806410
-75;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);450953920
-76;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);911361041
-77;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);378340135
-78;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);360377068
-79;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);154931826
-80;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);286748838
-81;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);318822599
-82;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);134361460
-83;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);681667269
-84;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);712502215
-85;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);6753465
-86;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);451693990
-87;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);98119686
-88;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);674285019
-89;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);60378459
-90;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);330562649
-91;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);214506082
-92;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);54291250
-93;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);325081022
-94;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);733836850
-95;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);816642203
-96;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);145003067
-97;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);861957363
-98;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);759714321
-99;MNIST;Linear(784,16)-Sigmoid-Linear(16,16)-Sigmoid-Linear(16,10)-LogSoftmax(1);269644774
\ No newline at end of file
diff --git a/tests/envs/test_sgd.py b/tests/envs/test_sgd.py
index 7d1d85df3..431c1adde 100644
--- a/tests/envs/test_sgd.py
+++ b/tests/envs/test_sgd.py
@@ -1,45 +1,146 @@
-
 import unittest
+import os
+
 import numpy as np
 from dacbench import AbstractEnv
+from dacbench.envs.sgd import SGDEnv, Reward
+from dacbench.abstract_benchmark import AbstractBenchmark, objdict
 from dacbench.benchmarks.sgd_benchmark import SGDBenchmark, SGD_DEFAULTS
+from dacbench.wrappers import ObservationWrapper
 
 
 class TestSGDEnv(unittest.TestCase):
-    def make_env(self):
+    def setUp(self):
         bench = SGDBenchmark()
-        env = bench.get_environment()
-        return env
+        self.env = bench.get_benchmark(seed=123)
+
+    @staticmethod
+    def data_path(path):
+        return os.path.join(os.path.dirname(__file__), 'data', path)
 
     def test_setup(self):
-        env = self.make_env()
-        self.assertTrue(issubclass(type(env), AbstractEnv))
-        self.assertTrue(env.no_cuda)
-        self.assertTrue(env.model is None)
-        self.assertTrue(env.current_training_loss is None)
-        self.assertTrue(env.batch_size == SGD_DEFAULTS["training_batch_size"])
-        self.assertTrue(env.initial_lr == env.current_lr)
+        self.assertTrue(issubclass(type(self.env), AbstractEnv))
+        self.assertFalse(self.env.no_cuda)
+        self.assertTrue(self.env.model is None)
+        self.assertTrue(self.env.current_training_loss is None)
+        self.assertTrue(self.env.batch_size == SGD_DEFAULTS["training_batch_size"])
+        self.assertTrue(self.env.initial_lr == self.env.current_lr)
+
+    def test_reward_type(self):
+        benchmark = SGDBenchmark()
+        benchmark.config = objdict(SGD_DEFAULTS.copy())
+        benchmark.read_instance_set()
+
+        env = SGDEnv(benchmark.config)
+        self.assertEqual(env.reward_type, SGD_DEFAULTS.reward_type)
+
+        benchmark.config.reward_type = SGD_DEFAULTS.reward_type.name
+        env = SGDEnv(benchmark.config)
+        self.assertEqual(env.reward_type, SGD_DEFAULTS.reward_type)
+
+        benchmark.config.reward_type = 'invalid_reward'
+        with self.assertRaises(ValueError):
+            env = SGDEnv(benchmark.config)
+
+        benchmark.config.reward_type = 0
+        with self.assertRaises(ValueError):
+            env = SGDEnv(benchmark.config)
 
     def test_reset(self):
-        env = self.make_env()
-        env.reset()
-        self.assertFalse(env.model is None)
-        self.assertFalse(env.train_dataset is None)
-        self.assertFalse(env.validation_dataset is None)
+        self.env.reset()
+        self.assertFalse(self.env.model is None)
+        self.assertFalse(self.env.train_dataset is None)
+        self.assertFalse(self.env.validation_dataset is None)
 
     def test_step(self):
-        env = self.make_env()
+        benchmark = SGDBenchmark()
+        benchmark.config = objdict(SGD_DEFAULTS.copy())
+        benchmark.read_instance_set()
+
+        for reward_type in Reward:
+            benchmark.config.reward_type = reward_type
+            env = SGDEnv(benchmark.config)
+            env = ObservationWrapper(env)
+            self.assertTrue(env.reward_range == reward_type.func.frange)
+
+            env.reset()
+            state, reward, done, meta = env.step(1.0)
+            self.assertTrue(reward >= env.reward_range[0])
+            self.assertTrue(reward <= env.reward_range[1])
+            self.assertFalse(done)
+            self.assertTrue(len(meta.keys()) == 0)
+
+    def test_crash(self):
+        env = ObservationWrapper(self.env)
         env.reset()
-        state, reward, done, meta = env.step(1.0)
-        self.assertTrue(reward >= env.reward_range[0])
-        self.assertTrue(reward <= env.reward_range[1])
-        self.assertFalse(done)
-        self.assertTrue(len(meta.keys()) == 0)
+        state, reward, done, _ = env.step(np.nan)
+        self.assertTrue(env.crashed)
+        self.assertFalse(any(np.isnan(state)))
+        self.assertTrue(reward == env.crash_penalty)
+
+    def test_stateless(self):
+        env = ObservationWrapper(self.env)
+        rng = np.random.default_rng(123)
+        mems = []
+        instance_idxs = []
+        for _ in range(3):
+            env.reset()
+            instance_idxs.append(env.instance_index)
+
+            done = False
+            mem = []
+            step = 0
+            while not done and step < 5:
+                action = np.exp(rng.integers(low=-10, high=1))
+                state, reward, done, _ = env.step(action)
+                mem.append(np.concatenate([state, [reward, int(done), action]]))
+                step += 1
+            mems.append(np.array(mem))
+
+        rng = np.random.default_rng(123)
+        for i, idx in enumerate(reversed(instance_idxs)):
+            env.instance_index = idx - 1
+            env.reset()
+            self.assertTrue(env.instance_index == idx)
+
+            done = False
+            mem = []
+            step = 0
+            while not done and step < 5:
+                action = mems[-(i + 1)][step][-1]
+                state, reward, done, _ = env.step(action)
+                mem.append(np.concatenate([state, [reward, int(done), action]]))
+                step += 1
+            np.testing.assert_allclose(mems[-(i + 1)], np.array(mem))
+
+    def test_reproducibility(self):
+        mems = []
+        instances = []
+        env = ObservationWrapper(self.env)
+        for _ in range(2):
+            rng = np.random.default_rng(123)
+            env.seed(123)
+            env.instance_index = 0
+            instances.append(env.get_instance_set())
+
+            env.reset()
+
+            done = False
+            mem = []
+            step = 0
+            while not done and step < 5:
+                action = np.exp(rng.integers(low=-10, high=1))
+                state, reward, done, _ = env.step(action)
+                mem.append(np.concatenate([state, [reward, int(done), action]]))
+                step += 1
+            mems.append(np.array(mem))
+        self.assertEqual(mems[0].size, mems[1].size)
+        self.assertEqual(instances[0], instances[1])
+        np.testing.assert_allclose(mems[0], mems[1])
 
     def test_get_default_state(self):
-        env = self.make_env()
-        env.reset()
-        state, _, _, _ = env.step(0.5)
+        self.env.reset()
+        state, _, _, _ = self.env.step(0.5)
         self.assertTrue(issubclass(type(state), dict))
         self.assertTrue(
             np.array_equal(
@@ -52,19 +153,20 @@ def test_get_default_state(self):
                     "currentLR",
                     "trainingLoss",
                     "validationLoss",
+                    "step",
+                    "alignment",
+                    "crashed"
                 ],
             )
         )
-        self.assertTrue(state["currentLR"] == 10 ** -0.5)
+        self.assertTrue(state["currentLR"] == 0.5)
         self.assertTrue(state["trainingLoss"] > 0)
         self.assertTrue(state["validationLoss"] > 0)
 
     def test_close(self):
-        env = self.make_env()
-        self.assertTrue(env.close())
+        self.assertTrue(self.env.close())
 
     def test_render(self):
-        env = self.make_env()
-        env.render("human")
+        self.env.render("human")
         with self.assertRaises(NotImplementedError):
-            env.render("random")
+            self.env.render("random")
diff --git a/tests/test_abstract_benchmark.py b/tests/test_abstract_benchmark.py
index 825035935..134176bab 100644
--- a/tests/test_abstract_benchmark.py
+++ b/tests/test_abstract_benchmark.py
@@ -115,15 +115,14 @@ def assert_restorable(space):
                 "box": Box(
                     low=np.array([0, 0]),
                     high=np.array([1, 1]),
+                ),
+             'discrete': Discrete(
+                    n=2
                 )
             }
         )
         assert_restorable(space)
 
-        with self.assertRaises(ValueError):
-            space = Dict({"discrete": Discrete(1)})
-            bench.space_to_list(space)
-
         space = MultiDiscrete([2, 3])
         assert_restorable(space)
 
diff --git a/tests/test_config.json b/tests/test_config.json
index 9c33eac53..da755af02 100644
--- a/tests/test_config.json
+++ b/tests/test_config.json
@@ -1,6 +1,6 @@
 {
-  "state_method": ["dacbench.challenge_benchmarks.state_space_challenge.random_states", "small_random_sigmoid_state"],
-  "reward_function": ["dacbench.challenge_benchmarks.reward_quality_challenge.reward_functions", "random_reward"],
+  "state_method": ["function", "dacbench.challenge_benchmarks.state_space_challenge.random_states", "small_random_sigmoid_state"],
+  "reward_function": ["function","dacbench.challenge_benchmarks.reward_quality_challenge.reward_functions", "random_reward"],
   "_observation_space_type": "None",
   "wrappers": [["RewardNoiseWrapper", []]]
 }
diff --git a/tests/test_runner.py b/tests/test_runner.py
index 302e1e261..d164cb613 100644
--- a/tests/test_runner.py
+++ b/tests/test_runner.py
@@ -40,7 +40,6 @@ def __init__(self, env):
 
             def act(self, reward, state):
                 action = np.ones(self.num_actions)
-                print(self.num_actions)
                 if self.num_actions == 1:
                     action = 1
                 return action