Repository: zj-zhang/AMBER (language: python)

Issue:
Revert "Lazy init of utils.sampler"
Reverts zj-zhang/AMBER#9


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index a2dfd79..29ccc49 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,7 @@
 .idea/
 data/
 outputs/
+slurm-*
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
diff --git a/amber/__init__.py b/amber/__init__.py
index cee2c37..130a4c6 100644
--- a/amber/__init__.py
+++ b/amber/__init__.py
@@ -1,4 +1,5 @@
 """Docstring"""
+from ._wrapper import Amber
 import tensorflow as tf
 import os
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
@@ -8,7 +9,6 @@
     print(e)
     tf.logging.set_verbosity(tf.logging.ERROR)
 
-from ._wrapper import Amber
 
 __all__ = [
     'Amber',
diff --git a/amber/_getter.py b/amber/_getter.py
index 946a1a3..fba1c09 100644
--- a/amber/_getter.py
+++ b/amber/_getter.py
@@ -73,13 +73,16 @@ def get_train_env(env_type, controller, manager, *args, **kwargs):
 
 
 # controller; needs model_space
-def get_controller(controller_type, model_space, session, *args, **kwargs):
+def get_controller(controller_type, model_space, session, **kwargs):
     if controller_type == 'General' or controller_type == 'GeneralController':
         from .architect import GeneralController
-        controller = GeneralController(model_space, session=session, *args, **kwargs)
+        controller = GeneralController(model_space=model_space, session=session, **kwargs)
     elif controller_type == 'Operation' or controller_type == 'OperationController':
         from .architect import OperationController
-        controller = OperationController(model_space, *args, **kwargs)
+        controller = OperationController(model_space=model_space, **kwargs)
+    elif controller_type == 'MultiIO' or controller_type == 'MultiIOController':
+        from .architect import MultiIOController
+        controller = MultiIOController(model_space=model_space, session=session, **kwargs)
     else:
         raise Exception('cannot understand controller type: %s' % controller_type)
     print("controller = %s" % controller_type)
@@ -146,9 +149,9 @@ def get_manager(manager_type, model_fn, reward_fn, data_dict, session, *args, **
 
 # model_fn
 def get_modeler(model_fn_type, model_space, session, *args, **kwargs):
+    from .architect.model_space import State
     if model_fn_type == 'DAG' or model_fn_type == 'DAGModelBuilder':
-        from .architect.model_space import State
-        from modeler import DAGModelBuilder
+        from .modeler import DAGModelBuilder
         assert 'inputs_op' in kwargs and 'outputs_op' in kwargs
         inp_op_list = kwargs.pop("inputs_op")
         inputs_op = [State(**x) if not isinstance(x, State) else x for x in inp_op_list]
@@ -162,8 +165,7 @@ def get_modeler(model_fn_type, model_space, session, *args, **kwargs):
             session=session,
             *args, **kwargs)
     elif model_fn_type == 'Enas' or model_fn_type == 'EnasAnnModelBuilder':
-        from .architect.model_space import State
-        from modeler import EnasAnnModelBuilder
+        from .modeler import EnasAnnModelBuilder
         inp_op_list = kwargs.pop("inputs_op")
         inputs_op = [State(**x) if not isinstance(x, State) else x for x in inp_op_list]
         out_op_list = kwargs.pop("outputs_op")
@@ -176,7 +178,6 @@ def get_modeler(model_fn_type, model_space, session, *args, **kwargs):
             session=session,
             *args, **kwargs)
     elif model_fn_type == 'EnasCnnModelBuilder':
-        from .architect.model_space import State
         from .modeler import EnasCnnModelBuilder
         inp_op_list = kwargs.pop("inputs_op")
         inputs_op = [State(**x) if not isinstance(x, State) else x for x in inp_op_list]
@@ -191,6 +192,19 @@ def get_modeler(model_fn_type, model_space, session, *args, **kwargs):
             session=session,
             controller=controller,
             *args, **kwargs)
+    elif model_fn_type == 'KerasMultiIOModelBuilder':
+        from .modeler import KerasMultiIOModelBuilder
+        inp_op_list = kwargs.pop("inputs_op")
+        inputs_op = [State(**x) if not isinstance(x, State) else x for x in inp_op_list]
+        out_op_list = kwargs.pop("outputs_op")
+        output_op = [State(**x) if not isinstance(x, State) else x for x in out_op_list]
+        model_fn = KerasMultiIOModelBuilder(
+            model_space=model_space,
+            inputs_op=inputs_op,
+            output_op=output_op,
+            session=session,
+            *args, **kwargs)
+ 
     else:
         raise Exception('cannot understand model_builder type: %s' % model_fn_type)
     print("modeler = %s" % model_fn_type)
@@ -211,9 +225,9 @@ def get_reward_fn(reward_fn_type, knowledge_fn, *args, **kwargs):
         reward_fn = MockReward(*args, **kwargs)
     elif reward_fn_type == 'LossAucReward':
         from .architect.reward import LossAucReward
-        assert knowledge_fn is None, \
-            "Incompatability: LossAucReward must have knownledge_fn=None; got %s" % knowledge_fn
-        reward_fn = LossAucReward(*args, **kwargs)
+        #assert knowledge_fn is None, \
+        #    "Incompatability: LossAucReward must have knownledge_fn=None; got %s" % knowledge_fn
+        reward_fn = LossAucReward(knowledge_function=knowledge_fn, *args, **kwargs)
     else:
         raise Exception("cannot understand reward_fn type: %s" % reward_fn_type)
     print("reward = %s" % reward_fn_type)
@@ -233,12 +247,16 @@ def get_knowledge_fn(knowledge_fn_type, knowledge_data_dict, *args, **kwargs):
     elif knowledge_fn_type == 'Motif':
         from .objective import MotifKLDivergence
         k_fn = MotifKLDivergence(*args, **kwargs)
+    elif knowledge_fn_type == 'AuxilaryAcc':
+        from .objective import AuxilaryAcc
+        k_fn = AuxilaryAcc(*args, **kwargs)
     elif knowledge_fn_type == 'None' or knowledge_fn_type == 'zero':
         k_fn = None
     else:
         raise Exception("cannot understand knowledge_fn type: %s" % knowledge_fn_type)
     if k_fn is not None:
-        k_fn.knowledge_encoder(**knowledge_data_dict)
+        if hasattr(k_fn, "knowledge_encoder"):
+            k_fn.knowledge_encoder(**knowledge_data_dict)
     print("knowledge = %s" % knowledge_fn_type)
     return k_fn
 
diff --git a/amber/_wrapper.py b/amber/_wrapper.py
index 03c3582..edae701 100644
--- a/amber/_wrapper.py
+++ b/amber/_wrapper.py
@@ -4,8 +4,14 @@
 Overall wrapper class for DeepAmbre
 """
 
-from tensorflow import Session
-
+import tensorflow as tf
+from keras import backend as K
+try:
+    from tensorflow import Session
+except ImportError:
+    from tensorflow.compat.v1 import Session
+    tf.compat.v1.disable_eager_execution()
+import os
 from . import _getter
 
 
@@ -47,6 +53,10 @@ def __init__(self, types, specs=None):
 
         # use one tf.Session throughout one DA instance
         self.session = Session()
+        try:
+          K.set_session(self.session)
+        except Exception as e:
+          print("Failed to set Keras backend becasue of %s"%e)
 
         if specs is not None:
             self.from_dict(specs)
@@ -92,3 +102,4 @@ def from_dict(self, d):
     def run(self):
         assert self.is_built
         self.env.train()
+        self.controller.save_weights(os.path.join(self.env.working_dir, "controller_weights.h5"))
diff --git a/amber/architect/_general_controller.py b/amber/architect/_general_controller.py
index cf35ab8..b85bde5 100644
--- a/amber/architect/_general_controller.py
+++ b/amber/architect/_general_controller.py
@@ -7,6 +7,9 @@
 
 import h5py
 import tensorflow as tf
+if tf.__version__.startswith("2"):
+    tf.compat.v1.disable_eager_execution()
+    import tensorflow.compat.v1 as tf
 
 from .buffer import get_buffer
 from .common_ops import get_keras_train_ops
@@ -92,7 +95,7 @@ class GeneralController(BaseController):
     def __init__(self, model_space, buffer_type='ordinal', with_skip_connection=True, share_embedding=None,
                  use_ppo_loss=False, kl_threshold=0.05, skip_connection_unique_connection=False, buffer_size=15,
                  batch_size=5, session=None, train_pi_iter=20, lstm_size=32, lstm_num_layers=2, lstm_keep_prob=1.0,
-                 tanh_constant=None, temperature=None, optim_algo="adam", skip_target=0.8, skip_weight=0.5,
+                 tanh_constant=None, temperature=None, optim_algo="adam", skip_target=0.8, skip_weight=None,
                  rescale_advantage_by_reward=False, name="controller", **kwargs):
         super().__init__(**kwargs)
 
@@ -109,7 +112,7 @@ def __init__(self, model_space, buffer_type='ordinal', with_skip_connection=True
 
         buffer_fn = get_buffer(buffer_type)
         self.buffer = buffer_fn(max_size=buffer_size,
-                                ewa_beta=max(1 - 1. / buffer_size, 0.9),
+                                #ewa_beta=max(1 - 1. / buffer_size, 0.9),
                                 discount_factor=0.,
                                 is_squeeze_dim=True,
                                 rescale_advantage_by_reward=rescale_advantage_by_reward)
@@ -130,9 +133,12 @@ def __init__(self, model_space, buffer_type='ordinal', with_skip_connection=True
 
         self.skip_target = skip_target
         self.skip_weight = skip_weight
+        if self.skip_weight is not None:
+            assert self.with_skip_connection, "If skip_weight is not None, must have with_skip_connection=True"
 
         self.optim_algo = optim_algo
         self.name = name
+        self.loss = 0
 
         with tf.variable_scope(self.name):
             self._create_weight()
@@ -188,8 +194,9 @@ def _create_weight(self):
                 for layer_id in range(self.num_layers):
                     if self.share_embedding:
                         if layer_id not in self.share_embedding:
-                            self.w_soft["start"].append(tf.get_variable(
-                                "w_start", [self.lstm_size, self.num_choices_per_layer[layer_id]]))
+                            with tf.variable_scope("layer_{}".format(layer_id)):
+                                self.w_soft["start"].append(tf.get_variable(
+                                    "w_start", [self.lstm_size, self.num_choices_per_layer[layer_id]]))
                         else:
                             shared_id = self.share_embedding[layer_id]
                             assert shared_id < layer_id, \
@@ -366,7 +373,11 @@ def _build_trainer(self):
                   range(self.lstm_num_layers)]
         prev_h = [tf.zeros([batch_size, self.lstm_size], tf.float32) for _ in
                   range(self.lstm_num_layers)]
-        inputs = tf.matmul(tf.ones((batch_size, 1)), self.g_emb)
+        # only expand `g_emb` if necessary
+        if self.g_emb.shape[0] is not None and self.g_emb.shape[0].value == 1:
+            inputs = tf.matmul(tf.ones((batch_size, 1)), self.g_emb)
+        else:
+            inputs = self.g_emb
         skip_targets = tf.constant([1.0 - self.skip_target, self.skip_target],
                                    dtype=tf.float32)
 
@@ -437,8 +448,9 @@ def _build_trainer(self):
                     skip = tf.to_int32(skip)
 
                     skip_prob = tf.sigmoid(logit)
-                    kl = skip_prob * tf.log(skip_prob / skip_targets)
-                    kl = tf.reduce_sum(kl)
+                    kl = skip_prob * tf.log(skip_prob / skip_targets)  # (batch_size*layer_id, 2)
+                    kl = tf.reduce_sum(kl, axis=1)    # (batch_size*layer_id,)
+                    kl = tf.reshape(kl, [batch_size, -1])  # (batch_size, layer_id)
                     skip_penaltys.append(kl)
 
                     log_prob3 = tf.nn.sparse_softmax_cross_entropy_with_logits(
@@ -484,8 +496,8 @@ def _build_trainer(self):
         self.onehot_log_prob = tf.reduce_sum(log_probs, axis=0)
         skip_count = tf.stack(skip_count)
         self.onehot_skip_count = tf.reduce_sum(skip_count, axis=0)
-        skip_penaltys = tf.stack(skip_penaltys)
-        self.onehot_skip_penaltys = tf.reduce_mean(skip_penaltys, axis=0)
+        skip_penaltys_flat = [tf.reduce_mean(x, axis=1) for x in skip_penaltys] # from (num_layer-1, batch_size, layer_id) to (num_layer-1, batch_size); layer_id makes each tensor of varying lengths in the list
+        self.onehot_skip_penaltys = tf.reduce_mean(skip_penaltys_flat, axis=0)  # (batch_size,)
 
     def _build_train_op(self):
         """build train_op
@@ -500,8 +512,10 @@ def _build_train_op(self):
         self.input_arc_onehot = self.convert_arc_to_onehot(self)
         self.old_probs = [tf.placeholder(shape=self.onehot_probs[i].shape, dtype=tf.float32, name="old_prob_%i" % i) for
                           i in range(len(self.onehot_probs))]
+        if self.skip_weight is not None:
+            self.loss += self.skip_weight * tf.reduce_mean(self.onehot_skip_penaltys)
         if self.use_ppo_loss:
-            self.loss = proximal_policy_optimization_loss(
+            self.loss += proximal_policy_optimization_loss(
                 curr_prediction=self.onehot_probs,
                 curr_onehot=self.input_arc_onehot,
                 old_prediction=self.old_probs,
@@ -510,9 +524,7 @@ def _build_train_op(self):
                 advantage=self.advantage,
                 clip_val=0.2)
         else:
-            self.loss = tf.reshape(tf.tensordot(self.onehot_log_prob, self.advantage, axes=1), [])
-            if self.skip_weight is not None:
-                self.loss += self.skip_weight * self.onehot_skip_penaltys
+            self.loss += tf.reshape(tf.tensordot(self.onehot_log_prob, self.advantage, axes=1), [])
 
         self.kl_div, self.ent = get_kl_divergence_n_entropy(curr_prediction=self.onehot_probs,
                                                             old_prediction=self.old_probs,
@@ -529,7 +541,7 @@ def _build_train_op(self):
             optim_algo=self.optim_algo
         )
 
-    def get_action(self, **kwargs):
+    def get_action(self, *args, **kwargs):
         probs, onehots = self.session.run([self.sample_probs, self.sample_arc])
         return onehots, probs
 
@@ -557,7 +569,6 @@ def train(self, episode, working_dir):
 
                 self.session.run(self.train_op, feed_dict=feed_dict)
                 curr_loss, curr_kl, curr_ent = self.session.run([self.loss, self.kl_div, self.ent], feed_dict=feed_dict)
-
                 aloss += curr_loss
                 kl_sum += curr_kl
                 ent_sum += curr_ent
@@ -577,8 +588,8 @@ def train(self, episode, working_dir):
 
         return aloss / g_t
 
-    def store(self, state, prob, action, reward):
-        self.buffer.store(state, prob, action, reward)
+    def store(self, state, prob, action, reward, *args, **kwargs):
+        self.buffer.store(state=state, prob=prob, action=action, reward=reward)
         return
 
     def remove_files(self, files, working_dir='.'):
diff --git a/amber/architect/_multiio_controller.py b/amber/architect/_multiio_controller.py
index 0775e27..3e2dfef 100644
--- a/amber/architect/_multiio_controller.py
+++ b/amber/architect/_multiio_controller.py
@@ -87,7 +87,7 @@ def __init__(self, model_space, buffer_type='ordinal', with_skip_connection=True
 
     def _create_weight(self):
         super()._create_weight()
-        if self.num_input_blocks > 1 and self.with_input_blocks:
+        if self.with_input_blocks:
             with tf.variable_scope("input", initializer=tf.random_uniform_initializer(minval=-0.1, maxval=0.1)):
                 # input_emb: embedding for input blocks, if present
                 self.input_emb = tf.get_variable("inp_emb", [self.num_input_blocks, self.lstm_size])
@@ -145,7 +145,7 @@ def _build_sampler(self):
             # END BLOCK 1
 
             # BLOCK 2 [optional]: sample input feature blocks
-            if self.num_input_blocks > 1 and self.with_input_blocks:
+            if self.with_input_blocks:
                 next_c, next_h = stack_lstm(inputs, prev_c, prev_h, self.w_lstm)
                 prev_c, prev_h = next_c, next_h
                 hidden_states.append(prev_h)
@@ -335,7 +335,7 @@ def _build_trainer(self):
             # inputs = tf.nn.embedding_lookup(self.w_emb["start"][branch_id], start)
             inputs = tf.nn.embedding_lookup(self.w_emb["start"][layer_id], start)
 
-            if self.num_input_blocks > 1 and self.with_input_blocks:
+            if self.with_input_blocks:
 
                 next_c, next_h = stack_lstm(inputs, prev_c, prev_h, self.w_lstm)
                 prev_c, prev_h = next_c, next_h
@@ -427,7 +427,7 @@ def _build_trainer(self):
                         logit = self.tanh_constant * tf.tanh(logit)
 
                     probs_.append(tf.reshape(tf.nn.softmax(logit), [batch_size, layer_id, 2]))
-                    if self.num_input_blocks > 1:
+                    if self.with_input_blocks:
                         skip = self.input_arc[(arc_pointer + ops_each_layer + self.num_input_blocks): (
                                 arc_pointer + ops_each_layer + self.num_input_blocks + layer_id)]
                     else:
@@ -436,9 +436,10 @@ def _build_trainer(self):
                     skip = tf.reshape(tf.transpose(skip), [batch_size * layer_id])
                     skip = tf.to_int32(skip)
 
-                    skip_prob = tf.sigmoid(logit)
+                    skip_prob = tf.sigmoid(logit)  # shape=(batch_size*layer_id, 2)
                     kl = skip_prob * tf.log(skip_prob / skip_targets)
-                    kl = tf.reduce_sum(kl)
+                    kl = tf.reduce_sum(kl, axis=1)  # shape=(batch_size*layer_id,)
+                    kl = tf.reshape(kl, [batch_size, -1])  # (batch_size, layer_id)
                     skip_penaltys.append(kl)
 
                     log_prob3 = tf.nn.sparse_softmax_cross_entropy_with_logits(
@@ -483,8 +484,10 @@ def _build_trainer(self):
         self.onehot_log_prob = tf.reduce_sum(log_probs, axis=0)
         skip_count = tf.stack(skip_count)
         self.onehot_skip_count = tf.reduce_sum(skip_count, axis=0)
-        skip_penaltys = tf.stack(skip_penaltys)
-        self.onehot_skip_penaltys = tf.reduce_mean(skip_penaltys, axis=0)
+        #skip_penaltys = tf.reduce_mean(tf.transpose(tf.stack(skip_penaltys), [1,0]), axis=1)
+        skip_penaltys_flat = [tf.reduce_mean(x, axis=1) for x in skip_penaltys] # reduce_mean of layer_id dim; each entry's shape is  (batch_size,); in total list len=num_layers-1
+        self.onehot_skip_penaltys = tf.reduce_mean(skip_penaltys_flat, axis=0)  # from [num_layers-1, batch_size] to [batch_size,]
+        #self.onehot_skip_penaltys = tf.reduce_mean(skip_penaltys, axis=0)
 
 
 class MultiIOController(MultiInputController):
@@ -522,16 +525,23 @@ def __init__(self,
                  num_output_blocks=2,
                  with_output_blocks=True,
                  output_block_unique_connection=True,
+                 output_block_diversity_weight=None,
                  **kwargs):
 
         # Attributes unique to the derived class:
-        self.with_skip_connection = True
-        self.with_input_blocks = True
+        #self.with_input_blocks = True
         self.with_output_blocks = with_output_blocks
         self.num_output_blocks = num_output_blocks
-        self.output_block_unique_connection = output_block_unique_connection
+        skip_weight = kwargs['skip_weight'] if 'skip_weight' in kwargs else None
+        if output_block_diversity_weight is not None:
+            assert skip_weight is not None, "Cannot use output_block_diversity_weight when skip_weight is None"
+            self.output_block_diversity_weight = output_block_diversity_weight / skip_weight
+        else:
+            self.output_block_diversity_weight = None
 
+        self.output_block_unique_connection = output_block_unique_connection
         super().__init__(**kwargs)
+        assert self.with_skip_connection is True, "Must have with_skip_connection=True for MultiIOController"
 
     # override
     def _create_weight(self):
@@ -546,7 +556,8 @@ def _create_weight(self):
     # override
     def _build_sampler(self):
         super()._build_sampler()
-        layer_hs = [self.sample_hidden_states[i][-1] for i in range(1, self.num_layers * 3, 3)]
+        step_size = 1 + int(self.with_input_blocks) + int(self.with_skip_connection)
+        layer_hs = [self.sample_hidden_states[i][-1] for i in range(0, self.num_layers*step_size-1, step_size)]
         layer_hs = tf.concat(layer_hs, axis=0)
         output_probs = []
         output_onehot = []
@@ -578,6 +589,10 @@ def _build_sampler(self):
         self.sample_probs.extend(output_probs)
         self.sample_arc = tf.concat([self.sample_arc, tf.reshape(output_onehot, [-1])], axis=0)
         self.sample_log_prob += tf.reduce_sum(output_log_probs)
+        #if self.output_block_diversity_weight is not None:
+        #    diversity = tf.math.reduce_std(output_probs, axis=0)
+        #    diversity = tf.reduce_mean(diversity)
+        #    self.skip_penaltys -= diversity * self.output_block_diversity_weight 
 
     # override
     def _build_trainer(self):
@@ -587,7 +602,8 @@ def _build_trainer(self):
                            for i in range(self.total_arc_len, self.total_arc_len + output_arc_len)]
         self.total_arc_len += output_arc_len
 
-        layer_hs = [self.train_hidden_states[i][-1] for i in range(1, self.num_layers * 3, 3)]
+        step_size = 1 + int(self.with_input_blocks) + int(self.with_skip_connection) 
+        layer_hs = [self.train_hidden_states[i][-1] for i in range(0, self.num_layers*step_size-1, step_size)]
         layer_hs = tf.transpose(tf.stack(layer_hs), [1, 0, 2])  # shape: batch, num_layers, lstm_size
         output_probs = []
         output_log_probs = []
@@ -620,3 +636,8 @@ def _build_trainer(self):
         # self.output_log_probs = output_log_probs
         output_log_probs = tf.squeeze(tf.transpose(tf.stack(output_log_probs), [1, 0, 2]), axis=-1)
         self.onehot_log_prob += tf.reduce_sum(output_log_probs, axis=1)
+        if self.output_block_diversity_weight is not None:
+            output_probs = tf.transpose(tf.stack(output_probs), [1, 0, 2]) # shape: (batch, num_out_blocks, num_layers)
+            diversity_penaltys = tf.math.reduce_std(output_probs, axis=1)  # std of probs. of out_blocks on each layer; connecting every out_block to one layer will be penalized
+            diversity_penaltys = tf.reduce_mean(diversity_penaltys, axis=1)
+            self.onehot_skip_penaltys -= diversity_penaltys * self.output_block_diversity_weight
diff --git a/amber/architect/_operation_controller.py b/amber/architect/_operation_controller.py
index 031ad07..9e8217d 100644
--- a/amber/architect/_operation_controller.py
+++ b/amber/architect/_operation_controller.py
@@ -17,7 +17,7 @@
 from keras.models import Model
 
 from .buffer import Buffer
-from .controller import BaseController
+from ._general_controller import BaseController
 from .common_ops import get_kl_divergence_n_entropy, proximal_policy_optimization_loss
 
 
diff --git a/amber/architect/_zero_shot_controller.py b/amber/architect/_zero_shot_controller.py
new file mode 100644
index 0000000..b848e57
--- /dev/null
+++ b/amber/architect/_zero_shot_controller.py
@@ -0,0 +1,128 @@
+"""
+Zero-shot controller takes in descriptive features from dataset when calling `.get_action()` method,
+then generate new architectures for new tasks using a trained controller based on the data-descriptive
+features for the new task.
+
+ZZ, May 13, 2020
+"""
+
+from ._general_controller import GeneralController
+from .common_ops import create_bias, create_weight
+from .buffer import MultiManagerBuffer
+import tensorflow as tf
+from tensorflow.keras.regularizers import L1L2
+if tf.__version__.startswith('2'):
+    tf.compat.v1.disable_eager_execution()
+    import tensorflow.compat.v1 as tf
+import sys
+from .common_ops import get_tf_layer
+
+
+class ZeroShotController(GeneralController):
+    def __init__(self, data_description_config, *args, **kwargs):
+        """
+        Args:
+            data_description_config: dict, must have key "length".
+                optional keys: "hidden_layer" (dict), "regularizer" (dict).
+        """
+        self.data_description_config = data_description_config
+        super().__init__(*args, **kwargs)
+        assert isinstance(self.buffer, MultiManagerBuffer), "ZeroShotController must have MultiManagerBuffer;" \
+                                                            " got %s" % self.buffer
+
+    # override
+    def _create_weight(self):
+        super()._create_weight()
+        with tf.variable_scope("description_features", initializer=tf.random_uniform_initializer(minval=-0.1, maxval=0.1)):
+            self.data_descriptive_feature = tf.placeholder(shape=(None, self.data_description_config['length']),
+                                                           dtype=tf.float32, name='description_features')
+            self.w_dd = []
+            self.b_dd = []
+            data_description_len = self.data_description_config['length']
+            if 'hidden_layer' in self.data_description_config:
+                try:
+                    hidden_units = self.data_description_config['hidden_layer']['units']
+                    hidden_actv = self.data_description_config['hidden_layer']['activation']
+                except KeyError:
+                    raise KeyError("Error in parsing data_description_config: missing keys units or activation")
+                w_dd_1 = create_weight(name="w_dd_1", shape=(data_description_len, hidden_units))
+                b_dd_1 = create_bias(name="b_dd_1", shape=(hidden_units,))
+                self.w_dd.append(w_dd_1)
+                self.b_dd.append(b_dd_1)
+                h = get_tf_layer(hidden_actv)(tf.matmul(self.data_descriptive_feature, w_dd_1) + b_dd_1)
+                input_dim = hidden_units
+            else:
+                h = self.data_descriptive_feature
+                input_dim = data_description_len
+            w_dd = create_weight(name="w_dd", shape=(input_dim, self.lstm_size))
+            b_dd = create_bias(name="b_dd", shape=(self.lstm_size,))
+            self.w_dd.append(w_dd)
+            self.b_dd.append(b_dd)
+            self.g_emb = tf.matmul(h, w_dd) + b_dd  # shape: none, lstm_size
+
+    # overwrite
+    def get_action(self, description_feature, *args, **kwargs):
+        feed_dict = {self.data_descriptive_feature: description_feature}
+        probs, onehots = self.session.run([self.sample_probs, self.sample_arc], feed_dict=feed_dict)
+        return onehots, probs
+
+    # ovewrite
+    def store(self, prob, action, reward, description, manager_index, *args, **kwargs):
+        self.buffer.store(prob=prob, action=action, reward=reward, description=description, 
+                manager_index=manager_index)
+
+    # overwrite
+    def train(self, episode, working_dir):
+        self.buffer.finish_path(self.model_space, episode, working_dir)
+        aloss = 0
+        g_t = 0
+
+        for epoch in range(self.train_pi_iter):
+            t = 0
+            kl_sum = 0
+            ent_sum = 0
+            # get data from buffer
+            for batch_data in self.buffer.get_data(self.batch_size):
+                p_batch, a_batch, ad_batch, nr_batch = \
+                    [batch_data[x] for x in ['prob', 'action', 'advantage', 'reward']]
+                desc_batch = batch_data['description']
+                feed_dict = {self.input_arc[i]: a_batch[:, [i]]
+                             for i in range(a_batch.shape[1])}
+                feed_dict.update({self.advantage: ad_batch})
+                feed_dict.update({self.old_probs[i]: p_batch[i]
+                                  for i in range(len(self.old_probs))})
+                feed_dict.update({self.reward: nr_batch})
+                feed_dict.update({self.data_descriptive_feature: desc_batch})
+
+                _ = self.session.run(self.train_op, feed_dict=feed_dict)
+                curr_loss, curr_kl, curr_ent = self.session.run([self.loss, self.kl_div, self.ent], feed_dict=feed_dict)
+                aloss += curr_loss
+                kl_sum += curr_kl
+                ent_sum += curr_ent
+                t += 1
+                g_t += 1
+
+                if kl_sum / t > self.kl_threshold and epoch > 0:
+                    print("     Early stopping at step {} as KL(old || new) = ".format(g_t), kl_sum / t)
+                    return aloss / g_t
+
+            if epoch % max(1, (self.train_pi_iter // 5)) == 0:
+                print("     Epoch: {} Actor Loss: {} KL(old || new): {} Entropy(new) = {}".format(
+                    epoch, aloss / g_t,
+                    kl_sum / t,
+                    ent_sum / t)
+                )
+
+        return aloss / g_t
+
+    def _build_train_op(self):
+        """add the L1/L2 regularizations to controller loss
+        """
+        if 'regularizer' in self.data_description_config:
+            l1 = self.data_description_config["regularizer"].pop('l1', 0)
+            l2 = self.data_description_config["regularizer"].pop('l2', 0)
+            l1l2_reg = L1L2(l1=l1, l2=l2)
+            dd_reg = tf.reduce_sum([ l1l2_reg(x) for x in self.w_dd])
+            self.loss += dd_reg
+
+        super()._build_train_op()
diff --git a/amber/architect/buffer.py b/amber/architect/buffer.py
index 1daf7b6..581b9c9 100644
--- a/amber/architect/buffer.py
+++ b/amber/architect/buffer.py
@@ -14,6 +14,8 @@ def get_buffer(arg):
             return Buffer
         elif arg.lower() == 'replay':
             return ReplayBuffer
+        elif arg.lower() == "multimanager":
+            return MultiManagerBuffer
     elif callable(arg):
         return arg
     else:
@@ -32,7 +34,7 @@ def parse_action_str_squeezed(action_onehot, state_space):
 class Buffer(object):
     def __init__(self, max_size,
                  discount_factor=0.,
-                 ewa_beta=0.95,
+                 ewa_beta=None,
                  is_squeeze_dim=False,
                  rescale_advantage_by_reward=False,
                  clip_advantage=10.):
@@ -41,32 +43,36 @@ def __init__(self, max_size,
         is_squeeze_dim: if True, controller samples a sequence of tokens, instead of one-hot arrays
         """
         self.max_size = max_size
-        self.ewa_beta = ewa_beta
+        self.ewa_beta = ewa_beta or float(1 - 1./self.max_size)
         self.discount_factor = discount_factor
         self.is_squeeze_dim = is_squeeze_dim
         self.rescale_advantage_by_reward = rescale_advantage_by_reward
         self.clip_advantage = clip_advantage
         self.r_bias = None
-        ## short term buffer storing single traj
+
+        # short term buffer storing single trajectory
         self.state_buffer = []
         self.action_buffer = []
         self.prob_buffer = []
         self.reward_buffer = []
 
-        ## long_term buffer
-        self.lt_sbuffer = []
-        self.lt_abuffer = []
-        self.lt_pbuffer = []
-        # self.lt_rbuffer = []
-        self.lt_adbuffer = []
-        self.lt_nrbuffer = []  ## lt buffer for non discounted reward
-        self.lt_rmbuffer = []  ## reward mean buffer
-
-    def store(self, state, prob, action, reward):  # , value, next_value):
-        self.state_buffer.append(state)
-        self.prob_buffer.append(prob)
-        self.reward_buffer.append(reward)
-        self.action_buffer.append(action)
+        # long_term buffer
+        self.lt_sbuffer = []     # state
+        self.lt_abuffer = []     # action
+        self.lt_pbuffer = []     # prob
+        self.lt_adbuffer = []    # advantage
+        self.lt_nrbuffer = []    # lt buffer for non discounted reward
+        self.lt_rmbuffer = []    # reward mean buffer
+
+    def store(self, state=None, prob=None, action=None, reward=None, *args, **kwargs):
+        if state is not None:
+            self.state_buffer.append(state)
+        if prob is not None:
+            self.prob_buffer.append(prob)
+        if action is not None:
+            self.action_buffer.append(action)
+        if reward is not None:
+            self.reward_buffer.append(reward)
 
     def discount_rewards(self):
         """
@@ -91,7 +97,7 @@ def discount_rewards(self):
         discounted_rewards.reverse()
         return discounted_rewards, self.reward_buffer
 
-    def finish_path(self, state_space, global_ep, working_dir):
+    def finish_path(self, state_space, global_ep, working_dir, *args, **kwargs):
         """dump buffers to file and return processed buffer
         return processed [state, ob_prob, ob_onehot, reward, advantage]
         """
@@ -100,7 +106,9 @@ def finish_path(self, state_space, global_ep, working_dir):
         # advantage = self.get_advantage()
 
         # get data from buffer
-        state = np.concatenate(self.state_buffer, axis=0)
+        # TODO: complete remove `state_buffer` as it's useless; ZZ 2020.5.15
+        if self.state_buffer:
+            state = np.concatenate(self.state_buffer, axis=0)
         old_prob = [np.concatenate(p, axis=0) for p in zip(*self.prob_buffer)]
         if self.is_squeeze_dim:
             action_onehot = np.array(self.action_buffer)
@@ -114,8 +122,6 @@ def finish_path(self, state_space, global_ep, working_dir):
 
         nr = np.array(reward)[:, np.newaxis]
 
-        # TODO: rescale the advantage as percentage change over mean, to
-        #  account for rewards of different scales
         if self.rescale_advantage_by_reward:
             ad = (r - self.r_bias) / np.abs(self.r_bias)
         else:
@@ -126,7 +132,8 @@ def finish_path(self, state_space, global_ep, working_dir):
         if self.clip_advantage:
             ad = np.clip(ad, -np.abs(self.clip_advantage), np.abs(self.clip_advantage))
 
-        self.lt_sbuffer.append(state)
+        if self.state_buffer:
+            self.lt_sbuffer.append(state)
         self.lt_pbuffer.append(old_prob)
         self.lt_abuffer.append(action_onehot)
         # self.lt_rbuffer.append(r)
@@ -187,16 +194,15 @@ def get_data(self, bs, shuffle=True):
         lt_nrbuffer = np.concatenate(lt_nrbuffer, axis=0)
 
         if shuffle:
-            slice = np.random.choice(lt_sbuffer.shape[0], size=lt_sbuffer.shape[0], replace=False)
-            # slice = np.random.choice(lt_sbuffer.shape[0], size=lt_sbuffer.shape[0])
-            lt_sbuffer = lt_sbuffer[slice]
-            lt_pbuffer = [p[slice] for p in lt_pbuffer]
+            slice_ = np.random.choice(lt_sbuffer.shape[0], size=lt_sbuffer.shape[0], replace=False)
+            lt_sbuffer = lt_sbuffer[slice_]
+            lt_pbuffer = [p[slice_] for p in lt_pbuffer]
             if self.is_squeeze_dim:
-                lt_abuffer = lt_abuffer[slice]
+                lt_abuffer = lt_abuffer[slice_]
             else:
-                lt_abuffer = [a[slice] for a in lt_abuffer]
-            lt_adbuffer = lt_adbuffer[slice]
-            lt_nrbuffer = lt_nrbuffer[slice]
+                lt_abuffer = [a[slice_] for a in lt_abuffer]
+            lt_adbuffer = lt_adbuffer[slice_]
+            lt_nrbuffer = lt_nrbuffer[slice_]
 
         for i in range(0, len(lt_sbuffer), bs):
             b = min(i + bs, len(lt_sbuffer))
@@ -205,8 +211,7 @@ def get_data(self, bs, shuffle=True):
                 a_batch = lt_abuffer[i:b]
             else:
                 a_batch = [a[i:b, :] for a in lt_abuffer]
-            yield lt_sbuffer[i:b, :, :], p_batch, a_batch, lt_adbuffer[i:b, :], \
-                  lt_nrbuffer[i:b, :]
+            yield lt_sbuffer[i:b, :, :], p_batch, a_batch, lt_adbuffer[i:b, :], lt_nrbuffer[i:b, :]
 
 
 class ReplayBuffer(Buffer):
@@ -237,16 +242,15 @@ def get_data(self, bs, shuffle=True):
         lt_nrbuffer = np.concatenate(lt_nrbuffer, axis=0)
 
         if shuffle:
-            slice = np.random.choice(lt_sbuffer.shape[0], size=lt_sbuffer.shape[0], replace=False)
-            # slice = np.random.choice(lt_sbuffer.shape[0], size=lt_sbuffer.shape[0])
-            lt_sbuffer = lt_sbuffer[slice]
-            lt_pbuffer = [p[slice] for p in lt_pbuffer]
+            slice_ = np.random.choice(lt_sbuffer.shape[0], size=lt_sbuffer.shape[0], replace=False)
+            lt_sbuffer = lt_sbuffer[slice_]
+            lt_pbuffer = [p[slice_] for p in lt_pbuffer]
             if self.is_squeeze_dim:
-                lt_abuffer = lt_abuffer[slice]
+                lt_abuffer = lt_abuffer[slice_]
             else:
-                lt_abuffer = [a[slice] for a in lt_abuffer]
-            lt_adbuffer = lt_adbuffer[slice]
-            lt_nrbuffer = lt_nrbuffer[slice]
+                lt_abuffer = [a[slice_] for a in lt_abuffer]
+            lt_adbuffer = lt_adbuffer[slice_]
+            lt_nrbuffer = lt_nrbuffer[slice_]
 
         for i in range(0, len(lt_sbuffer), bs):
             b = min(i + bs, len(lt_sbuffer))
@@ -259,3 +263,202 @@ def get_data(self, bs, shuffle=True):
             nr_batch = lt_nrbuffer[i:b, :]
             ad_batch = lt_adbuffer[i:b, :]
             yield s_batch, p_batch, a_batch, ad_batch, nr_batch
+
+
+class MultiManagerBuffer:
+    def __init__(self, max_size,
+                 ewa_beta=None,
+                 is_squeeze_dim=False,
+                 rescale_advantage_by_reward=False,
+                 clip_advantage=10.,
+                 **kwargs
+                 ):
+        self.max_size = max_size
+        self.ewa_beta = ewa_beta or float(1 - 1./max_size)
+        self.is_squeeze_dim = is_squeeze_dim
+        self.rescale_advantage_by_reward = rescale_advantage_by_reward
+        self.clip_advantage = clip_advantage
+        self.r_bias = None
+
+        # short term buffer storing single trajectory
+        self.action_buffer = []
+        self.prob_buffer = []
+        self.reward_buffer = []
+
+        # long_term buffer
+        self.lt_action = []         # action
+        self.lt_prob = []           # prob
+        self.lt_adv = []            # advantage
+        self.lt_reward = []         # lt buffer for non discounted reward
+        self.lt_reward_mean = []    # reward mean buffer
+        # unique to this class
+        self.description_buffer = []   # short-term descriptive feature buffer
+        self.lt_desc_buffer = []       # long-term descriptive feature buffer
+        self.manager_index_buffer = []
+        self.lt_manager_idx_buffer = []
+
+    def store(self, prob, action, reward, description, manager_index):
+        """
+        Parameters
+        ----------
+        prob : list
+            A list of np.array of different sizes corresponding to each layer
+        action : list
+            A list of np.array of one-hot actions
+        reward : float
+            Float number of reward to be maximized
+        description : 1D-array like
+            A vector of data descriptor
+        manager_index: int
+            Integer number showing the index for which the reward is coming from
+        """
+        self.prob_buffer.append(prob)
+        self.action_buffer.append(action)
+        self.reward_buffer.append(reward)
+        self.description_buffer.append(description)
+        self.manager_index_buffer.append(manager_index)
+
+    def reset_short_term(self):
+        self.prob_buffer = []
+        self.action_buffer = []
+        self.reward_buffer = []
+        self.description_buffer = []
+        self.manager_index_buffer = []
+
+    def dump_buffer(self, model_space, global_ep, working_dir):
+        with open(os.path.join(working_dir, 'buffers.txt'), mode='a+') as f:
+            f.write("-" * 80 + "\n")
+            # write out one example for each manager
+            for j in self.r_bias:
+                this = np.where(np.array(self.manager_index_buffer)==j)[0]
+                max_idx = np.argmax(np.array(self.reward_buffer)[this]).tolist()
+                action_onehot = self.action_buffer[this[max_idx]]
+                if self.is_squeeze_dim:
+                    action_readable_str = ','.join([str(x) for x in parse_action_str_squeezed(action_onehot, model_space)])
+                else:
+                    action_readable_str = ','.join([str(x) for x in parse_action_str(action_onehot, model_space)])
+                f.write(
+                    "Episode:%d\tManager:%i\tReward:%.4f\tR_bias:%.4f\tAdvantage:%s\n" %
+                    (
+                        global_ep,
+                        j,
+                        self.reward_buffer[this[max_idx]],
+                        self.r_bias[j],
+                        np.round(self.lt_adv[-1][this].flatten(), 2),
+                    ) +
+                    "\tAction:%s\n\tProb:%s\n" % (
+                        action_readable_str,
+                        ' || '.join([str(np.round(x, 2)).replace("\n ", ";") for x in self.prob_buffer[this[max_idx]]])
+                    )
+                )
+            print("Saved buffers to file `buffers.txt` !")
+
+    def stratified_reward_mean(self):
+        """
+        Returns
+        ----------
+        dict
+            for each reward (key), the average reward (value) computed in current short-term
+            buffer
+        """
+        stratified_mean = {}
+        # Assumes manager_index is starting from 0
+        for idx in range(0, max(self.manager_index_buffer)+1):
+            this = np.where(np.array(self.manager_index_buffer)==idx)[0]
+            r_mean = np.mean(np.array(self.reward_buffer)[this])
+            stratified_mean[idx] = r_mean
+        return stratified_mean
+
+    def finish_path(self, model_space, global_ep, working_dir, *args, **kwargs):
+        old_prob = [np.concatenate(p, axis=0) for p in zip(*self.prob_buffer)]
+        if self.is_squeeze_dim:
+            action_onehot = np.array(self.action_buffer)
+        else:  # action squeezed into sequence
+            action_onehot = [np.concatenate(onehot, axis=0) for onehot in zip(*self.action_buffer)]
+
+        reward = np.array(self.reward_buffer)
+        if not self.lt_reward_mean:
+            self.r_bias = self.stratified_reward_mean()
+        else:
+            self.r_bias = {
+                    j: self.r_bias[j] * self.ewa_beta +
+                       self.lt_reward_mean[-1][j] * (1 - self.ewa_beta)
+                            for j in self.r_bias
+            }
+
+        # unroll the r_bias to fill in corresponding locations in short-term buffer
+        expanded_r_bias = np.array([ self.r_bias[j] for j in self.manager_index_buffer  ])
+        assert expanded_r_bias.shape == reward.shape
+        if self.rescale_advantage_by_reward:
+            ad = (reward - expanded_r_bias) / np.abs(expanded_r_bias)
+        else:
+            ad = (reward - expanded_r_bias)
+
+        if self.clip_advantage:
+            ad = np.clip(ad, -np.abs(self.clip_advantage), np.abs(self.clip_advantage))
+
+        self.lt_prob.append(old_prob)
+        self.lt_action.append(action_onehot)
+        self.lt_adv.append(ad)
+        self.lt_reward.append(reward)
+        self.lt_reward_mean.append(self.stratified_reward_mean())
+
+        self.dump_buffer(model_space=model_space, global_ep=global_ep, working_dir=working_dir)
+
+        # NOTE: this will only keep the `max_size` number of short-term buffers;
+        # whereas each short-term buffer could have multiple samples
+        if len(self.lt_prob) > self.max_size:
+            self.lt_prob = self.lt_prob[-self.max_size:]
+            self.lt_adv = self.lt_adv[-self.max_size:]
+            self.lt_action = self.lt_action[-self.max_size:]
+            self.lt_reward = self.lt_reward[-self.max_size:]
+            self.lt_reward_mean = self.lt_reward_mean[-self.max_size:]
+
+        description = np.concatenate(self.description_buffer, axis=0)
+        self.lt_desc_buffer.append(description)
+        self.lt_manager_idx_buffer.append(np.array(self.manager_index_buffer))
+        if len(self.lt_desc_buffer) > self.max_size:
+            self.lt_desc_buffer = self.lt_desc_buffer[-self.max_size:]
+            self.lt_manager_idx_buffer = self.lt_manager_idx_buffer[-self.max_size:]
+        self.reset_short_term()
+
+    def get_data(self, bs, shuffle=True):
+        lt_prob, lt_action, lt_adv, lt_reward = self.lt_prob, self.lt_action, self.lt_adv, self.lt_reward
+        lt_desc = self.lt_desc_buffer
+
+        lt_prob = [np.concatenate(p, axis=0) for p in zip(*lt_prob)]
+        if self.is_squeeze_dim:
+            lt_action = np.concatenate(lt_action, axis=0)
+        else:
+            lt_action = [np.concatenate(a, axis=0) for a in zip(*lt_action)]
+        lt_adv = np.concatenate(lt_adv, axis=0)
+        lt_reward = np.concatenate(lt_reward, axis=0)
+        lt_desc = np.concatenate(lt_desc, axis=0)
+
+        if shuffle:
+            slice_ = np.random.choice(lt_adv.shape[0], size=lt_adv.shape[0], replace=False)
+            lt_prob = [p[slice_] for p in lt_prob]
+            if self.is_squeeze_dim:
+                lt_action = lt_action[slice_]
+            else:
+                lt_action = [a[slice_] for a in lt_action]
+            lt_adv = lt_adv[slice_]
+            lt_reward = lt_reward[slice_]
+            lt_desc = lt_desc[slice_]
+
+        for i in range(0, len(lt_prob), bs):
+            b = min(i + bs, len(lt_prob))
+            p_batch = [p[i:b, :] for p in lt_prob]
+            if self.is_squeeze_dim:
+                a_batch = lt_action[i:b]
+            else:
+                a_batch = [a[i:b, :] for a in lt_action]
+            batch_data = {
+                "prob": p_batch,
+                "action": a_batch,
+                "advantage": np.expand_dims(lt_adv[i:b], axis=-1),
+                "reward": np.expand_dims(lt_reward[i:b], axis=-1),
+                "description": lt_desc[i:b]
+            }
+            yield batch_data
+
diff --git a/amber/architect/common_ops.py b/amber/architect/common_ops.py
index 0078177..b9ad9dd 100644
--- a/amber/architect/common_ops.py
+++ b/amber/architect/common_ops.py
@@ -3,9 +3,43 @@
 import keras.backend as K
 import numpy as np
 import tensorflow as tf
+if tf.__version__.startswith("2"):
+    tf.compat.v1.disable_eager_execution()
+    import tensorflow.compat.v1 as tf
 from tensorflow.python.training import moving_averages
 
 
+def unpack_data(data, unroll_generator_x=False, unroll_generator_y=False, callable_kwargs=None):
+    is_generator = False
+    unroll_generator = unroll_generator_x or unroll_generator_y
+    if type(data) in (tuple, list):
+        x, y = data[0], data[1]
+    elif isinstance(data, tf.keras.utils.Sequence):
+        x = data
+        y = None
+        is_generator = True
+    elif hasattr(data, '__next__'):
+        x = data
+        y = None
+        is_generator = True
+    elif callable(data):
+        callable_kwargs = callable_kwargs or {}
+        x, y = unpack_data(data=data(**callable_kwargs),
+                unroll_generator_x=unroll_generator_x,
+                unroll_generator_y=unroll_generator_y)
+    else:
+        raise Exception("cannot unpack data of type: %s"%type(data))
+    if is_generator and unroll_generator:
+        gen = data if hasattr(data, '__next__') else iter(data)
+        d_ = [d for d in zip(*gen)]
+        if unroll_generator_x ^ unroll_generator_y:
+            if hasattr(data, "shuffle"):
+                assert data.shuffle == False
+        x = np.concatenate(d_[0], axis=0) if unroll_generator_x else data
+        y = np.concatenate(d_[1], axis=0) if unroll_generator_y else None
+    return x, y
+
+
 def batchify(x, y=None, batch_size=None, shuffle=True, drop_remainder=True):
     if not type(x) is list: x = [x]
     if y is not None and type(y) is not list: y = [y]
@@ -88,7 +122,10 @@ def get_tf_layer(fn_str):
 
 def create_weight(name, shape, initializer=None, trainable=True, seed=None):
     if initializer is None:
-        initializer = tf.contrib.keras.initializers.he_normal(seed=seed)
+        try:
+            initializer = tf.contrib.keras.initializers.he_normal(seed=seed)
+        except AttributeError:
+            initializer = tf.keras.initializers.he_normal(seed=seed)
     return tf.get_variable(name, shape, initializer=initializer, trainable=trainable)
 
 
@@ -155,18 +192,24 @@ def get_keras_train_ops(loss, tf_variables, optim_algo, **kwargs):
     for g, v in zip(grads, tf_variables):
         if g is None:
             # get sub-scope name; if is optimizer-related, ignore
-            if v.name.split('/')[1] == 'compile':
+            if 'compile' in v.name.split('/'):
                 continue
             no_grad_var.append(v)
         else:
             grad_var.append(v)
     if no_grad_var:
         warnings.warn(
-            "=" * 80 + "\n\nWarning: the following tf.variables have no gradients"
-                       " and have been discarded: \n %s" % no_grad_var)
+            "\n" + "=" * 80 + "\nWarning: the following tf.variables have no gradients"
+                       " and have been discarded: \n %s" % no_grad_var, stacklevel=2)
     train_op = opt.get_updates(loss, grad_var)
-    config = opt.get_config()
-    learning_rate = config['lr']
+    try:
+        config = opt.get_config()
+    except NotImplementedError:  # if cannot get learning-rate when eager-execution is disableed
+        config = {'lr':None}
+    try:
+        learning_rate = config['lr']
+    except:  # for newer version of keras
+        learning_rate = config['learning_rate']
     return train_op, learning_rate, None, opt
 
 
diff --git a/amber/architect/complexity_calculator.py b/amber/architect/complexity_calculator.py
index 59aba58..e77f067 100644
--- a/amber/architect/complexity_calculator.py
+++ b/amber/architect/complexity_calculator.py
@@ -14,6 +14,7 @@ def child_model_params(num_features, num_layers, max_units):
 
 
 def controller_search_space(input_blocks, output_blocks, num_layers, num_choices_per_layer):
-    s = np.log10(num_choices_per_layer) * num_layers + np.log10(np.prod(np.arange(1, num_layers + 1)))
+    s = np.log10(num_choices_per_layer) * num_layers
+    s += np.log10(2) * (num_layers-1)*num_layers/2
     s += np.log10(input_blocks) * num_layers + np.log10(output_blocks) * num_layers
     return s
diff --git a/amber/architect/controller.py b/amber/architect/controller.py
index 6c5c109..5f2c9de 100644
--- a/amber/architect/controller.py
+++ b/amber/architect/controller.py
@@ -18,11 +18,13 @@
 from ._general_controller import BaseController, GeneralController
 from ._multiio_controller import MultiInputController, MultiIOController
 from ._operation_controller import OperationController
+from ._zero_shot_controller import ZeroShotController
 
 
 __all__ = [
     'GeneralController',
     'MultiInputController',
     'MultiIOController',
-    'OperationController'
+    'OperationController',
+    'ZeroShotController'
 ]
diff --git a/amber/architect/manager.py b/amber/architect/manager.py
index a338af5..e361db8 100644
--- a/amber/architect/manager.py
+++ b/amber/architect/manager.py
@@ -1,16 +1,29 @@
 # -*- coding: UTF-8 -*-
 
 import gc
-import os
+import os, sys
 import warnings
 
 import numpy as np
-import tensorflow as tf
-from keras.callbacks import ModelCheckpoint, EarlyStopping
-from keras.models import Model
-
+import tensorflow.keras as keras
+from ..utils import corrected_tf as tf
+import tensorflow as tf2
+from keras import backend as K
+from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping
+from tensorflow.keras.models import Model
+import time
+from datetime import datetime
+from collections import defaultdict
+
+from .common_ops import unpack_data
 from .store import get_store_fn
 
+__all__ = [
+    'BaseNetworkManager',
+    'NetworkManager',
+    'GeneralManager'
+]
+
 
 class BaseNetworkManager:
     def __init__(self, *args, **kwargs):
@@ -38,14 +51,21 @@ def __init__(self,
                  epochs=5,
                  child_batchsize=128,
                  verbose=0,
+                 fit_kwargs=None,
                  **kwargs):
         super(GeneralManager, self).__init__(**kwargs)
         self.train_data = train_data
         self.validation_data = validation_data
         self.working_dir = working_dir
+        self.fit_kwargs = fit_kwargs or {}
         if not os.path.exists(self.working_dir):
             os.makedirs(self.working_dir)
-        self.model_compile_dict = model_fn.model_compile_dict
+        self.model_compile_dict = kwargs.pop("model_compile_dict", None)
+        if self.model_compile_dict is None:
+            self.model_compile_dict = model_fn.model_compile_dict
+
+        # added 2020.5.19: parse model_space to manager for compatibility with newer versions of controllers
+        self.model_space = kwargs.pop("model_space", None)
 
         self.save_full_model = save_full_model
         self.epochs = epochs
@@ -56,64 +76,219 @@ def __init__(self,
         self.reward_fn = reward_fn
         self.store_fn = get_store_fn(store_fn)
 
-    def get_rewards(self, trial, model_arc):
+
+    def get_rewards(self, trial, model_arc, **kwargs):
         # print('-'*80, model_arc, '-'*80)
         train_graph = tf.Graph()
         train_sess = tf.Session(graph=train_graph)
         with train_graph.as_default(), train_sess.as_default():
+            try:
+                K.set_session(train_sess)
+            except RuntimeError: # keras 2.3.1 `set_session` not available for tf2.0
+                pass
             model = self.model_fn(model_arc)  # a compiled keras Model
+            if model is None:
+                assert hasattr(self.reward_fn, "min"), "model_fn of type %s returned a non-valid model, but the given reward_fn of type %s does not have .min() method" % (type(self.model_fn), type(self.reward_fn))
+                hist = None
+                this_reward, loss_and_metrics, reward_metrics = self.reward_fn.min(data=self.validation_data)
+                loss = loss_and_metrics.pop(0)
+                loss_and_metrics = {str(self.model_compile_dict['metrics'][i]): loss_and_metrics[i] for i in
+                                    range(len(loss_and_metrics))}
+                loss_and_metrics['loss'] = loss
+                if reward_metrics:
+                    loss_and_metrics.update(reward_metrics)
 
-            # unpack the dataset
-            X_val, y_val = self.validation_data[0:2]
-            X_train, y_train = self.train_data
+            else:
+                # train the model using Keras methods
+                print(" Trial %i: Start training model..." % trial)
+                train_x, train_y = unpack_data(self.train_data)
+                hist = model.fit(x=train_x,
+                                 y=train_y,
+                                 batch_size=self.batchsize if train_y else None,
+                                 epochs=self.epochs,
+                                 verbose=self.verbose,
+                                 #shuffle=True,
+                                 validation_data=self.validation_data,
+                                 callbacks=[ModelCheckpoint(os.path.join(self.working_dir, 'temp_network.h5'),
+                                                            monitor='val_loss', verbose=self.verbose,
+                                                            save_best_only=True),
+                                            EarlyStopping(monitor='val_loss', patience=self.fit_kwargs.pop("earlystop_patience", 5), verbose=self.verbose)],
+                                 **self.fit_kwargs
+                                 )
+                # load best performance epoch in this training session
+                model.load_weights(os.path.join(self.working_dir, 'temp_network.h5'))
+
+                # evaluate the model by `reward_fn`
+                this_reward, loss_and_metrics, reward_metrics = \
+                    self.reward_fn(model, (self.validation_data),
+                                   session=train_sess,
+                                   graph=train_graph)
+                loss = loss_and_metrics.pop(0)
+                loss_and_metrics = {str(self.model_compile_dict['metrics'][i]): loss_and_metrics[i] for i in
+                                    range(len(loss_and_metrics))}
+                loss_and_metrics['loss'] = loss
+                if reward_metrics:
+                    loss_and_metrics.update(reward_metrics)
+
+                # do any post processing,
+                # e.g. save child net, plot training history, plot scattered prediction.
+                if self.store_fn:
+                    val_pred = model.predict(self.validation_data) 
+                    self.store_fn(
+                        trial=trial,
+                        model=model,
+                        hist=hist,
+                        data=self.validation_data,
+                        pred=val_pred,
+                        loss_and_metrics=loss_and_metrics,
+                        working_dir=self.working_dir,
+                        save_full_model=self.save_full_model,
+                        knowledge_func=self.reward_fn.knowledge_function
+                    )
 
-            # train the model using Keras methods
-            print(" Trial %i: Start training model..." % trial)
-            hist = model.fit(X_train, y_train,
-                             batch_size=self.batchsize,
-                             epochs=self.epochs,
-                             verbose=self.verbose,
-                             validation_data=(X_val, y_val),
-                             callbacks=[ModelCheckpoint(os.path.join(self.working_dir, 'temp_network.h5'),
-                                                        monitor='val_loss', verbose=self.verbose,
-                                                        save_best_only=True),
-                                        EarlyStopping(monitor='val_loss', patience=5, verbose=self.verbose)]
-                             )
-            # load best performance epoch in this training session
-            model.load_weights(os.path.join(self.working_dir, 'temp_network.h5'))
+        # clean up resources and GPU memory
+        del model
+        del hist
+        gc.collect()
+        return this_reward, loss_and_metrics
 
-            # evaluate the model by `reward_fn`
-            this_reward, loss_and_metrics, reward_metrics = \
-                self.reward_fn(model, (X_val, y_val),
-                               session=train_sess,
-                               graph=train_graph)
-            loss = loss_and_metrics.pop(0)
-            loss_and_metrics = {str(self.model_compile_dict['metrics'][i]): loss_and_metrics[i] for i in
-                                range(len(loss_and_metrics))}
-            loss_and_metrics['loss'] = loss
-            if reward_metrics:
-                loss_and_metrics.update(reward_metrics)
 
-            # do any post processing,
-            # e.g. save child net, plot training history, plot scattered prediction.
-            if self.store_fn:
-                val_pred = model.predict(X_val)
-                self.store_fn(
-                    trial=trial,
-                    model=model,
-                    hist=hist,
-                    data=self.validation_data,
-                    pred=val_pred,
-                    loss_and_metrics=loss_and_metrics,
-                    working_dir=self.working_dir,
-                    save_full_model=self.save_full_model,
-                    knowledge_func=self.reward_fn.knowledge_function
-                )
+class DistributedGeneralManager(GeneralManager):
+    def __init__(self, devices, train_data_kwargs, validate_data_kwargs, *args, **kwargs):
+        self.devices = devices
+        super().__init__(*args, **kwargs)
+        assert devices is None or len(self.devices) == 1, "Only supports one GPU device currently"
+        # For keeping & closing file connection at multi-processing
+        self.train_data_kwargs = train_data_kwargs or {}
+        self.validate_data_kwargs = validate_data_kwargs or {}
+        self.train_x = None
+        self.train_y = None
+        self.file_connected = False
+        # For resampling; TODO: how to implement a Bayesian version of this?
+        self.arc_records = defaultdict(dict)
+
+    def close_handler(self):
+        if self.file_connected:
+            self.train_x.close()
+            if self.train_y:
+                self.train_y.close()
+            self.train_x = None
+            self.train_y = None
+            self.file_connected = False
+
+    def get_rewards(self, trial, model_arc, **kwargs):
+        # TODO: use tensorflow distributed strategy
+        #strategy = tf2.distribute.MirroredStrategy(devices=self.devices)
+        #print('Number of devices: {} - {}'.format(strategy.num_replicas_in_sync, self.devices))
+        #with strategy.scope():
+        pid = os.getpid()
+        sys.stderr.write("[%s][%s] Preprocessing.."%(pid, datetime.now().strftime("%H:%M:%S") ))
+        start_time = time.time()
+        train_graph = tf.Graph()
+        config = tf.ConfigProto()
+        config.gpu_options.allow_growth = True
+        train_sess = tf.Session(graph=train_graph, config=config)
+        if self.devices is None:
+            from ..utils.gpu_query import get_idle_gpus
+            idle_gpus = get_idle_gpus()
+            target_device = idle_gpus[0] 
+            target_device = "/device:GPU:%i"%target_device
+            self.devices = [target_device]
+            sys.stderr.write("[%s] Auto-assign device: %s"% (pid, target_device) )
+        else:
+            target_device = self.devices[0]
+        with train_graph.as_default(), train_sess.as_default():
+            with tf.device(target_device):
+                try:
+                    K.set_session(train_sess)
+                except RuntimeError: # keras 2.3.1 `set_session` not available for tf2.0
+                    pass
+                model = self.model_fn(model_arc)  # a compiled keras Model
+
+                # unpack the dataset
+                if not self.file_connected:
+                    X_train, y_train = unpack_data(self.train_data, callable_kwargs=self.train_data_kwargs)
+                    self.train_x = X_train
+                    self.train_y = y_train
+                    assert callable(self.validation_data)
+                    self.validation_data = self.validation_data(**self.validate_data_kwargs)
+                    self.file_connected = True
+                elapse_time = time.time() - start_time
+                sys.stderr.write("  %.3f sec\n"%elapse_time)
+                model_arc_ = tuple(model_arc)
+                if model_arc_ in self.arc_records:
+                    this_reward = self.arc_records[model_arc_]['reward'] 
+                    old_trial = self.arc_records[model_arc_]['trial'] 
+                    loss_and_metrics = self.arc_records[model_arc_]['loss_and_metrics'] 
+                    sys.stderr.write("[%s][%s] Trial %i: Re-sampled from history %i\n" % (pid, datetime.now().strftime("%H:%M:%S"), old_trial))
+                else:
+                    # train the model using Keras methods
+                    start_time = time.time()
+                    sys.stderr.write("[%s][%s] Trial %i: Start training model.." % (pid, datetime.now().strftime("%H:%M:%S"), trial))
+                    hist = model.fit(self.train_x, self.train_y,
+                                     batch_size=self.batchsize,
+                                     epochs=self.epochs,
+                                     verbose=self.verbose,
+                                     validation_data=self.validation_data,
+                                     callbacks=[ModelCheckpoint(os.path.join(self.working_dir, 'temp_network.h5'),
+                                                                monitor='val_loss', verbose=self.verbose,
+                                                                save_best_only=True),
+                                                EarlyStopping(monitor='val_loss', patience=self.fit_kwargs.pop("earlystop_patience", 10), verbose=self.verbose)],
+                                     **self.fit_kwargs
+                                     )
+
+                    # load best performance epoch in this training session
+                    model.load_weights(os.path.join(self.working_dir, 'temp_network.h5'))
+                    elapse_time = time.time() - start_time
+                    sys.stderr.write("  %.3f sec\n"%elapse_time)
+
+                    start_time = time.time()
+                    sys.stderr.write("[%s] Postprocessing.."% pid )
+                    # evaluate the model by `reward_fn`
+                    this_reward, loss_and_metrics, reward_metrics = \
+                        self.reward_fn(model, self.validation_data,
+                                       session=train_sess,
+                                       )
+                    loss = loss_and_metrics.pop(0)
+                    loss_and_metrics = {str(self.model_compile_dict['metrics'][i]): loss_and_metrics[i] for i in
+                                        range(len(loss_and_metrics))}
+                    loss_and_metrics['loss'] = loss
+                    if reward_metrics:
+                        loss_and_metrics.update(reward_metrics)
+
+                    # do any post processing,
+                    # e.g. save child net, plot training history, plot scattered prediction.
+                    if self.store_fn:
+                        val_pred = model.predict(self.validation_data)
+                        self.store_fn(
+                            trial=trial,
+                            model=model,
+                            hist=hist,
+                            data=self.validation_data,
+                            pred=val_pred,
+                            loss_and_metrics=loss_and_metrics,
+                            working_dir=self.working_dir,
+                            save_full_model=self.save_full_model,
+                            knowledge_func=self.reward_fn.knowledge_function
+                        )
+                    elapse_time = time.time() - start_time
+                    sys.stderr.write("  %.3f sec\n"%elapse_time)
+
+                    # store the rewards in records
+                    self.arc_records[model_arc_]['trial'] = trial
+                    self.arc_records[model_arc_]['reward'] = this_reward
+                    self.arc_records[model_arc_]['loss_and_metrics'] = loss_and_metrics
 
         # clean up resources and GPU memory
+        start_time = time.time()
+        sys.stderr.write("[%s] Cleaning up.."%pid)
         del model
         del hist
+        del train_sess
+        del train_graph
         gc.collect()
+        elapse_time = time.time() - start_time
+        sys.stderr.write("  %.3f sec\n"%elapse_time)
         return this_reward, loss_and_metrics
 
 
diff --git a/amber/architect/model_space.py b/amber/architect/model_space.py
index 21465fa..035b4bb 100644
--- a/amber/architect/model_space.py
+++ b/amber/architect/model_space.py
@@ -146,3 +146,7 @@ def from_dict(d):
 
             ms.add_layer(layer_id=i, layer_states=[State(**d[i][j]) for j in range(len(d[i]))])
         return ms
+
+
+# alias
+Operation = State
diff --git a/amber/architect/reward.py b/amber/architect/reward.py
index c1890df..2013ad7 100644
--- a/amber/architect/reward.py
+++ b/amber/architect/reward.py
@@ -6,7 +6,8 @@
 '''
 
 import numpy as np
-
+import copy
+from .common_ops import unpack_data
 from ..utils.io import read_history
 
 
@@ -54,14 +55,19 @@ def __init__(self, *args, **kwargs):
         self.c = 1.
 
     def __call__(self, model, data, *args, **kwargs):
-        X, y = data
+        X, y = unpack_data(data)
         loss_and_metrics = model.evaluate(X, y)
         # Loss function values will always be the first value
         if type(loss_and_metrics) is list:
             L = loss_and_metrics[0]
-        else:
+        elif type(loss_and_metrics) is dict:
+            L = loss_and_metrics['val_loss']
+            loss_and_metrics = [L]
+        elif isinstance(loss_and_metrics, float):
             L = loss_and_metrics
             loss_and_metrics = [loss_and_metrics]
+        else:
+            raise Exception("Cannot understand return type of model.evaluate; got %s" % type(loss_and_metrics))
         return -L, loss_and_metrics, None
     # return self.c/L, loss_and_metrics, None
 
@@ -74,6 +80,8 @@ def __init__(self, method='auc', knowledge_function=None, Lambda=1, loss_c=None,
         elif method == 'aupr' or method == 'auprc':
             from sklearn.metrics import average_precision_score
             self.scorer = average_precision_score
+        elif callable(method):
+            self.scorer = method
         else:
             raise Exception("cannot understand scorer method: %s" % method)
         self.knowledge_function = knowledge_function
@@ -83,7 +91,7 @@ def __init__(self, method='auc', knowledge_function=None, Lambda=1, loss_c=None,
         self.pred = kwargs.pop('pred', None)
 
     def __call__(self, model, data, *args, **kwargs):
-        X, y = data
+        X, y = unpack_data(data, unroll_generator_y=True)
         if self.pred is None:
             pred = model.predict(X)
         else:
@@ -95,6 +103,8 @@ def __call__(self, model, data, *args, **kwargs):
             pred = [pred]
         for i in range(len(y)):
             tmp = []
+            if len(y[i].shape) == 1: y[i] = np.expand_dims(y[i], axis=-1)
+            if len(pred[i].shape) == 1: pred[i] = np.expand_dims(pred[i], axis=-1)
             for j in range(y[i].shape[1]):
                 try:
                     score = self.scorer(y_true=y[i][:, j], y_score=pred[i][:, j])
@@ -120,6 +130,17 @@ def __call__(self, model, data, *args, **kwargs):
         loss_and_metrics = [L]
         reward_metrics = {'knowledge': K}
         return reward, loss_and_metrics, reward_metrics
+    
+    def min(self, data):
+        """For dealing with non-valid model"""
+        X, y = unpack_data(data, unroll_generator_y=True)
+        if type(y) is not list: y = [y]
+        pred = copy.deepcopy(y)
+        _ = [np.random.shuffle(p) for p in pred]
+        self.pred = pred
+        reward, loss_and_metrics, reward_metrics = self.__call__(None, (X,y))
+        self.pred = None  # release self.pred after use
+        return reward, loss_and_metrics, reward_metrics
 
 
 def MockReward(train_history_list, metric, stringify_states, metric_name_dict, Lambda=1.):
diff --git a/amber/architect/store.py b/amber/architect/store.py
index 8761ab9..5520551 100644
--- a/amber/architect/store.py
+++ b/amber/architect/store.py
@@ -5,6 +5,7 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from ..plots import plot_hessian, plot_training_history
+from .common_ops import unpack_data
 
 
 def get_store_fn(arg):
@@ -38,6 +39,7 @@ def store_with_model_plot(
         *args, **kwargs
 ):
     par_dir = os.path.join(working_dir, 'weights', 'trial_%s' % trial)
+    os.makedirs(par_dir, exist_ok=True)
     store_general(trial=trial,
                   model=model,
                   hist=hist,
@@ -47,8 +49,8 @@ def store_with_model_plot(
                   working_dir=working_dir,
                   save_full_model=save_full_model
                   )
-    from keras.utils import plot_model
-    plot_model(model, to_file=os.path.join(par_dir, "model_arc.png"))
+    from tensorflow.keras.utils import plot_model
+    plot_model(model, to_file=os.path.join(par_dir, "model_arc.png"), show_shapes=True, show_layer_names=True)
 
 
 def store_with_hessian(
@@ -101,9 +103,15 @@ def store_general(
     os.mkdir(par_dir)
     if save_full_model:
         model.save(os.path.join(working_dir, 'weights', 'trial_%s' % trial, 'full_bestmodel.h5'))
-    plot_training_history(hist, par_dir)
+    try:
+        plot_training_history(hist, par_dir)
+    except:
+        # TODO: this still not working for non-keras models, e.g. EnasAnn/Cnn
+        pass
     if os.path.isfile(os.path.join(working_dir, 'temp_network.h5')):
         shutil.move(os.path.join(working_dir, 'temp_network.h5'), os.path.join(par_dir, 'bestmodel.h5'))
+    # TODO: REVAMP THIS. if data is a generator/keras.utils.Sequence
+    data = unpack_data(data, unroll_generator_y=True)
     metadata = data[2] if len(data) > 2 else None
     obs = data[1]
     write_pred_to_disk(
@@ -140,7 +148,7 @@ def write_pred_to_disk(fn, y_pred, y_obs, metadata=None, metrics=None):
         if len(np.unique(y_obs)) < 10:  # is categorical
             str_format = "%i"
         else:
-            str_format = "%.3f'"
+            str_format = "%.3f"
 
         f.write('pred\tobs\tmetadata\n')
         for i in range(len(y_pred)):
diff --git a/amber/architect/train_env.py b/amber/architect/train_env.py
index 1ecd27b..b47dd72 100644
--- a/amber/architect/train_env.py
+++ b/amber/architect/train_env.py
@@ -1,38 +1,18 @@
 # -*- coding: UTF-8 -*-
-
+import os, sys
 import csv
-import datetime
-import logging
+import datetime, time
 import shutil
 import warnings
-
-from keras.optimizers import *
+from collections import defaultdict
 
 from .buffer import parse_action_str, parse_action_str_squeezed
 from ._operation_controller import *
 from ..utils.io import save_action_weights, save_stats
 from ..plots import plot_stats2, plot_environment_entropy, plot_controller_performance, \
     plot_action_weights, plot_wiring_weights
-
-
-def setup_logger(working_dir='.', verbose_level=logging.INFO):
-    # setup logger
-    logger = logging.getLogger('AMBER')
-    logger.setLevel(verbose_level)
-    # create file handler which logs even debug messages
-    fh = logging.FileHandler(os.path.join(working_dir, 'log.AMBER.txt'))
-    fh.setLevel(verbose_level)
-    # create console handler with a higher log level
-    ch = logging.StreamHandler()
-    ch.setLevel(verbose_level)
-    # create formatter and add it to the handlers
-    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s -\n %(message)s')
-    fh.setFormatter(formatter)
-    ch.setFormatter(formatter)
-    # add the handlers to the logger
-    logger.addHandler(fh)
-    logger.addHandler(ch)
-    return logger
+from ..utils.logging import setup_logger
+from .manager import BaseNetworkManager, EnasManager
 
 
 def get_controller_states(model):
@@ -80,7 +60,7 @@ def __init__(self,
                  initial_buffering_queue=15,
                  working_dir='.', entropy_converge_epsilon=0.01,
                  squeezed_action=True,
-                 with_input_blocks=True,
+                 with_input_blocks=False,
                  with_skip_connection=True,
                  save_controller=False,
                  continuous_run=False,
@@ -110,17 +90,21 @@ def __init__(self,
         try:
             self.last_actionState_size = len(self.controller.state_space[-1])
         except Exception as e:
-            warnings.warn("DEPRECATED Exception in ControllerTrainEnv: %s" % e)
+            warnings.warn("DEPRECATED Exception in ControllerTrainEnv: %s" % e, stacklevel=2)
             self.last_actionState_size = 1
 
+        self.resume_prev_run = resume_prev_run
+        self.logger = logger if logger else setup_logger(working_dir)
+        if issubclass(type(manager), BaseNetworkManager):
+            if os.path.realpath(manager.working_dir) != os.path.realpath(self.working_dir):
+                warnings.warn("manager working dir and environment working dir are different.", stacklevel=2)
+        else:
+            warnings.warn("ControllerTrainEnvironment: input manager is not a subclass of BaseNetworkManager; "
+                          "make sure this is intended", stacklevel=2)
         if resume_prev_run:
             self.restore()
         else:
             self.clean()
-        self.resume_prev_run = resume_prev_run
-        self.logger = logger if logger else setup_logger(working_dir)
-        if os.path.realpath(manager.working_dir) != os.path.realpath(self.working_dir):
-            warnings.warn("manager working dir and environment working dir are different.")
 
     def __str__(self):
         s = 'ControllerTrainEnv for %i max steps, %i child mod. each step' % (self.max_episode, self.max_step_per_ep)
@@ -330,8 +314,12 @@ def __init__(self, *args, **kwargs):
         self.save_controller_every = kwargs.pop('save_controller_every', None)
         super().__init__(*args, **kwargs)
         self.initial_buffering_queue = 0
-        if self.manager.model_fn.controller is None:
-            self.manager.model_fn.set_controller(self.controller)
+        if issubclass(type(self.manager), BaseNetworkManager):
+            if self.manager.model_fn.controller is None:
+                self.manager.model_fn.set_controller(self.controller)
+        else:
+            warnings.warn("EnasTrainEnv: input manager is not a subclass of BaseNetworkManager; "
+                          "make sure this is intended", stacklevel=2)
         if self.time_budget is None:
             pass
         elif type(self.time_budget) is str:
@@ -392,7 +380,7 @@ def train(self):
                         loss_and_metrics_ep[x] += loss_and_metrics[x]
 
                     # save the arc_seq and reward
-                    self.controller.store(state, probs, arc_seq, reward)
+                    self.controller.store(state=state, prob=probs, action=arc_seq, reward=reward)
 
                     # write the results of this trial into a file
                     data = [controller_step, [loss_and_metrics[x] for x in sorted(loss_and_metrics.keys())],
@@ -423,6 +411,7 @@ def train(self):
                     LOGGER.info("-" * 10)
 
                 if self.save_controller_every is not None and child_step % self.save_controller_every == 0:
+                    self.logger.info("Saving controller weights for epoch %d" % child_step)
                     self.controller.save_weights(
                         os.path.join(self.working_dir, "controller_weights-epoch-%i.h5" % child_step))
 
@@ -464,3 +453,357 @@ def train(self):
         for p in ep_p:
             act_idx.append(np.argmax(p))
         return act_idx
+
+
+class MultiManagerEnvironment(EnasTrainEnv):
+    """
+    MultiManagerEnvironment is an environment that allows one controller to interact with multiple EnasManagers
+    """
+
+    def __init__(self, data_descriptive_features, is_enas='auto', *args, **kwargs):
+        super(MultiManagerEnvironment, self).__init__(*args, **kwargs)
+        assert type(self.manager) is list, \
+            "MultiManagerEnasEnvironment must have a List of manager instances, got %s" % type(self.manager)
+
+        self.manager_cnt = len(self.manager)
+        self.is_enas = is_enas
+        for i in range(self.manager_cnt):
+            assert issubclass(type(self.manager[i]), BaseNetworkManager), \
+                "MultiManagerEnasEnvironment expects a List of Manager instances, " \
+                "got %s for %i-th element" % (type(self.manager[i]), i)
+
+        self.data_descriptive_features = data_descriptive_features
+        assert len(self.data_descriptive_features) == self.manager_cnt, \
+            "data descriptive features must match the number of managers; " \
+            "got %i description, %i managers" % ( len(self.data_descriptive_features), self.manager_cnt )
+
+        if self.is_enas == "auto":
+            if all([isinstance(self.manager[i], EnasManager) for i in range(self.manager_cnt)]):
+                self.is_enas = True
+            else:
+                self.is_enas = False
+
+    def _warmup(self):
+        assert self.is_enas, "You can only set warm_up_epochs>0 if is_enas=True"
+        self.logger.info("warm-up for child model: %i epochs" % self.child_warm_up_epochs)
+        warmup_nsteps = None
+        for i in range(1, self.child_warm_up_epochs + 1):
+            self.logger.info("warm-up : %i epoch" % i)
+            for j in range(self.manager_cnt):
+                self.manager[j].get_rewards(trial=-i, model_arc=None, nsteps=warmup_nsteps)
+
+    def _train_loop(self):
+        if self.resume_prev_run:
+            f = open(os.path.join(self.working_dir, 'train_history.csv'), mode='a+')
+        else:
+            f = open(os.path.join(self.working_dir, 'train_history.csv'), mode='w')
+        writer = csv.writer(f)
+        starttime = datetime.datetime.now()
+        action_probs_record = []
+        loss_and_metrics_list = []
+        controller_step = self.start_ep * self.max_step_per_ep
+        for child_step in range(self.start_ep, self.max_episode):
+            try:
+                if self.is_enas:
+                    # train child parameters w, if is_enas
+                    for j in range(self.manager_cnt):
+                        self.logger.info("sampling with manager %i" % j)
+                        self.manager[j].get_rewards(child_step, None, nsteps=self.child_train_steps)
+
+                # train controller parameters theta
+                ep_reward = 0
+                loss_and_metrics_ep = defaultdict(float)
+                for step in range(self.max_step_per_ep):
+                    for j in range(self.manager_cnt):
+                        ep_probs = []
+                        arc_seq, probs = self.controller.get_action(
+                            description_feature=self.data_descriptive_features[[j]])    # preserve the first dimension
+                        self.entropy_record.append(compute_entropy(probs))
+                        ep_probs.append(probs)
+                        # LOGGER.debug the action probabilities
+                        action_list = parse_action_str_squeezed(arc_seq, self.controller.state_space)
+                        self.logger.debug("Manager {}, Predicted actions : {}".format(j, [str(x) for x in action_list]))
+
+                        # build a model, train and get reward and accuracy from the network manager
+                        reward, loss_and_metrics = self.manager[j].get_rewards(
+                            trial=controller_step, model_arc=arc_seq, nsteps=self.child_train_steps)
+                        self.logger.debug("Rewards : " + str(reward) + " Metrics : " + str(loss_and_metrics))
+
+                        ep_reward += reward
+                        for x in loss_and_metrics.keys():
+                            loss_and_metrics_ep[x] += loss_and_metrics[x]
+                        # save the arc_seq and reward
+                        self.controller.store(prob=probs, action=arc_seq, reward=reward,
+                                              description=self.data_descriptive_features[[j]],
+                                              manager_index=j)
+                        # write the results of this trial into a file
+                        data = ["%i-%i"%(j,controller_step), [loss_and_metrics[x] for x in sorted(loss_and_metrics.keys())],
+                                reward]
+                        if self.squeezed_action:
+                            data.extend(arc_seq)
+                        else:
+                            data.extend(action_list)
+                        writer.writerow(data)
+                        f.flush()
+                    # update trial
+                    controller_step += 1
+                loss_and_metrics_list.append({x: (v / self.max_step_per_ep) for x, v in loss_and_metrics_ep.items()})
+                # average probs over trajectory
+                ep_p = [sum(p) / len(p) for p in zip(*ep_probs)]
+                action_probs_record.append(ep_p)
+                if child_step >= self.initial_buffering_queue - 1:
+                    # train the controller on the saved state and the discounted rewards
+                    loss = self.controller.train(child_step, self.working_dir)
+                    self.total_reward += np.sum(np.array(self.controller.buffer.lt_adv[-1]).flatten())
+                    self.logger.info("Total reward : " + str(self.total_reward))
+                    self.logger.info("END episode %d: Controller loss : %0.6f" % (child_step, loss))
+                    self.logger.info("-" * 10)
+                else:
+                    self.logger.info("END episode %d: Buffering" % child_step)
+                    self.logger.info("-" * 10)
+
+                if self.save_controller_every is not None and child_step % self.save_controller_every == 0:
+                    self.logger.info("Saving controller weights for epoch %d" % child_step)
+                    self.controller.save_weights(
+                        os.path.join(self.working_dir, "controller_weights-epoch-%i.h5" % child_step))
+            except KeyboardInterrupt:
+                self.logger.info("User disrupted training")
+                break
+
+            consumed_time = (datetime.datetime.now() - starttime).total_seconds()
+            self.logger.info("used time: %.2f %%" % (consumed_time / self.time_budget * 100))
+            if consumed_time >= self.time_budget:
+                self.logger.info("training ceased because run out of time budget")
+                break
+
+        self.logger.debug("Total Reward : %s" % self.total_reward)
+        f.close()
+        return action_probs_record, loss_and_metrics_list
+
+    def train(self):
+        if self.child_warm_up_epochs > 0:
+            self._warmup()
+
+        action_probs_record, loss_and_metrics_list = self._train_loop()
+
+        metrics_dict = {k: v for k, v in zip(sorted(loss_and_metrics_list[0].keys()),
+                                             range(len(loss_and_metrics_list[0])))}
+        plot_controller_performance(os.path.join(self.working_dir, 'train_history.csv'),
+                                    metrics_dict=metrics_dict,
+                                    save_fn=os.path.join(self.working_dir, 'train_history.png'), N_sma=5)
+        plot_environment_entropy(self.entropy_record,
+                                 os.path.join(self.working_dir, 'entropy.png'))
+        save_kwargs = {}
+        if self.with_input_blocks:
+            save_kwargs['input_nodes'] = self.manager.model_fn.inputs_op
+        self.action_probs_record = action_probs_record
+        save_action_weights(action_probs_record, self.controller.state_space, self.working_dir,
+                            with_input_blocks=self.with_input_blocks, with_skip_connection=self.with_skip_connection,
+                            **save_kwargs)
+        save_stats(loss_and_metrics_list, self.working_dir)
+
+        if self.should_plot:
+            plot_action_weights(self.working_dir)
+            plot_wiring_weights(self.working_dir, self.with_input_blocks, self.with_skip_connection)
+            plot_stats2(self.working_dir)
+
+        # return converged config idx
+        act_idx = []
+        for p in action_probs_record[-1]:
+            act_idx.append(np.argmax(p))
+        return act_idx
+
+
+class ParallelMultiManagerEnvironment(MultiManagerEnvironment):
+    def __init__(self, processes=2, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.processes = processes
+        assert self.processes <= self.manager_cnt, "Cannot have more processes than managers"
+        assert self.processes >= 1
+
+    @staticmethod
+    def _reward_getter(args):
+        pid = os.getpid()
+        res = []
+        for i in range(len(args)):
+            try:
+                if hasattr(args[i]['manager'], 'devices'):
+                    devices = args[i]['manager'].devices
+                else:
+                    devices = "NoAttribute"
+                sys.stderr.write("PID %i: %i/%i run; devices=%s\n" % (pid, i, len(args), devices))
+                reward, loss_and_metrics = args[i]['manager'].get_rewards(
+                    trial=args[i]['trial'], model_arc=args[i]['model_arc'], nsteps=args[i]['nsteps']
+                    )
+            except Exception as e:
+                raise Exception("child pid %i has exception %s" % (pid, e))
+            res.append({'reward': reward, 'loss_and_metrics': loss_and_metrics})
+        # close all handlers opened in this thread
+        for i in range(len(args)):
+            if hasattr(args[i]['manager'], "close_handler"):
+                args[i]['manager'].close_handler()
+        return res
+
+    # overwrite
+    def _train_loop(self):
+        if self.resume_prev_run:
+            f = open(os.path.join(self.working_dir, 'train_history.csv'), mode='a+')
+        else:
+            f = open(os.path.join(self.working_dir, 'train_history.csv'), mode='w')
+        writer = csv.writer(f)
+        starttime = datetime.datetime.now()
+        action_probs_record = []
+        loss_and_metrics_list = []
+        controller_step = self.start_ep * self.max_step_per_ep
+
+        from multiprocessing import set_start_method, get_context
+        set_start_method("spawn")
+        for child_step in range(self.start_ep, self.max_episode):
+            try:
+                if self.is_enas:
+                    # train child parameters w, if is_enas
+                    pool_args = []
+                    for j in range(self.manager_cnt):
+                        pool_args.append(
+                            [{
+                                'manager': self.manager[j],
+                                'trial': child_step,
+                                'model_arc': None,
+                                'nsteps': self.child_train_steps
+                            }]
+                        )
+                    self.logger.info("sampling of n=%i managers executed in parallel.." % self.manager_cnt)
+                    _ = pool.map(self._reward_getter, pool_args)
+
+                # train controller parameters theta
+                ep_reward = 0
+                loss_and_metrics_ep = defaultdict(float)
+
+                # pool_args is a list of list; each element is list of args (dict)
+                # for a spawned process handling one manager
+                pool_args = []
+                # store_args is a list of list; each element is list of args (dict)
+                # for controller.store; ordered the same as pool_args
+                store_args = []
+
+                for j in range(self.manager_cnt):
+                    this_pool = []
+                    this_store = []
+                    controller_step_ = controller_step
+                    for step in range(self.max_step_per_ep):
+                        ep_probs = []
+                        arc_seq, probs = self.controller.get_action(
+                            description_feature=self.data_descriptive_features[[j]])  # preserve the first dimension
+                        self.entropy_record.append(compute_entropy(probs))
+                        ep_probs.append(probs)
+                        # LOGGER.debug the action probabilities
+                        action_list = parse_action_str_squeezed(arc_seq, self.controller.state_space)
+                        self.logger.debug("Manager {}, Predicted actions : {}".format(j, [str(x) for x in action_list]))
+
+                        this_pool.append(
+                            {
+                                'manager': self.manager[j],
+                                'trial': controller_step_,
+                                'model_arc': arc_seq,
+                                'nsteps': self.child_train_steps
+                            }
+                        )
+
+                        this_store.append(
+                            {
+                                'prob': probs,
+                                'action': arc_seq,
+                                'description': self.data_descriptive_features[[j]]
+                            }
+                        )
+
+                        controller_step_ += 1
+
+                    # append the j-th manager's args to the pool
+                    pool_args.append(this_pool)
+                    store_args.append(this_store)
+
+                # update trial
+                controller_step += self.max_step_per_ep
+
+                # Plain `Pool` won't work; see here for a very nice explanation:
+                # https://pythonspeed.com/articles/python-multiprocessing/
+                a_time = time.time()
+                with get_context('spawn').Pool(processes=self.processes) as pool:
+                    elp = time.time() - a_time
+                    # distribute args to the pool of workers
+                    if self.processes > 1:
+                        self.logger.info("distributing reward-getter to a pool of %i workers.. %.3f sec" % (self.processes, elp))
+                        res_list = pool.map(self._reward_getter, pool_args)
+                    else:
+                        self.logger.info("Only %i worker, running sequentially.. %.3f" % (self.processes, elp))
+                        res_list = []
+                        for x in pool_args:
+                            res_list.append(self._reward_getter(x))
+
+                for m, (store_, res_) in enumerate(zip(store_args, res_list)):  # manager level
+                    for t, (store, res) in enumerate(zip(store_, res_)):        # trial level
+
+                        reward, loss_and_metrics = res['reward'], res['loss_and_metrics']
+                        probs, arc_seq, description = store['prob'], store['action'], store['description']
+                        ep_reward += reward
+                        trial = pool_args[m][t]['trial']
+                        for x in loss_and_metrics.keys():
+                            loss_and_metrics_ep[x] += loss_and_metrics[x]
+                        # save the arc_seq and reward
+                        self.controller.store(prob=probs, action=arc_seq, reward=reward,
+                                              description=self.data_descriptive_features[[j]],
+                                              manager_index=m)
+                        # write the results of this trial into a file
+                        data = ["%i-%i"%(m,trial), [loss_and_metrics[x] for x in sorted(loss_and_metrics.keys())],
+                                reward]
+                        if self.squeezed_action:
+                            data.extend(arc_seq)
+                        else:
+                            data.extend(action_list)
+                        writer.writerow(data)
+                        f.flush()
+                        loss_and_metrics_list.append({x: (v / self.max_step_per_ep) for x, v in loss_and_metrics_ep.items()})
+                        # average probs over trajectory
+                        ep_p = [sum(p) / len(p) for p in zip(*ep_probs)]
+                        action_probs_record.append(ep_p)
+
+                if child_step >= self.initial_buffering_queue - 1:
+                    # train the controller on the saved state and the rewards
+                    loss = self.controller.train(child_step, self.working_dir)
+                    self.total_reward += np.sum(np.array(self.controller.buffer.lt_adv[-1]).flatten())
+                    self.logger.info("Total reward : " + str(self.total_reward))
+                    self.logger.info("END episode %d: Controller loss : %0.6f" % (child_step, loss))
+                    self.logger.info("-" * 10)
+                else:
+                    self.logger.info("END episode %d: Buffering" % child_step)
+                    self.logger.info("-" * 10)
+
+                # save controller weights every N epoches, if turned on
+                if self.save_controller_every is not None and child_step % self.save_controller_every == 0:
+                    self.logger.info("Saving controller weights for epoch %d" % child_step)
+                    self.controller.save_weights(
+                        os.path.join(self.working_dir, "controller_weights-epoch-%i.h5" % child_step))
+            except KeyboardInterrupt:
+                self.logger.info("User disrupted training")
+                break
+
+            consumed_time = (datetime.datetime.now() - starttime).total_seconds()
+            self.logger.info("used time: %.2f %%" % (consumed_time / self.time_budget * 100))
+            if consumed_time >= self.time_budget:
+                self.logger.info("training ceased because run out of time budget")
+                break
+
+        self.logger.debug("Total Reward : %s" % self.total_reward)
+        f.close()
+        return action_probs_record, loss_and_metrics_list
+
+    def restore(self):
+        self.controller.load_weights(os.path.join(self.working_dir, "controller_weights.h5"))
+        self.logger.info("Loaded existing weights")
+        finished_records = 0
+        with open(os.path.join(self.working_dir, "train_history.csv"), "r") as f:
+            for _ in f:
+                finished_records += 1
+        self.start_ep = finished_records // (len(self.manager)*self.max_step_per_ep)
+        self.logger.info("Loaded existing history; starting from ep %i"%self.start_ep)
diff --git a/amber/bootstrap/grid_search.py b/amber/bootstrap/grid_search.py
index 460d6a8..18b4f83 100644
--- a/amber/bootstrap/grid_search.py
+++ b/amber/bootstrap/grid_search.py
@@ -92,16 +92,16 @@ def rewrite_train_hist(working_dir, model_fn, knowledge_fn, data, suffix='new',
 
 
 def grid_search(model_space_generator, manager, working_dir, B=10, resume_prev_run=True):
-    # for backward comparability
-    if getattr(model_space_generator, "__next__", None):
-        model_space_generator_ = model_space_generator
-    else:
-        model_space_generator_ = itertools.product(*model_space_generator)
     write_mode = "a" if resume_prev_run else "w"
     fh = open(os.path.join(working_dir, 'train_history.csv'), write_mode)
     writer = csv.writer(fh)
     i = 0
     for b in range(B):
+        # for backward comparability
+        if getattr(model_space_generator, "__next__", None):
+            model_space_generator_ = model_space_generator
+        else:
+            model_space_generator_ = itertools.product(*model_space_generator)
         for model_states in model_space_generator_:
             i += 1
             print("B={} i={} arc={}".format(b, i, ','.join([str(x) for x in model_states])))
diff --git a/amber/bootstrap/mock_manager.py b/amber/bootstrap/mock_manager.py
index 2aab999..9f862ce 100644
--- a/amber/bootstrap/mock_manager.py
+++ b/amber/bootstrap/mock_manager.py
@@ -29,9 +29,8 @@ def get_mock_reward(model_states, train_history_df, metric, stringify_states=Tru
         return train_history_df[metric].iloc[np.random.choice(np.where(index)[0])]
 
 
-def get_default_mock_reward_fn(model_states, train_history_df, lbd=1.):
+def get_default_mock_reward_fn(model_states, train_history_df, lbd=1.0, metric=['loss', 'knowledge', 'acc']):
     Lambda = lbd
-    metric = ['loss', 'knowledge', 'acc']
     mock_reward = get_mock_reward(model_states, train_history_df, metric)
     this_reward = -(mock_reward['loss'] + Lambda * mock_reward['knowledge'])
     loss_and_metrics = [mock_reward['loss'], mock_reward['acc']]
@@ -90,7 +89,9 @@ def __str__(self):
         s = "MockManager with %i records" % self.train_history_df.shape[0]
         return s
 
-    def get_rewards(self, trial, model_states):
+    def get_rewards(self, trial, model_states=None, **kwargs):
+        if model_states is None:
+            model_states = kwargs.pop('model_arc', None)
         # evaluate the model by `reward_fn`
         if self.reward_fn:
             this_reward, loss_and_metrics, reward_metrics = self.reward_fn(model_states, self.train_history_df,
diff --git a/amber/modeler/__init__.py b/amber/modeler/__init__.py
index 4fb3a73..672654c 100644
--- a/amber/modeler/__init__.py
+++ b/amber/modeler/__init__.py
@@ -1,7 +1,5 @@
 """modeler Module Docstring
 """
 
-from ._enas_modeler import DAGModelBuilder, EnasAnnModelBuilder, EnasCnnModelBuilder
-from ._keras_modeler import build_sequential_model, build_multi_gpu_sequential_model, \
-    build_multi_gpu_sequential_model_from_string, build_sequential_model_from_string
-
+from .modeler import DAGModelBuilder, EnasAnnModelBuilder, EnasCnnModelBuilder
+from .modeler import KerasModelBuilder, KerasMultiIOModelBuilder
diff --git a/amber/modeler/_enas_modeler.py b/amber/modeler/_enas_modeler.py
index 98fe641..e504054 100755
--- a/amber/modeler/_enas_modeler.py
+++ b/amber/modeler/_enas_modeler.py
@@ -2,7 +2,7 @@
 
 from .dag import get_dag, get_layer
 from .dag import ComputationNode
-
+from .dag import EnasConv1dDAG
 
 class ModelBuilder:
     """Scaffold of Model Builder
@@ -166,6 +166,7 @@ def __init__(self, session=None, controller=None, dag_func='enas', l1_reg=0.0, l
                  with_output_blocks=False,
                  use_node_dag=True,
                  feature_model=None,
+                 dag_kwargs=None,
                  *args,
                  **kwargs):
         """
@@ -191,6 +192,7 @@ def __init__(self, session=None, controller=None, dag_func='enas', l1_reg=0.0, l
         self.node_dag = None
         self.use_node_dag = use_node_dag
         self.feature_model = feature_model
+        self.dag_kwargs = dag_kwargs or {}
         self.with_output_blocks = with_output_blocks
         assert not (
                     self.with_output_blocks and self.use_node_dag), "Currently `use_node_dag` is incompatible with `with_output_blocks`"
@@ -210,13 +212,16 @@ def _build_dag(self):
         self.dag = self.dag_func(model_space=self.model_space,
                                  input_node=self.inputs_op,
                                  output_node=self.output_op,
+                                 with_input_blocks=self.with_input_blocks,
+                                 with_skip_connection=self.with_skip_connection,
                                  with_output_blocks=self.with_output_blocks,
                                  session=self.session,
                                  model_compile_dict=self.model_compile_dict,
                                  l1_reg=self.l1_reg,
                                  l2_reg=self.l2_reg,
                                  controller=self.controller,
-                                 feature_model=self.feature_model)
+                                 feature_model=self.feature_model,
+                                 **self.dag_kwargs)
 
     def __call__(self, arc_seq=None, *args, **kwargs):
         input_nodes = self._get_input_nodes()
@@ -267,7 +272,7 @@ def __init__(self, session=None, controller=None, dag_func='enas', l1_reg=0.0, l
             **kwargs:
         """
         super().__init__(dag_func=dag_func, *args, **kwargs)
-        assert dag_func.lower() in ('enascnndag', 'enasconv1ddag'), "EnasModelBuilder only support enasDAG."
+        #assert dag_func.lower() in ('enascnndag', 'enasconv1ddag'), "EnasModelBuilder only support enasDAG."
         self.session = session
         self.controller = controller
         self.l1_reg = float(l1_reg)
@@ -279,6 +284,7 @@ def __init__(self, session=None, controller=None, dag_func='enas', l1_reg=0.0, l
         # END NEW ARGS
         # -----
         self._build_dag()
+        assert issubclass(type(self.dag), EnasConv1dDAG), "EnasModelBuilder only support enasDAG and its derivatives"
 
     def _build_dag(self):
         """
diff --git a/amber/modeler/_keras_modeler.py b/amber/modeler/_keras_modeler.py
index b8f0f31..40fcf25 100644
--- a/amber/modeler/_keras_modeler.py
+++ b/amber/modeler/_keras_modeler.py
@@ -1,7 +1,146 @@
-from keras.models import Model
+import tensorflow.keras as keras
+from ..architect import Operation
+from .dag import get_layer
+import numpy as np
+from ._enas_modeler import ModelBuilder
+import tensorflow as tf
+from .architecture_decoder import MultiIOArchitecture
+from tensorflow.keras.layers import Concatenate
+from tensorflow.keras.models import Model
+import copy
+import numpy as np
 
 
-def build_sequential_model(model_states, input_state, output_state, model_compile_dict):
+class KerasModelBuilder(ModelBuilder):
+    def __init__(self, inputs_op, output_op, model_compile_dict, model_space=None, gpus=None, **kwargs):
+        self.model_compile_dict = model_compile_dict
+        self.input_node = inputs_op
+        self.output_node = output_op
+        self.model_space = model_space
+        self.gpus = gpus
+
+    def __call__(self, model_states):
+        if self.gpus is None or self.gpus == 1:
+            model = build_sequential_model(
+                        model_states=model_states,
+                        input_state=self.input_node,
+                        output_state=self.output_node,
+                        model_compile_dict=self.model_compile_dict,
+                        model_space=self.model_space
+                        )
+        elif type(self.gpus) is int:
+            model = build_multi_gpu_sequential_model(
+                        model_states=model_states,
+                        input_state=self.input_node,
+                        output_state=self.output_node,
+                        model_compile_dict=self.model_compile_dict,
+                        model_space=self.model_space,
+                        gpus=self.gpus
+                        )
+        elif type(self.gpus) is list:
+            mirrored_strategy = tf.distribute.MirroredStrategy(devices=self.gpus)
+            with mirrored_strategy.scope():
+                model = build_sequential_model(
+                            model_states=model_states,
+                            input_state=self.input_node,
+                            output_state=self.output_node,
+                            model_compile_dict=self.model_compile_dict,
+                            model_space=self.model_space
+                            )
+        return model
+
+
+class KerasMultiIOModelBuilder(ModelBuilder):
+    """
+    Note:
+        Still not working if num_outputs=0
+    """
+    def __init__(self, inputs_op, output_op, model_compile_dict, model_space, with_input_blocks, with_output_blocks, dropout_rate=0.2, wsf=1, **kwargs):
+        self.model_compile_dict = model_compile_dict
+        self.inputs = inputs_op
+        self.outputs = output_op
+        self.model_space = model_space
+        self.num_inputs = len(inputs_op) if type(inputs_op) in (list, tuple) else 0
+        self.num_outputs = len(output_op) if type(output_op) in (list, tuple) else 0
+        assert not (self.num_inputs==0 & self.num_outputs==0), "MultiIO cannot have single input and single output at the same time"
+        self.with_input_blocks = with_input_blocks
+        self.with_output_blocks = with_output_blocks
+        if self.with_input_blocks: assert self.num_inputs > 0, "you specified with_input_blocks=True for KerasMultiIOModelBuilder, but only provided 1 num_inputs"
+        self.decoder = MultiIOArchitecture(num_layers=len(self.model_space), num_inputs=self.num_inputs*self.with_input_blocks, num_outputs=self.num_outputs*self.with_output_blocks)
+
+    def __call__(self, model_states):
+        model = self._convert(model_states)
+        if model is not None:
+            model.compile(**self.model_compile_dict)
+        return model
+
+    def _convert(self, arc_seq, with_bn=True, wsf=1):
+        inputs = [get_layer(x=None, state=x) for x in self.inputs] if self.num_inputs>0 else [get_layer(x=None, state=self.inputs)]
+        op, inp, skp, out = self.decoder.decode(arc_seq)
+        out_rowsum = np.apply_along_axis(np.sum, 1, out)
+        out_colsum = np.apply_along_axis(np.sum, 0, out)
+        skp_rowsum = np.array([1] + [sum(x) for x in skp])
+        with_input_blocks = self.with_input_blocks
+        # missing output connection
+        if any(out_rowsum==0):
+            print("invalid model: unconnected output")
+            return None
+        # missing output with skip connection
+        if self.with_input_blocks is False and any( (skp_rowsum==0)&(out_colsum!=0) ):
+            print("invalid model: output connected to layer with no input")
+            return None
+
+        # Build the model until outputs
+        prev_layers = []
+        for layer_id in range(len(self.model_space)):
+            this_op = op[layer_id]
+            # Prepare the inputs
+            if with_input_blocks:
+                this_inputs = [inputs[i] for i in np.where(inp[layer_id])[0]]
+            else:
+                this_inputs = inputs if layer_id == 0 else []
+            if layer_id > 0:
+                this_inputs += [ prev_layers[i] for i in np.where(skp[layer_id-1])[0] if prev_layers[i] is not None ]
+
+            # Connect tensors
+            model_op = copy.deepcopy(self.model_space[layer_id][this_op])
+            if 'units' in model_op.Layer_attributes:
+                model_op.Layer_attributes['units'] *= wsf
+            elif 'filters' in model_op.Layer_attributes:
+                model_op.Layer_attributes['filters'] *= wsf
+            else:
+                raise Exception("Cannot use wsf")
+ 
+            if len(this_inputs) > 1:
+                input_tensor = Concatenate()(this_inputs)
+                layer = get_layer(x=input_tensor, state=model_op, with_bn=with_bn)
+                prev_layers.append(layer)
+            elif len(this_inputs) == 1:
+                input_tensor = this_inputs[0]
+                layer = get_layer(x=input_tensor, state=model_op, with_bn=with_bn)
+                prev_layers.append(layer)
+            else:
+                prev_layers.append(None)  # skipped a layer
+
+        # Build the outputs
+        outputs_inputs = []
+        for m, o in enumerate(out):
+            idx = [i for i in np.where(o)[0] if prev_layers[i] is not None]
+            if len(idx) > 1:
+                outputs_inputs.append( Concatenate()([prev_layers[i] for i in idx])  )
+            elif len(idx) == 1:
+                outputs_inputs.append(prev_layers[idx[0]] )
+            else:
+                #raise Exception("Unconnected output %i"%m)
+                print("Secondary unconnected output %i"%m)
+                return None
+        outputs = [get_layer(x=outputs_inputs[i], state=self.outputs[i]) for i in range(self.num_outputs)  ]
+        model = Model(inputs=inputs, outputs=outputs)
+        return model
+     
+
+
+def build_sequential_model(model_states, input_state, output_state, model_compile_dict, **kwargs):
     """
     Args:
         model_states: a list of operators sampled from operator space
@@ -13,21 +152,30 @@ def build_sequential_model(model_states, input_state, output_state, model_compil
     """
     inp = get_layer(None, input_state)
     x = inp
-    for state in model_states:
-        x = get_layer(x, state)
+    model_space = kwargs.pop("model_space", None)
+    for i, state in enumerate(model_states):
+        if issubclass(type(state), Operation):
+            x = get_layer(x, state)
+        elif issubclass(type(state), int) or np.issubclass_(type(state), np.integer):
+            assert model_space is not None, "if provided integer model_arc, must provide model_space in kwargs"
+            x = get_layer(x, model_space[i][state])
+        else:
+            raise Exception("cannot understand %s of type %s" % (state, type(state)))
     out = get_layer(x, output_state)
     model = Model(inputs=inp, outputs=out)
-    model.compile(**model_compile_dict)
+    if not kwargs.pop('stop_compile', False):
+        model.compile(**model_compile_dict)
     return model
 
 
 def build_multi_gpu_sequential_model(model_states, input_state, output_state, model_compile_dict, gpus=4, **kwargs):
     try:
-        from keras.utils import multi_gpu_model
+        from tensorflow.keras.utils import multi_gpu_model
     except Exception as e:
         raise Exception("multi gpu not supported in keras. check your version. Error: %s" % e)
-    vanilla_model = build_sequential_model(model_states, input_state, output_state, model_compile_dict)
-    model = multi_gpu_model(vanilla_model, gpus=gpus, **kwargs)
+    with tf.device('/cpu:0'):
+        vanilla_model = build_sequential_model(model_states, input_state, output_state, model_compile_dict, stop_compile=True, **kwargs)
+    model = multi_gpu_model(vanilla_model, gpus=gpus)
     model.compile(**model_compile_dict)
     return model
 
diff --git a/amber/modeler/architecture_decoder.py b/amber/modeler/architecture_decoder.py
new file mode 100644
index 0000000..efb9f5f
--- /dev/null
+++ b/amber/modeler/architecture_decoder.py
@@ -0,0 +1,88 @@
+
+"""
+Classes for breaking down an architecture sequence into a more structured format for later use
+"""
+
+
+# Author : zzj
+# Date : 2020.8.2; revised 2020-10-06
+
+import numpy as np
+
+
+class MultiIOArchitecture:
+    def __init__(self, num_layers, num_inputs, num_outputs):
+        #self.model_space = model_space
+        #self.num_layers = len(model_space)
+        self.num_layers = num_layers
+        self.num_inputs = num_inputs
+        self.num_outputs = num_outputs
+
+    def decode(self, arc_seq):
+        start_idx = 0
+        operations = []
+        inputs = []
+        skips = []
+        for layer_id in range(self.num_layers):
+            operation = arc_seq[start_idx]
+            #operation = self.model_space[layer_id][operation]
+            start_idx += 1
+            inp = arc_seq[start_idx : (start_idx + self.num_inputs)]
+            if layer_id > 0:
+                skip = arc_seq[(start_idx + self.num_inputs) : (start_idx + self.num_inputs + layer_id)]
+                skips.append(skip)
+            operations.append(operation)
+            inputs.append(inp)
+            start_idx +=  self.num_inputs + layer_id
+        inputs = np.asarray(inputs)
+        outputs = np.asarray(arc_seq[start_idx::]).reshape((-1, self.num_layers))
+        return operations, inputs, skips, outputs
+
+
+class ResConvNetArchitecture:
+    def __init__(self, model_space):
+        """ResConvNetArchitecture is a class for decoding and encoding neural architectures of convolutional neural
+        networks with residual connections
+
+        Parameters
+        ----------
+        model_space : amber.architect.ModelSpace
+            The model space which architectures are being sampled from
+        """
+        self.model_space = model_space
+        self._num_layers = len(self.model_space)
+
+    def decode(self, arc_seq):
+        """Decode a sequence of architecture tokens into operations and res-connections
+        """
+        start_idx = 0
+        operations = []
+        res_con = []
+        for layer_id in range(self._num_layers):
+            operations.append(arc_seq[start_idx])
+            if layer_id > 0:
+                res_con.append(arc_seq[(start_idx+1) : (start_idx + layer_id + 1)])
+            start_idx += layer_id + 1
+        return operations, res_con
+
+    def encode(self, operations, res_con):
+        """Encode operations and residual connections to a sequence of architecture tokens
+
+        This is the inverse function for `decode`
+
+        Parameters
+        ----------
+        operations : list
+            A list of integers for categorically-encoded operations
+        res_con : list
+            A list of list where each entry is a binary-encoded residual connections
+        """
+        operations_ = list(operations)
+        arc_seq = [operations_.pop(0)]
+        for op, res in zip(operations_, res_con):
+            arc_seq.append(op)
+            arc_seq.extend(res)
+        return arc_seq
+
+
+
diff --git a/amber/modeler/child.py b/amber/modeler/child.py
index cf858c5..79e4672 100644
--- a/amber/modeler/child.py
+++ b/amber/modeler/child.py
@@ -472,7 +472,7 @@ def predict_pipe(self, x, batch_size=None, verbose=0):
 
     def evaluate_ph(self, x, y, batch_size=None):
         if batch_size is None:
-            batch_size = min(1000, len(x[0]))
+            batch_size = min(100, x.shape[0])
         loss_and_metrics = []
         seen = 0
         for x_, y_ in batchify(x, y, batch_size=batch_size, shuffle=False):
@@ -571,13 +571,15 @@ class EnasCnnModel:
           methods though
     """
 
-    def __init__(self, inputs, outputs, labels, arc_seq, dag, session, dropouts=None, use_pipe=None, name='EnasModel'):
+    def __init__(self, inputs, outputs, labels, arc_seq, dag, session, dropouts=None, use_pipe=None, name='EnasModel',
+                 **kwargs):
         assert type(inputs) in (tf.Tensor, list), "get unexpected inputs types: %s" % type(inputs)
         assert type(outputs) in (tf.Tensor, list), "get unexpected outputs types: %s" % type(outputs)
         self.arc_seq = arc_seq
         self.dag = dag
         self.inputs = [inputs] if type(inputs) is tf.Tensor else inputs
         self.outputs = [outputs] if type(outputs) is tf.Tensor else outputs
+        self.callbacks = None
         self.labels = [labels] if type(labels) is tf.Tensor else labels
         self.session = session
         self.dropouts = dropouts
@@ -598,6 +600,9 @@ def __init__(self, inputs, outputs, labels, arc_seq, dag, session, dropouts=None
         self.use_pipe = use_pipe or False
         self.reinitialize_train_pipe = None
 
+        # added 2020.5.17: add default feed-dict for sampled architecture, to account for data description NAS
+        self.sample_dag_feed_dict = kwargs.pop("sample_dag_feed_dict", {})
+
         # for Keras
         self.stop_training = False
         self.metrics_name = []
@@ -646,7 +651,7 @@ def _make_feed_dict(self, x=None, y=None, is_training_phase=False):
         assert x is None or type(x) is list, "x arg for _make_feed_dict must be List, got %s" % type(x)
         assert y is None or type(y) is list, "x arg for _make_feed_dict must be List, got %s" % type(y)
         if self.arc_seq is None:
-            feed_dict = {}
+            feed_dict = self.sample_dag_feed_dict
         else:
             feed_dict = {self.dag.input_arc[i]: self.arc_seq[i] for i in range(len(self.arc_seq))}
         if x is not None:
@@ -709,8 +714,9 @@ def fit(self, x, y, batch_size=None, nsteps=None, epochs=1, verbose=1, callbacks
                 if verbose == 2:
                     if it % 1000 == 0:
                         print(
-                        "%s %i/%i, loss=%.5f" % (datetime.datetime.now().strftime("%H:%M:%S"), it, nsteps, curr_loss),
-                        flush=True)
+                            "%s %i/%i, loss=%.5f" %
+                            (datetime.datetime.now().strftime("%H:%M:%S"), it, nsteps, curr_loss),
+                            flush=True)
 
             if validation_data:
                 val_logs = self.evaluate(validation_data[0], validation_data[1])
@@ -791,7 +797,7 @@ def save(self, *args, **kwargs):
         Returns:
 
         """
-        warnings.warn("Not implemented yet; rolling back to `save_weights`")
+        warnings.warn("Not implemented yet; rolling back to `save_weights`", stacklevel=2)
         self.save_weights(*args, **kwargs)
         return
 
diff --git a/amber/modeler/dag.py b/amber/modeler/dag.py
index 63a7aac..31e0a78 100644
--- a/amber/modeler/dag.py
+++ b/amber/modeler/dag.py
@@ -9,9 +9,11 @@
     6.12.2019
 """
 import numpy as np
-import tensorflow as tf
-from keras.layers import Concatenate
-from keras.models import Model
+import warnings
+#import tensorflow as tf
+from ..utils import corrected_tf as tf
+from tensorflow.keras.layers import Concatenate
+from tensorflow.keras.models import Model
 # TODO: need to clean up `State` as `Operation`
 from ..architect.model_space import State
 
@@ -21,11 +23,11 @@
     create_bias, batch_norm1d
 # for get layers
 from keras import backend as K
-from keras.layers import GlobalAveragePooling1D, GlobalMaxPooling1D, GaussianNoise
-from keras.layers.core import Dense, Dropout, Flatten
-from keras.layers.convolutional import Conv1D, MaxPooling1D, AveragePooling1D
-from keras.layers import Input, Lambda, Permute
-from keras.layers.recurrent import LSTM
+from tensorflow.keras.layers import GlobalAveragePooling1D, GlobalMaxPooling1D, GaussianNoise
+from tensorflow.keras.layers import Dense, Dropout, Flatten
+from tensorflow.keras.layers import Conv1D, MaxPooling1D, AveragePooling1D
+from tensorflow.keras.layers import Input, Lambda, Permute, BatchNormalization, Activation
+from tensorflow.keras.layers import LSTM
 from ..modeler.operators import Layer_deNovo, SeparableFC, sparsek_vec
 from ..architect.model_space import get_layer_shortname
 
@@ -44,15 +46,24 @@ def get_dag(arg):
             return EnasAnnDAG
         elif arg.lower() == 'enascnndag' or arg.lower() == 'enasconv1ddag':
             return EnasConv1dDAG
+        elif arg == 'EnasConv1DwDataDescrption':
+            return EnasConv1DwDataDescrption
     elif callable(arg):
         return arg
     else:
         raise ValueError("Could not understand the DAG func:", arg)
 
 
-def get_layer(x, state):
+def get_layer(x, state, with_bn=False):
     if state.Layer_type == 'dense':
-        return Dense(**state.Layer_attributes)(x)
+        if with_bn is True:
+            actv_fn = state.Layer_attributes.pop('activation', 'linear')
+            x = Dense(**state.Layer_attributes)(x)
+            x = BatchNormalization()(x)
+            x = Activation(actv_fn)(x)
+            return x
+        else:
+            return Dense(**state.Layer_attributes)(x)
 
     elif state.Layer_type == 'sfc':
         return SeparableFC(**state.Layer_attributes)(x)
@@ -61,7 +72,14 @@ def get_layer(x, state):
         return Input(**state.Layer_attributes)
 
     elif state.Layer_type == 'conv1d':
-        return Conv1D(**state.Layer_attributes)(x)
+        if with_bn is True:
+            actv_fn = state.Layer_attributes.pop('activation', 'linear')
+            x = Conv1D(**state.Layer_attributes)(x)
+            x = BatchNormalization()(x)
+            x = Activation(actv_fn)(x)
+            return x
+        else:
+            return Conv1D(**state.Layer_attributes)(x)
 
     elif state.Layer_type == 'denovo':
         x = Lambda(lambda t: K.expand_dims(t))(x)
@@ -96,7 +114,8 @@ def get_layer(x, state):
         return Dropout(**state.Layer_attributes)(x)
 
     elif state.Layer_type == 'identity':
-        return Lambda(lambda t: t, **state.Layer_attributes)(x)
+        #return Lambda(lambda t: t, **state.Layer_attributes)(x)
+        return Lambda(lambda t: t)(x)
 
     elif state.Layer_type == 'gaussian_noise':
         return GaussianNoise(**state.Layer_attributes)(x)
@@ -528,8 +547,8 @@ def __init__(self,
             with_input_blocks:
             name:
         """
-        assert with_skip_connection == with_input_blocks == True, \
-            "EnasAnnDAG must have with_input_blocks and with_skip_connection"
+        #assert with_skip_connection == with_input_blocks == True, \
+        #    "EnasAnnDAG must have with_input_blocks and with_skip_connection"
         self.model_space = model_space
         if not type(input_node) in (tuple, list):
             self.input_node = [input_node]
@@ -595,6 +614,9 @@ def set_controller(self, controller):
     def _verify_args(self):
         """verify vanilla ANN model space, input nodes, etc.,
          and configure internal attr. like masking steps"""
+        # check the consistency of with_output_blocks and output_op
+        if not self.with_output_blocks and len(self.output_node)>1:
+            warnings.warn("You specified `with_output_blocks=False`, but gave a List of output operations of length %i"%len(self.output_node), stacklevel=2)
         # model space
         assert len(set([tuple(self.model_space[i]) for i in
                         range(self.num_layers)])) == 1, "model_space for EnasDAG must be identical for all layers"
@@ -767,29 +789,42 @@ def _build_dag(self, arc_seq):
             start_idx += 1
 
             # input masking for with_input_blocks
-            inp_mask = arc_seq[start_idx: start_idx + self.num_input_blocks]
-            inp_mask = tf.boolean_mask(self._input_block_map, tf.squeeze(inp_mask))
-            new_range = tf.range(0, limit=self._feature_max_size, dtype=tf.int32)
-            inp_mask = tf.map_fn(lambda x: tf.cast(tf.logical_and(x[0] <= new_range, new_range < x[1]), dtype=tf.int32),
-                                 inp_mask)
-            inp_mask = tf.reduce_sum(inp_mask, axis=0)
-            start_idx += self.num_input_blocks * self.with_input_blocks
+            if self.with_input_blocks:
+                inp_mask = arc_seq[start_idx: start_idx + self.num_input_blocks]
+                inp_mask = tf.boolean_mask(self._input_block_map, tf.squeeze(inp_mask))
+                new_range = tf.range(0, limit=self._feature_max_size, dtype=tf.int32)
+                inp_mask = tf.map_fn(lambda x: tf.cast(tf.logical_and(x[0] <= new_range, new_range < x[1]), dtype=tf.int32),
+                                     inp_mask)
+                inp_mask = tf.reduce_sum(inp_mask, axis=0)
+                start_idx += self.num_input_blocks * self.with_input_blocks
+            else:
+                # get all inputs if layer_id=0, else mask all
+                inp_mask = tf.ones(shape=(self._feature_max_size), dtype=tf.int32) if layer_id == 0 else \
+                           tf.zeros(shape=(self._feature_max_size), dtype=tf.int32)
 
             # hidden layer masking for with_skip_connection
-            if layer_id > 0:
-                layer_mask = arc_seq[
-                             start_idx: start_idx + layer_id]
-                layer_mask = tf.boolean_mask(self._skip_conn_map[layer_id], layer_mask)
-                new_range2 = tf.range(0, limit=layer_id * self._weight_max_units, delta=1, dtype=tf.int32)
-                layer_mask = tf.map_fn(
-                    lambda t: tf.cast(tf.logical_and(t[0] <= new_range2, new_range2 < t[1]), dtype=tf.int32),
-                    layer_mask)
-                layer_mask = tf.reduce_sum(layer_mask, axis=0)
+            if self.with_skip_connection:
+                if layer_id > 0:
+                    layer_mask = arc_seq[
+                                 start_idx: start_idx + layer_id]
+                    layer_mask = tf.boolean_mask(self._skip_conn_map[layer_id], layer_mask)
+                    new_range2 = tf.range(0, limit=layer_id * self._weight_max_units, delta=1, dtype=tf.int32)
+                    layer_mask = tf.map_fn(
+                        lambda t: tf.cast(tf.logical_and(t[0] <= new_range2, new_range2 < t[1]), dtype=tf.int32),
+                        layer_mask)
+                    layer_mask = tf.reduce_sum(layer_mask, axis=0)
+                    start_idx += layer_id * self.with_skip_connection
+                else:
+                    layer_mask = []
                 row_mask = tf.concat([inp_mask, layer_mask], axis=0)
-                start_idx += layer_id * self.with_skip_connection
             else:
-                row_mask = inp_mask
+                if layer_id > 0:
+                    # keep last/closest layer, mask all others
+                    layer_masks = [tf.zeros(shape=(self._weight_max_units*(layer_id-1)), dtype=tf.int32), tf.ones(shape=(self._weight_max_units), dtype=tf.int32)]
+                else:
+                    layer_masks = []
 
+                row_mask = tf.concat([inp_mask] + layer_masks, axis=0)
             w_mask = tf.matmul(tf.expand_dims(row_mask, -1), tf.expand_dims(col_mask, 0))
 
             # get the TF layer
@@ -887,103 +922,104 @@ def _compile(self, w_masks, model_output, use_pipe=True, var_scope=None):
         metrics = self.model_compile_dict['metrics'] if 'metrics' in self.model_compile_dict else None
         var_scope = var_scope or self.name
 
-        if self.feature_model is None or self.feature_model.pseudo_inputs_pipe is None:
-            labels = self.child_model_label
-        else:
-            # TODO: for now, only use data_pipe for sample_arc
-            if use_pipe:
-                labels = self.child_model_label_pipe
-            else:
+        with tf.compat.v1.variable_scope("compile", reuse=tf.AUTO_REUSE):
+            if self.feature_model is None or self.feature_model.pseudo_inputs_pipe is None:
                 labels = self.child_model_label
+            else:
+                # TODO: for now, only use data_pipe for sample_arc
+                if use_pipe:
+                    labels = self.child_model_label_pipe
+                else:
+                    labels = self.child_model_label
 
-        # TODO: process loss_weights
-        loss_weights = self.model_compile_dict['loss_weights'] if 'loss_weights' in self.model_compile_dict else None
-        if type(loss) is str:
-            loss_ = [get_tf_loss(loss, labels[i], model_output[i]) for i in range(len(model_output))]
-            # loss_ should originally all have the batch_size dim, then reduce_mean to 1-unit sample
-            # if the loss is homogeneous, then take the mean
-            loss_ = tf.reduce_mean(loss_)
-        elif type(loss) is list:
-            loss_ = []
-            for i in range(len(loss)):
-                loss_.append(get_tf_loss(
-                    loss[i],
-                    labels[i],
-                    model_output[i]
-                ))
-            # if the loss are provided by a list, assume its heterogeneous, and return the sum of
-            # each individual loss components
-            loss_ = tf.reduce_sum(loss_)
-        elif callable(loss):
-            loss_ = tf.reduce_sum([loss(labels[i], model_output[i]) for i in range(len(model_output))])
-        else:
-            raise Exception("expect loss to be str, list, dict or callable; got %s" % loss)
-        trainable_var = [var for var in tf.trainable_variables() if var.name.startswith(var_scope)]
-        if self.feature_model_trainable:
-            feature_model_trainable_var = [var for var in tf.trainable_variables() if
-                                           var.name.startswith(self.feature_model.name)]
-            assert len(feature_model_trainable_var) > 0, "You asked me to train featureModel but there is no trainable " \
-                                                         "variables in featureModel"
-            trainable_var += feature_model_trainable_var
-        regularization_penalty = 0.
-        if self.l1_reg > 0:
-            l1_regularizer = tf.contrib.layers.l1_regularizer(
-                scale=self.l1_reg, scope=self.name
-            )
-            l1_regularization_penalty = tf.contrib.layers.apply_regularization(l1_regularizer,
-                                                                               [var[0] for var in w_masks])
-            loss_ += l1_regularization_penalty
-        else:
-            l1_regularization_penalty = 0.
-
-        if self.l2_reg > 0:
-            l2_regularizer = tf.contrib.layers.l2_regularizer(
-                scale=self.l2_reg, scope=self.name
-            )
-            l2_regularization_penalty = tf.contrib.layers.apply_regularization(l2_regularizer,
-                                                                               [var[0] for var in w_masks])
-            loss_ += l2_regularization_penalty
-        else:
-            l2_regularization_penalty = 0.
-
-        regularization_penalty += l1_regularization_penalty + l2_regularization_penalty
-
-        # default settings used from enas
-        if self.child_train_op_kwargs is None:
-            # more sensible default values
-            train_op, lr, grad_norm, optimizer_ = get_keras_train_ops(
-                # get_train_ops(
-                loss=loss_,
-                tf_variables=trainable_var,
-                optim_algo=optimizer,
-                train_step=self.train_step)
-        # user specific settings; useful when training the final searched arc
-        else:
-            train_op, lr, grad_norm, optimizer_ = get_keras_train_ops(
-                # get_train_ops
-                loss=loss_,
-                tf_variables=trainable_var,
-                train_step=self.train_step,
-                optim_algo=optimizer,
-                **self.child_train_op_kwargs)
-        if metrics is None:
-            metrics = []
-        else:
-            metrics = [get_tf_metrics(x) for x in metrics]
-        # TODO: this needs fixing to be more generic;
-        # TODO: for example, the squeeze op is not usable for
-        # TODO: other metrics such as Acc
-        metrics_ = [f(tf.squeeze(self.child_model_label[i]), tf.squeeze(model_output[i]))
-                    for i in range(len(model_output)) for f in metrics]
-        ops = {'train_op': train_op,
-               'lr': lr,
-               'grad_norm': grad_norm,
-               'optimizer': optimizer,
-               'loss': loss_,
-               'metrics': metrics_,
-               'reg_cost': regularization_penalty
-               }
-        return ops
+            # TODO: process loss_weights
+            loss_weights = self.model_compile_dict['loss_weights'] if 'loss_weights' in self.model_compile_dict else None
+            if type(loss) is str:
+                loss_ = [get_tf_loss(loss, labels[i], model_output[i]) for i in range(len(model_output))]
+                # loss_ should originally all have the batch_size dim, then reduce_mean to 1-unit sample
+                # if the loss is homogeneous, then take the mean
+                loss_ = tf.reduce_mean(loss_)
+            elif type(loss) is list:
+                loss_ = []
+                for i in range(len(loss)):
+                    loss_.append(get_tf_loss(
+                        loss[i],
+                        labels[i],
+                        model_output[i]
+                    ))
+                # if the loss are provided by a list, assume its heterogeneous, and return the sum of
+                # each individual loss components
+                loss_ = tf.reduce_sum(loss_)
+            elif callable(loss):
+                loss_ = tf.reduce_sum([loss(labels[i], model_output[i]) for i in range(len(model_output))])
+            else:
+                raise Exception("expect loss to be str, list, dict or callable; got %s" % loss)
+            trainable_var = [var for var in tf.trainable_variables() if var.name.startswith(var_scope)]
+            if self.feature_model_trainable:
+                feature_model_trainable_var = [var for var in tf.trainable_variables() if
+                                               var.name.startswith(self.feature_model.name)]
+                assert len(feature_model_trainable_var) > 0, "You asked me to train featureModel but there is no trainable " \
+                                                             "variables in featureModel"
+                trainable_var += feature_model_trainable_var
+            regularization_penalty = 0.
+            if self.l1_reg > 0:
+                l1_regularizer = tf.contrib.layers.l1_regularizer(
+                    scale=self.l1_reg, scope=self.name
+                )
+                l1_regularization_penalty = tf.contrib.layers.apply_regularization(l1_regularizer,
+                                                                                   [var[0] for var in w_masks])
+                loss_ += l1_regularization_penalty
+            else:
+                l1_regularization_penalty = 0.
+
+            if self.l2_reg > 0:
+                l2_regularizer = tf.contrib.layers.l2_regularizer(
+                    scale=self.l2_reg, scope=self.name
+                )
+                l2_regularization_penalty = tf.contrib.layers.apply_regularization(l2_regularizer,
+                                                                                   [var[0] for var in w_masks])
+                loss_ += l2_regularization_penalty
+            else:
+                l2_regularization_penalty = 0.
+
+            regularization_penalty += l1_regularization_penalty + l2_regularization_penalty
+
+            # default settings used from enas
+            if self.child_train_op_kwargs is None:
+                # more sensible default values
+                train_op, lr, grad_norm, optimizer_ = get_keras_train_ops(
+                    # get_train_ops(
+                    loss=loss_,
+                    tf_variables=trainable_var,
+                    optim_algo=optimizer,
+                    train_step=self.train_step)
+            # user specific settings; useful when training the final searched arc
+            else:
+                train_op, lr, grad_norm, optimizer_ = get_keras_train_ops(
+                    # get_train_ops
+                    loss=loss_,
+                    tf_variables=trainable_var,
+                    train_step=self.train_step,
+                    optim_algo=optimizer,
+                    **self.child_train_op_kwargs)
+            if metrics is None:
+                metrics = []
+            else:
+                metrics = [get_tf_metrics(x) for x in metrics]
+            # TODO: this needs fixing to be more generic;
+            # TODO: for example, the squeeze op is not usable for
+            # TODO: other metrics such as Acc
+            metrics_ = [f(tf.squeeze(self.child_model_label[i]), tf.squeeze(model_output[i]))
+                        for i in range(len(model_output)) for f in metrics]
+            ops = {'train_op': train_op,
+                   'lr': lr,
+                   'grad_norm': grad_norm,
+                   'optimizer': optimizer,
+                   'loss': loss_,
+                   'metrics': metrics_,
+                   'reg_cost': regularization_penalty
+                   }
+            return ops
 
     def connect_controller(self, controller):
         self.sample_arc = controller.sample_arc
@@ -998,6 +1034,7 @@ def __init__(self,
                  output_node,
                  model_compile_dict,
                  session,
+                 with_skip_connection=True,
                  batch_size=128,
                  keep_prob=0.9,
                  l1_reg=0.0,
@@ -1037,6 +1074,7 @@ def __init__(self,
         self.session = session
         self.l1_reg = l1_reg
         self.l2_reg = l2_reg
+        self.with_skip_connection = with_skip_connection
         self.controller = None
         if controller is not None:
             self.set_controller(controller)
@@ -1200,10 +1238,10 @@ def _build_sample_arc(self, input_tensor=None, label_tensor=None, **kwargs):
             self.sample_metrics = ops['metrics']
             self.sample_ops = ops
         # if not self.is_initialized:
-        vars = [v for v in tf.all_variables() if v.name.startswith(var_scope) and v not in self.vars]
-        if len(vars):
-            self.session.run(tf.initialize_variables(vars))
-            self.vars += vars
+        vars_ = [v for v in tf.all_variables() if v.name.startswith(var_scope) and v not in self.vars]
+        if len(vars_):
+            self.session.run(tf.initialize_variables(vars_))
+            self.vars += vars_
             # self.is_initialized = True
 
     def _build_fixed_arc(self, input_tensor=None, label_tensor=None, **kwargs):
@@ -1366,14 +1404,18 @@ def _build_dag(self, arc_seq, input_tensor=None, is_training=True, reuse=True):
             # input = self.input_node
             input = self.input_ph if input_tensor is None else input_tensor
             layers = []
-            with tf.variable_scope("stem_conv"):
-                stem_kernel_size = self.stem_config.pop('stem_kernel_size', 8)
-                stem_filters = out_filters[0]
-                w = create_weight("w", [stem_kernel_size, 4, stem_filters])
-                x = tf.nn.conv1d(input, w, 1, "SAME", data_format=self.data_format)
-                x = batch_norm1d(x, is_training, data_format=self.data_format)
-                layers.append(x)
-                self.layers.append(x)
+            has_stem_conv = self.stem_config['has_stem_conv'] if 'has_stem_conv' in self.stem_config else True
+            if has_stem_conv:
+                with tf.variable_scope("stem_conv"):
+                    stem_kernel_size = self.stem_config.pop('stem_kernel_size', 8)
+                    stem_filters = out_filters[0]
+                    w = create_weight("w", [stem_kernel_size, 4, stem_filters])
+                    x = tf.nn.conv1d(input, w, 1, "SAME", data_format=self.data_format)
+                    x = batch_norm1d(x, is_training, data_format=self.data_format)
+                    layers.append(x)
+                    self.layers.append(x)
+            else:
+                layers = [input]
 
             start_idx = 0
             for layer_id in range(self.num_layers):
@@ -1389,7 +1431,7 @@ def _build_dag(self, arc_seq, input_tensor=None, is_training=True, reuse=True):
                     else:
                         layers.append(x)
                     self.layers.append(x)
-                    if layer_id in self.pool_layers:
+                    if (self.with_skip_connection is True) and (layer_id in self.pool_layers):
                         with tf.variable_scope("pool_at_{0}".format(layer_id)):
                             pooled_layers = []
                             for i, layer in enumerate(layers):
@@ -1403,7 +1445,7 @@ def _build_dag(self, arc_seq, input_tensor=None, is_training=True, reuse=True):
                                             layer, out_filters[layer_id + 1], is_training)
                                 pooled_layers.append(x)
                             layers = pooled_layers
-                start_idx += 1 + layer_id
+                start_idx += 1 + layer_id*int(self.with_skip_connection)
 
             flatten_op = self.stem_config['flatten_op'] if 'flatten_op' in self.stem_config else 'flatten'
             if flatten_op == 'global_avg_pool' or flatten_op == 'gap':
@@ -1423,12 +1465,20 @@ def _build_dag(self, arc_seq, input_tensor=None, is_training=True, reuse=True):
                 x = tf.nn.dropout(x, rate=dropout_placeholders[-1])
             with tf.variable_scope("fc"):
                 fc_units = self.stem_config['fc_units'] if 'fc_units' in self.stem_config else 1000
-                if flatten_op == 'global_avg_pool':
-                    inp_c = x.get_shape()[-1].value
+                if flatten_op == 'global_avg_pool' or flatten_op == 'gap':
+                    try:
+                        inp_c = x.get_shape()[-1].value
+                    except AttributeError:
+                        inp_c = x.get_shape()[-1]
                     w = create_weight("w_fc", [inp_c, fc_units])
                 elif flatten_op == 'flatten':
-                    inp_c = np.prod(x.get_shape()[1:]).value
+                    try:
+                        inp_c = np.prod(x.get_shape()[1:]).value
+                    except AttributeError:
+                        inp_c = np.prod(x.get_shape()[1:])
                     w = create_weight("w_fc", [inp_c, fc_units])
+                else:
+                    raise Exception("Unknown fc string: %s" % flatten_op)
                 b = create_bias("b_fc", shape=[fc_units])
                 x = tf.matmul(x, w) + b
                 x = tf.nn.relu(x)
@@ -1447,10 +1497,16 @@ def _build_dag(self, arc_seq, input_tensor=None, is_training=True, reuse=True):
     def _refactorized_channels_for_skipcon(self, layer, out_filters, is_training):
         """for dealing with mismatch-dimensions in skip connections: use a linear transformation"""
         if self.data_format == 'NWC':
-            inp_c = layer.get_shape()[-1].value
+            try:
+                inp_c = layer.get_shape()[-1].value
+            except AttributeError:
+                inp_c = layer.get_shape()[-1]
             actual_data_format = 'channels_last'
         elif self.data_format == 'NCW':
-            inp_c = layer.get_shape()[1].value
+            try:    
+                inp_c = layer.get_shape()[1].value
+            except AttributeError:
+                inp_c = layer.get_shape()[1]
             actual_data_format = 'channels_first'
 
         with tf.variable_scope("path1_conv"):
@@ -1463,15 +1519,26 @@ def _refactorized_channels_for_skipcon(self, layer, out_filters, is_training):
     def _layer(self, arc_seq, layer_id, prev_layers, start_idx, out_filters, is_training):
         inputs = prev_layers[-1]
         if self.data_format == "NWC":
-            inp_w = inputs.get_shape()[1].value
-            inp_c = inputs.get_shape()[2].value
+            try:
+                inp_w = inputs.get_shape()[1].value
+                inp_c = inputs.get_shape()[2].value
+            except AttributeError:    # for newer tf2
+                inp_w = inputs.get_shape()[1]
+                inp_c = inputs.get_shape()[2]
+
         elif self.data_format == "NCW":
-            inp_c = inputs.get_shape()[1].value
-            inp_w = inputs.get_shape()[2].value
+            try:
+                inp_c = inputs.get_shape()[1].value
+                inp_w = inputs.get_shape()[2].value
+            except AttributeError:
+                inp_c = inputs.get_shape()[1]
+                inp_w = inputs.get_shape()[2]
+
         else:
             raise Exception("cannot understand data format: %s" % self.data_format)
         count = arc_seq[start_idx]
         branches = {}
+        strides = []
         for i in range(len(self.model_space[layer_id])):
             if self.train_fixed_arc and i != count:
                 continue
@@ -1488,12 +1555,14 @@ def _layer(self, arc_seq, layer_id, prev_layers, start_idx, out_filters, is_trai
                                           layer_attr=self.model_space[layer_id][i].Layer_attributes,
                                           is_training=is_training)
                     branches[tf.equal(count, i)] = y
+                    strides.append(self.model_space[layer_id][i].Layer_attributes['strides'])
                 elif self.model_space[layer_id][i].Layer_type == 'avgpool1d':
                     # print('%i, avgpool1d' % layer_id)
                     y = self._pool_branch(inputs, "avg",
                                           layer_attr=self.model_space[layer_id][i].Layer_attributes,
                                           is_training=is_training)
                     branches[tf.equal(count, i)] = y
+                    strides.append(self.model_space[layer_id][i].Layer_attributes['strides'])
                 elif self.model_space[layer_id][i].Layer_type == 'identity':
                     y = self._identity_branch(inputs)
                     branches[tf.equal(count, i)] = y
@@ -1501,6 +1570,9 @@ def _layer(self, arc_seq, layer_id, prev_layers, start_idx, out_filters, is_trai
                     raise Exception("Unknown layer: %s" % self.model_space[layer_id][i])
 
         self.branches.append(branches)
+        if len(strides) > 0:
+            assert len(set(strides)) == 1, "If you set strides!=1 (i.e. a reduction layer), then all candidate operations must have the same strides to keep the shape identical; got %s" % strides
+            inp_w = int(np.ceil(inp_w / strides[0]))
         if self.train_fixed_arc:
             ks = list(branches.keys())
             assert len(ks) == 1
@@ -1514,7 +1586,7 @@ def _layer(self, arc_seq, layer_id, prev_layers, start_idx, out_filters, is_trai
             out.set_shape([None, inp_w, out_filters])
         elif self.data_format == "NCW":
             out.set_shape([None, out_filters, inp_w])
-        if layer_id > 0:
+        if self.with_skip_connection is True and layer_id > 0:
             skip_start = start_idx + 1
             skip = arc_seq[skip_start: skip_start + layer_id]
             with tf.variable_scope("skip"):
@@ -1539,10 +1611,15 @@ def _conv_branch(self, inputs, layer_attr, is_training):
         dilation = layer_attr['dilation'] if 'dilation' in layer_attr else 1
         filters = layer_attr['filters']
         if self.data_format == "NWC":
-            inp_c = inputs.get_shape()[-1].value
+            try:
+                inp_c = inputs.get_shape()[-1].value
+            except AttributeError:
+                inp_c = inputs.get_shape()[-1]
         elif self.data_format == "NCW":
-            inp_c = inputs.get_shape()[1].value
-
+            try:
+                inp_c = inputs.get_shape()[1].value
+            except AttributeError:
+                inp_c = inputs.get_shape()[1]
         w = create_weight("w", [kernel_size, inp_c, filters])
         x = tf.nn.conv1d(inputs, filters=w, stride=1, padding="SAME", dilations=dilation)
         x = batch_norm1d(x, is_training, data_format=self.data_format)
@@ -1555,10 +1632,16 @@ def _pool_branch(self, inputs, avg_or_max, layer_attr, is_training):
         strides = layer_attr['strides']
         filters = layer_attr['filters']
         if self.data_format == "NWC":
-            inp_c = inputs.get_shape()[-1].value
+            try:
+                inp_c = inputs.get_shape()[-1].value
+            except AttributeError:
+                inp_c = inputs.get_shape()[-1]
             actual_data_format = "channels_last"
         elif self.data_format == "NCW":
-            inp_c = inputs.get_shape()[1].value
+            try:
+                inp_c = inputs.get_shape()[1].value
+            except AttributeError:
+                inp_c = inputs.get_shape()[1]
             actual_data_format = "channels_first"
         else:
             raise Exception("Unknown data format: %s" % self.data_format)
@@ -1584,3 +1667,55 @@ def _pool_branch(self, inputs, avg_or_max, layer_attr, is_training):
 
     def _identity_branch(self, inputs):
         return lambda: inputs
+
+
+class EnasConv1DwDataDescrption(EnasConv1dDAG):
+    """This is a modeler that specifiied for convolution network with data description features
+    Date: 2020.5.17
+    """
+    def __init__(self, data_description, *args, **kwargs):
+        self.data_description = data_description
+        super().__init__(*args, **kwargs)
+        if len(self.data_description.shape) < 2:
+            self.data_description = np.expand_dims(self.data_description, axis=0)
+
+    # overwrite
+    def _model(self, arc, **kwargs):
+        """
+        Overwrite the parent `_model` method to feed the description to controller when sampling architectures
+        :param arc:
+        :param kwargs:
+        :return:
+        """
+        if self.train_fixed_arc:
+            assert arc == self.fixed_arc or arc is None, "This DAG instance is built to train fixed arc, hence you can only provide arc=None or arc=self.fixed_arc; check the initialization of this instances"
+        if arc is None:
+            if self.train_fixed_arc:
+                model = EnasCnnModel(inputs=self.fixed_model_input,
+                                     outputs=self.fixed_model_output,
+                                     labels=self.fixed_model_label,
+                                     arc_seq=arc,
+                                     dag=self,
+                                     session=self.session,
+                                     name=self.name)
+
+            else:
+                model = EnasCnnModel(inputs=self.sample_model_input,
+                                     outputs=self.sample_model_output,
+                                     labels=self.sample_model_label,
+                                     arc_seq=arc,
+                                     dag=self,
+                                     session=self.session,
+                                     name=self.name,
+                                     sample_dag_feed_dict={
+                                         self.controller.data_descriptive_feature: self.data_description}
+                                     )
+        else:
+            model = EnasCnnModel(inputs=self.fixed_model_input,
+                                 outputs=self.fixed_model_output,
+                                 labels=self.fixed_model_label,
+                                 arc_seq=arc,
+                                 dag=self,
+                                 session=self.session,
+                                 name=self.name)
+        return model
diff --git a/amber/modeler/modeler.py b/amber/modeler/modeler.py
new file mode 100644
index 0000000..019d306
--- /dev/null
+++ b/amber/modeler/modeler.py
@@ -0,0 +1,23 @@
+"""
+Modeler is an interface class that interacts outside with manager, and inside coordinates with dag and child.
+- Dag builds the underlying tensors
+- child facilitates training and evaluating
+"""
+
+from ._enas_modeler import DAGModelBuilder, EnasAnnModelBuilder, EnasCnnModelBuilder
+from ._keras_modeler import KerasModelBuilder, KerasMultiIOModelBuilder, \
+    build_sequential_model, build_multi_gpu_sequential_model, \
+    build_multi_gpu_sequential_model_from_string, build_sequential_model_from_string
+
+
+__all__ = [
+    'DAGModelBuilder',
+    'EnasCnnModelBuilder',
+    'EnasAnnModelBuilder',
+    'KerasModelBuilder',
+    'KerasMultiIOModelBuilder'
+    #'build_sequential_model',
+    #'build_sequential_model_from_string',
+    #'build_multi_gpu_sequential_model_from_string',
+    #'build_multi_gpu_sequential_model'
+]
diff --git a/amber/objective/__init__.py b/amber/objective/__init__.py
index 0803781..fea8950 100644
--- a/amber/objective/__init__.py
+++ b/amber/objective/__init__.py
@@ -1,6 +1,7 @@
 """Module Docstring: objective
 """
 
+from .general_obj_func import AuxilaryAcc
 from .motif import MotifSaliency, MotifKLDivergence
 from .graph import GraphHierarchyTreeAuxLoss, GraphHierarchyTree, GraphKnowledgeHessBias, GraphKnowledgeHessFunc
 
diff --git a/amber/objective/general_obj_func.py b/amber/objective/general_obj_func.py
index ad48ec9..cedbfe0 100644
--- a/amber/objective/general_obj_func.py
+++ b/amber/objective/general_obj_func.py
@@ -45,4 +45,49 @@ def get_obj_val(self, **kwargs):
         return self.obj_fn(self.W_model, self.W_knowledge, **kwargs)
 
 
+class AuxilaryAcc(GeneralKnowledgeObjectiveFunction):
+    """Perform the accuracy measurement (defined by method argument)  on an 
+    auxilary dataset
+    """
+    def __init__(self, method='auc', *args, **kwargs):
+        if method == 'auc' or method == 'auroc':
+            from sklearn.metrics import roc_auc_score
+            self.scorer = roc_auc_score
+        elif method == 'aupr' or method == 'auprc':
+            from sklearn.metrics import average_precision_score
+            self.scorer = average_precision_score
+        elif callable(method):
+            self.scorer = method
+        else:
+            raise Exception("cannot understand scorer method: %s" % method)
+        super().__init__(*args, **kwargs)
+
+    def model_encoder(self, model, *args, **kwargs):
+        self.W_model = model.predict(self.W_knowledge[0])
+
+    def knowledge_encoder(self, data, *args, **kwargs):
+        self.W_knowledge = data
 
+    def _build_obj_func(self, **kwargs):
+        def obj_fn(W_model, W_knowledge, **kwargs):
+            y = W_knowledge[1]
+            pred = W_model
+            auc_list = []
+            if type(y) is not list:
+                y = [y]
+            if type(pred) is not list:
+                pred = [pred]
+            for i in range(len(y)):
+                tmp = []
+                if len(y[i].shape) == 1: y[i] = np.expand_dims(y[i], axis=-1)
+                if len(pred[i].shape) == 1: pred[i] = np.expand_dims(pred[i], axis=-1)
+                for j in range(y[i].shape[1]):
+                    try:
+                        score = self.scorer(y_true=y[i][:, j], y_score=pred[i][:, j])
+                        tmp.append(score)
+                    except ValueError:  # only one-class present
+                        pass
+                auc_list.append(tmp)
+            avg_auc = np.nanmean(np.concatenate(auc_list, axis=0))
+            return avg_auc
+        self.obj_fn = obj_fn
diff --git a/amber/objective/motif.py b/amber/objective/motif.py
index 96aab9c..e06184a 100644
--- a/amber/objective/motif.py
+++ b/amber/objective/motif.py
@@ -57,7 +57,7 @@ def model_encoder(self, model, data, **kwargs):
         return self
 
     def knowledge_encoder(self, motif_name_list, motif_file, is_log_motif, autoremove_dup=False):
-        from plots import motif
+        from ..utils import motif
         motif_dict = motif.load_binding_motif_pssm(motif_file, is_log_motif)
         self.W_knowledge = {motif_name: motif_dict[motif_name] for motif_name in motif_name_list}
         if autoremove_dup:
diff --git a/amber/plots/_plotsV1.py b/amber/plots/_plotsV1.py
index 99153c2..69cf488 100644
--- a/amber/plots/_plotsV1.py
+++ b/amber/plots/_plotsV1.py
@@ -9,7 +9,6 @@
 import numpy as np
 import os
 import json
-from .heatmap2 import heatmap2
 
 matplotlib.use('Agg')
 
@@ -384,6 +383,8 @@ def plot_hessian(gkf, save_fn):
     f, ax = plt.subplots(1, 1, figsize=figsize)
     cbar_ax = f.add_axes([0.9, 0.1, cbar_width, 0.8])
     cbar_ax.tick_params(length=0.5, labelsize='x-small')
+    
+    from .heatmap2 import heatmap2
     heatmap2(g_ratio, ax=ax,
              cbar_ax=cbar_ax,
              vmax=vmax, vmin=vmin,
diff --git a/amber/plots/heatmap2.py b/amber/plots/heatmap2.py
index 5ade687..c9b3526 100644
--- a/amber/plots/heatmap2.py
+++ b/amber/plots/heatmap2.py
@@ -15,7 +15,11 @@
 import numpy as np
 import pandas as pd
 from seaborn import cm
-from seaborn.external.six import string_types
+try:
+    from seaborn.external.six import string_types
+except ModuleNotFoundError:
+    # patch on 2020.6.8
+    from six import string_types
 from seaborn.utils import (despine, axis_ticklabels_overlap, relative_luminance, to_utf8)
 
 
diff --git a/amber/utils/__init__.py b/amber/utils/__init__.py
index e69de29..4409959 100644
--- a/amber/utils/__init__.py
+++ b/amber/utils/__init__.py
@@ -0,0 +1,23 @@
+import tensorflow as corrected_tf
+if corrected_tf.__version__.startswith("2"):
+    corrected_tf.compat.v1.disable_eager_execution()
+    import tensorflow.compat.v1 as corrected_tf
+
+
+def run_from_ipython():
+    try:
+        __IPYTHON__
+        return True
+    except NameError:
+        return False
+
+
+def get_available_gpus():
+    from tensorflow.python.client import device_lib
+    local_device_protos = device_lib.list_local_devices()
+    return [x.name for x in local_device_protos if x.device_type == 'GPU']
+
+
+
+
+
diff --git a/amber/utils/data_parser.py b/amber/utils/data_parser.py
index f536672..d6ff65d 100644
--- a/amber/utils/data_parser.py
+++ b/amber/utils/data_parser.py
@@ -170,12 +170,16 @@ def simdata_reader(fn, targets):
         a tuple of (X, y) in np.array format
     """
     assert len(targets) == len(set(targets)), Exception('List `targets` must be unique, found duplication')
-    f = gzip.GzipFile(fn, 'rb') if fn.endswith('gz') else open(fn, 'r')
+    GZIPPED = True if fn.endswith("gz") else False
+    f = gzip.GzipFile(fn, 'rb') if GZIPPED else open(fn, 'r')
     f.readline()  # skip the header
     tmp_X = []
     tmp_y = []
     for line in f:
-        line = line.decode('UTF-8').strip()
+        if GZIPPED:
+            line = line.decode('UTF-8').strip()
+        else:
+            line = line.strip()
         ele = line.split('\t')
         seqName, seq = ele[0], ele[1]
         embeddings = ele[2] if len(ele) > 2 else ''
diff --git a/amber/utils/io.py b/amber/utils/io.py
index d545d14..d389957 100644
--- a/amber/utils/io.py
+++ b/amber/utils/io.py
@@ -32,6 +32,7 @@ def read_history(fn_list,
                 t = t[1:]
             elif idx == len(tmp) - 1:
                 t = t[:-1]
+            t = t.strip("[]")
             t = float(t)
             metrics[name].append(t)
     for name in metrics:
@@ -185,10 +186,13 @@ def save_stats(loss_and_metrics_list, working_dir):
     data = [list(loss_and_metrics.values()) for loss_and_metrics in loss_and_metrics_list]
     data_per_cat = list(zip(*data))
     k_data = data_per_cat[keys.index('knowledge')]
-    acc_data = data_per_cat[keys.index('acc')]
     loss_data = data_per_cat[keys.index('loss')]
     df['Knowledge'].append(list(k_data))
-    df['Accuracy'].append(list(acc_data))
     df['Loss'].append(list(loss_data))
-
+    # modified 2020.6.7 by ZZ: not every case will have the Accuracy metric..
+    try:
+        acc_data = data_per_cat[keys.index('acc')]
+        df['Accuracy'].append(list(acc_data))
+    except ValueError:
+        pass
     json.dump(df, open(save_path, 'w'))
diff --git a/amber/utils/logging.py b/amber/utils/logging.py
new file mode 100644
index 0000000..b50e967
--- /dev/null
+++ b/amber/utils/logging.py
@@ -0,0 +1,23 @@
+import logging
+import os
+
+
+def setup_logger(working_dir='.', verbose_level=logging.INFO):
+    # setup logger
+    logger = logging.getLogger('AMBER')
+    logger.setLevel(verbose_level)
+    # create file handler which logs even debug messages
+    fh = logging.FileHandler(os.path.join(working_dir, 'log.AMBER.txt'))
+    fh.setLevel(verbose_level)
+    # create console handler with a higher log level
+    ch = logging.StreamHandler()
+    ch.setLevel(verbose_level)
+    # create formatter and add it to the handlers
+    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s -\n %(message)s')
+    fh.setFormatter(formatter)
+    ch.setFormatter(formatter)
+    # add the handlers to the logger
+    logger.addHandler(fh)
+    logger.addHandler(ch)
+    return logger
+
diff --git a/amber/plots/motif.py b/amber/utils/motif.py
similarity index 100%
rename from amber/plots/motif.py
rename to amber/utils/motif.py
diff --git a/amber/utils/sampler.py b/amber/utils/sampler.py
new file mode 100644
index 0000000..542463a
--- /dev/null
+++ b/amber/utils/sampler.py
@@ -0,0 +1,548 @@
+# Author: Evan M. Cofer
+# Created on June 5, 2020
+"""
+This module provides the `BioIntervalSource` class and its children.
+These are essentially wrappers for sets of sequence intervals and
+associated labels.
+"""
+#import keras
+import tensorflow as tf
+import numpy
+from .sequences import EncodedHDF5Genome
+import h5py
+
+
+class BioIntervalSource(object):
+    """A generic class for labeled examples of biological intervals.
+    The amount of padding added to the end of the intervals is able to
+    be changed during runtime. This allows these functions to be passed
+    to objects such as a model controller.
+
+    Parameters
+    ----------
+    example_file : str
+        A path to a file that contains the examples in BED-like format.
+        Specifically, this file will have one example per line, with
+        the chromosome, start, end, and label for the example. Each
+        column is separated by tabs.
+    reference_sequence : Sequence
+        The reference sequence used to generate the input sequences
+        from the example coordinates.
+    seed : int, optional
+        Default is `1337`. The value used to seed random number generation.
+    n_examples : int, optional
+        Default is `None`. The number of examples. If left as `None`, will
+        use all of the examples in the file. If fewer than `n_examples`
+        are found, and error will be thrown.
+
+    Attributes
+    ----------
+    reference_sequence : Sequence
+        The reference sequence used to generate the input sequences
+        from the example coordinates.
+    examples : list
+        A list of the example coordinates.
+    labels : list
+        A list of the labels for the examples.
+    left_pad : int
+        The length of padding added to the left side of the interval.
+    right_pad : int
+        The length of padding added to the right side of the interval.
+    random_state : numpy.random.RandomState
+        A random number generator to use.
+    seed : int
+        The value used to seed the random number generator.
+    """
+    def __init__(self, example_file, reference_sequence, n_examples=None, seed=1337, pad=400):
+        if type(reference_sequence) is str:
+            self.reference_sequence = EncodedHDF5Genome(input_path=reference_sequence, in_memory=False)
+        else:
+            self.reference_sequence = reference_sequence
+        self.left_pad = 0
+        self.right_pad = 0
+
+        # Setup RNG.
+        self.seed = seed
+        self.random_state = numpy.random.RandomState(seed=self.seed)
+
+        # Load examples.
+        self.labels = list()
+        self.examples = list()
+        with open(example_file, "r") as read_file:
+            for line in read_file:
+                line = line.strip()
+                if not line.startswith("#"):
+                    if line:
+                        line = line.split("\t")
+                        chrom, start, end, strand = line[:4]
+                        label = [int(x) for x in line[4:]]
+                        self.labels.append(numpy.array(label))
+                        self.examples.append((chrom, int(start), int(end), strand))
+        # TODO: Consider using separate random states for index shuffling and this part?
+        if n_examples is not None:
+            if len(self.examples) < n_examples:
+                s = ("Specified value of examples was {}".format(n_examples) +
+                     ", but only {} were found in \"{}\".".format(len(self.examples),
+                                                                  example_file))
+                raise RuntimeError(s)
+            elif len(self.examples) > n_examples:
+                idx = self.random_state.choice(len(self.examples),
+                                               n_examples,
+                                               replace=False)
+                idx.sort()
+                self.examples = numpy.array(self.examples, dtype='O')[idx].tolist()
+                self.labels = numpy.array(self.labels, dtype='O')[idx].tolist()
+                self.labels = [numpy.array(x) for x in self.labels]
+            else:
+                # Ensure random state not affected by using input with length of n_examples.
+                idx = self.random_state.choice(2, 1, replace=False)
+                del idx
+        else: # Ensure random state not affected by not using n_examples.
+            idx = self.random.state.choice(2, 1, replace=False)
+            del idx
+
+        self.set_pad(pad)
+
+    def padding_is_valid(self, value):
+        """Determine if the specified value is a valid value for padding
+        intervals.
+
+        Parameters
+        ----------
+        value : int
+            Proposed amount of padding.
+
+        Returns
+        -------
+        bool
+            Whether the input value is valid.
+        """
+        if value < 0:
+            return False
+        else:
+            return True
+
+    def _test_padding(self, value):
+        """Tests if padding is valid or not. If invalid, raises an error.
+
+
+        Parameters
+        ----------
+        value : int
+            Amount of padding to test.
+
+        Returns
+        -------
+        None
+
+        Raises
+        ------
+        ValueError
+            This method throws an error if the proposed amount of padding is
+            not a valid amount.
+        """
+        if not self.padding_is_valid(value):
+            s = "Invalid padding amount : {}".format(value)
+            raise ValueError(s)
+
+    def set_left_pad(self, value):
+        """Sets the length of the padding added to the left
+        side of the input sequence.
+
+        Parameters
+        ----------
+        value : int
+            The length of the padding to add to the left side of an example
+            interval.
+        """
+        self._test_padding(value)
+        self.left_pad = value
+
+    def set_right_pad(self, value):
+        """Sets the length of the padding added to the right side of an
+        example interval.
+
+
+        Parameters
+        ----------
+        value : int
+            The length of the padding to add to the right side of an example
+            interval.
+        """
+        self._test_padding(value)
+        self.right_pad = value
+
+    def set_pad(self, value):
+        """Sets the length of padding added to both the left and right sides of
+        example intervals.
+
+        Parameters
+        ----------
+        value : int
+            The length of the padding to add to the left and right sides of
+            input example intervals.
+        """
+        self._test_padding(value)
+        self.left_pad = value
+        self.right_pad = value
+
+    def __len__(self):
+        """Number of examples available.
+
+        Returns
+        -------
+        int
+            The number of examples available.
+        """
+        return len(self.examples)
+
+    def _load_unshuffled(self, item):
+        """Loads example `item` from the unshuffled list of examples.
+
+        Parameters
+        ----------
+        item : int
+            The index of the example to load.
+
+        Returns
+        -------
+        tuple(numpy.ndarray, numpy.ndarray)
+        """
+        chrom, start, end, strand = self.examples[item]
+        x = self.reference_sequence.get_sequence_from_coords(chrom, start - self.left_pad, end + self.right_pad, strand)
+        y = self.labels[item]
+        return x, y
+
+
+class BioIntervalSequence(BioIntervalSource, tf.keras.utils.Sequence):
+    """This data sequence type holds intervals in a genome and a
+    label associated with each interval. Unlike a generator, this
+    is based off of `keras.utils.Sequence`, which shifts things like
+    shuffling elsewhere. The amount of padding added to the end of
+    the intervals is able to be changed during runtime. This allows
+    these functions to be passed to objects such as a model
+    controller.
+
+    Parameters
+    ----------
+    example_file : str
+        A path to a file that contains the examples in BED-like format.
+        Specifically, this file will have one example per line, with
+        the chromosome, start, end, and label for the example. Each
+        column is separated by tabs.
+    reference_sequence : Sequence
+        The reference sequence used to generate the input sequences
+        from the example coordinates.
+    n_examples : int, optional
+        Default is `None`. The number of examples. If left as `None`, will
+        use all of the examples in the file. If fewer than `n_examples` are
+        found, an error will be thrown.
+    seed : int, optional
+        Default is `1337`. The value used to seed random number generation.
+
+    Attributes
+    ----------
+    reference_sequence : Sequence
+        The reference sequence used to generate the input sequences
+        from the example coordinates.
+    examples : list
+        A list of the example coordinates.
+    labels : list
+        A list of the labels for the examples.
+    left_pad : int
+        The length of padding added to the left side of the interval.
+    right_pad : int
+        The length of padding added to the right side of the interval.
+    random_state : numpy.random.RandomState
+        A random number generator to use.
+    seed : int
+        The value used to seed the random number generator.
+    """
+    def __init__(self, example_file, reference_sequence, n_examples=None, seed=1337):
+        super(BioIntervalSequence, self).__init__(
+            example_file=example_file,
+            reference_sequence=reference_sequence,
+            n_examples=n_examples, seed=seed)
+
+    def __getitem__(self, item):
+        """
+        Indexes into the set of examples and labels.
+
+        Parameters
+        ----------
+        item : int
+            The index in the example/label pairs to fetch.
+
+        Returns
+        -------
+        tuple(numpy.ndarray, numpy.ndarray)
+            A tuple consisting of the example and the target label.
+
+        """
+        return self._load_unshuffled(item)
+
+
+class BioIntervalGenerator(BioIntervalSource):
+    """This data generator type holds intervals in a genome and a
+    label associated with each interval. This essentially acts as
+    an iterator over the inputs examples. This approach is useful
+    and preferable to `BioIntervalSequence` when there are a very
+    large number of examples in the input. The amount of padding
+    added to the end of the intervals is able to be changed during
+    runtime. This allows these functions to be passed to objects
+    such as a model controller.
+
+    Parameters
+    ----------
+    example_file : str
+        A path to a file that contains the examples in BED-like format.
+        Specifically, this file will have one example per line, with
+        the chromosome, start, end, and label for the example. Each
+        column is separated by tabs.
+    reference_sequence : Sequence
+        The reference sequence used to generate the input sequences
+        from the example coordinates.
+    n_examples : int, optional
+        Default is `None`. The number of examples. If left as `None`, will
+        use all of the examples in the file. If fewer than `n_examples` are
+        found, an error will be thrown.
+    seed : int, optional
+        Default is `1337`. The value used to seed random number generation.
+
+    Attributes
+    ----------
+    reference_sequence : Sequence
+        The reference sequence used to generate the input sequences
+        from the example coordinates.
+    examples : list
+        A list of the example coordinates.
+    labels : list
+        A list of the labels for the examples.
+    left_pad : int
+        The length of padding added to the left side of the interval.
+    right_pad : int
+        The length of padding added to the right side of the interval.
+    random_state : numpy.random.RandomState
+        A random number generator to use.
+    seed : int
+        The value used to seed random number generation.
+
+    """
+    def __init__(self, example_file, reference_sequence, n_examples=None, seed=1337):
+        super(BioIntervalGenerator, self).__init__(
+            example_file=example_file,
+            reference_sequence=reference_sequence,
+            n_examples=n_examples,
+            seed=seed)
+        raise NotImplementedError
+
+
+class BatchedBioIntervalSequence(BioIntervalSource, tf.keras.utils.Sequence):
+    """This data sequence type holds intervals in a genome and a
+    label associated with each interval. Unlike a generator, this
+    is based off of `keras.utils.Sequence`, which shifts things like
+    shuffling elsewhere. The amount of padding added to the end of
+    the intervals is able to be changed during runtime. This allows
+    these functions to be passed to objects such as a model
+    controller. Examples are divided into batches.
+
+    Parameters
+    ----------
+    example_file : str
+        A path to a file that contains the examples in BED-like format.
+        Specifically, this file will have one example per line, with
+        the chromosome, start, end, and label for the example. Each
+        column is separated by tabs.
+    reference_sequence : Sequence or str
+        The reference sequence used to generate the input sequences
+        from the example coordinates; could be a Sequence instance or a 
+        filepath to reference sequence.
+    batch_size : int
+        Specifies size of the mini-batches.
+    shuffle : bool
+        Specifies whether to shuffle the mini-batches.
+    n_examples : int, optional
+        Default is `None`. The number of examples. If left as `None`, will
+        use all of the examples in the file. If fewer than `n_examples` are
+        found, an error will be thrown.
+    seed : int, optional
+        Default is `1337`. The value used to seed random number generation.
+
+    Attributes
+    ----------
+    reference_sequence : Sequence
+        The reference sequence used to generate the input sequences
+        from the example coordinates.
+    examples : list
+        A list of the example coordinates.
+    labels : list
+        A list of the labels for the examples.
+    left_pad : int
+        The length of padding added to the left side of the interval.
+    right_pad : int
+        The length of padding added to the right side of the interval.
+    batch_size : int
+        Specifies size of the mini-batches.
+    shuffle : bool
+        Specifies whether to shuffle the mini-batches.
+    random_state : numpy.random.RandomState
+        A random number generator to use.
+    seed : int
+        The value used to seed the random number generator.
+    """
+    def __init__(self, example_file, reference_sequence,
+                 batch_size, shuffle=True, n_examples=None, seed=1337, pad=0):
+        super(BatchedBioIntervalSequence, self).__init__(
+            example_file=example_file,
+            reference_sequence=reference_sequence,
+            n_examples=n_examples,
+            seed=seed,
+            pad=pad
+            )
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.index = numpy.arange(len(self.examples))
+        self.total_batches = len(self)
+
+    def __len__(self):
+        """Number of examples available.
+
+        Returns
+        -------
+        int
+            The number of examples available.
+        """
+        l = super(BatchedBioIntervalSequence, self).__len__()
+        return l // self.batch_size
+
+    def __getitem__(self, item):
+        """
+        Indexes into the set of examples and labels.
+
+        Parameters
+        ----------
+        item : int
+            The index in the example/label pairs to fetch.
+
+        Returns
+        -------
+        tuple(numpy.ndarray, numpy.ndarray)
+            A tuple consisting of the example and the target label.
+
+        """
+        x = list()
+        y = list()
+        #for i in range(self.batch_size):
+        #    cur_x, cur_y = self._load_unshuffled(self.index[item + i])
+        for i in range(item*self.batch_size, (item+1)*self.batch_size):
+            cur_x, cur_y = self._load_unshuffled(self.index[i])
+            x.append(cur_x)
+            y.append(cur_y)
+        x = numpy.stack(x)
+        y = numpy.stack(y)
+        return x, y
+
+    def on_epoch_end(self):
+        """
+        If applicable, shuffle the examples at the end of an epoch.
+        """
+        if self.shuffle:
+            self.index = self.random_state.choice(len(self.examples),
+                                                  len(self.examples),
+                                                  replace=False)
+
+    def close(self):
+        """
+        Close the file connection of Sequence
+        """
+        self.reference_sequence.close()
+
+
+class BatchedBioIntervalSequenceGenerator(BatchedBioIntervalSequence):
+    """This class modifies on top of BatchedBioIntervalSequence by performing the 
+    generator loop infinitely
+    """
+    def __init__(self, *args, **kwargs):
+        super(BatchedBioIntervalSequenceGenerator, self).__init__(*args, **kwargs)
+        self.step = 0
+
+    def __getitem__(self, item):
+        x = list()
+        y = list()
+        if self.step >= self.total_batches:
+            #print(self.step)
+            if self.shuffle: self._shuffle()
+            self.step = 0
+        for i in range(self.step*self.batch_size, (self.step+1)*self.batch_size):
+            cur_x, cur_y = self._load_unshuffled(self.index[i])
+            x.append(cur_x)
+            y.append(cur_y)
+        x = numpy.stack(x)
+        y = numpy.stack(y)
+        self.step += 1
+        return x, y
+
+    def _shuffle(self):
+        #print("Shuffled")
+        self.index = self.random_state.choice(len(self.examples),
+                                              len(self.examples),
+                                              replace=False)
+
+    def on_epoch_end(self):
+        pass
+
+
+class BatchedHDF5Generator(tf.keras.utils.Sequence):
+    def __init__(self, hdf5_fp, batch_size, shuffle=True, in_memory=False, seed=None, x_selector=None, y_selector=None):
+        super(BatchedHDF5Generator, self).__init__()
+        self.in_memory = in_memory
+        if self.in_memory is True:
+            with h5py.File(hdf5_fp, "r") as f:
+                self.h5py_store = (f['x'][:], f['y'][:])
+        else:
+            self.hdf5_store = h5py.File(hdf5_fp, "r")
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.seed = seed
+        self.random_state = numpy.random.RandomState(seed=self.seed)
+        self.x_selector = x_selector
+        self.y_selector = y_selector
+        self.n_total_samp = (self.hdf5_store['x'].shape[0] // self.batch_size)*self.batch_size
+        #self.n_total_samp = self.hdf5_store['x'].shape[0]
+        self.n_total_batch = len(self)
+        # Method 2: only shuffle batch order, not batch composition, allowing chunk storage in hdf5
+        self.index = numpy.arange(self.n_total_samp).reshape((self.n_total_batch, -1))
+        # Method 1: column-wise shuffle 
+        #self.index = numpy.arange(self.n_total_samp).reshape((-1, self.n_total_batch)).T
+        if self.shuffle is True:
+            self._shuffle()
+        self.step = 0
+
+    def __len__(self):
+        return int(numpy.ceil(self.n_total_samp / self.batch_size))
+
+    def __getitem__(self, index):
+        if self.step >= self.n_total_batch:
+            if self.shuffle: self._shuffle()
+            self.step = 0
+        samp_idx = self.index[self.step].tolist()
+        x_batch = self.hdf5_store['x'][samp_idx]
+        y_batch = self.hdf5_store['y'][samp_idx]
+        if self.x_selector is not None:
+            x_batch = x_batch[:, self.x_selector]
+        if self.y_selector is not None:
+            y_batch = y_batch[:, self.y_selector]
+        self.step += 1
+        return x_batch, y_batch
+
+    def _shuffle(self):
+        # for method 1
+        #_ = numpy.apply_along_axis(
+        #        self.random_state.shuffle,
+        #        0,
+        #        self.index)
+        _ = self.random_state.shuffle(self.index)
+
+    def close(self):
+        if self.in_memory is False:
+            self.hdf5_store.close()
diff --git a/amber/utils/sequences.py b/amber/utils/sequences.py
new file mode 100644
index 0000000..fd16232
--- /dev/null
+++ b/amber/utils/sequences.py
@@ -0,0 +1,498 @@
+# Author: Evan M. Cofer
+# Created on June 7, 2020
+"""
+This module provides the `Sequence` class, which is an abstract class
+that defines the interface for loading biological sequence data.
+
+TODO
+----
+- Abstract away the PYFAIDX interface, so that inheritence and extension of the
+    genome class is more straightforward.
+"""
+import abc
+
+import h5py
+import numpy
+import pyfaidx
+
+_DNA_COMP_TABLE = str.maketrans("ATCGN", "TAGCN")
+_STORE_TYPE = numpy.half
+
+class Sequence(metaclass=abc.ABCMeta):
+    """This class represents a source of sequence data, which can be
+    fetched by querying different coordinates.
+    """
+    @abc.abstractmethod
+    def __len__(self):
+        """Number of queryable positions in the sequence.
+
+        Returns
+        -------
+        int
+            The number of queryable positions.
+
+        """
+        pass
+
+    @abc.abstractmethod
+    def coords_are_valid(self, *args, **kwargs):
+        """Checks if the queried coordinates are valid.
+
+        Returns
+        -------
+        bool
+            `True` if the coordinates are valid, otherwise `False`.
+        """
+        pass
+
+    @abc.abstractmethod
+    def get_sequence_from_coords(self, *args, **kwargs):
+        """Fetches a string representation of a sequence at
+        the specified coordinates.
+
+        Returns
+        -------
+        str
+            The sequence of bases occuring at the queried
+            coordinates. Behavior is undefined for invalid
+            coordinates.
+        """
+        pass
+
+
+class Genome(Sequence):
+    """This class allows the user to a query a potentially file-backed
+    genome by coordinate. It is essentially a wrapper around the
+    `pyfaidx.Fasta` class.
+
+    Parameters
+    ----------
+    input_path : str
+        Path to an indexed FASTA file.
+    in_memory : bool
+        Specifies whether the genome should be loaded from
+        disk and stored in memory.
+
+
+    Attributes
+    ----------
+    data : pyfaidx.Fasta or dict
+        The FASTA file containing the genome sequence. Alternatively,
+        this can be a `dict` object mapping chromosomes to sequences that
+        stores the file in memory.
+    in_memory : bool
+        Specified whether the genomic data is being stored in memory.
+    chrom_len_dict : dict
+        A dictionary mapping the chromosome names to their lengths.
+    """
+    def __init__(self, input_path, in_memory=False):
+        """
+        Constructs a new `Genome` object.
+        """
+        super(Genome, self).__init__()
+        self.in_memory = in_memory
+        if in_memory is True:
+            fasta = pyfaidx.Fasta(input_path)
+            self.data = {k: str(fasta[k][:].seq).upper() for k in fasta.keys()}
+            fasta.close()
+        else:
+            self.data = pyfaidx.Fasta(input_path)
+        self.chrom_len_dict = {k: len(self.data[k]) for k in self.data.keys()}
+
+    def __len__(self):
+        """Number of queryable positions in the genome.
+
+        Returns
+        -------
+        int
+            The number of queryable positions.
+        """
+        return sum(self.chrom_len_dict.values())
+
+    def coords_are_valid(self, chrom, start, end, strand="+"):
+        """Checks if the queried coordinates are valid.
+
+        Parameters
+        ----------
+        chrom : str
+            The chromosome to query from.
+        start : int
+            The first position in the queried corodinates.
+        end : int
+            One past the last position in the queried coordinates.
+        strand : str
+            Strand of sequence to draw from.
+
+        Returns
+        -------
+        bool
+            `True` if the coordinates are valid, otherwise `False`.
+        """
+        if chrom not in self.chrom_len_dict:
+            return False
+        elif start < 0 or end <= 0:
+            return False
+        elif start >= end:
+            return False
+        elif end > self.chrom_len_dict[chrom]:
+            return False
+        elif start >= self.chrom_len_dict[chrom]:
+            return False
+        elif strand not in {"+", "-"}:
+            return False
+        else:
+            return True
+
+    def get_sequence_from_coords(self, chrom, start, end, strand="+"):
+        """Fetches a string representation of a sequence at
+        the specified coordinates.
+
+
+        Parameters
+        ----------
+        chrom : str
+            Chromosome to query from.
+        start : int
+            First position in queried sequence.
+        end : int
+            One past the last position in the queried sequence.
+        strand : str
+            The strand to sample from.
+
+        Returns
+        -------
+        str
+            The sequence of bases occuring at the queried
+            coordinates.
+
+        Raises
+        ------
+        IndexError
+            If the coordinates are not valid.
+        """
+        if self.coords_are_valid(chrom, start, end, strand):
+            x = self.data[chrom][start:end]
+            if not self.in_memory:
+                x = str(x.seq).upper()
+            if strand == "-":
+                x = x.translate(_DNA_COMP_TABLE)[::-1]
+            return x
+        else:
+            s = "Specified coordinates ({} to {} on \"{}\", strand of \"{}\") are invalid!".format(
+                start, end, chrom, strand)
+            raise IndexError(s)
+
+
+class Encoding(metaclass=abc.ABCMeta):
+    """This class is a mostly-abstract class used to represent some dataset that
+    should be transformed with an encoding.
+    """
+    @abc.abstractmethod
+    def encode(self, *args, **kwargs):
+        """
+        Method to encode some input.
+        """
+        pass
+
+
+class EncodedSequence(Encoding, Sequence):
+    """Mixin of `Encoding` and `Sequence` to define the approach for
+    encoding biological sequence data.
+    """
+    @property
+    @abc.abstractmethod
+    def ALPHABET_TO_ARRAY(self):
+        """
+        The alphabet used to encode the input sequence.
+        """
+        pass
+
+    def encode(self, s):
+        """Encodes a string with a numpy array.
+
+        Parameters
+        ----------
+        s : str
+            The string to encode.
+
+        Returns
+        -------
+        numpy.ndarray
+            An array with the encoded string.
+        """
+        ret = list()
+        for i in range(len(s)):
+            ret.append(self.ALPHABET_TO_ARRAY[s[i]].copy())
+        ret = numpy.stack(ret)
+        return ret
+
+    def get_sequence_from_coords(self, *args, **kwargs):
+        """Fetches an encoded sequence at the specified coordinates.
+
+        Returns
+        -------
+        numpy.ndarray
+            The numpy array encoding the queried sequence.
+        """
+        return self.encode(
+            super(EncodedSequence, self).get_sequence_from_coords(
+                *args, **kwargs))
+
+
+class EncodedGenome(EncodedSequence, Genome):
+    """This class allows the user to a query a potentially file-backed
+    genome by coordinate. It is essentially a wrapper around the
+    `pyfaidx.Fasta` class. The returned values have been encoded as numpy
+    arrays.
+
+    Parameters
+    ----------
+    input_path : str
+        Path to an indexed FASTA file.
+    in_memory : bool
+        Specifies whether the genome should be loaded from
+        disk and stored in memory.
+
+
+    Attributes
+    ----------
+    data : pyfaidx.Fasta or dict
+        The FASTA file containing the genome sequence. Alternatively,
+        this can be a `dict` object mapping chromosomes to sequences that
+        stores the file in memory.
+    in_memory : bool
+        Specified whether the genomic data is being stored in memory.
+    chrom_len_dict : dict
+        A dictionary mapping the chromosome names to their lengths.
+    ALPHABET_TO_ARRAY : dict
+        A mapping from characters in the genome to their
+        `numpy.ndarray` representations.
+    """
+    ALPHABET_TO_ARRAY = dict(A=numpy.array([1, 0, 0, 0], dtype=_STORE_TYPE),
+                             C=numpy.array([0, 1, 0, 0], dtype=_STORE_TYPE),
+                             G=numpy.array([0, 0, 1, 0], dtype=_STORE_TYPE),
+                             T=numpy.array([0, 0, 0, 1], dtype=_STORE_TYPE),
+                             N=numpy.array([.25, .25, .25, .25], dtype=_STORE_TYPE))
+    """
+    A dictionary mapping possible characters in the genome to
+    their `numpy.ndarray` representations.
+    """
+
+    def __init__(self, *args, **kwargs):
+        """
+        Constructs a new `EncodedGenome` object.
+        """
+        super(EncodedGenome, self).__init__(*args, **kwargs)
+        if self.in_memory is True: # Pre-encode the genome if storing in memory.
+            for k, v in self.data.items():
+                print(k, flush=True)
+                self.data[k] = numpy.zeros((len(v), 4), dtype=_STORE_TYPE)
+                for i in range(self.data[k].shape[0]):
+                    self.data[k][i, :] = self.ALPHABET_TO_ARRAY[v[i]]
+                #self.data[k] = self.encode(self.data[k])#.astype(_STORE_TYPE)
+
+    def get_sequence_from_coords(self, chrom, start, end, strand="+"):
+        """Fetches a string representation of a sequence at
+        the specified coordinates.
+
+
+        Parameters
+        ----------
+        chrom : str
+            Chromosome to query from.
+        start : int
+            First position in queried sequence.
+        end : int
+            One past the last position in the queried sequence.
+        strand : str
+            The strand to sample from.
+
+        Returns
+        -------
+        str
+            The sequence of bases occuring at the queried
+            coordinates.
+
+        Raises
+        ------
+        IndexError
+            If the coordinates are not valid.
+        """
+        if self.in_memory is True:
+            if self.coords_are_valid(chrom, start, end, strand):
+                x = self.data[chrom][start:end].copy()
+                if strand == "-":
+                    x = numpy.flip(numpy.flip(x, 0), 1)
+                return x.astype(numpy.float64)
+            else:
+                s = "Specified coordinates ({} to {} on \"{}\", strand of \"{}\") are invalid!".format(
+                     start, end, chrom, strand)
+                raise IndexError(s)
+        else:
+            return self.encode(
+                super(EncodedSequence, self).get_sequence_from_coords(
+                    chrom=chrom, start=start, end=end, strand=strand)).astype(numpy.float64)
+
+
+class HDF5Genome(Genome):
+    """This class allows the user to query a Genome stored in an HDF5 file.
+
+    Parameters
+    ----------
+    input_path : str
+        Path to an HDF5 file.
+    in_memory : bool
+        Specifies whether the genome should be loaded from
+        disk and stored in memory.
+
+
+    Attributes
+    ----------
+    data : h5py.File
+        The HDF5 file containing the genome sequence. Alternatively,
+        this can be a `dict` object mapping chromosomes to sequences that
+        stores the file in memory.
+    in_memory : bool
+        Specified whether the genomic data is being stored in memory.
+    chrom_len_dict : dict
+        A dictionary mapping the chromosome names to their lengths.
+    """
+    def __init__(self, input_path, in_memory=False):
+        """
+        Constructs a new `HDF5Genome` object.
+        """
+        self.in_memory = in_memory
+        if in_memory is True:
+            f = h5py.File(input_path, "r")
+            self.data = dict()
+            for k in f.keys():
+                self.data[k] = f[k][:].tostring().decode("UTF-8")
+            f.close()
+        else:
+            self.data = h5py.File(input_path, "r")
+        self.chrom_len_dict = {k: len(self.data[k]) for k in self.data.keys()}
+
+    def get_sequence_from_coords(self, chrom, start, end, strand="+"):
+        """Fetches a string representation of a sequence at
+        the specified coordinates.
+
+
+        Parameters
+        ----------
+        chrom : str
+            Chromosome to query from.
+        start : int
+            First position in queried sequence.
+        end : int
+            One past the last position in the queried sequence.
+        strand : str
+            The strand to sample from.
+
+        Returns
+        -------
+        str
+            The sequence of bases occuring at the queried
+            coordinates.
+
+        Raises
+        ------
+        IndexError
+            If the coordinates are not valid.
+        """
+        if self.coords_are_valid(chrom, start, end, strand):
+            x = self.data[chrom][start:end]
+            if not self.in_memory:
+                x = x.tostring().decode("UTF-8")
+            if strand == "-":
+                x = x.translate(_DNA_COMP_TABLE)[::-1]
+            return x
+        else:
+            s = "Specified coordinates ({} to {} on \"{}\", strand of \"{}\") are invalid!".format(
+                start, end, chrom, strand)
+            raise IndexError(s)
+
+
+class EncodedHDF5Genome(EncodedGenome):
+    """
+    This class allows the user to specify an HDF5-backed encoded genome.
+
+    Parameters
+    ----------
+    input_path : str
+        Path to the HDF5 file.
+    in_memory : bool
+        Specifies whether the genome should be loaded from
+        disk and stored in memory.
+
+
+    Attributes
+    ----------
+    data : h5py.File or dict
+        The HDF5 file pointer or a dict containing the sequences in memory.
+    in_memory : bool
+        Specified whether the genomic data is being stored in memory.
+    chrom_len_dict : dict
+        A dictionary mapping the chromosome names to their lengths.
+    """
+    def __init__(self, input_path, in_memory=False):
+        """
+        Constructs a new `EncodedHDF5Genome` object.
+        """
+        self.in_memory = in_memory
+        if in_memory is True:
+            f = h5py.File(input_path, "r")
+            self.data = dict()
+            for k in f.keys():
+                self.data[k] = f[k][()]
+            f.close()
+        else:
+            self.data = h5py.File(input_path, "r")
+        self.chrom_len_dict = {k: self.data[k].shape[0] for k
+                               in self.data.keys()}
+
+    def get_sequence_from_coords(self, chrom, start, end, strand="+"):
+        """Fetches an array representation of a sequence at
+        the specified coordinates.
+
+
+        Parameters
+        ----------
+        chrom : str
+            Chromosome to query from.
+        start : int
+            First position in queried sequence.
+        end : int
+            One past the last position in the queried sequence.
+        strand : str
+            The strand to sample from.
+
+        Returns
+        -------
+        str
+            The sequence of bases occuring at the queried
+            coordinates.
+
+        Raises
+        ------
+        IndexError
+            If the coordinates are not valid.
+        """
+        if self.coords_are_valid(chrom, start, end, strand):
+            if self.in_memory:
+                x = self.data[chrom][start:end].copy()
+            else:
+                x = numpy.zeros((end - start, 4), dtype=_STORE_TYPE)
+                self.data[chrom].read_direct(x, numpy.s_[start:end], numpy.s_[0:x.shape[0]])
+            if strand == "-":
+                x = numpy.flip(numpy.flip(x, 0), 1)
+            return x.astype(numpy.float64)
+        else:
+            s = "Specified coordinates ({} to {} on \"{}\", strand of \"{}\") are invalid!".format(start, end, chrom, strand)
+            raise IndexError(s)
+
+    def close(self):
+        """
+        Close the file connection to HDF5
+        """
+        self.data.close()
+
diff --git a/bin/convert_fasta_to_encoded_h5py.py b/bin/convert_fasta_to_encoded_h5py.py
new file mode 100644
index 0000000..4cd5cad
--- /dev/null
+++ b/bin/convert_fasta_to_encoded_h5py.py
@@ -0,0 +1,58 @@
+"""
+Author: Evan M. Cofer
+Created on: June 29, 2020
+"""
+import argparse
+import os
+
+import h5py
+import numpy
+import pyfaidx
+
+from amber.utils.sequences import EncodedGenome
+
+
+def convert_file(input_file, output_file):
+    g = EncodedGenome(input_file, in_memory=False)
+    h5 = h5py.File(output_file, "w")
+
+    # Convert.
+    for chrom, chrom_len in g.chrom_len_dict.items():
+        ds = h5.create_dataset(chrom, (chrom_len, 4), dtype=g.ALPHABET_TO_ARRAY['N'].dtype)
+        s = str(g.data[chrom]).upper()
+        for i in range(chrom_len):
+            if s[i] == 'A':
+                ds[i, 0] = 1
+            elif s[i] == 'C':
+                ds[i, 1] = 1
+            elif s[i] == 'T':
+                ds[i, 2] = 1
+            elif s[i] == 'G':
+                ds[i, 3] = 1
+            elif s[i] == 'N':
+                ds[i, :] = 0.25
+            else:
+                s = "Failed on chrom {} in position {} with character \"{}\"".format(chrom, i, s[i])
+                raise ValueError(s)
+
+    # Close.
+    h5.close()
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert fasta file to H5PY file with array for genome.")
+    parser.add_argument("--input-file", type=str, required=True, help="Path to FASTA input file.")
+    parser.add_argument("--output-file", type=str, required=True, help="Path to store H5PY file at")
+    args = parser.parse_args()
+
+
+    # Validate arguments.
+    for x in [args.input_file]:
+        if not os.path.exists(x):
+            s = 'Found no input file at {}'.format(x)
+            raise ValueError(s)
+    for x in [args.output_file]:
+        if os.path.exists(x):
+            s = 'Found output file already exists at {}'.format(x)
+            raise ValueError(s)
+    # Run command.
+    convert_file(args.input_file, args.output_file)
diff --git a/bin/convert_fasta_to_h5py.py b/bin/convert_fasta_to_h5py.py
new file mode 100644
index 0000000..56dbcfc
--- /dev/null
+++ b/bin/convert_fasta_to_h5py.py
@@ -0,0 +1,44 @@
+"""
+Author: Evan M. Cofer
+Created on: June 29, 2020
+"""
+import argparse
+import os
+
+import h5py
+import numpy
+import pyfaidx
+
+
+def convert_file(input_file, output_file):
+    fasta = pyfaidx.Fasta(input_file)
+    h5 = h5py.File(output_file, "w")
+
+    # Convert.
+    for k in fasta.keys():
+        s = str(fasta[k][:].seq).upper()
+        ds = h5.create_dataset(k, (len(s),), dtype='S1')
+        for i in range(len(s)):
+            ds[i] = numpy.string_(s[i])
+
+    # Close.
+    h5.close()
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert fasta file to H5PY file with array for genome.")
+    parser.add_argument("--input-file", type=str, required=True, help="Path to FASTA input file.")
+    parser.add_argument("--output-file", type=str, required=True, help="Path to store H5PY file at")
+    args = parser.parse_args()
+
+
+    # Validate arguments.
+    for x in [args.input_file]:
+        if not os.path.exists(x):
+            s = 'Found no input file at {}'.format(x)
+            raise ValueError(s)
+    for x in [args.output_file]:
+        if os.path.exists(x):
+            s = 'Found output file already exists at {}'.format(x)
+            raise ValueError(s)
+    # Run command.
+    convert_file(args.input_file, args.output_file)
diff --git a/bin/draw_samples.py b/bin/draw_samples.py
new file mode 100644
index 0000000..e274fb2
--- /dev/null
+++ b/bin/draw_samples.py
@@ -0,0 +1,233 @@
+"""
+Evan Cofer, 2020
+"""
+import argparse
+import collections
+import os
+import re
+import sys
+
+import bx.intervals.intersection
+import gzip
+import numpy
+import pyfaidx
+
+def draw_samples(genome_file, bed_file, output_file, feature_name_file,
+                 bin_size, cvg_frac, n_examples,
+                 chrom_pad, chrom_pattern, max_unk, interval_file):
+    # Read feature names from file.
+    feature_name_set = set()
+    i_to_feature_name = list()
+    feature_name_to_i = dict()
+    i = 0
+    with open(feature_name_file, "r") as read_file:
+        for line in read_file:
+            line = line.strip()
+            if line:
+                feature_name_set.add(line)
+                i_to_feature_name.append(line)
+                feature_name_to_i[line] = i
+                i += 1
+    n_feats = len(i_to_feature_name)
+    print("Loaded features", file=sys.stderr)
+
+    # Load genome and get estimate of chromosome weights.
+    genome = pyfaidx.Fasta(genome_file)
+    chroms = list()
+    chrom_lens = list()
+    max_examples = 0
+    for k in genome.keys():
+        if chrom_pattern.match(k) is not None:
+            l = len(genome[k])
+            if l > chrom_pad * 3:
+                for s in ["+", "-"]: # Different mass on strands.
+                    chroms.append((s, k))
+                    chrom_lens.append(l - 2 * chrom_pad - bin_size)
+    n_chrom = len(chroms)
+    chrom_to_i = {k: i for (i, k) in enumerate(chroms)}
+    print("Loaded chroms", file=sys.stderr)
+
+    # Get intervals.
+    chrom_bound_ivals = {k: list() for k in chrom_to_i.values()}
+    chrom_bound_ival_weights = {k: list() for k in chrom_to_i.values()}
+    chrom_lens = numpy.array(chrom_lens)
+    if interval_file is None:
+        chrom_weighting_lens = chrom_lens.copy()
+        for x, chrom_len in zip(chroms, chrom_lens.tolist()):
+            c_i = chrom_to_i[x]
+            chrom_bound_ivals[c_i].append((chrom_pad, chrom_len + chrom_pad))
+            chrom_bound_ival_weights[c_i].append(chrom_len)
+    else:
+        chrom_weighting_lens = numpy.zeros_like(chrom_lens)
+        with open(interval_file, "r") as read_file:
+            for line_i, line in enumerate(read_file):
+                line = line.strip()
+                if line:
+                    if not line.startswith("#"):
+                        line = line.split("\t")
+                        if len(line) != 6:
+                            s = "Found that line #{} has {} elements and not 6".format(line_i, len(line))
+                            raise ValueError(s)
+                        chrom, start, end, _, _, strand = line
+                        start = int(start)
+                        end = int(end)
+                        start, end = min(start, end - 1), max(start + 1, end)
+                        if (strand, chrom) in chrom_to_i:
+                            i = chrom_to_i[(strand, chrom)]
+                            start = max(start, chrom_pad)
+                            end = min(end, chrom_lens[i] - chrom_pad - bin_size)
+                            if end - start <= 0:
+                                continue
+                            chrom_bound_ivals[i].append((start, end))
+                            dist = abs(end - start)
+                            chrom_weighting_lens[i] += dist
+                            chrom_bound_ival_weights[i].append(dist)
+    print("Loaded chrom ivals", file=sys.stderr)
+
+    # Calculate weighting and number of examples.
+    chrom_weights = chrom_weighting_lens / chrom_weighting_lens.sum()
+    max_examples = numpy.sum(chrom_weighting_lens)
+    if max_examples < n_examples:
+        msg = "Got {} max examples possible, but need {} examples".format(
+            max_examples, n_examples)
+        raise ValueError(msg)
+
+    # Create interval tree for fast label query.
+    ivt = {k : {kk: bx.intervals.intersection.IntervalTree() for kk in feature_name_to_i.values()} for k in chrom_to_i.values()}
+    with gzip.open(bed_file, "rt") as read_file:
+        for line in read_file:
+            line = line.strip()
+            if not line.startswith("#"):
+                line = line.split("\t")
+                chrom, start, end, name = line[:4]
+                if len(line) >= 6:
+                    if line[5] == ".":
+                        strand = ["+", "-"]
+                    else:
+                        strand = [line[5]]
+                else:
+                    strand = ["+", "-"]
+                start = int(start)
+                end = int(end)
+                if name in feature_name_set:
+                    for x in strand:
+                        i = chrom_to_i[(x, chrom)]
+                        ivt[i][feature_name_to_i[name]].insert_interval(
+                            bx.intervals.intersection.Interval(start, end))
+    print("Loaded labels", file=sys.stderr)
+
+    # Create outputs.
+    outputs = list()
+    i = 0
+    while i < n_examples:
+        c_i = numpy.random.choice(n_chrom, p=chrom_weights)
+        strand, chrom = chroms[c_i]
+        try:
+            ival_bin_i = numpy.random.choice(len(chrom_bound_ivals[c_i]),
+                                         p=numpy.array(chrom_bound_ival_weights[c_i]).flatten() / chrom_weighting_lens[c_i])
+        except:
+            print("***FAILED ON RANDOM CHOICE",
+                  c_i,
+                  chrom_bound_ivals[c_i],
+                  chrom_weights,
+                  chrom_weighting_lens, sep="\n", flush=True, file=sys.stderr)
+            raise RuntimeError()
+        start, end = chrom_bound_ivals[c_i].pop(ival_bin_i)
+        cur_weight = chrom_bound_ival_weights[c_i].pop(ival_bin_i) # Remove weight.
+        try:
+            pos = numpy.random.choice(end - start) + start
+        except:
+            print("***FAILED ON RANDOM CHOICE",
+                  c_i,
+                  chrom_bound_ivals[c_i],
+                  chrom_weights,
+                  chrom_weighting_lens, sep="\n", flush=True, file=sys.stderr)
+            raise RuntimeError()
+        if end - start > 1:
+            if pos == start:
+                chrom_bound_ival_weights[c_i].append(cur_weight - 1)
+                chrom_bound_ivals[c_i].append((start + 1, end))
+            elif pos == end - 1:
+                chrom_bound_ival_weights[c_i].append(cur_weight - 1)
+                chrom_bound_ivals[c_i].append((start, end - 1))
+            else:
+                chrom_bound_ivals[c_i].append((start, pos))
+                chrom_bound_ival_weights[c_i].append(pos - start)
+                chrom_bound_ivals[c_i].append((pos + 1, end))
+                chrom_bound_ival_weights[c_i].append(end - (pos + 1))
+        start = pos
+        end = pos + bin_size
+
+        # Adjust chrom weights.
+        chrom_weighting_lens[c_i] -= 1
+        chrom_weights = chrom_weighting_lens / numpy.sum(chrom_weighting_lens)
+
+        # Determine label etc w/ ivt.
+        cvg = numpy.zeros(n_feats)
+        for feat_i in feature_name_to_i.values():
+            for x in ivt[c_i][feat_i].find(start, end):
+                cvg[feat_i] += min(x.end, end) - max(x.start, start)
+        cvg /= bin_size
+        cvg = (cvg > cvg_frac).astype(int).tolist()
+        outputs.append((chrom, start, end, strand, *cvg))
+        i += 1
+        if i % 10000 == 0:
+            print(i)
+    print("Loaded all examples.", file=sys.stderr)
+
+    # write outputs to file.
+    with open(output_file, "w") as write_file:
+        for x in sorted(outputs):
+            x = [str(y) for y in list(x)]
+            write_file.write("\t".join(x) + "\n")
+    print("Wrote examples.", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    # Get command line arguments.
+    parser = argparse.ArgumentParser( description="sampling data for tf genomics models")
+    parser.add_argument("--output", type=str, required=True, help="Path to the output file.")
+    parser.add_argument("--genome", type=str, required=True, help="Path to the indexed fasta file.")
+    parser.add_argument("--bed", type=str, required=True, help="Path to the input bed file.")
+    parser.add_argument("--feature-name-file", type=str, required=True, help="Name of feature in bed file to look for.")
+    parser.add_argument("--bin-size", type=int, required=True, help="Size of the bin composing minimal examples.")
+    parser.add_argument("--cvg-frac", type=float, required=True, help="Fraction of bin that must be covered to be positive example.")
+    parser.add_argument("--n-examples", type=int, required=True, help="Number of examples to include.")
+    parser.add_argument("--chrom-pad", type=int, required=True, help="Length of region to ignore at the start and end of chromosomes.")
+    parser.add_argument("--seed", type=int, required=True, help="Seed for RNG.")
+    parser.add_argument("--include-chroms", type=str, required=True, help="Regex for chromosomes to include.")
+    parser.add_argument("--max-n", type=int, required=True, help="Maximum N chars in sequences.")
+    parser.add_argument("--interval-file", type=str, required=False, default=None, help="Path to file with intervals to draw positives from")
+    args = parser.parse_args()
+
+    # Validate arguments.
+    for x in [args.genome, args.bed, args.feature_name_file] + ([] if args.interval_file is None else [args.interval_file]):
+        if not os.path.exists(x):
+            raise ValueError(x + " does not exist")
+
+    if args.n_examples <= 0:
+        raise ValueError("--n-examples must be > 0")
+
+    if args.chrom_pad < 0:
+        raise ValueError("--chrom-pad must be >= 0")
+
+    if args.bin_size <= 0:
+        raise ValueError("--bin-size must be > 0")
+
+    if args.cvg_frac < 0:
+        raise ValueError("--cvg-frac must be >= 0")
+    elif args.cvg_frac > 1:
+        raise ValueError("--cvg-frac must be < 1")
+
+    if args.max_n < 0:
+        raise ValueError("--max-n must be >= 0")
+
+
+    # Prepare for function.
+    numpy.random.seed(args.seed)
+    pattern = re.compile("^" + args.include_chroms + "$")
+
+    # Run function.
+    draw_samples(args.genome, args.bed, args.output, args.feature_name_file,
+                 args.bin_size, args.cvg_frac, args.n_examples,
+                 args.chrom_pad, pattern, args.max_n, args.interval_file)
diff --git a/bin/split_samples.py b/bin/split_samples.py
new file mode 100644
index 0000000..c5b93d4
--- /dev/null
+++ b/bin/split_samples.py
@@ -0,0 +1,60 @@
+"""
+Evan Cofer June 29, 2020
+"""
+import argparse
+import os
+
+
+def split_samples(input_file, feature_name_file, output_dir):
+    # Read in feature names.
+    i_to_feat = list()
+    with open(feature_name_file, "r") as read_file:
+        for line in read_file:
+            line = line.strip()
+            if line:
+                i_to_feat.append(line)
+    feat_to_i = dict()
+    for i, feat in enumerate(i_to_feat):
+        feat_to_i[feat] = i
+
+    # Create file handles for features.
+    i_to_fp = list()
+    for feat in i_to_feat:
+        f = "{}.txt".format(feat)
+        f = os.path.join(output_dir, f)
+        i_to_fp.append(open(f, "w"))
+
+    # Read in data and write out as we go.
+    with open(input_file, "r") as read_file:
+        for line in read_file:
+            line = line.strip()
+            if line:
+                line = line.split("\t")
+                stub = "\t".join(line[:4]) + "\t"
+                for i, fp in enumerate(i_to_fp):
+                    fp.write(stub + line[4 + i] + "\n")
+
+    # Close files.
+    for x in i_to_fp:
+        x.close()
+
+if __name__ == "__main__":
+    # Get command line arguments.
+    parser = argparse.ArgumentParser( description="split input training examples for several models")
+    parser.add_argument("--input-file", type=str, required=True, help="Path to input samples to split.")
+    parser.add_argument("--feature-name-file", type=str, required=True, help="Path to file with feature names listed in order.")
+    parser.add_argument("--output-dir", type=str, required=True, help="Path to directory to write outputs to.")
+    args = parser.parse_args()
+
+    # Validate arguments.
+    for x in [args.input_file, args.feature_name_file]:
+        if not os.path.exists(x):
+            s = "Found no file at {}".format(x)
+            raise ValueError(s)
+
+    if not os.path.isdir(args.output_dir):
+        s = "Found no directory at {}".format(args.output_dir)
+        raise ValueError(s)
+
+    # Run main function.
+    split_samples(args.input_file, args.feature_name_file, args.output_dir)
diff --git a/.DS_Store b/examples/.DS_Store
similarity index 73%
rename from .DS_Store
rename to examples/.DS_Store
index edda34d..f732f76 100644
Binary files a/.DS_Store and b/examples/.DS_Store differ
diff --git a/templates/AmberDeepSea.py b/examples/amber_deepsea.py
similarity index 97%
rename from templates/AmberDeepSea.py
rename to examples/amber_deepsea.py
index 38d8b68..e3f5ae8 100644
--- a/templates/AmberDeepSea.py
+++ b/examples/amber_deepsea.py
@@ -96,7 +96,7 @@ def get_model_space(out_filters=64, num_layers=9):
     'manager': {
         'data': {
             'train_data': read_train_data('/mnt/ceph/users/zzhang/DeepSEA/data/deepsea_train/train.mat'),
-            'validation_data': read_val_data('./mnt/ceph/users/zzhang/DeepSEA/data/deepsea_train/valid.mat')
+            'validation_data': read_val_data('/mnt/ceph/users/zzhang/DeepSEA/data/deepsea_train/valid.mat')
         },
         'params': {
             'epochs': 1,
diff --git a/examples/amber_env.yml b/examples/amber_env.yml
new file mode 100644
index 0000000..5ac2f74
--- /dev/null
+++ b/examples/amber_env.yml
@@ -0,0 +1,474 @@
+name: amber
+channels:
+  - bioconda
+  - defaults
+  - conda-forge
+dependencies:
+  - _anaconda_depends=2019.03=py37_0
+  - _libgcc_mutex=0.1=main
+  - _r-mutex=1.0.0=anacondar_1
+  - _tflow_select=2.1.0=gpu
+  - absl-py=0.7.1=py37_0
+  - alabaster=0.7.12=py37_0
+  - anaconda=custom=py37_1
+  - anaconda-client=1.7.2=py37_0
+  - anaconda-project=0.8.3=py_0
+  - asn1crypto=0.24.0=py37_0
+  - astor=0.8.0=py37_0
+  - astroid=2.2.5=py37_0
+  - astropy=3.2.1=py37h7b6447c_0
+  - atomicwrites=1.3.0=py37_1
+  - attrs=19.1.0=py37_1
+  - babel=2.7.0=py_0
+  - backcall=0.1.0=py37_0
+  - backports=1.0=py_2
+  - backports.os=0.1.1=py37_0
+  - backports.shutil_get_terminal_size=1.0.0=py37_2
+  - beautifulsoup4=4.8.0=py37_0
+  - binutils_impl_linux-64=2.31.1=h6176602_1
+  - binutils_linux-64=2.31.1=h6176602_8
+  - bitarray=1.0.1=py37h7b6447c_0
+  - bkcharts=0.2=py37_0
+  - blas=1.0=mkl
+  - bleach=3.1.0=py37_0
+  - blosc=1.16.3=hd408876_0
+  - bokeh=1.3.4=py37_0
+  - boto=2.49.0=py37_0
+  - bottleneck=1.2.1=py37h035aef0_1
+  - bwidget=1.9.11=1
+  - bzip2=1.0.8=h7b6447c_0
+  - c-ares=1.15.0=h7b6447c_1001
+  - ca-certificates=2020.1.1=0
+  - cairo=1.14.12=h8948797_3
+  - certifi=2020.4.5.1=py37_0
+  - cffi=1.12.3=py37h2e261b9_0
+  - chardet=3.0.4=py37_1003
+  - click=7.0=py37_0
+  - cloudpickle=1.2.1=py_0
+  - clyent=1.2.2=py37_1
+  - colorama=0.4.1=py37_0
+  - contextlib2=0.5.5=py37_0
+  - cryptography=2.7=py37h1ba5d50_0
+  - cudatoolkit=10.1.168=0
+  - cudnn=7.6.0=cuda10.1_0
+  - cupti=10.1.168=0
+  - curl=7.65.3=hbc83047_0
+  - cycler=0.10.0=py37_0
+  - cython=0.29.13=py37he6710b0_0
+  - cytoolz=0.10.0=py37h7b6447c_0
+  - dask=2.3.0=py_0
+  - dask-core=2.3.0=py_0
+  - dbus=1.13.6=h746ee38_0
+  - decorator=4.4.0=py37_1
+  - defusedxml=0.6.0=py_0
+  - distributed=2.3.2=py_1
+  - docutils=0.15.2=py37_0
+  - entrypoints=0.3=py37_0
+  - et_xmlfile=1.0.1=py37_0
+  - expat=2.2.6=he6710b0_0
+  - fastcache=1.1.0=py37h7b6447c_0
+  - filelock=3.0.12=py_0
+  - flask=1.1.1=py_0
+  - fontconfig=2.13.0=h9420a91_0
+  - freetype=2.9.1=h8a8886c_1
+  - fribidi=1.0.5=h7b6447c_0
+  - fsspec=0.4.0=py_0
+  - gast=0.2.2=py37_0
+  - gcc_impl_linux-64=7.3.0=habb00fd_1
+  - gcc_linux-64=7.3.0=h553295d_8
+  - get_terminal_size=1.0.0=haa9412d_0
+  - gevent=1.4.0=py37h7b6447c_0
+  - gfortran_impl_linux-64=7.3.0=hdf63c60_1
+  - gfortran_linux-64=7.3.0=h553295d_8
+  - glib=2.56.2=hd408876_0
+  - glob2=0.7=py_0
+  - gmp=6.1.2=h6c8ec71_1
+  - gmpy2=2.0.8=py37h10f8cd9_2
+  - google-pasta=0.1.7=py_0
+  - graphite2=1.3.13=h23475e2_0
+  - graphviz=2.40.1=h21bd128_2
+  - greenlet=0.4.15=py37h7b6447c_0
+  - grpcio=1.16.1=py37hf8bcb03_1
+  - gsl=2.4=h14c3975_4
+  - gst-plugins-base=1.14.0=hbbd80ab_1
+  - gstreamer=1.14.0=hb453b48_1
+  - gxx_impl_linux-64=7.3.0=hdf63c60_1
+  - gxx_linux-64=7.3.0=h553295d_8
+  - h5py=2.9.0=py37h7918eee_0
+  - harfbuzz=1.8.8=hffaf4a1_0
+  - hdf5=1.10.4=hb1b8bf9_0
+  - heapdict=1.0.0=py37_2
+  - html5lib=1.0.1=py37_0
+  - icu=58.2=h9c2bf20_1
+  - idna=2.8=py37_0
+  - imageio=2.5.0=py37_0
+  - imagesize=1.1.0=py37_0
+  - importlib_metadata=0.19=py37_0
+  - intel-openmp=2019.4=243
+  - ipykernel=5.1.2=py37h39e3cac_0
+  - ipython=7.8.0=py37h39e3cac_0
+  - ipython_genutils=0.2.0=py37_0
+  - ipywidgets=7.5.1=py_0
+  - isort=4.3.21=py37_0
+  - itsdangerous=1.1.0=py37_0
+  - jbig=2.1=hdba287a_0
+  - jdcal=1.4.1=py_0
+  - jedi=0.15.1=py37_0
+  - jeepney=0.4=py37_0
+  - jinja2=2.10.1=py37_0
+  - joblib=0.13.2=py37_0
+  - jpeg=9b=h024ee3a_2
+  - json5=0.8.5=py_0
+  - jsonschema=3.0.2=py37_0
+  - jupyter=1.0.0=py37_7
+  - jupyter_client=5.3.1=py_0
+  - jupyter_console=6.0.0=py37_0
+  - jupyter_core=4.5.0=py_0
+  - jupyterlab=1.0.2=py37hf63ae98_0
+  - jupyterlab_server=1.0.0=py_1
+  - keras-applications=1.0.8=py_0
+  - keras-preprocessing=1.1.0=py_1
+  - keyring=18.0.0=py37_0
+  - kiwisolver=1.1.0=py37he6710b0_0
+  - krb5=1.16.1=h173b8e3_7
+  - lazy-object-proxy=1.4.2=py37h7b6447c_0
+  - libarchive=3.3.3=h5d8350f_5
+  - libcurl=7.65.3=h20c2e04_0
+  - libdeflate=1.0=h14c3975_1
+  - libedit=3.1.20181209=hc058e9b_0
+  - libffi=3.2.1=hd88cf55_4
+  - libgcc-ng=9.1.0=hdf63c60_0
+  - libgfortran-ng=7.3.0=hdf63c60_0
+  - liblief=0.9.0=h7725739_2
+  - libpng=1.6.37=hbc83047_0
+  - libprotobuf=3.8.0=hd408876_0
+  - libsodium=1.0.16=h1bed415_0
+  - libssh2=1.8.2=h1ba5d50_0
+  - libstdcxx-ng=9.1.0=hdf63c60_0
+  - libtiff=4.0.10=h2733197_2
+  - libtool=2.4.6=h7b6447c_5
+  - libuuid=1.0.3=h1bed415_2
+  - libxcb=1.13=h1bed415_1
+  - libxml2=2.9.9=hea5a465_1
+  - libxslt=1.1.33=h7d1a2b0_0
+  - llvmlite=0.29.0=py37hd408876_0
+  - locket=0.2.0=py37_1
+  - lxml=4.4.1=py37hefd8a0e_0
+  - lz4-c=1.8.1.2=h14c3975_0
+  - lzo=2.10=h49e0be7_2
+  - make=4.2.1=h1bed415_1
+  - markdown=3.1.1=py37_0
+  - markupsafe=1.1.1=py37h7b6447c_0
+  - matplotlib=3.1.3=py37_0
+  - matplotlib-base=3.1.3=py37hef1b27d_0
+  - mccabe=0.6.1=py37_1
+  - mistune=0.8.4=py37h7b6447c_0
+  - mkl=2019.4=243
+  - mkl-service=2.3.0=py37he904b0f_0
+  - mkl_fft=1.0.14=py37ha843d7b_0
+  - mkl_random=1.0.2=py37hd81dba3_0
+  - mock=3.0.5=py37_0
+  - more-itertools=7.2.0=py37_0
+  - mpc=1.1.0=h10f8cd9_1
+  - mpfr=4.0.1=hdf1c602_3
+  - mpmath=1.1.0=py37_0
+  - msgpack-python=0.6.1=py37hfd86e86_1
+  - multipledispatch=0.6.0=py37_0
+  - nbconvert=5.5.0=py_0
+  - nbformat=4.4.0=py37_0
+  - ncurses=6.1=he6710b0_1
+  - networkx=2.3=py_0
+  - nltk=3.4.5=py37_0
+  - nose=1.3.7=py37_2
+  - notebook=6.0.0=py37_0
+  - numba=0.45.1=py37h962f231_0
+  - numexpr=2.7.0=py37h9e4a6bb_0
+  - numpy=1.16.4=py37h7e9f1db_0
+  - numpy-base=1.16.4=py37hde5b4d6_0
+  - numpydoc=0.9.1=py_0
+  - olefile=0.46=py37_0
+  - openjdk=8.0.152=h7b6447c_3
+  - openpyxl=2.6.2=py_0
+  - openssl=1.1.1g=h7b6447c_0
+  - packaging=19.1=py37_0
+  - pandas=0.25.1=py37he6710b0_0
+  - pandoc=2.2.3.2=0
+  - pandocfilters=1.4.2=py37_1
+  - pango=1.42.4=h049681c_0
+  - parso=0.5.1=py_0
+  - partd=1.0.0=py_0
+  - patchelf=0.9=he6710b0_3
+  - path.py=12.0.1=py_0
+  - pathlib2=2.3.4=py37_0
+  - patsy=0.5.1=py37_0
+  - pcre=8.43=he6710b0_0
+  - pep8=1.7.1=py37_0
+  - pexpect=4.7.0=py37_0
+  - picard=2.22.2=0
+  - pickleshare=0.7.5=py37_0
+  - pillow=6.1.0=py37h34e0f95_0
+  - pip=19.2.2=py37_0
+  - pixman=0.38.0=h7b6447c_0
+  - pkginfo=1.5.0.1=py37_0
+  - pluggy=0.12.0=py_0
+  - ply=3.11=py37_0
+  - prometheus_client=0.7.1=py_0
+  - prompt_toolkit=2.0.9=py37_0
+  - protobuf=3.8.0=py37he6710b0_0
+  - psutil=5.6.3=py37h7b6447c_0
+  - ptyprocess=0.6.0=py37_0
+  - py=1.8.0=py37_0
+  - py-lief=0.9.0=py37h7725739_2
+  - pycodestyle=2.5.0=py37_0
+  - pycosat=0.6.3=py37h14c3975_0
+  - pycparser=2.19=py37_0
+  - pycrypto=2.6.1=py37h14c3975_9
+  - pycurl=7.43.0.3=py37h1ba5d50_0
+  - pydot=1.4.1=py37_0
+  - pyfaidx=0.5.8=py_1
+  - pyflakes=2.1.1=py37_0
+  - pygments=2.4.2=py_0
+  - pylint=2.3.1=py37_0
+  - pyodbc=4.0.27=py37he6710b0_0
+  - pyopenssl=19.0.0=py37_0
+  - pyparsing=2.4.2=py_0
+  - pyqt=5.9.2=py37h05f1152_2
+  - pyrsistent=0.14.11=py37h7b6447c_0
+  - pysam=0.15.3=py37hda2845c_1
+  - pysocks=1.7.0=py37_0
+  - pytables=3.5.2=py37h71ec239_1
+  - pytest=5.0.1=py37_0
+  - pytest-arraydiff=0.3=py37h39e3cac_0
+  - pytest-astropy=0.5.0=py37_0
+  - pytest-doctestplus=0.3.0=py37_0
+  - pytest-openfiles=0.3.2=py37_0
+  - pytest-remotedata=0.3.2=py37_0
+  - python=3.7.4=h265db76_1
+  - python-dateutil=2.8.0=py37_0
+  - python-libarchive-c=2.8=py37_13
+  - pytz=2019.2=py_0
+  - pywavelets=1.0.3=py37hdd07704_1
+  - pyyaml=5.1.2=py37h7b6447c_0
+  - pyzmq=17.1.2=py37h14c3975_0
+  - qt=5.9.7=h5867ecd_1
+  - qtawesome=0.5.7=py37_1
+  - qtconsole=4.5.4=py_0
+  - qtpy=1.9.0=py_0
+  - r-askpass=1.0=r36h14c3975_0
+  - r-assertthat=0.2.1=r36h6115d3f_0
+  - r-backports=1.1.4=r36h96ca727_0
+  - r-base=3.6.1=hce969dd_0
+  - r-base64enc=0.1_3=r36h96ca727_4
+  - r-bh=1.69.0_1=r36h6115d3f_0
+  - r-boot=1.3_20=r36h6115d3f_0
+  - r-broom=0.5.2=r36h6115d3f_0
+  - r-callr=3.2.0=r36h6115d3f_0
+  - r-caret=6.0_83=r36h96ca727_0
+  - r-cellranger=1.1.0=r36h6115d3f_0
+  - r-class=7.3_15=r36h96ca727_0
+  - r-cli=1.1.0=r36h6115d3f_0
+  - r-clipr=0.6.0=r36h6115d3f_0
+  - r-cluster=2.0.8=r36ha65eedd_0
+  - r-codetools=0.2_16=r36h6115d3f_0
+  - r-colorspace=1.4_1=r36h96ca727_0
+  - r-crayon=1.3.4=r36h6115d3f_0
+  - r-curl=3.3=r36h96ca727_0
+  - r-data.table=1.12.2=r36h96ca727_0
+  - r-dbi=1.0.0=r36h6115d3f_0
+  - r-dbplyr=1.4.0=r36h6115d3f_0
+  - r-dichromat=2.0_0=r36h6115d3f_4
+  - r-digest=0.6.18=r36h96ca727_0
+  - r-dplyr=0.8.0.1=r36h29659fb_0
+  - r-ellipsis=0.1.0=r36h96ca727_0
+  - r-essentials=3.6.0=r36_0
+  - r-evaluate=0.13=r36h6115d3f_0
+  - r-fansi=0.4.0=r36h96ca727_0
+  - r-forcats=0.4.0=r36h6115d3f_0
+  - r-foreach=1.4.4=r36h6115d3f_0
+  - r-foreign=0.8_71=r36h96ca727_0
+  - r-formatr=1.6=r36h6115d3f_0
+  - r-fs=1.2.7=r36h29659fb_0
+  - r-generics=0.0.2=r36h6115d3f_0
+  - r-ggplot2=3.1.1=r36h6115d3f_0
+  - r-glmnet=2.0_16=r36ha65eedd_0
+  - r-glue=1.3.1=r36h96ca727_0
+  - r-gower=0.2.0=r36h96ca727_0
+  - r-gtable=0.3.0=r36h6115d3f_0
+  - r-haven=2.1.0=r36h29659fb_0
+  - r-hexbin=1.27.2=r36ha65eedd_0
+  - r-highr=0.8=r36h6115d3f_0
+  - r-hms=0.4.2=r36h6115d3f_0
+  - r-htmltools=0.3.6=r36h29659fb_0
+  - r-htmlwidgets=1.3=r36h6115d3f_0
+  - r-httpuv=1.5.1=r36h29659fb_0
+  - r-httr=1.4.0=r36h6115d3f_0
+  - r-ipred=0.9_8=r36h96ca727_0
+  - r-irdisplay=0.7.0=r36h6115d3f_0
+  - r-irkernel=0.8.15=r36_0
+  - r-iterators=1.0.10=r36h6115d3f_0
+  - r-jsonlite=1.6=r36h96ca727_0
+  - r-kernsmooth=2.23_15=r36ha65eedd_4
+  - r-knitr=1.22=r36h6115d3f_0
+  - r-labeling=0.3=r36h6115d3f_4
+  - r-later=0.8.0=r36h29659fb_0
+  - r-lattice=0.20_38=r36h96ca727_0
+  - r-lava=1.6.5=r36h6115d3f_0
+  - r-lazyeval=0.2.2=r36h96ca727_0
+  - r-lubridate=1.7.4=r36h29659fb_0
+  - r-magrittr=1.5=r36h6115d3f_4
+  - r-maps=3.3.0=r36h96ca727_0
+  - r-markdown=0.9=r36h96ca727_0
+  - r-mass=7.3_51.3=r36h96ca727_0
+  - r-matrix=1.2_17=r36h96ca727_0
+  - r-mgcv=1.8_28=r36h96ca727_0
+  - r-mime=0.6=r36h96ca727_0
+  - r-modelmetrics=1.2.2=r36h29659fb_0
+  - r-modelr=0.1.4=r36h6115d3f_0
+  - r-munsell=0.5.0=r36h6115d3f_0
+  - r-nlme=3.1_139=r36ha65eedd_0
+  - r-nnet=7.3_12=r36h96ca727_0
+  - r-numderiv=2016.8_1=r36h6115d3f_0
+  - r-openssl=1.3=r36h96ca727_0
+  - r-pbdzmq=0.3_3=r36h29659fb_0
+  - r-pillar=1.3.1=r36h6115d3f_0
+  - r-pkgconfig=2.0.2=r36h6115d3f_0
+  - r-plogr=0.2.0=r36h6115d3f_0
+  - r-plyr=1.8.4=r36h29659fb_0
+  - r-prettyunits=1.0.2=r36h6115d3f_0
+  - r-processx=3.3.0=r36h96ca727_0
+  - r-prodlim=2018.04.18=r36h29659fb_0
+  - r-progress=1.2.0=r36h6115d3f_0
+  - r-promises=1.0.1=r36h29659fb_0
+  - r-ps=1.3.0=r36h96ca727_0
+  - r-purrr=0.3.2=r36h96ca727_0
+  - r-quantmod=0.4_14=r36h6115d3f_0
+  - r-r6=2.4.0=r36h6115d3f_0
+  - r-randomforest=4.6_14=r36ha65eedd_0
+  - r-rbokeh=0.6.3=r36_0
+  - r-rcolorbrewer=1.1_2=r36h6115d3f_0
+  - r-rcpp=1.0.1=r36h29659fb_0
+  - r-rcpproll=0.3.0=r36h29659fb_0
+  - r-readr=1.3.1=r36h29659fb_0
+  - r-readxl=1.3.1=r36h29659fb_0
+  - r-recipes=0.1.5=r36h6115d3f_0
+  - r-recommended=3.6.0=r36_0
+  - r-rematch=1.0.1=r36h6115d3f_0
+  - r-repr=0.19.2=r36h6115d3f_0
+  - r-reprex=0.2.1=r36h6115d3f_0
+  - r-reshape2=1.4.3=r36h29659fb_0
+  - r-rlang=0.3.4=r36h96ca727_0
+  - r-rmarkdown=1.12=r36h6115d3f_0
+  - r-rpart=4.1_15=r36h96ca727_0
+  - r-rstudioapi=0.10=r36h6115d3f_0
+  - r-rvest=0.3.3=r36h6115d3f_0
+  - r-scales=1.0.0=r36h29659fb_0
+  - r-selectr=0.4_1=r36h6115d3f_0
+  - r-shiny=1.3.2=r36h6115d3f_0
+  - r-sourcetools=0.1.7=r36h29659fb_0
+  - r-spatial=7.3_11=r36h96ca727_4
+  - r-squarem=2017.10_1=r36h6115d3f_0
+  - r-stringi=1.4.3=r36h29659fb_0
+  - r-stringr=1.4.0=r36h6115d3f_0
+  - r-survival=2.44_1.1=r36h96ca727_0
+  - r-sys=3.2=r36h96ca727_0
+  - r-tibble=2.1.1=r36h96ca727_0
+  - r-tidyr=0.8.3=r36h29659fb_0
+  - r-tidyselect=0.2.5=r36h29659fb_0
+  - r-tidyverse=1.2.1=r36h6115d3f_0
+  - r-timedate=3043.102=r36h6115d3f_0
+  - r-tinytex=0.12=r36h6115d3f_0
+  - r-ttr=0.23_4=r36ha65eedd_0
+  - r-utf8=1.1.4=r36h96ca727_0
+  - r-uuid=0.1_2=r36h96ca727_4
+  - r-viridislite=0.3.0=r36h6115d3f_0
+  - r-whisker=0.3_2=r36h6115d3f_4
+  - r-withr=2.1.2=r36h6115d3f_0
+  - r-xfun=0.6=r36h6115d3f_0
+  - r-xml2=1.2.0=r36h29659fb_0
+  - r-xtable=1.8_4=r36h6115d3f_0
+  - r-xts=0.11_2=r36h96ca727_0
+  - r-yaml=2.2.0=r36h96ca727_0
+  - r-zoo=1.8_6=r36h96ca727_0
+  - readline=7.0=h7b6447c_5
+  - requests=2.22.0=py37_0
+  - rope=0.14.0=py_0
+  - ruamel_yaml=0.15.46=py37h14c3975_0
+  - scikit-image=0.15.0=py37he6710b0_0
+  - scikit-learn=0.21.2=py37hd81dba3_0
+  - scipy=1.4.1=py37h0b6359f_0
+  - seaborn=0.9.0=py37_0
+  - secretstorage=3.1.1=py37_0
+  - send2trash=1.5.0=py37_0
+  - setuptools=41.0.1=py37_0
+  - simplegeneric=0.8.1=py37_2
+  - singledispatch=3.4.0.3=py37_0
+  - sip=4.19.8=py37hf484d3e_0
+  - six=1.12.0=py37_0
+  - snappy=1.1.7=hbae5bb6_3
+  - snowballstemmer=1.9.0=py_0
+  - sortedcollections=1.1.2=py37_0
+  - sortedcontainers=2.1.0=py37_0
+  - soupsieve=1.9.2=py37_0
+  - sphinx=2.1.2=py_0
+  - sphinxcontrib=1.0=py37_1
+  - sphinxcontrib-applehelp=1.0.1=py_0
+  - sphinxcontrib-devhelp=1.0.1=py_0
+  - sphinxcontrib-htmlhelp=1.0.2=py_0
+  - sphinxcontrib-jsmath=1.0.1=py_0
+  - sphinxcontrib-qthelp=1.0.2=py_0
+  - sphinxcontrib-serializinghtml=1.1.3=py_0
+  - sphinxcontrib-websupport=1.1.2=py_0
+  - spyder=3.3.6=py37_0
+  - spyder-kernels=0.5.1=py37_0
+  - sqlalchemy=1.3.7=py37h7b6447c_0
+  - sqlite=3.29.0=h7b6447c_0
+  - statsmodels=0.10.1=py37hdd07704_0
+  - sympy=1.4=py37_0
+  - tbb=2020.0=hfd86e86_0
+  - tblib=1.4.0=py_0
+  - tensorboard=1.14.0=py37hf484d3e_0
+  - tensorflow=1.14.0=gpu_py37h74c33d7_0
+  - tensorflow-base=1.14.0=gpu_py37he45bfe2_0
+  - tensorflow-estimator=1.14.0=py_0
+  - tensorflow-gpu=1.14.0=h0d30ee6_0
+  - termcolor=1.1.0=py37_1
+  - terminado=0.8.2=py37_0
+  - testpath=0.4.2=py37_0
+  - tk=8.6.8=hbc83047_0
+  - tktable=2.10=h14c3975_0
+  - toolz=0.10.0=py_0
+  - tornado=6.0.3=py37h7b6447c_0
+  - tqdm=4.32.1=py_0
+  - traitlets=4.3.2=py37_0
+  - unicodecsv=0.14.1=py37_0
+  - unixodbc=2.3.7=h14c3975_0
+  - urllib3=1.24.2=py37_0
+  - wcwidth=0.1.7=py37_0
+  - webencodings=0.5.1=py37_1
+  - werkzeug=0.15.5=py_0
+  - wheel=0.33.4=py37_0
+  - widgetsnbextension=3.5.1=py37_0
+  - wrapt=1.11.2=py37h7b6447c_0
+  - wurlitzer=1.0.3=py37_0
+  - xlrd=1.2.0=py37_0
+  - xlsxwriter=1.1.8=py_0
+  - xlwt=1.3.0=py37_0
+  - xz=5.2.4=h14c3975_4
+  - yaml=0.1.7=had09818_2
+  - zeromq=4.2.5=hf484d3e_1
+  - zict=1.0.0=py_0
+  - zipp=0.5.2=py_0
+  - zlib=1.2.11=h7b6447c_3
+  - zstd=1.3.7=h0b5b093_0
+  - pip:
+    - appdirs==1.4.3
+    - blessings==1.7
+    - configargparse==1.1
+    - datrie==0.8
+    - gitdb==4.0.2
+    - gitpython==3.1.0
+    - gpustat==0.6.0
+    - keras==2.2.4
+    - nvidia-ml-py3==7.352.0
+    - smmap==3.0.1
+
diff --git a/examples/amber_tf2_env.yml b/examples/amber_tf2_env.yml
new file mode 100644
index 0000000..d702a48
--- /dev/null
+++ b/examples/amber_tf2_env.yml
@@ -0,0 +1,354 @@
+name: tf2
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - _anaconda_depends=2019.03=py37_0
+  - _ipyw_jlab_nb_ext_conf=0.1.0=py37_0
+  - _libgcc_mutex=0.1=main
+  - alabaster=0.7.12=py37_0
+  - anaconda=custom=py37_1
+  - anaconda-client=1.7.2=py37_0
+  - anaconda-navigator=1.9.7=py37_0
+  - anaconda-project=0.8.4=py_0
+  - argh=0.26.2=py37_0
+  - asn1crypto=1.3.0=py37_0
+  - astroid=2.3.3=py37_0
+  - astropy=4.0=py37h7b6447c_0
+  - atomicwrites=1.3.0=py37_1
+  - attrs=19.3.0=py_0
+  - autopep8=1.4.4=py_0
+  - babel=2.8.0=py_0
+  - backcall=0.1.0=py37_0
+  - backports=1.0=py_2
+  - backports.functools_lru_cache=1.6.1=py_0
+  - backports.os=0.1.1=py37_0
+  - backports.shutil_get_terminal_size=1.0.0=py37_2
+  - backports.tempfile=1.0=py_1
+  - backports.weakref=1.0.post1=py_1
+  - beautifulsoup4=4.8.2=py37_0
+  - bitarray=1.2.1=py37h7b6447c_0
+  - bkcharts=0.2=py37_0
+  - blas=1.0=mkl
+  - bleach=3.1.0=py37_0
+  - blosc=1.16.3=hd408876_0
+  - bokeh=1.4.0=py37_0
+  - boto=2.49.0=py37_0
+  - bottleneck=1.3.1=py37hdd07704_0
+  - bzip2=1.0.8=h7b6447c_0
+  - ca-certificates=2020.1.1=0
+  - cairo=1.14.12=h8948797_3
+  - certifi=2019.11.28=py37_0
+  - cffi=1.13.2=py37h2e261b9_0
+  - chardet=3.0.4=py37_1003
+  - click=7.0=py37_0
+  - cloudpickle=1.2.2=py_0
+  - clyent=1.2.2=py37_1
+  - colorama=0.4.3=py_0
+  - conda-package-handling=1.6.0=py37h7b6447c_0
+  - conda-verify=3.4.2=py_1
+  - contextlib2=0.6.0.post1=py_0
+  - cryptography=2.8=py37h1ba5d50_0
+  - curl=7.68.0=hbc83047_0
+  - cycler=0.10.0=py37_0
+  - cython=0.29.14=py37he6710b0_0
+  - cytoolz=0.10.1=py37h7b6447c_0
+  - dask=2.10.1=py_0
+  - dask-core=2.10.1=py_0
+  - dbus=1.13.12=h746ee38_0
+  - decorator=4.4.1=py_0
+  - defusedxml=0.6.0=py_0
+  - diff-match-patch=20181111=py_0
+  - distributed=2.10.0=py_0
+  - docutils=0.16=py37_0
+  - entrypoints=0.3=py37_0
+  - et_xmlfile=1.0.1=py37_0
+  - expat=2.2.6=he6710b0_0
+  - fastcache=1.1.0=py37h7b6447c_0
+  - filelock=3.0.12=py_0
+  - flake8=3.7.9=py37_0
+  - flask=1.1.1=py_0
+  - fontconfig=2.13.0=h9420a91_0
+  - freetype=2.9.1=h8a8886c_1
+  - fribidi=1.0.5=h7b6447c_0
+  - fsspec=0.6.2=py_0
+  - future=0.18.2=py37_0
+  - get_terminal_size=1.0.0=haa9412d_0
+  - gettext=0.19.8.1=hc5be6a0_1002
+  - gevent=1.4.0=py37h7b6447c_0
+  - glib=2.63.1=h5a9c865_0
+  - glob2=0.7=py_0
+  - gmp=6.1.2=h6c8ec71_1
+  - gmpy2=2.0.8=py37h10f8cd9_2
+  - graphite2=1.3.13=h23475e2_0
+  - greenlet=0.4.15=py37h7b6447c_0
+  - gst-plugins-base=1.14.0=hbbd80ab_1
+  - gstreamer=1.14.0=hb453b48_1
+  - h5py=2.10.0=py37h7918eee_0
+  - harfbuzz=1.8.8=hffaf4a1_0
+  - hdf5=1.10.4=hb1b8bf9_0
+  - heapdict=1.0.1=py_0
+  - html5lib=1.0.1=py37_0
+  - hypothesis=5.4.1=py_0
+  - icu=58.2=h9c2bf20_1
+  - idna=2.8=py37_0
+  - imageio=2.6.1=py37_0
+  - imagesize=1.2.0=py_0
+  - importlib_metadata=1.5.0=py37_0
+  - inflect=4.1.0=py37_0
+  - intel-openmp=2019.4=243
+  - intervaltree=3.0.2=py_0
+  - ipykernel=5.1.4=py37h39e3cac_0
+  - ipython=7.12.0=py37h5ca1d4c_0
+  - ipython_genutils=0.2.0=py37_0
+  - ipywidgets=7.5.1=py_0
+  - isort=4.3.21=py37_0
+  - itsdangerous=1.1.0=py37_0
+  - jaraco.itertools=5.0.0=py_0
+  - jbig=2.1=hdba287a_0
+  - jdcal=1.4.1=py_0
+  - jedi=0.14.1=py37_0
+  - jeepney=0.4.2=py_0
+  - jinja2=2.11.1=py_0
+  - joblib=0.14.1=py_0
+  - jpeg=9b=h024ee3a_2
+  - json5=0.9.0=py_0
+  - jsonschema=3.2.0=py37_0
+  - jupyter=1.0.0=py37_7
+  - jupyter_client=5.3.4=py37_0
+  - jupyter_console=6.1.0=py_0
+  - jupyter_core=4.6.1=py37_0
+  - jupyterlab=1.2.6=pyhf63ae98_0
+  - jupyterlab_server=1.0.6=py_0
+  - keyring=21.1.0=py37_0
+  - kiwisolver=1.1.0=py37he6710b0_0
+  - krb5=1.17.1=h173b8e3_0
+  - lazy-object-proxy=1.4.3=py37h7b6447c_0
+  - libarchive=3.3.3=h5d8350f_5
+  - libcurl=7.68.0=h20c2e04_0
+  - libedit=3.1.20181209=hc058e9b_0
+  - libffi=3.2.1=hd88cf55_4
+  - libgcc-ng=9.1.0=hdf63c60_0
+  - libgfortran-ng=7.3.0=hdf63c60_0
+  - libiconv=1.15=h516909a_1005
+  - liblief=0.9.0=h7725739_2
+  - libpng=1.6.37=hbc83047_0
+  - libsodium=1.0.16=h1bed415_0
+  - libspatialindex=1.9.3=he6710b0_0
+  - libssh2=1.8.2=h1ba5d50_0
+  - libstdcxx-ng=9.1.0=hdf63c60_0
+  - libtiff=4.1.0=h2733197_0
+  - libtool=2.4.6=h7b6447c_5
+  - libuuid=1.0.3=h1bed415_2
+  - libxcb=1.13=h1bed415_1
+  - libxml2=2.9.9=hea5a465_1
+  - libxslt=1.1.33=h7d1a2b0_0
+  - llvmlite=0.29.0=py37hd408876_0
+  - locket=0.2.0=py37_1
+  - lxml=4.5.0=py37hefd8a0e_0
+  - lz4-c=1.8.1.2=h14c3975_0
+  - lzo=2.10=h49e0be7_2
+  - markupsafe=1.1.1=py37h7b6447c_0
+  - matplotlib=3.1.3=py37_0
+  - matplotlib-base=3.1.3=py37hef1b27d_0
+  - mccabe=0.6.1=py37_1
+  - mistune=0.8.4=py37h7b6447c_0
+  - mkl=2019.4=243
+  - mkl-service=2.3.0=py37he904b0f_0
+  - mkl_fft=1.0.15=py37ha843d7b_0
+  - mkl_random=1.1.0=py37hd6b4f25_0
+  - mock=3.0.5=py37_0
+  - more-itertools=8.2.0=py_0
+  - mpc=1.1.0=h10f8cd9_1
+  - mpfr=4.0.1=hdf1c602_3
+  - mpmath=1.1.0=py37_0
+  - msgpack-python=0.6.1=py37hfd86e86_1
+  - multipledispatch=0.6.0=py37_0
+  - navigator-updater=0.2.1=py37_0
+  - nbconvert=5.6.1=py37_0
+  - nbformat=5.0.4=py_0
+  - ncurses=6.1=he6710b0_1
+  - networkx=2.4=py_0
+  - nltk=3.4.5=py37_0
+  - nose=1.3.7=py37_2
+  - notebook=6.0.3=py37_0
+  - numba=0.44.1=py37h962f231_0
+  - numexpr=2.7.1=py37h423224d_0
+  - numpy=1.18.1=py37h4f9e942_0
+  - numpy-base=1.18.1=py37hde5b4d6_1
+  - numpydoc=0.9.2=py_0
+  - olefile=0.46=py37_0
+  - openpyxl=3.0.3=py_0
+  - openssl=1.1.1d=h7b6447c_3
+  - packaging=20.1=py_0
+  - pandas=1.0.0=py37h0573a6f_0
+  - pandoc=2.2.3.2=0
+  - pandocfilters=1.4.2=py37_1
+  - pango=1.42.4=h049681c_0
+  - parso=0.6.0=py_0
+  - partd=1.1.0=py_0
+  - patchelf=0.10=he6710b0_0
+  - path=13.1.0=py37_0
+  - path.py=12.4.0=0
+  - pathlib2=2.3.5=py37_0
+  - pathtools=0.1.2=py_1
+  - patsy=0.5.1=py37_0
+  - pcre=8.43=he6710b0_0
+  - pep8=1.7.1=py37_0
+  - pexpect=4.8.0=py37_0
+  - pickleshare=0.7.5=py37_0
+  - pillow=7.0.0=py37hb39fc2d_0
+  - pixman=0.38.0=h7b6447c_0
+  - pkginfo=1.5.0.1=py37_0
+  - pluggy=0.13.1=py37_0
+  - ply=3.11=py37_0
+  - prometheus_client=0.7.1=py_0
+  - prompt_toolkit=3.0.3=py_0
+  - psutil=5.6.7=py37h7b6447c_0
+  - ptyprocess=0.6.0=py37_0
+  - py=1.8.1=py_0
+  - py-lief=0.9.0=py37h7725739_2
+  - pycodestyle=2.5.0=py37_0
+  - pycosat=0.6.3=py37h7b6447c_0
+  - pycparser=2.19=py37_0
+  - pycrypto=2.6.1=py37h14c3975_9
+  - pycurl=7.43.0.5=py37h1ba5d50_0
+  - pydocstyle=4.0.1=py_0
+  - pyflakes=2.1.1=py37_0
+  - pygments=2.5.2=py_0
+  - pylint=2.4.4=py37_0
+  - pyodbc=4.0.28=py37he6710b0_0
+  - pyopenssl=19.1.0=py37_0
+  - pyparsing=2.4.6=py_0
+  - pyqt=5.9.2=py37h05f1152_2
+  - pyrsistent=0.15.7=py37h7b6447c_0
+  - pysocks=1.7.1=py37_0
+  - pytables=3.6.1=py37h71ec239_0
+  - pytest=5.3.5=py37_0
+  - pytest-arraydiff=0.3=py37h39e3cac_0
+  - pytest-astropy=0.7.0=py_0
+  - pytest-astropy-header=0.1.2=py_0
+  - pytest-doctestplus=0.5.0=py_0
+  - pytest-openfiles=0.4.0=py_0
+  - pytest-remotedata=0.3.2=py37_0
+  - python=3.7.3=h0371630_0
+  - python-dateutil=2.8.1=py_0
+  - python-jsonrpc-server=0.3.4=py_0
+  - python-language-server=0.31.7=py37_0
+  - python-libarchive-c=2.8=py37_13
+  - pytz=2019.3=py_0
+  - pywavelets=1.1.1=py37h7b6447c_0
+  - pyxdg=0.26=py_0
+  - pyyaml=5.3=py37h7b6447c_0
+  - pyzmq=18.1.1=py37he6710b0_0
+  - qdarkstyle=2.8=py_0
+  - qt=5.9.7=h5867ecd_1
+  - qtawesome=0.6.1=py_0
+  - qtconsole=4.6.0=py_1
+  - qtpy=1.9.0=py_0
+  - readline=7.0=h7b6447c_5
+  - requests=2.22.0=py37_1
+  - rope=0.16.0=py_0
+  - rtree=0.9.3=py37_0
+  - ruamel_yaml=0.15.87=py37h7b6447c_0
+  - scikit-image=0.16.2=py37h0573a6f_0
+  - scikit-learn=0.22.1=py37hd81dba3_0
+  - scipy=1.4.1=py37h0b6359f_0
+  - seaborn=0.10.0=py_0
+  - secretstorage=3.1.2=py37_0
+  - send2trash=1.5.0=py37_0
+  - setuptools=45.1.0=py37_0
+  - simplegeneric=0.8.1=py37_2
+  - singledispatch=3.4.0.3=py37_0
+  - sip=4.19.8=py37hf484d3e_0
+  - six=1.14.0=py37_0
+  - snappy=1.1.7=hbae5bb6_3
+  - snowballstemmer=2.0.0=py_0
+  - sortedcollections=1.1.2=py37_0
+  - sortedcontainers=2.1.0=py37_0
+  - soupsieve=1.9.5=py37_0
+  - sphinx=2.3.1=py_0
+  - sphinxcontrib=1.0=py37_1
+  - sphinxcontrib-applehelp=1.0.1=py_0
+  - sphinxcontrib-devhelp=1.0.1=py_0
+  - sphinxcontrib-htmlhelp=1.0.2=py_0
+  - sphinxcontrib-jsmath=1.0.1=py_0
+  - sphinxcontrib-qthelp=1.0.2=py_0
+  - sphinxcontrib-serializinghtml=1.1.3=py_0
+  - sphinxcontrib-websupport=1.1.2=py_0
+  - spyder=4.0.1=py37_0
+  - spyder-kernels=1.8.1=py37_0
+  - sqlalchemy=1.3.13=py37h7b6447c_0
+  - sqlite=3.31.1=h7b6447c_0
+  - statsmodels=0.11.0=py37h7b6447c_0
+  - sympy=1.5.1=py37_0
+  - tbb=2020.0=hfd86e86_0
+  - tblib=1.6.0=py_0
+  - terminado=0.8.3=py37_0
+  - testpath=0.4.4=py_0
+  - tk=8.6.8=hbc83047_0
+  - toolz=0.10.0=py_0
+  - tornado=6.0.3=py37h7b6447c_0
+  - tqdm=4.42.0=py_0
+  - traitlets=4.3.3=py37_0
+  - ujson=1.35=py37h14c3975_0
+  - unicodecsv=0.14.1=py37_0
+  - unixodbc=2.3.7=h14c3975_0
+  - urllib3=1.25.8=py37_0
+  - watchdog=0.10.1=py37_0
+  - wcwidth=0.1.8=py_0
+  - webencodings=0.5.1=py37_1
+  - werkzeug=0.16.1=py_0
+  - wheel=0.34.2=py37_0
+  - widgetsnbextension=3.5.1=py37_0
+  - wrapt=1.11.2=py37h7b6447c_0
+  - wurlitzer=2.0.0=py37_0
+  - xlrd=1.2.0=py37_0
+  - xlsxwriter=1.2.7=py_0
+  - xlwt=1.3.0=py37_0
+  - xz=5.2.4=h14c3975_4
+  - yaml=0.1.7=had09818_2
+  - yapf=0.28.0=py_0
+  - zeromq=4.3.1=he6710b0_3
+  - zict=1.0.0=py_0
+  - zipp=2.1.0=py_0
+  - zlib=1.2.11=h7b6447c_3
+  - zstd=1.3.7=h0b5b093_0
+  - pip:
+    - absl-py==0.9.0
+    - astunparse==1.6.3
+    - blessings==1.7
+    - cachetools==4.1.0
+    - gast==0.3.3
+    - google-auth==1.16.1
+    - google-auth-oauthlib==0.4.1
+    - google-pasta==0.2.0
+    - gpustat==0.6.0
+    - grpcio==1.29.0
+    - keras==2.3.1
+    - keras-applications==1.0.8
+    - keras-preprocessing==1.1.2
+    - markdown==3.2.2
+    - nvidia-ml-py3==7.352.0
+    - oauthlib==3.1.0
+    - opt-einsum==3.2.1
+    - pip==20.1.1
+    - plotly==4.4.1
+    - protobuf==3.12.2
+    - pyasn1==0.4.8
+    - pyasn1-modules==0.2.8
+    - pybedtools==0.8.1
+    - pydot==1.4.1
+    - pysam==0.15.4
+    - pytabix==0.1
+    - requests-oauthlib==1.3.0
+    - rsa==4.0
+    - tensorboard==2.2.2
+    - tensorboard-plugin-wit==1.6.0.post3
+    - tensorflow-estimator==2.2.0
+    - tensorflow-gpu==2.2.0
+    - termcolor==1.1.0
+    - torch==1.4.0
+    - torchvision==0.5.0
+prefix: /mnt/home/zzhang/workspace/programs/anaconda3/envs/tf2
+
diff --git a/examples/mock_nas.py b/examples/mock_nas.py
new file mode 100644
index 0000000..cdc2b2b
--- /dev/null
+++ b/examples/mock_nas.py
@@ -0,0 +1,108 @@
+import os
+import sys
+import logging
+from amber.architect.controller import GeneralController
+from amber.architect.model_space import State, ModelSpace
+from amber.architect.train_env import ControllerTrainEnvironment
+from amber.bootstrap.mock_manager import MockManager
+from amber.bootstrap.simple_conv1d_space import get_state_space
+from amber.utils.logging import setup_logger
+import tensorflow as tf
+
+
+def get_controller(state_space, sess):
+	'''Test function for building controller network. A controller is a LSTM cell that predicts the next
+	layer given the previous layer and all previous layers (as stored in the hidden cell states). The 
+	controller model is trained by policy gradients as in reinforcement learning.
+	'''
+	with tf.device("/cpu:0"):
+		controller = GeneralController(
+			model_space=state_space,
+			lstm_size=50,
+			lstm_num_layers=1,
+			with_skip_connection=False,
+			kl_threshold=0.01,
+			train_pi_iter=100,
+			optim_algo='adam',
+			#tanh_constant=1.5,
+			buffer_size=5, ## num of episodes saved
+			batch_size=5,
+			session=sess,
+			#use_ppo_loss=True
+			)
+	return controller
+
+
+
+def get_mock_manager(history_fn_list, Lambda=1., wd='./tmp_mock'):
+	'''Test function for building a mock manager. A mock manager 
+	returns a loss and knowledge instantly based on previous 
+	training history.
+	Args:
+		train_history_fn_list: a list of 
+	'''
+	manager = MockManager(
+			history_fn_list = history_fn_list,
+			model_compile_dict = { 'loss':'binary_crossentropy', 'optimizer':'adam', 'metrics':['acc']},
+			working_dir=wd,
+			Lambda=Lambda,
+ 			verbose=0
+ 		)
+	return manager
+
+
+def get_environment(controller, manager, should_plot, logger=None, wd='./tmp_mock/'):
+	'''Test function for getting a training environment for controller.
+	Args:
+		controller: a built controller net
+		manager: a manager is a function that manages child-networks. Manager is built upon `model_fn` and `reward_fn`.
+		max_trials: maximum number of child-net generated
+		working_dir: specifies the working dir where all files will be generated in.
+		resume_prev_run: restore the controller parameters and states from a preivous run
+		logger: a Logging file handler; creates a new logger if None.
+	'''
+	env = ControllerTrainEnvironment(
+		controller,
+		manager,
+		max_episode=200,
+		max_step_per_ep=3,
+		logger=logger,
+		resume_prev_run=False,
+		should_plot=should_plot,
+		working_dir=wd,
+		with_skip_connection=False,
+		with_input_blocks=False
+	)
+	return env
+
+def train_simple_controller(should_plot=False, logger=None, Lambda=1., wd='./outputs/mock_nas/'):
+	sess = tf.Session()
+
+	# first get state_space
+	state_space = get_state_space()
+
+
+	# init network manager
+	hist_file_list = ["./data/mock_black_box/tmp_%i/train_history.csv"%i for i in range(1,21) ]
+	#print(hist_file_list)
+	manager = get_mock_manager(hist_file_list, Lambda=Lambda, wd=wd)
+
+	# get controller
+	controller = get_controller(state_space, sess)
+
+	# get the training environment
+	env = get_environment(controller, manager, should_plot, logger, wd=wd)
+
+	# train one step
+	env.train()
+
+	return
+
+if __name__ == '__main__':
+    wd = "./outputs/mock_nas/"
+    os.makedirs(wd, exist_ok=True)
+    logger = setup_logger(working_dir=wd, verbose_level=logging.DEBUG)
+    B = 1
+    Lambda = 1.
+    for t in range(B):
+        train_simple_controller(t==(B-1), logger, Lambda)
diff --git a/examples/sbatch_chipseq.sh b/examples/sbatch_chipseq.sh
new file mode 100644
index 0000000..fc6a95a
--- /dev/null
+++ b/examples/sbatch_chipseq.sh
@@ -0,0 +1,93 @@
+#!/usr/bin/env bash
+#SBATCH --time=7-0:00:00
+#SBATCH --gres=gpu:v100-32gb:4
+#SBATCH --partition=gpu
+#SBATCH --mem 160g
+#SBATCH -c 32
+
+echo 'Just plain run.'
+
+# Load bash config.
+source "${HOME}"'/.bashrc'
+if [ $? != 0 ]; then
+    echo 'Failed sourcing '"${HOME}"'/.bashrc'
+    exit 1
+fi
+
+# Start up conda.
+conda activate
+if [ $? != 0 ]; then
+    echo 'Failed to activate conda.'
+    exit 1
+fi
+
+# Start up anaconda.
+conda activate 'amber'
+if [ $? != 0 ]; then
+    echo 'Failed to activate conda environment.'
+    exit 1
+fi
+
+# Change directories.
+#SRC_DIR='/mnt/ceph/users/ecofer/AMBER/7_initial_parallel_debug'
+SRC_DIR='.'
+cd "${SRC_DIR}"
+if [ $? != 0 ]; then
+    echo 'Failed changing directories to '"${SRC_DIR}"
+    exit 1
+fi
+
+# Remove things on ram disk.
+function cleanup {
+    echo 'Cleaning up data on ram disk.'
+    rm -rf '/dev/shm/'"${SLURM_JOB_ID}"
+    if [ $? != 0 ]; then
+        echo 'Failed to cleanup after exit.'
+        exit 1
+    else
+        echo 'Finished cleaning up data on ram disk.'
+    fi
+}
+
+# Set trap for exit to ensure cleanup.
+trap cleanup EXIT
+
+# Make ram disk.
+echo 'Making ram disk dir.'
+mkdir '/dev/shm/'"${SLURM_JOB_ID}"
+if [ $? != 0 ]; then
+    echo 'Failed to make dir /dev/shm/'"${SLURM_JOB_ID}"
+    exit 1
+else
+    echo 'Finished making ram disk dir.'
+fi
+
+# Copy genome to ram disk.
+echo 'Copying genome to ram disk.'
+cp './data/zero_shot_chipseq/hg19.encoded.h5' '/dev/shm/'"${SLURM_JOB_ID}"'/hg19.encoded.h5'
+if [ $? != 0 ]; then
+    echo 'Failed to copy genome to /dev/shm/'"${SLURM_JOB_ID}"'/hg19.encoded.h5'
+    exit 1
+else
+    echo 'Finished copying genome.'
+fi
+
+# Run GPU logging.
+nvidia-smi -l 600 >"${SRC_DIR}"'/slurm-'"${SLURM_JOB_ID}"'.nvidia-smi.out' 2>&1  &
+
+# Run train script.
+/usr/bin/time -v python -u "${SRC_DIR}"'/zero_shot_nas.real_data.py' \
+    --analysis 'train' \
+    --wd "${SRC_DIR}"'/outputs/zero_shot_8chipseq' \
+    --config-file "${SRC_DIR}"'/data/zero_shot_chipseq/train_feats.config_file.8_random_feats.tsv' \
+    --genome-file '/dev/shm/'"${SLURM_JOB_ID}"'/hg19.encoded.h5' \
+    --dfeature-name-file "${SRC_DIR}"'/data/zero_shot_chipseq/dfeatures_ordered_list.txt' \
+    --n-test 1 \
+    --n-train 4000000 \
+    --n-validate 64000 \
+    --parallel 
+echo $?
+
+# Deactivate conda.
+conda deactivate
+
diff --git a/examples/sbatch_deepsea.sh b/examples/sbatch_deepsea.sh
new file mode 100644
index 0000000..35ef6a5
--- /dev/null
+++ b/examples/sbatch_deepsea.sh
@@ -0,0 +1,90 @@
+#!/usr/bin/env bash
+#SBATCH --time=7-0:00:00
+#SBATCH --gres=gpu:v100-32gb:4
+#SBATCH --partition=gpu
+#SBATCH --mem 128g
+#SBATCH -c 32
+
+echo 'Just plain run.'
+
+# Load bash config.
+source "${HOME}"'/.bashrc'
+if [ $? != 0 ]; then
+    echo 'Failed sourcing '"${HOME}"'/.bashrc'
+    exit 1
+fi
+
+# Start up conda.
+conda activate
+if [ $? != 0 ]; then
+    echo 'Failed to activate conda.'
+    exit 1
+fi
+
+# Start up anaconda.
+conda activate 'amber'
+if [ $? != 0 ]; then
+    echo 'Failed to activate conda environment.'
+    exit 1
+fi
+
+# Change directories.
+SRC_DIR='.'
+cd "${SRC_DIR}"
+if [ $? != 0 ]; then
+    echo 'Failed changing directories to '"${SRC_DIR}"
+    exit 1
+fi
+
+# Remove things on ram disk.
+function cleanup {
+    echo 'Cleaning up data on ram disk.'
+    rm -rf '/dev/shm/'"${SLURM_JOB_ID}"
+    if [ $? != 0 ]; then
+        echo 'Failed to cleanup after exit.'
+        exit 1
+    else
+        echo 'Finished cleaning up data on ram disk.'
+    fi
+}
+
+# Set trap for exit to ensure cleanup.
+trap cleanup EXIT
+
+# Make ram disk.
+echo 'Making ram disk dir.'
+#SLURM_JOB_ID='test-amber'
+mkdir '/dev/shm/'"${SLURM_JOB_ID}"
+if [ $? != 0 ]; then
+    echo 'Failed to make dir /dev/shm/'"${SLURM_JOB_ID}"
+    exit 1
+else
+    echo 'Finished making ram disk dir.'
+fi
+
+# Copy genome to ram disk.
+echo 'Copying hdf5 data to ram disk.'
+cp './data/zero_shot_deepsea/train.h5' '/dev/shm/'"${SLURM_JOB_ID}"'/train.h5'
+cp './data/zero_shot_deepsea/val.h5' '/dev/shm/'"${SLURM_JOB_ID}"'/val.h5'
+if [ $? != 0 ]; then
+    echo 'Failed to copy genome to /dev/shm/'"${SLURM_JOB_ID}"'/hg19.encoded.h5'
+    exit 1
+else
+    echo 'Finished copying hdf5.'
+fi
+
+# Run train script.
+/usr/bin/time -v python -u "${SRC_DIR}"'/zero_shot_nas.deepsea.py' \
+    --analysis 'train' \
+    --wd "${SRC_DIR}"'/outputs/zero_shot_deepsea' \
+    --config-file "${SRC_DIR}"'/data/zero_shot_deepsea/debug_feats.config.4_cats.tsv' \
+    --train-file '/dev/shm/'"${SLURM_JOB_ID}"'/train.h5' \
+    --val-file '/dev/shm/'"${SLURM_JOB_ID}"'/val.h5' \
+    --dfeature-name-file "${SRC_DIR}"'/data/zero_shot_deepsea/dfeatures_ordered_list.txt' \
+    --parallel  #\
+#    --resume
+echo $?
+
+# Deactivate conda.
+conda deactivate
+
diff --git a/examples/sbatch_grid.sh b/examples/sbatch_grid.sh
new file mode 100644
index 0000000..25198e6
--- /dev/null
+++ b/examples/sbatch_grid.sh
@@ -0,0 +1,12 @@
+#!/bin/bash
+#SBATCH -p gpu --gres=gpu:v100-32gb:1
+#SBATCH --mem=16G
+#SBATCH --time=1-1
+
+dataset=$1
+parDir=$2
+if [ ! -d $parDir ]; then
+	  mkdir -p $parDir
+fi
+
+python zeroshot_grid_searcher.py --dataset $dataset --wd $parDir --B 50
diff --git a/examples/simulate_dna.sh b/examples/simulate_dna.sh
new file mode 100644
index 0000000..78eab39
--- /dev/null
+++ b/examples/simulate_dna.sh
@@ -0,0 +1,11 @@
+# This script simulated DNA sequences using `simdna` package
+# ZZ, 2020.5.16
+
+SIMDNA_FOREGROUND="../../simdna/scripts/densityMotifSimulation.py"
+SIMDNA_BACKGROUND="../../simdna/scripts/emptyBackground.py"
+
+python $SIMDNA_FOREGROUND --prefix MYC_known10 --motifNames MYC_known10 --max-motifs 10 --min-motifs 1 --mean-motifs 5 --seqLength 1000 --numSeqs 10000 --seed 777
+
+python $SIMDNA_FOREGROUND --prefix CTCF_known1 --motifNames CTCF_known1 --max-motifs 1 --min-motifs 1 --mean-motifs 1 --seqLength 1000 --numSeqs 10000 --seed 777
+
+python $SIMDNA_BACKGROUND --prefix empty_bg --seqLength 1000 --numSeqs 10000
diff --git a/examples/zero_shot_nas.deepsea.py b/examples/zero_shot_nas.deepsea.py
new file mode 100644
index 0000000..1c48d18
--- /dev/null
+++ b/examples/zero_shot_nas.deepsea.py
@@ -0,0 +1,427 @@
+#!/usr/bin/env python
+
+"""
+Testing for ZeroShotNAS, aka data description driven nas
+ZZ
+May 14, 2020
+"""
+
+
+import tensorflow as tf
+tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+import numpy as np
+import os
+import copy
+import logging
+import pickle
+import pandas as pd
+import argparse
+
+from amber.modeler import EnasCnnModelBuilder
+from amber.architect.controller import ZeroShotController
+from amber.architect.model_space import State, ModelSpace
+from amber.architect.common_ops import count_model_params, unpack_data
+
+from amber.architect.train_env import MultiManagerEnvironment, ParallelMultiManagerEnvironment
+from amber.architect.reward import LossAucReward, LossReward
+from amber.plots import plot_controller_hidden_states
+from amber.utils import run_from_ipython, get_available_gpus
+from amber.utils.logging import setup_logger
+from amber.utils.data_parser import get_data_from_simdata
+
+from amber.modeler.modeler import build_sequential_model, KerasModelBuilder
+from amber.architect.manager import GeneralManager, DistributedGeneralManager
+from amber.architect.model_space import get_layer_shortname
+
+from amber.utils.sequences import EncodedGenome
+from amber.utils.sequences import EncodedHDF5Genome
+from amber.utils.sampler import BatchedHDF5Generator
+
+
+def get_controller(model_space, session, data_description_len=3, layer_embedding_sharing=None):
+    with tf.device("/cpu:0"):
+        controller = ZeroShotController(
+            data_description_config={
+                "length": data_description_len,
+                "hidden_layer": {"units":8, "activation": "relu"},
+                "regularizer": {"l1":1e-8 }
+                },
+            share_embedding=layer_embedding_sharing,
+            model_space=model_space,
+            session=session,
+            with_skip_connection=False,
+            skip_weight=None,
+            skip_target=0.2,
+            lstm_size=128,
+            lstm_num_layers=2,
+            kl_threshold=0.05,
+            train_pi_iter=100,
+            optim_algo='adam',
+            temperature=1.5,
+            tanh_constant=2,
+            buffer_type="MultiManager",
+            buffer_size=5,
+            batch_size=10,
+            use_ppo_loss=True,
+            rescale_advantage_by_reward=False
+        )
+    return controller
+
+
+def get_model_space_common():
+
+    # Setup and params.
+    state_space = ModelSpace()
+    default_params = {"kernel_initializer": "glorot_uniform",
+                      "activation": "relu"}
+    param_list = [
+            # Block 1:
+            [
+                {"filters": 64, "kernel_size": 8},
+                {"filters": 64, "kernel_size": 14},
+                {"filters": 64, "kernel_size": 20}
+            ],
+            # Block 2:
+            [
+                {"filters": 128, "kernel_size": 8},
+                {"filters": 128, "kernel_size": 14},
+                {"filters": 128, "kernel_size": 20}
+            ],
+            # Block 3:
+            [
+                {"filters": 256, "kernel_size": 8},
+                {"filters": 256, "kernel_size": 14},
+                {"filters": 256, "kernel_size": 20}
+            ],
+        ]
+
+    # Build state space.
+    layer_embedding_sharing = {}
+    conv_seen = 0
+    for i in range(len(param_list)):
+        # Build conv states for this layer.
+        conv_states = [State("Identity")]
+        for j in range(len(param_list[i])):
+            d = copy.deepcopy(default_params)
+            for k, v in param_list[i][j].items():
+                d[k] = v
+            conv_states.append(State('conv1d', name="conv{}".format(conv_seen), **d))
+        state_space.add_layer(conv_seen*3, conv_states)
+        if i > 0:
+            layer_embedding_sharing[conv_seen*3] = 0
+        conv_seen += 1
+
+        # Add pooling states.
+        if i < len(param_list) - 1:
+            pool_states = [State('Identity'),
+                           State('maxpool1d', pool_size=4, strides=4),
+                           State('avgpool1d', pool_size=4, strides=4)]
+            if i > 0:
+                layer_embedding_sharing[conv_seen*3-2] = 1
+        else:
+            pool_states = [
+                    State('Flatten'),
+                    State('GlobalMaxPool1D'),
+                    State('GlobalAvgPool1D')
+                ]
+        state_space.add_layer(conv_seen*3-2, pool_states)
+
+
+        # Add dropout
+        state_space.add_layer(conv_seen*3-1, [
+            State('Identity'),
+            State('Dropout', rate=0.1),
+            State('Dropout', rate=0.3),
+            State('Dropout', rate=0.5)
+            ])
+        if i > 0:
+            layer_embedding_sharing[conv_seen*3-1] = 2
+
+    # Add final classifier layer.
+    state_space.add_layer(conv_seen*3, [
+            State('Dense', units=30, activation='relu'),
+            State('Dense', units=100, activation='relu'),
+            State('Identity')
+        ])
+    return state_space, layer_embedding_sharing
+
+
+def get_manager_distributed(train_data, val_data, controller, model_space, wd, data_description, verbose=0,
+                            devices=None, train_data_kwargs=None, validate_data_kwargs=None, **kwargs):
+    reward_fn = LossAucReward(method='auc')
+    input_node = State('input', shape=(1000, 4), name="input", dtype='float32')
+    output_node = State('dense', units=1, activation='sigmoid')
+    model_compile_dict = {
+        'loss': 'binary_crossentropy',
+        'optimizer': 'adam',
+        'metrics': ['acc']
+    }
+    mb = KerasModelBuilder(inputs=input_node, outputs=output_node, model_compile_dict=model_compile_dict, model_space=model_space)
+    manager = DistributedGeneralManager(
+        devices=devices,
+        train_data_kwargs=train_data_kwargs,
+        train_data=train_data,
+        validate_data_kwargs=validate_data_kwargs,
+        validation_data=val_data,
+        epochs=1000,
+        child_batchsize=1000,
+        reward_fn=reward_fn,
+        model_fn=mb,
+        store_fn='model_plot',
+        model_compile_dict=model_compile_dict,
+        working_dir=wd,
+        verbose=verbose,
+        save_full_model=True,
+        model_space=model_space,
+        fit_kwargs={
+            'steps_per_epoch': 50,
+            'workers': 3, 'max_queue_size': 50,
+            'earlystop_patience': 10}
+    )
+    return manager
+
+
+def get_manager_common(train_data, val_data, controller, model_space, wd, data_description, verbose=2, n_feats=1, **kwargs):
+    input_node = State('input', shape=(1000, 4), name="input", dtype='float32')
+    output_node = State('dense', units=n_feats, activation='sigmoid')
+    model_compile_dict = {
+        'loss': 'binary_crossentropy',
+        'optimizer': 'adam',
+        'metrics': ['acc']
+    }
+    session = controller.session
+
+    reward_fn = LossAucReward(method='auc')
+    gpus = get_available_gpus()
+    num_gpus = len(gpus)
+    mb = KerasModelBuilder(inputs=input_node, outputs=output_node, model_compile_dict=model_compile_dict, model_space=model_space, gpus=num_gpus)
+
+    # TODO: batch_size here is not effective because it's been set at generator init
+    child_batch_size = 1000*num_gpus
+    manager = GeneralManager(
+        train_data=train_data,
+        validation_data=val_data,
+        epochs=1000,
+        child_batchsize=child_batch_size,
+        reward_fn=reward_fn,
+        model_fn=mb,
+        store_fn='model_plot',
+        model_compile_dict=model_compile_dict,
+        working_dir=wd,
+        verbose=verbose,
+        save_full_model=True,
+        model_space=model_space,
+        fit_kwargs={
+            'steps_per_epoch': 50,
+            'workers': 8, 'max_queue_size': 50,
+            'earlystop_patience': 20}
+    )
+    return manager
+
+
+def get_samples_controller(dfeatures, controller, model_space, T=100):
+    res = []
+    for i in range(len(dfeatures)):
+        dfeature = dfeatures[[i]]
+        prob_arr = [ np.zeros((T, len(model_space[i]))) for i in range(len(model_space)) ]
+        for t in range(T):
+            a, p = controller.get_action(dfeature)
+            for i, p in enumerate(p):
+                prob_arr[i][t] = p.flatten()
+        res.append(prob_arr)
+    return res
+
+
+def convert_to_dataframe(res, model_space, data_names):
+    probs = []
+    layer = []
+    description = []
+    operation = []
+    for i in range(len(res)):
+        for j in range(len(model_space)):
+            for k in range(len(model_space[j])):
+                o = get_layer_shortname(model_space[j][k])
+                p = res[i][j][:, k]
+                # extend the array
+                T = p.shape[0]
+                probs.extend(p)
+                description.extend([data_names[i]]*T)
+                layer.extend([j]*T)
+                operation.extend([o]*T)
+    df = pd.DataFrame({
+        'description':description,
+        'layer': layer,
+        'operation': operation,
+        'prob': probs
+        })
+    return df
+
+
+def reload_trained_controller(arg):
+    wd = arg.wd #wd = "./outputs/zero_shot/"
+    configs, config_keys, controller, model_space = read_configs(arg)
+    dfeatures = np.stack([configs[k]["dfeatures"] for k in config_keys])
+    res = get_samples_controller(dfeatures, controller, model_space, T=1000)
+   
+    import seaborn as sns
+    import matplotlib.pyplot as plt
+    df = convert_to_dataframe(res, model_space, data_names=[configs[k]['feat_name'] for k in config_keys])
+
+    for i in range(len(model_space)):
+        sub_df = df.loc[ (df.layer==i) ]
+        plt.clf()
+        fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10,5))
+        sns.boxplot(x="operation", y="prob",
+            hue="description",
+            data=sub_df,
+            ax=ax
+            )
+        ax.set_xticklabels(ax.get_xticklabels(), rotation=45)
+        plt.legend(bbox_to_anchor=(1.05, 1.0), loc='upper left')
+        plt.tight_layout()
+        plt.savefig(os.path.join(wd, "layer_%i.png"%i), bbox_inches="tight")
+    return res
+
+
+def read_configs(arg):
+    dfeature_names = list()
+    with open(arg.dfeature_name_file, "r") as read_file:
+        for line in read_file:
+            line = line.strip()
+            if line:
+                dfeature_names.append(line)
+ 
+    wd = arg.wd
+    model_space, layer_embedding_sharing = get_model_space_common()
+    print(layer_embedding_sharing)
+    try:
+        session = tf.Session()
+    except AttributeError:
+        session = tf.compat.v1.Session()
+
+    controller = get_controller(
+            model_space=model_space,
+            session=session,
+            data_description_len=len(dfeature_names),
+            layer_embedding_sharing=layer_embedding_sharing)
+    # Load in datasets and configurations for them.
+    if arg.config_file.endswith("tsv"):
+        sep = "\t"
+    else:
+        sep = ","
+    configs = pd.read_csv(arg.config_file, sep=sep)
+    tmp = list(configs.columns) # Because pandas doesn't have infer quotes...
+    if any(["\"" in x for x in tmp]):
+        configs = pd.read_csv(arg.config_file, sep=sep, quoting=2)
+        print("Re-read with quoting")
+    configs = configs.to_dict(orient='index')
+    # Get available gpus for parsing to DistributedManager
+    gpus = get_available_gpus()
+    gpus_ = gpus * len(configs)
+
+    manager_getter = get_manager_distributed if arg.parallel else get_manager_common
+    config_keys = list()
+    seed_generator = np.random.RandomState(seed=1337)
+    for i, k in enumerate(configs.keys()):
+        # Build datasets for train/test/validate splits.
+        for x in ["train", "validate"]:
+            if arg.lockstep_sampling is False and x == "train":
+                cur_seed = seed_generator.randint(0, np.iinfo(np.uint32).max)
+            else:
+                cur_seed = 1337
+            d = {
+                        'hdf5_fp':  arg.train_file if x=='train' else arg.val_file,
+                        'y_selector': configs[k]['ds_col'],
+                        'batch_size': 512 if arg.parallel else 512*len(gpus),
+                        'shuffle': x=='train',
+                }
+            if arg.parallel is True:
+                configs[k][x] = BatchedHDF5Generator
+                configs[k]['%s_data_kwargs'%x] = d
+            else:
+                configs[k][x] = BatchedHDF5Generator(**d)
+
+        # Build covariates and manager.
+        configs[k]["dfeatures"] = np.array(
+            [configs[k][x] for x in dfeature_names]) # TODO: Make cols dynamic.
+
+        tmp = dict() if arg.parallel is False else dict(
+                train_data_kwargs=configs[k]['train_data_kwargs'],
+                validate_data_kwargs=configs[k]["validate_data_kwargs"]
+                )
+
+        configs[k]["manager"] = manager_getter(
+            devices=[gpus_[i]],
+            train_data=configs[k]["train"],
+            val_data=configs[k]["validate"],
+            controller=controller,
+            model_space=model_space,
+            wd=os.path.join(wd, "manager_%s"%k),
+            data_description=configs[k]["dfeatures"],
+            dag_name="AmberDAG{}".format(k),
+            verbose=0 if arg.parallel is True else 2,
+            n_feats=configs[k]["n_feats"],
+            **tmp
+            )
+        config_keys.append(k)
+    return configs, config_keys, controller, model_space
+
+
+
+def train_nas(arg):
+    wd=arg.wd
+    logger = setup_logger(wd, verbose_level=logging.INFO)
+
+    gpus = get_available_gpus()
+    configs, config_keys, controller, model_space = read_configs(arg)
+    # Setup env kwargs.
+    tmp = dict(data_descriptive_features=np.stack([configs[k]["dfeatures"] for k in config_keys]),
+               controller=controller,
+               manager=[configs[k]["manager"] for k in config_keys],
+               logger=logger,
+               max_episode=200,
+               max_step_per_ep=3,
+               working_dir=wd,
+               time_budget="150:00:00",
+               with_input_blocks=False,
+               with_skip_connection=False,
+               save_controller_every=1,
+               resume_prev_run=arg.resume
+           )
+
+    if arg.parallel is True:
+        env = ParallelMultiManagerEnvironment(
+                    processes=len(gpus),
+                    **tmp)
+    else:
+        env = MultiManagerEnvironment(**tmp)
+
+    try:
+        env.train()
+    except KeyboardInterrupt:
+        print("user interrupted training")
+        pass
+    controller.save_weights(os.path.join(wd, "controller_weights.h5"))
+
+
+if __name__ == "__main__":
+    if not run_from_ipython():
+        parser = argparse.ArgumentParser(description="experimental zero-shot nas")
+        parser.add_argument("--analysis", type=str, choices=['train', 'reload'], required=True, help="analysis type")
+        parser.add_argument("--wd", type=str, default="./outputs/zero_shot/", help="working dir")
+        parser.add_argument("--parallel", default=False, action="store_true", help="Use parallel")
+        parser.add_argument("--resume", default=False, action="store_true", help="resume previous run")
+        parser.add_argument("--config-file", type=str, required=True, help="Path to the config file to use.")
+        parser.add_argument("--dfeature-name-file", type=str, required=True, help="Path to file with dataset feature names listed one per line.")
+        parser.add_argument("--train-file", type=str, required=True, help="Path to the hdf5 file of training data.")
+        parser.add_argument("--val-file", type=str, required=True, help="Path to the hdf5 file of validation data.")
+        parser.add_argument("--lockstep-sampling", default=False, action="store_true", help="Ensure same training samples used for all models.")
+
+        arg = parser.parse_args()
+
+        if arg.analysis == "train":
+            train_nas(arg)
+        elif arg.analysis == "reload":
+            reload_trained_controller(arg)
+        else:
+            raise Exception("Unknown analysis type: %s"% arg.analysis)
diff --git a/examples/zero_shot_nas.real_data.py b/examples/zero_shot_nas.real_data.py
new file mode 100644
index 0000000..8da110f
--- /dev/null
+++ b/examples/zero_shot_nas.real_data.py
@@ -0,0 +1,637 @@
+#!/usr/bin/env python
+
+"""
+Testing for ZeroShotNAS, aka data description driven nas
+ZZ
+May 14, 2020
+"""
+
+
+import tensorflow as tf
+tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+import numpy as np
+import os
+import copy
+import logging
+import pickle
+import pandas as pd
+import argparse
+
+from amber.modeler import EnasCnnModelBuilder
+from amber.architect.controller import ZeroShotController
+from amber.architect.model_space import State, ModelSpace
+from amber.architect.common_ops import count_model_params, unpack_data
+
+from amber.architect.train_env import MultiManagerEnvironment, ParallelMultiManagerEnvironment
+from amber.architect.reward import LossAucReward, LossReward
+from amber.plots import plot_controller_hidden_states
+from amber.utils import run_from_ipython, get_available_gpus
+from amber.utils.logging import setup_logger
+from amber.utils.data_parser import get_data_from_simdata
+
+from amber.modeler.modeler import build_sequential_model, KerasModelBuilder
+from amber.architect.manager import GeneralManager, DistributedGeneralManager
+from amber.architect.model_space import get_layer_shortname
+
+from amber.utils.sequences import EncodedGenome
+from amber.utils.sequences import EncodedHDF5Genome
+from amber.utils.sampler import BatchedBioIntervalSequence, BatchedBioIntervalSequenceGenerator
+
+
+def get_controller(model_space, session, data_description_len=3, layer_embedding_sharing=None):
+    with tf.device("/cpu:0"):
+        controller = ZeroShotController(
+            data_description_config={
+                "length": data_description_len,
+                "hidden_layer": {"units":8, "activation": "relu"},
+                "regularizer": {"l1":1e-8 }
+                },
+            share_embedding=layer_embedding_sharing,
+            model_space=model_space,
+            session=session,
+            with_skip_connection=False,
+            skip_weight=None,
+            skip_target=0.2,
+            lstm_size=64,
+            lstm_num_layers=1,
+            kl_threshold=0.01,
+            train_pi_iter=100,
+            optim_algo='adam',
+            temperature=2.,
+            tanh_constant=1.5,
+            buffer_type="MultiManager",
+            buffer_size=5,
+            batch_size=5,
+            use_ppo_loss=True,
+            rescale_advantage_by_reward=False
+        )
+    return controller
+
+
+def get_model_space_common():
+
+    # Setup and params.
+    state_space = ModelSpace()
+    default_params = {"kernel_initializer": "glorot_uniform",
+                      "activation": "relu"}
+    param_list = [
+            # Block 1:
+            [
+                {"filters": 64, "kernel_size": 8},
+                {"filters": 64, "kernel_size": 14},
+                {"filters": 64, "kernel_size": 20}
+            ],
+            # Block 2:
+            [
+                {"filters": 128, "kernel_size": 8},
+                {"filters": 128, "kernel_size": 14},
+                {"filters": 128, "kernel_size": 20}
+            ],
+            # Block 3:
+            [
+                {"filters": 256, "kernel_size": 8},
+                {"filters": 256, "kernel_size": 14},
+                {"filters": 256, "kernel_size": 20}
+            ],
+        ]
+
+    # Build state space.
+    layer_embedding_sharing = {}
+    conv_seen = 0
+    for i in range(len(param_list)):
+        # Build conv states for this layer.
+        conv_states = [State("Identity")]
+        for j in range(len(param_list[i])):
+            d = copy.deepcopy(default_params)
+            for k, v in param_list[i][j].items():
+                d[k] = v
+            conv_states.append(State('conv1d', name="conv{}".format(conv_seen), **d))
+        state_space.add_layer(conv_seen*3, conv_states)
+        if i > 0:
+            layer_embedding_sharing[conv_seen*3] = 0
+        conv_seen += 1
+
+        # Add pooling states.
+        if i < len(param_list) - 1:
+            pool_states = [State('Identity'),
+                           State('maxpool1d', pool_size=4, strides=4),
+                           State('avgpool1d', pool_size=4, strides=4)]
+            if i > 0:
+                layer_embedding_sharing[conv_seen*3-2] = 1
+        else:
+            pool_states = [
+                    State('Flatten'),
+                    State('GlobalMaxPool1D'),
+                    State('GlobalAvgPool1D')
+                ]
+        state_space.add_layer(conv_seen*3-2, pool_states)
+
+
+        # Add dropout
+        state_space.add_layer(conv_seen*3-1, [
+            State('Identity'),
+            State('Dropout', rate=0.1),
+            State('Dropout', rate=0.3),
+            State('Dropout', rate=0.5)
+            ])
+        if i > 0:
+            layer_embedding_sharing[conv_seen*3-1] = 2
+
+    # Add final classifier layer.
+    state_space.add_layer(conv_seen*3, [
+            State('Dense', units=30, activation='relu'),
+            State('Dense', units=100, activation='relu'),
+            State('Identity')
+        ])
+    return state_space, layer_embedding_sharing
+
+
+def get_model_space_with_long_model():
+    # Allow option for convolution to be smaller but have two layers instead of just 1.
+
+    # Setup and params.
+    state_space = ModelSpace()
+    default_params = {"kernel_initializer": "glorot_uniform",
+                      "activation": "relu"}
+    param_list = [
+            # Block 1:
+            [{"filters": 100, "kernel_size": 8},
+             {"filters": 100, "kernel_size": 14},
+             {"filters": 100, "kernel_size": 20}],
+            # Block 2:
+            [{"filters": 200, "kernel_size": 8},
+             {"filters": 200, "kernel_size": 14},
+             {"filters": 200, "kernel_size": 20}],
+            # Block 3:
+            [{"filters": 300, "kernel_size": 8},
+             {"filters": 300, "kernel_size": 14},
+             {"filters": 300, "kernel_size": 20}]
+        ]
+
+    # Build state space.
+    conv_seen = 0
+    for i in range(len(param_list)):
+        # Build conv states for this layer.
+        conv_states = [State("Identity")]
+        second_conv_states = [State("Identity")]
+        for j in range(len(param_list[i])):
+            d = copy.deepcopy(default_params)
+            for k, v in param_list[i][j].items():
+                d[k] = v
+            conv_states.append(State('conv1d', name="conv{}.1".format(conv_seen), **d))
+            d["kernel_size"] = d["kernel_size"] // 2
+            conv_states.append(State('conv1d', name="conv{}.1".format(conv_seen), **d))
+            second_conv_states.append(State("conv1d", name="conv{}.2".format(conv_seen), **d))
+        state_space.add_layer(conv_seen * 3, conv_states)
+        state_space.add_layer(conv_seen * 3 + 1, second_conv_states)
+        conv_seen += 1
+
+        # Add pooling states.
+        if i < len(param_list) - 1:
+            pool_states = [State('Identity'),
+                           State('maxpool1d', pool_size=4, strides=4),
+                           State('avgpool1d', pool_size=4, strides=4)]
+        else:
+            pool_states = [State('Flatten'),
+                           State('GlobalMaxPool1D'),
+                           State('GlobalAvgPool1D')]
+        state_space.add_layer(conv_seen * 3 - 1, pool_states)
+
+    # Add final classifier layer.
+    state_space.add_layer(conv_seen * 2, [
+            State('Dense', units=30, activation='relu'),
+            State('Dense', units=100, activation='relu'),
+            State('Identity')
+        ])
+    return state_space
+
+
+def get_model_space_with_dilation():
+
+    # Setup and params.
+    state_space = ModelSpace()
+    default_params = {"kernel_initializer": "glorot_uniform",
+                      "activation": "relu"}
+    param_list = [
+            # Block 1:
+            [{"filters": 100, "kernel_size": 8},
+             {"filters": 100, "kernel_size": 14},
+             {"filters": 100, "kernel_size": 20},
+             {"filters": 100, "kernel_size": 8, "dilation_rate": 2},
+             {"filters": 100, "kernel_size": 14, "dilation_rate": 2},
+             {"filters": 100, "kernel_size": 20, "dilation_rate": 2}],
+            # Block 2:
+            [{"filters": 200, "kernel_size": 8},
+             {"filters": 200, "kernel_size": 14},
+             {"filters": 200, "kernel_size": 20},
+             {"filters": 200, "kernel_size": 8, "dilation_rate": 2},
+             {"filters": 200, "kernel_size": 14, "dilation_rate": 2},
+             {"filters": 200, "kernel_size": 20, "dilation_rate": 2}],
+            # Block 3:
+            [{"filters": 300, "kernel_size": 8},
+             {"filters": 300, "kernel_size": 14},
+             {"filters": 300, "kernel_size": 20},
+             {"filters": 300, "kernel_size": 8, "dilation_rate": 2},
+             {"filters": 300, "kernel_size": 14, "dilation_rate": 2},
+             {"filters": 300, "kernel_size": 20, "dilation_rate": 2}],
+        ]
+
+    # Build state space.
+    conv_seen = 0
+    for i in range(len(param_list)):
+        # Build conv states for this layer.
+        conv_states = [State("Identity")]
+        for j in range(len(param_list[i])):
+            d = copy.deepcopy(default_params)
+            for k, v in param_list[i][j].items():
+                d[k] = v
+            conv_states.append(State('conv1d', name="conv{}".format(conv_seen), **d))
+        state_space.add_layer(conv_seen * 2, conv_states)
+        conv_seen += 1
+
+        # Add pooling states.
+        if i < len(param_list) - 1:
+            pool_states = [State('Identity'),
+                           State('maxpool1d', pool_size=4, strides=4),
+                           State('avgpool1d', pool_size=4, strides=4)]
+        else:
+            pool_states = [State('Flatten'),
+                           State('GlobalMaxPool1D'),
+                           State('GlobalAvgPool1D')]
+        state_space.add_layer(conv_seen * 2 - 1, pool_states)
+
+    # Add final classifier layer.
+    state_space.add_layer(conv_seen * 2, [
+            State('Dense', units=30, activation='relu'),
+            State('Dense', units=100, activation='relu'),
+            State('Identity')
+        ])
+    return state_space
+
+
+def get_model_space_with_long_model_and_dilation():
+    # Allow option for convolution to be smaller but have two layers instead of just 1.
+
+    # Setup and params.
+    state_space = ModelSpace()
+    default_params = {"kernel_initializer": "glorot_uniform",
+                      "activation": "relu"}
+    param_list = [
+            # Block 1:
+            [{"filters": 100, "kernel_size": 8},
+             {"filters": 100, "kernel_size": 14},
+             {"filters": 100, "kernel_size": 20},
+             {"filters": 100, "kernel_size": 8, "dilation_rate": 2},
+             {"filters": 100, "kernel_size": 14, "dilation_rate": 2},
+             {"filters": 100, "kernel_size": 20, "dilation_rate": 2}],
+            # Block 2:
+            [{"filters": 200, "kernel_size": 8},
+             {"filters": 200, "kernel_size": 14},
+             {"filters": 200, "kernel_size": 20},
+             {"filters": 200, "kernel_size": 8, "dilation_rate": 2},
+             {"filters": 200, "kernel_size": 14, "dilation_rate": 2},
+             {"filters": 200, "kernel_size": 20, "dilation_rate": 2}],
+            # Block 3:
+            [{"filters": 300, "kernel_size": 8},
+             {"filters": 300, "kernel_size": 14},
+             {"filters": 300, "kernel_size": 20},
+             {"filters": 300, "kernel_size": 8, "dilation_rate": 2},
+             {"filters": 300, "kernel_size": 14, "dilation_rate": 2},
+             {"filters": 300, "kernel_size": 20, "dilation_rate": 2}],
+        ]
+
+    # Build state space.
+    conv_seen = 0
+    for i in range(len(param_list)):
+        # Build conv states for this layer.
+        conv_states = [State("Identity")]
+        second_conv_states = [State("Identity")]
+        for j in range(len(param_list[i])):
+            d = copy.deepcopy(default_params)
+            for k, v in param_list[i][j].items():
+                d[k] = v
+            conv_states.append(State('conv1d', name="conv{}.1".format(conv_seen), **d))
+            d["kernel_size"] = d["kernel_size"] // 2
+            conv_states.append(State('conv1d', name="conv{}.1".format(conv_seen), **d))
+            second_conv_states.append(State("conv1d", name="conv{}.2".format(conv_seen), **d))
+        state_space.add_layer(conv_seen * 3, conv_states)
+        state_space.add_layer(conv_seen * 3 + 1, second_conv_states)
+        conv_seen += 1
+
+        # Add pooling states.
+        if i < len(param_list) - 1:
+            pool_states = [State('Identity'),
+                           State('maxpool1d', pool_size=4, strides=4),
+                           State('avgpool1d', pool_size=4, strides=4)]
+        else:
+            pool_states = [State('Flatten'),
+                           State('GlobalMaxPool1D'),
+                           State('GlobalAvgPool1D')]
+        state_space.add_layer(conv_seen * 3 - 1, pool_states)
+
+    # Add final classifier layer.
+    state_space.add_layer(conv_seen * 2, [
+            State('Dense', units=30, activation='relu'),
+            State('Dense', units=100, activation='relu'),
+            State('Identity')
+        ])
+    return state_space
+
+
+def get_manager_distributed(train_data, val_data, controller, model_space, wd, data_description, verbose=0,
+                            devices=None, train_data_kwargs=None, validate_data_kwargs=None, **kwargs):
+    reward_fn = LossAucReward(method='auc')
+    input_node = State('input', shape=(1000, 4), name="input", dtype='float32')
+    output_node = State('dense', units=1, activation='sigmoid')
+    model_compile_dict = {
+        'loss': 'binary_crossentropy',
+        'optimizer': 'adam',
+        'metrics': ['acc']
+    }
+    mb = KerasModelBuilder(inputs=input_node, outputs=output_node, model_compile_dict=model_compile_dict, model_space=model_space)
+    manager = DistributedGeneralManager(
+        devices=devices,
+        train_data_kwargs=train_data_kwargs,
+        train_data=train_data,
+        validate_data_kwargs=validate_data_kwargs,
+        validation_data=val_data,
+        epochs=100,
+        child_batchsize=1000,
+        reward_fn=reward_fn,
+        model_fn=mb,
+        store_fn='model_plot',
+        model_compile_dict=model_compile_dict,
+        working_dir=wd,
+        verbose=verbose,
+        save_full_model=True,
+        model_space=model_space,
+        fit_kwargs={
+            'steps_per_epoch': 50,
+            'workers': 3, 'max_queue_size': 50,
+            'earlystop_patience': 10}
+    )
+    return manager
+
+
+def get_manager_common(train_data, val_data, controller, model_space, wd, data_description, verbose=2, n_feats=1, **kwargs):
+    input_node = State('input', shape=(1000, 4), name="input", dtype='float32')
+    output_node = State('dense', units=n_feats, activation='sigmoid')
+    model_compile_dict = {
+        'loss': 'binary_crossentropy',
+        'optimizer': 'adam',
+        'metrics': ['acc']
+    }
+    session = controller.session
+
+    reward_fn = LossAucReward(method='auc')
+    gpus = get_available_gpus()
+    num_gpus = len(gpus)
+    mb = KerasModelBuilder(inputs=input_node, outputs=output_node, model_compile_dict=model_compile_dict, model_space=model_space, gpus=num_gpus)
+
+    # TODO: batch_size here is not effective because it's been set at generator init
+    child_batch_size = 1000*num_gpus
+    manager = GeneralManager(
+        train_data=train_data,
+        validation_data=val_data,
+        epochs=1000,
+        child_batchsize=child_batch_size,
+        reward_fn=reward_fn,
+        model_fn=mb,
+        store_fn='model_plot',
+        model_compile_dict=model_compile_dict,
+        working_dir=wd,
+        verbose=verbose,
+        save_full_model=True,
+        model_space=model_space,
+        fit_kwargs={
+            'steps_per_epoch': 50,
+            'workers': 8, 'max_queue_size': 50,
+            'earlystop_patience': 20}
+    )
+    return manager
+
+
+def get_samples_controller(dfeatures, controller, model_space, T=100):
+    res = []
+    for i in range(len(dfeatures)):
+        dfeature = dfeatures[[i]]
+        prob_arr = [ np.zeros((T, len(model_space[i]))) for i in range(len(model_space)) ]
+        for t in range(T):
+            a, p = controller.get_action(dfeature)
+            for i, p in enumerate(p):
+                prob_arr[i][t] = p.flatten()
+        res.append(prob_arr)
+    return res
+
+
+def convert_to_dataframe(res, model_space, data_names):
+    probs = []
+    layer = []
+    description = []
+    operation = []
+    for i in range(len(res)):
+        for j in range(len(model_space)):
+            for k in range(len(model_space[j])):
+                o = get_layer_shortname(model_space[j][k])
+                p = res[i][j][:, k]
+                # extend the array
+                T = p.shape[0]
+                probs.extend(p)
+                description.extend([data_names[i]]*T)
+                layer.extend([j]*T)
+                operation.extend([o]*T)
+    df = pd.DataFrame({
+        'description':description,
+        'layer': layer,
+        'operation': operation,
+        'prob': probs
+        })
+    return df
+
+
+def reload_trained_controller(arg):
+    wd = arg.wd #wd = "./outputs/zero_shot/"
+    model_space = get_model_space_with_long_model_and_dilation()
+    try:
+        session = tf.Session()
+    except AttributeError:
+        session = tf.compat.v1.Session()
+    controller = get_controller(model_space=model_space, session=session, data_description_len=2)
+    controller.load_weights(os.path.join(wd, "controller_weights.h5"))
+
+    dfeatures = np.array([[1,0], [0,1]])  # one-hot encoding
+    res = get_samples_controller(dfeatures, controller, model_space, T=1000)
+   
+    import seaborn as sns
+    import matplotlib.pyplot as plt
+    df = convert_to_dataframe(res, model_space, data_names=['MYC_known10', 'CTCF_known1'])
+
+    for i in range(len(model_space)):
+        sub_df = df.loc[ (df.layer==i) ]
+        plt.clf()
+        plt.tight_layout()
+        ax = sns.boxplot(x="operation", y="prob",
+            hue="description", palette=["m", "g"],
+            data=sub_df)
+        ax.set_xticklabels(ax.get_xticklabels(), rotation=45)
+        plt.savefig(os.path.join(wd, "layer_%i.png"%i), bbox_inches="tight")
+    return res
+
+
+def train_nas(arg):
+    dfeature_names = list()
+    with open(arg.dfeature_name_file, "r") as read_file:
+        for line in read_file:
+            line = line.strip()
+            if line:
+                dfeature_names.append(line)
+    wd = arg.wd
+    verbose = 1
+    model_space, layer_embedding_sharing = get_model_space_common()
+    print(layer_embedding_sharing)
+    try:
+        session = tf.Session()
+    except AttributeError:
+        session = tf.compat.v1.Session()
+
+    controller = get_controller(
+            model_space=model_space,
+            session=session,
+            data_description_len=len(dfeature_names),
+            layer_embedding_sharing=layer_embedding_sharing)
+    # Re-load previously saved weights, if specified
+    if arg.resume:
+        try:
+            controller.load_weights(os.path.join(wd, "controller_weights.h5"))
+            print("loaded existing weights")
+        except Exception as e:
+            print("cannot load controller weights because of %s"%e)
+    # Load in datasets and configurations for them.
+    if arg.config_file.endswith("tsv"):
+        sep = "\t"
+    else:
+        sep = ","
+    configs = pd.read_csv(arg.config_file, sep=sep)
+    tmp = list(configs.columns) # Because pandas doesn't have infer quotes...
+    if any(["\"" in x for x in tmp]):
+        configs = pd.read_csv(arg.config_file, sep=sep, quoting=2)
+        print("Re-read with quoting")
+    configs = configs.to_dict(orient='index')
+    # Get available gpus for parsing to DistributedManager
+    gpus = get_available_gpus()
+    gpus_ = gpus * len(configs)
+
+    # Build genome. This only works under the assumption that all configs use same genome.
+    genome = EncodedHDF5Genome(input_path=arg.genome_file, in_memory=False)
+
+    manager_getter = get_manager_distributed if arg.parallel else get_manager_common
+    config_keys = list()
+    seed_generator = np.random.RandomState(seed=1337)
+    for i, k in enumerate(configs.keys()):
+        # Build datasets for train/test/validate splits.
+        for x in ["train", "validate"]:
+            if arg.lockstep_sampling is False and x == "train":
+                cur_seed = seed_generator.randint(0, np.iinfo(np.uint32).max)
+            else:
+                cur_seed = 1337
+            if x == "train":
+                n = arg.n_train
+            elif x == "test":
+                n = arg.n_test
+            elif x == "validate":
+                n = arg.n_validate
+            else:
+                s = "Unknown mode: {}".format(x)
+                raise ValueError(s)
+            d = {
+                        'example_file': configs[k][x + "_file"],
+                        'reference_sequence': arg.genome_file,
+                        'batch_size': 512 if arg.parallel else 512*len(gpus),
+                        'seed': cur_seed,
+                        'shuffle': x=='train',
+                        'n_examples': n,
+                        'pad': 400
+                    }
+            if arg.parallel is True:
+                configs[k][x] = BatchedBioIntervalSequenceGenerator if x=='train' else BatchedBioIntervalSequence
+                configs[k]['%s_data_kwargs'%x] = d
+            else:
+                configs[k][x] = BatchedBioIntervalSequenceGenerator(**d) if x=='train' else BatchedBioIntervalSequence(**d)
+
+        # Build covariates and manager.
+        configs[k]["dfeatures"] = np.array(
+            [configs[k][x] for x in dfeature_names]) # TODO: Make cols dynamic.
+
+        tmp = dict() if arg.parallel is False else dict(
+                train_data_kwargs=configs[k]['train_data_kwargs'],
+                validate_data_kwargs=configs[k]["validate_data_kwargs"]
+                )
+        configs[k]["manager"] = manager_getter(
+            devices=[gpus_[i]],
+            train_data=configs[k]["train"],
+            val_data=configs[k]["validate"],
+            controller=controller,
+            model_space=model_space,
+            wd=os.path.join(wd, "manager_%s"%k),
+            data_description=configs[k]["dfeatures"],
+            dag_name="AmberDAG{}".format(k),
+            verbose=0 if arg.parallel is True else 2,
+            n_feats=configs[k]["n_feats"],
+            **tmp
+            )
+        config_keys.append(k)
+
+    logger = setup_logger(wd, verbose_level=logging.INFO)
+
+    # Setup env kwargs.
+    tmp = dict(data_descriptive_features=np.stack([configs[k]["dfeatures"] for k in config_keys]),
+               controller=controller,
+               manager=[configs[k]["manager"] for k in config_keys],
+               logger=logger,
+               max_episode=200,
+               max_step_per_ep=3,
+               working_dir=wd,
+               time_budget="150:00:00",
+               with_input_blocks=False,
+               with_skip_connection=False,
+               save_controller_every=1
+           )
+
+    if arg.parallel is True:
+        env = ParallelMultiManagerEnvironment(
+                    processes=len(gpus) if arg.parallel else 1,
+                    **tmp)
+    else:
+        env = MultiManagerEnvironment(**tmp)
+
+    try:
+        env.train()
+    except KeyboardInterrupt:
+        print("user interrupted training")
+        pass
+    controller.save_weights(os.path.join(wd, "controller_weights.h5"))
+
+
+if __name__ == "__main__":
+    if not run_from_ipython():
+        parser = argparse.ArgumentParser(description="experimental zero-shot nas")
+        parser.add_argument("--analysis", type=str, choices=['train', 'reload'], required=True, help="analysis type")
+        parser.add_argument("--wd", type=str, default="./outputs/zero_shot/", help="working dir")
+        parser.add_argument("--parallel", default=False, action="store_true", help="Use parallel")
+        parser.add_argument("--resume", default=False, action="store_true", help="resume previous run")
+        parser.add_argument("--config-file", type=str, required=True, help="Path to the config file to use.")
+        parser.add_argument("--genome-file", type=str, required=True, help="Path to genome file to use.")
+        parser.add_argument("--dfeature-name-file", type=str, required=True, help="Path to file with dataset feature names listed one per line.")
+        parser.add_argument("--n-test", type=int, required=True, help="Number of test examples.")
+        parser.add_argument("--n-train", type=int, required=True, help="Number of train examples.")
+        parser.add_argument("--n-validate", type=int, required=True, help="Number of validation examples.")
+        parser.add_argument("--lockstep-sampling", default=False, action="store_true", help="Ensure same training samples used for all models.")
+
+        arg = parser.parse_args()
+
+        if arg.analysis == "train":
+            train_nas(arg)
+        elif arg.analysis == "reload":
+            reload_trained_controller(arg)
+        else:
+            raise Exception("Unknown analysis type: %s"% arg.analysis)
diff --git a/examples/zero_shot_nas.sim_data.py b/examples/zero_shot_nas.sim_data.py
new file mode 100644
index 0000000..830ee06
--- /dev/null
+++ b/examples/zero_shot_nas.sim_data.py
@@ -0,0 +1,431 @@
+#!/usr/bin/env python
+
+"""
+Testing for ZeroShotNAS, aka data description driven nas
+ZZ
+May 14, 2020
+"""
+
+import tensorflow as tf
+import numpy as np
+import os
+import logging
+import pickle
+import pandas as pd
+import argparse
+
+from amber.modeler import EnasCnnModelBuilder
+from amber.architect.controller import ZeroShotController
+from amber.architect.model_space import State, ModelSpace
+from amber.architect.common_ops import count_model_params
+
+from amber.architect.manager import EnasManager
+from amber.architect.train_env import MultiManagerEnvironment
+from amber.architect.reward import LossAucReward
+from amber.plots import plot_controller_hidden_states
+from amber.utils import run_from_ipython
+from amber.utils.logging import setup_logger
+from amber.utils.data_parser import get_data_from_simdata
+
+from amber.modeler.modeler import build_sequential_model
+from amber.architect.manager import GeneralManager
+from amber.architect.model_space import get_layer_shortname
+
+from amber.bootstrap.mock_manager import MockManager
+
+import keras.backend as K
+from keras.optimizers import SGD, Adam
+
+
+def get_controller(model_space, session, layer_embedding_sharing, data_description_len=3, is_enas=True):
+    with tf.device("/cpu:0"):
+        controller = ZeroShotController(
+            data_description_config={
+                "length": 2,
+                "hidden_layer": {"units":8, "activation": "relu"},
+                "regularizer": {"l1":1e-8 }
+                },
+            model_space=model_space,
+            session=session,
+            share_embedding=layer_embedding_sharing,
+            with_skip_connection=False,
+            skip_weight=None,
+            skip_target=0.2,
+            lstm_size=32,
+            lstm_num_layers=1,
+            kl_threshold=0.05,
+            train_pi_iter=100,
+            optim_algo='adam',
+            temperature=1.5,
+            tanh_constant=2,
+            buffer_type="MultiManager",
+            buffer_size=1 if is_enas else 5,
+            batch_size=10,
+            use_ppo_loss=True,
+            rescale_advantage_by_reward=False
+        )
+    return controller
+
+
+def get_model_space_enas(out_filters=16, num_layers=4, num_pool=2):
+    state_space = ModelSpace()
+    expand_layers = [num_layers//num_pool*i-1 for i in range(num_pool)]
+    assert num_layers%2 == 0
+    layer_embedding_sharing = {}
+    for i in range(0, num_layers, 2):
+        state_space.add_layer(i, [
+            State('conv1d', filters=out_filters, kernel_size=8, activation='relu'),
+            State('conv1d', filters=out_filters, kernel_size=14, activation='relu'),
+            State('conv1d', filters=out_filters, kernel_size=20, activation='relu'),
+        ])
+        state_space.add_layer(i+1, [
+            State('maxpool1d', filters=out_filters, pool_size=8, strides=8),
+            State('avgpool1d', filters=out_filters, pool_size=8, strides=8)
+        ])
+        if i in expand_layers:
+            out_filters *= 2
+        if i>0:
+            layer_embedding_sharing[i] = 0
+            layer_embedding_sharing[i+1] = 1
+    return state_space, layer_embedding_sharing
+
+
+def get_model_space_common():
+    state_space = ModelSpace()
+    state_space.add_layer(0, [
+        State('conv1d', filters=3, kernel_size=8, kernel_initializer='glorot_uniform', activation='relu',
+              name="conv1"),
+        State('conv1d', filters=3, kernel_size=14, kernel_initializer='glorot_uniform', activation='relu',
+              name="conv1"),
+        State('conv1d', filters=3, kernel_size=20, kernel_initializer='glorot_uniform', activation='relu',
+              name="conv1"),
+    ])
+    state_space.add_layer(1, [
+        State('Identity'),
+        State('maxpool1d', pool_size=8, strides=8),
+        State('avgpool1d', pool_size=8, strides=8),
+
+    ])
+    state_space.add_layer(2, [
+        State('Flatten'),
+        State('GlobalMaxPool1D'),
+        State('GlobalAvgPool1D'),
+    ])
+    state_space.add_layer(3, [
+        State('Dense', units=3, activation='relu'),
+        State('Dense', units=10, activation='relu'),
+        State('Identity')
+    ])
+    return state_space, None
+
+
+
+def read_data():
+    dataset1 = get_data_from_simdata(
+        positive_file="./data/zero_shot_sim/DensityEmbedding_prefix-MYC_known10_motifs-MYC_known10_min-1_max-10_mean-5_zeroProb-0_seqLength-1000_numSeqs-10000.simdata",
+        negative_file="./data/zero_shot_sim/EmptyBackground_prefix-empty_bg_seqLength-1000_numSeqs-10000.simdata",
+        targets=["MYC"])
+    dataset2 = get_data_from_simdata(
+        positive_file="./data/zero_shot_sim/DensityEmbedding_prefix-CTCF_known1_motifs-CTCF_known1_min-1_max-1_mean-1_zeroProb-0_seqLength-1000_numSeqs-10000.simdata",
+        negative_file="./data/zero_shot_sim/EmptyBackground_prefix-empty_bg_seqLength-1000_numSeqs-10000.simdata",
+        targets=["CTCF"])
+
+    num_seqs = 20000
+    train_idx = np.arange(0, int(num_seqs*0.8))
+    val_idx = np.arange(int(num_seqs*0.8), int(num_seqs*0.9) )
+    test_idx = np.arange(int(num_seqs*0.9), num_seqs )
+    dictionarize = lambda dataset: {
+            "train": (dataset[0][train_idx], dataset[1][train_idx]),
+            "val": (dataset[0][val_idx], dataset[1][val_idx]),
+            "test": (dataset[0][test_idx], dataset[1][test_idx]),
+            }
+
+    dataset1 = dictionarize(dataset1)
+    dataset2 = dictionarize(dataset2)
+
+    return dataset1, dataset2
+
+
+def get_manager_enas(train_data, val_data, controller, model_space, wd, data_description, dag_name, verbose=2):
+    input_node = State('input', shape=(1000, 4), name="input", dtype='float32')
+    output_node = State('dense', units=1, activation='sigmoid')
+    model_compile_dict = {
+        'loss': 'binary_crossentropy',
+        'optimizer': SGD(lr=0.01, momentum=0.9, decay=1e-5),
+    }
+    session = controller.session
+
+    reward_fn = LossAucReward(method='auc')
+    
+    child_batch_size = 512
+    model_fn = EnasCnnModelBuilder(
+        dag_func='EnasConv1DwDataDescrption',
+        batch_size=child_batch_size,
+        session=session,
+        model_space=model_space,
+        inputs_op=[input_node],
+        output_op=[output_node],
+        num_layers=len(model_space),
+        l1_reg=1e-8,
+        l2_reg=5e-7,
+        model_compile_dict=model_compile_dict,
+        controller=controller,
+        dag_kwargs={
+            'with_skip_connection': False,
+            'add_conv1_under_pool': False,
+            'stem_config':{
+                'has_stem_conv': False,
+                'flatten_op': 'flatten',
+                'fc_units': 30
+                },
+            'name': dag_name,
+            'data_description': data_description
+            }
+    )
+    
+    manager = EnasManager(
+        train_data=train_data,
+        validation_data=val_data,
+        epochs=1,
+        child_batchsize=child_batch_size,
+        reward_fn=reward_fn,
+        model_fn=model_fn,
+        store_fn='minimal',
+        model_compile_dict=model_compile_dict,
+        working_dir=wd,
+        verbose=verbose
+        )
+    return manager
+
+
+def get_manager_common(train_data, val_data, controller, model_space, wd, data_description, verbose=2, **kwargs):
+    input_node = State('input', shape=(1000, 4), name="input", dtype='float32')
+    output_node = State('dense', units=1, activation='sigmoid')
+    model_compile_dict = {
+        'loss': 'binary_crossentropy',
+        'optimizer': 'adam',
+        'metrics': ['acc']
+    }
+
+    reward_fn = LossAucReward(method='auc')
+
+    child_batch_size = 500
+    # TODO: convert functions in `_keras_modeler.py` to classes, and wrap up this Lambda function step
+    model_fn = lambda model_arc: build_sequential_model(
+        model_states=model_arc, input_state=input_node, output_state=output_node, model_compile_dict=model_compile_dict,
+        model_space=model_space)
+    manager = GeneralManager(
+        train_data=train_data,
+        validation_data=val_data,
+        epochs=30,
+        child_batchsize=child_batch_size,
+        reward_fn=reward_fn,
+        model_fn=model_fn,
+        store_fn='model_plot',
+        model_compile_dict=model_compile_dict,
+        working_dir=wd,
+        verbose=0,
+        save_full_model=False,
+        model_space=model_space
+    )
+    return manager
+
+
+def get_manager_mock(dataset, model_space, **kwargs):
+    if dataset in [1,2]:
+        model_compile_dict = {
+            'loss': 'binary_crossentropy',
+            'optimizer': 'adam',
+            'metrics': ['acc']
+        }
+        
+        def mock_reward_fn(model_arc, train_history_df, *args, **kwargs):
+            model_states_ = [str(model_space[i][model_arc[i]]) for i in range(len(model_arc))]
+            idx_bool = np.array([train_history_df['L%i' % (i + 1)] == model_states_[i] for i in range(len(model_states_))])
+            index = np.apply_along_axis(func1d=lambda x: all(x), axis=0, arr=idx_bool)
+            l, k = train_history_df[['loss', 'knowledge']].iloc[np.random.choice(np.where(index)[0])]
+            this_reward =  l 
+            loss_and_metrics = [l, k]
+            reward_metrics = {'knowledge': k}
+            return this_reward, loss_and_metrics, reward_metrics
+
+        manager = MockManager(
+                #history_fn_list=['./outputs/zs_grid_history/dataset%i/zs_grid_%i/train_history.csv'%(dataset, i)
+                #    for i in range(1,3)],
+                history_fn_list = ['./outputs/zs_hist/dataset%i/train_history.csv'%dataset ],
+                model_compile_dict=model_compile_dict,
+                #metric_name_dict={"acc":0, "knowledge":1, "loss":2},
+                metric_name_dict={"knowledge":0, "loss":1},
+                reward_fn=mock_reward_fn
+                )
+    else:
+        raise Exception("Unknown dataset: %s"%dataset)
+    return manager
+
+
+def get_samples_controller(dfeatures, controller, model_space, T=100, is_enas=True):
+    res = []
+    for i in range(len(dfeatures)):
+        dfeature = dfeatures[[i]]
+        prob_arr = [ np.zeros((T, len(model_space[i]))) for i in range(len(model_space)) ]
+        if is_enas:
+            skip_layers = [i*2 for i in range(1, len(model_space))]
+        else:
+            skip_layers = []
+        for t in range(T):
+            a, p = controller.get_action(dfeature)
+            layer_id = 0
+            for i, p in enumerate(p):
+                if i in skip_layers:
+                    continue
+                prob_arr[layer_id][t] = p.flatten()
+                layer_id += 1
+        res.append(prob_arr)
+    return res
+
+
+def convert_to_dataframe(res, model_space, data_names):
+    probs = []
+    layer = []
+    description = []
+    operation = []
+    for i in range(len(res)):
+        for j in range(len(model_space)):
+            for k in range(len(model_space[j])):
+                o = get_layer_shortname(model_space[j][k])
+                p = res[i][j][:, k]
+                # extend the array
+                T = p.shape[0]
+                probs.extend(p)
+                description.extend([data_names[i]]*T)
+                layer.extend([j]*T)
+                operation.extend([o]*T)
+    df = pd.DataFrame({
+        'description':description, 
+        'layer': layer,
+        'operation': operation,
+        'prob': probs
+        })
+    return df
+
+
+def reload_trained_controller(arg):
+    wd = arg.wd #wd = "./outputs/zero_shot/"
+    model_space, layer_embedding_sharing = get_model_space_common() if arg.mode!="enas" else \
+            get_model_space_enas(out_filters=4, num_layers=3, num_pool=3) 
+    try:
+        session = tf.Session()
+    except:
+        session = tf.compat.v1.Session()
+
+    controller = get_controller(model_space=model_space, session=session, data_description_len=2,
+            layer_embedding_sharing=layer_embedding_sharing,
+            is_enas=arg.mode=='enas')
+    controller.load_weights(os.path.join(wd, "controller_weights.h5"))
+
+    dfeatures = np.array([[1,0], [0,1]])  # one-hot encoding
+    res = get_samples_controller(dfeatures, controller, model_space, T=1000, is_enas=arg.mode=='enas')
+   
+    import seaborn as sns
+    import matplotlib.pyplot as plt
+    df = convert_to_dataframe(res, model_space, data_names=['MYC_known10', 'CTCF_known1'])
+
+    for i in range(len(model_space)):
+        sub_df = df.loc[ (df.layer==i) ]
+        plt.clf()
+        plt.tight_layout()
+        ax = sns.boxplot(x="operation", y="prob",
+            hue="description", palette=["m", "g"],
+            data=sub_df)
+        ax.set_xticklabels(ax.get_xticklabels(), rotation=45)
+        plt.savefig(os.path.join(wd, "layer_%i.png"%i), bbox_inches="tight")
+    return res
+
+
+def train_nas(arg):
+    wd = arg.wd #wd = "./outputs/zero_shot/"
+    logger = setup_logger(wd, verbose_level=logging.INFO)
+    verbose = 2
+    model_space, layer_embedding_sharing = get_model_space_enas(out_filters=16, num_layers=4, num_pool=2) if arg.mode == 'enas' else \
+            get_model_space_common()
+    try:
+        session = tf.Session()
+    except:
+        session = tf.compat.v1.Session()
+    K.set_session(session)
+    controller = get_controller(model_space=model_space, session=session, layer_embedding_sharing=layer_embedding_sharing,
+            data_description_len=2,
+            is_enas=arg.mode=='enas')
+
+    dfeatures = np.array([[1,0], [0,1]])  # one-hot encoding
+    if arg.mode == 'enas':
+        dataset1, dataset2 = read_data()
+        manager1 = get_manager_enas(train_data=dataset1['train'], val_data=dataset1['val'], controller=controller,
+                model_space=model_space, wd=wd, data_description=dfeatures[[0]], dag_name="EnasDAG1", verbose=verbose)
+        manager2 = get_manager_enas(train_data=dataset2['train'], val_data=dataset2['val'], controller=controller,
+                model_space=model_space, wd=wd, data_description=dfeatures[[1]], dag_name="EnasDAG2", verbose=verbose)
+        # only for enas
+        vars_list1 = [v for v in tf.trainable_variables() if v.name.startswith(manager1.model_fn.dag.name)]
+        vars_list2 = [v for v in tf.trainable_variables() if v.name.startswith(manager2.model_fn.dag.name)]
+        # remove optimizer related vars (e.g. momentum, rms)
+        vars_list1 = [v for v in vars_list1 if not v.name.startswith("%s/compile"%manager1.model_fn.dag.name)]
+        vars_list2 = [v for v in vars_list2 if not v.name.startswith("%s/compile"%manager2.model_fn.dag.name)]
+        logger.info("total model1 params: %i" % count_model_params(vars_list1))
+        logger.info("total model2 params: %i" % count_model_params(vars_list2))
+        with open(os.path.join(wd,"tensor_vars.txt"), "w") as f:
+            for v in vars_list1 + vars_list2:
+                f.write("%s\t%i\n"%(v.name, int(np.prod(v.shape).value) ))
+
+    elif arg.mode == 'vanilla':
+        dataset1, dataset2 = read_data()
+        manager1 = get_manager_common(train_data=dataset1['train'], val_data=dataset1['val'], controller=controller,
+                model_space=model_space, wd=wd, data_description=dfeatures[[0]], dag_name="EnasDAG1", verbose=verbose)
+        manager2 = get_manager_common(train_data=dataset2['train'], val_data=dataset2['val'], controller=controller,
+                model_space=model_space, wd=wd, data_description=dfeatures[[1]], dag_name="EnasDAG2", verbose=verbose)
+    elif arg.mode == 'bootstrap':
+        manager1 = get_manager_mock(dataset=1, model_space=model_space)
+        manager2 = get_manager_mock(dataset=2, model_space=model_space)
+    
+    else:
+        raise Exception("Cannot understand mode: %s"% arg.mode)
+
+
+    env = MultiManagerEnvironment(
+        data_descriptive_features= dfeatures,
+        controller=controller,
+        manager=[manager1, manager2],
+        logger=logger,
+        max_episode=200,
+        max_step_per_ep=15,
+        working_dir=wd,
+        time_budget="8:00:00",
+        with_input_blocks=False,
+        with_skip_connection=False,
+        child_warm_up_epochs=2 if arg.mode=='enas' else 0
+    )
+
+    try:
+        env.train()
+    except KeyboardInterrupt:
+        print("user interrupted training")
+        pass
+    controller.save_weights(os.path.join(wd, "controller_weights.h5"))
+
+
+if __name__ == "__main__":
+    if not run_from_ipython():
+        parser = argparse.ArgumentParser(description="experimental zero-shot nas")
+        parser.add_argument("--analysis", type=str, choices=['train', 'reload'], required=True, help="analysis type")
+        parser.add_argument("--mode", type=str, choices=['enas', 'vanilla', 'bootstrap'], required=True, help="amber mode")
+        parser.add_argument("--wd", type=str, default="./outputs/zero_shot/", help="working dir")
+
+        arg = parser.parse_args()
+
+        os.makedirs(arg.wd, exist_ok=True)
+        if arg.analysis == "train":
+            train_nas(arg)
+        elif arg.analysis == "reload":
+            reload_trained_controller(arg)
+        else:
+            raise Exception("Unknown analysis type: %s"% arg.analysis)
+
diff --git a/examples/zeroshot_grid_searcher.py b/examples/zeroshot_grid_searcher.py
new file mode 100644
index 0000000..a653b78
--- /dev/null
+++ b/examples/zeroshot_grid_searcher.py
@@ -0,0 +1,96 @@
+# -*- coding: UTF-8 -*-
+
+from amber.utils import run_from_ipython
+from amber.bootstrap.grid_search import grid_search
+from amber.objective.motif import MotifKLDivergence
+from amber.modeler import KerasModelBuilder
+from amber.architect.reward import KnowledgeReward, LossAucReward
+from amber.architect.model_space import State
+from amber.architect.manager import GeneralManager
+from zero_shot_nas import get_model_space_common as get_model_space
+#from zero_shot_nas import get_manager_common as get_manager
+from zero_shot_nas import read_data
+import pickle
+import sys
+
+
+def get_manager(train_data, val_data, controller, model_space, wd, motif_name, verbose=2, **kwargs):
+    input_node = State('input', shape=(1000, 4), name="input", dtype='float32')
+    output_node = State('dense', units=1, activation='sigmoid')
+    model_compile_dict = {
+        'loss': 'binary_crossentropy',
+        'optimizer': 'adam',
+        'metrics': ['acc']
+    }
+
+    knowledge_fn = MotifKLDivergence(temperature=0.1, Lambda_regularizer=0.0)
+    knowledge_fn.knowledge_encoder(
+            motif_name_list=[motif_name],
+            motif_file="/mnt/home/zzhang/workspace/src/BioNAS/BioNAS/resources/rbp_motif/encode_motifs.txt.gz",
+            is_log_motif=False
+            )
+    reward_fn = LossAucReward(method='auc', knowledge_function=knowledge_fn)
+    #reward_fn = KnowledgeReward(knowledge_function=knowledge_fn, Lambda=2)
+
+    child_batch_size = 500
+    #model_fn = lambda model_arc: build_sequential_model(
+    #    model_states=model_arc, input_state=input_node, output_state=output_node, model_compile_dict=model_compile_dict,
+    #    model_space=model_space)
+    model_fn = KerasModelBuilder(inputs=input_node, outputs=output_node, model_compile_dict=model_compile_dict,
+            model_space=model_space)
+    manager = GeneralManager(
+        train_data=train_data,
+        validation_data=val_data,
+        epochs=200,
+        child_batchsize=child_batch_size,
+        reward_fn=reward_fn,
+        model_fn=model_fn,
+        store_fn='model_plot',
+        model_compile_dict=model_compile_dict,
+        working_dir=wd,
+        verbose=0,
+        save_full_model=True,
+        model_space=model_space
+    )
+    return manager
+
+
+def main(arg):
+    model_space = get_model_space()
+
+    # read the data
+    dataset1, dataset2 = read_data()
+    if arg.dataset == 1:
+        train_data, validation_data = dataset1['train'], dataset1['val']
+    elif arg.dataset == 2:
+        train_data, validation_data = dataset2['train'], dataset2['val']
+    else:
+        raise Exception("Unknown dataset: %s" % arg.dataset)
+
+    # init network manager
+    manager = get_manager(
+            train_data=train_data,
+            val_data=validation_data,
+            controller=None,
+            model_space=model_space,
+            motif_name="MYC_known10" if arg.dataset==1 else "CTCF_known1",
+            wd=arg.wd)
+
+    # grid search
+    grid_search(model_space, manager, arg.wd, B=arg.B)
+
+
+if __name__ == '__main__':
+    if not run_from_ipython():
+        import argparse
+        parser = argparse.ArgumentParser()
+        parser.add_argument("--wd", help="working dir")
+        parser.add_argument("--dataset", type=int, choices=[1,2], help="dataset choice")
+        parser.add_argument("--B", type=int, default=10, help="run each model architecture for B times")
+
+        arg = parser.parse_args()
+
+        main(arg)
+
+    
+
diff --git a/setup.py b/setup.py
index 8149967..818c03d 100644
--- a/setup.py
+++ b/setup.py
@@ -9,10 +9,10 @@
     'package_data': {'deepambre': ['./amber/resources/gui/*']},
     'setup_requires': [],
     'install_requires': [
-        'numpy', 
+        'numpy',
+        'tensorflow',
         'matplotlib', 
         'scipy', 
-        'tensorflow>=1.9.0,<2.0.0', 
         'keras>=2.0.0', 
         'seaborn>=0.9.0',
         'networkx'


Test patch:
diff --git a/amber/tests/test_Controller_Manager.py b/amber/tests/test_Controller_Manager.py
deleted file mode 100644
index ebabb4f..0000000
--- a/amber/tests/test_Controller_Manager.py
+++ /dev/null
@@ -1,5 +0,0 @@
-# -*- coding: UTF-8 -*-
-
-from BioNAS.Controller.manager import *
-from BioNAS.Controller.model import *
-from BioNAS.Controller.model_space import *
diff --git a/amber/tests/test_KFunctions.py b/amber/tests/test_KFunctions.py
deleted file mode 100644
index 3b15705..0000000
--- a/amber/tests/test_KFunctions.py
+++ /dev/null
@@ -1,29 +0,0 @@
-# -*- coding: UTF-8 -*-
-import unittest
-
-import matplotlib.pyplot as plt
-from BioNAS.KFunctions.K_func import *
-from BioNAS.resources.eclip import cnnModel
-from BioNAS.utils import motif
-
-
-class MotifKfunction_Test(unittest.TestCase):
-
-    def test_MKF(self):
-        mkf = Motif_K_function(temperature=0.1, Lambda_regularizer=0.01)
-        mkf.knowledge_encoder(['RBFOX2_v1', 'RBFOX2_v2', 'RBFOX2_v3'], '../resources/rbp_motif/human_rbp_pwm.txt', True)
-
-        model = cnnModel.build_model()
-        print(mkf(model, None))
-        model.load_weights('../resources/eclip/bestmodel.h5')
-        print(mkf(model, None))
-
-        score_dict, weight_dict = mkf.get_matched_model_weights()
-        motif.draw_dnalogo_matplot(weight_dict['RBFOX2_v1'])
-        plt.close()
-        motif.draw_dnalogo_matplot(mkf.W_knowledge['RBFOX2_v1'])
-        plt.close()
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/amber/tests/test_SFC.py b/amber/tests/test_SFC.py
deleted file mode 100644
index c050e9d..0000000
--- a/amber/tests/test_SFC.py
+++ /dev/null
@@ -1,16 +0,0 @@
-# -*- coding: UTF-8 -*-
-import unittest
-
-from BioNAS.operators.convolutional import SeparableFC
-from keras.layers import Dense, Conv1D
-from keras.models import Sequential
-
-
-class SFC_Test(unittest.TestCase):
-
-    def test_SFC_build(self):
-        model = Sequential()
-        model.add(Conv1D(filters=3, kernel_size=4, input_shape=(200, 4)))
-        model.add(SeparableFC(output_dim=10, symmetric=True))
-        model.add(Dense(1, activation='sigmoid'))
-        model.compile(loss='binary_crossentropy', metrics=['acc'])
diff --git a/amber/tests/test_StateSpace.py b/amber/tests/test_StateSpace.py
deleted file mode 100644
index 4a8c8d0..0000000
--- a/amber/tests/test_StateSpace.py
+++ /dev/null
@@ -1,48 +0,0 @@
-# -*- coding: UTF-8 -*-
-import unittest
-
-from BioNAS.Controller.model import *
-from BioNAS.Controller.model_space import *
-
-
-class StateSpace_Test(unittest.TestCase):
-
-    def test_StateSpace(self):
-        state_space = ModelSpace()
-
-        # add layer by individual states
-        state_space.add_layer(0)
-        state_space.add_state(0, State('conv1d', filters=10, kernel_size=10))
-        s1 = State('conv1d', filters=3, kernel_size=10)
-        state_space.add_state(0, s1)
-        s2 = State('conv1d', filters=3, kernel_size=15)
-        state_space.add_state(0, s2)
-
-        # add layer by entire list
-        state_space.add_layer(1, [State('maxpool1d', pool_size=10, strides=10),
-                                  State('avgpool1d', pool_size=10, strides=10), ])
-        state_space[2] = [State('Flatten'), State('GlobalMaxPool1D')]
-
-        # testing prints
-        print(state_space)
-        print('length=', len(state_space))
-        state_space.print_state_space()
-        state_space.get_random_model_states()
-
-    def test_ModelBuild_from_StateSpace(self):
-        state_space = ModelSpace()
-        state_space.add_layer(0, [State('conv1d', filters=3, kernel_size=10),
-                                  State('conv1d', filters=10, kernel_size=10), ])
-        state_space.add_layer(1, [State('maxpool1d', pool_size=10, strides=10),
-                                  State('avgpool1d', pool_size=10, strides=10), ])
-        model_states = state_space.get_random_model_states()
-
-        input_state = State('Input', shape=(200, 4))
-        output_state = State('Dense', units=1, activation='sigmoid')
-
-        model_compile_dict = {'loss': 'binary_crossentropy', 'optimizer': 'adam', 'metrics': ['acc']}
-        model = build_sequential_model(model_states, input_state, output_state, model_compile_dict)
-
-
-if __name__ == '__main__':
-    unittest.main()
diff --git a/amber/utils/tests/files/small_genome.encoded.h5 b/amber/utils/tests/files/small_genome.encoded.h5
new file mode 100644
index 0000000..162d64e
Binary files /dev/null and b/amber/utils/tests/files/small_genome.encoded.h5 differ
diff --git a/amber/utils/tests/files/small_genome.fa b/amber/utils/tests/files/small_genome.fa
new file mode 100644
index 0000000..ef39046
--- /dev/null
+++ b/amber/utils/tests/files/small_genome.fa
@@ -0,0 +1,6 @@
+>seq0
+aAnNcCtTgG
+>seq1
+AAAAACCCCCTTTTTGGGGGNNNNN
+>seq2
+ACTGACTGACTGACTGACTGACTG
diff --git a/amber/utils/tests/files/small_genome.fa.fai b/amber/utils/tests/files/small_genome.fa.fai
new file mode 100644
index 0000000..4cc9ce7
--- /dev/null
+++ b/amber/utils/tests/files/small_genome.fa.fai
@@ -0,0 +1,3 @@
+seq0	10	6	10	11
+seq1	25	23	25	26
+seq2	24	55	24	25
diff --git a/amber/utils/tests/files/small_genome.h5 b/amber/utils/tests/files/small_genome.h5
new file mode 100644
index 0000000..d8bfdfa
Binary files /dev/null and b/amber/utils/tests/files/small_genome.h5 differ
diff --git a/amber/utils/tests/test_encoded_genome.py b/amber/utils/tests/test_encoded_genome.py
new file mode 100644
index 0000000..c03098a
--- /dev/null
+++ b/amber/utils/tests/test_encoded_genome.py
@@ -0,0 +1,144 @@
+import numpy
+import unittest
+
+from amber.utils.sequences import EncodedGenome
+
+
+class TestEncodedGenome(unittest.TestCase):
+    def setUp(self):
+        self.bases = ['A', 'C', 'G', 'T']
+        self.bases_to_arr = dict(A=numpy.array([1, 0, 0, 0]),
+                                 C=numpy.array([0, 1, 0, 0]),
+                                 G=numpy.array([0, 0, 1, 0]),
+                                 T=numpy.array([0, 0, 0, 1]),
+                                 N=numpy.array([.25, .25, .25, .25]))
+
+        self.chrom_to_lens = {"seq0": 10,
+                              "seq1": 25,
+                              "seq2": 24}
+        self.in_memory = False
+
+    def _get_small_genome(self):
+        return EncodedGenome("amber/utils/tests/files/small_genome.fa", in_memory=self.in_memory)
+
+    def test_load_mixed_case_sequence(self):
+        expected = [
+                [1, 0, 0, 0],
+                [1, 0, 0, 0],
+                [.25, .25, .25, .25],
+                [.25, .25, .25, .25],
+                [0, 1, 0, 0],
+                [0, 1, 0, 0],
+                [0, 0, 0, 1],
+                [0, 0, 0, 1],
+                [0, 0, 1, 0],
+                [0, 0, 1, 0]
+            ]
+        expected = numpy.array(expected)
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq0", 0, 10)
+        self.assertSequenceEqual(expected.tolist(), observed.tolist())
+
+
+    def test_load_sequence0(self):
+        expected = [[0, 1, 0, 0]]
+        expected = numpy.array(expected)
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq0", 5, 6)
+        self.assertSequenceEqual(expected.tolist(), observed.tolist())
+
+    def test_load_sequence1(self):
+        expected = [[0, 0, 0, 1],
+                    [0, 0, 1, 0]]
+        expected = numpy.array(expected)
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq2", 22, 24)
+        self.assertSequenceEqual(expected.tolist(), observed.tolist())
+
+    def test_load_rc(self):
+        expected = "CCAAGGNNTT"
+        expected = [[0, 1, 0, 0],
+                    [0, 1, 0, 0],
+                    [1, 0, 0, 0],
+                    [1, 0, 0, 0],
+                    [0, 0, 1, 0],
+                    [0, 0, 1, 0],
+                    [.25, .25, .25, .25],
+                    [.25, .25, .25, .25],
+                    [0, 0, 0, 1],
+                    [0, 0, 0, 1]]
+        expected = numpy.array(expected)
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq0", 0, 10, "-")
+        self.assertSequenceEqual(expected.tolist(), observed.tolist())
+
+    def test_coords_flipped(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 2, 1)
+        self.assertFalse(observed)
+
+    def test_end_coord_too_large(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 9, 11)
+        self.assertFalse(observed)
+
+    def test_both_coords_too_large(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 11, 14)
+        self.assertFalse(observed)
+
+    def test_coords_too_large_negative(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", -100, -99)
+        self.assertFalse(observed)
+
+    def test_coords_negative_flipped(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", -1, -2)
+        self.assertFalse(observed)
+
+    def test_coords_negative(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", -2, -1)
+        self.assertFalse(observed)
+
+    def test_bad_chrom(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq-1", 1, 2)
+        self.assertFalse(observed)
+
+    def test_too_small(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 1, 1)
+        self.assertFalse(observed)
+
+    def test_first_coord(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 0, 1)
+        self.assertTrue(observed)
+
+    def test_last_coord(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 9, 10)
+        self.assertTrue(observed)
+
+    def test_in_memory(self):
+        self.assertFalse(self.in_memory)
+
+    def test_genome_in_memory(self):
+        g = self._get_small_genome()
+        self.assertFalse(g.in_memory)
+
+
+class TestEncodedGenomeInMemory(TestEncodedGenome):
+    def setUp(self):
+        super(TestEncodedGenomeInMemory, self).setUp()
+        self.in_memory = True
+
+    def test_in_memory(self):
+        self.assertTrue(self.in_memory)
+
+    def test_genome_in_memory(self):
+        g = self._get_small_genome()
+        self.assertTrue(g.in_memory)
+
diff --git a/amber/utils/tests/test_encoded_hdf5_genome.py b/amber/utils/tests/test_encoded_hdf5_genome.py
new file mode 100644
index 0000000..afb5338
--- /dev/null
+++ b/amber/utils/tests/test_encoded_hdf5_genome.py
@@ -0,0 +1,144 @@
+import numpy
+import unittest
+
+from amber.utils.sequences import EncodedHDF5Genome
+
+
+class TestEncodedHDF5Genome(unittest.TestCase):
+    def setUp(self):
+        self.bases = ['A', 'C', 'G', 'T']
+        self.bases_to_arr = dict(A=numpy.array([1, 0, 0, 0]),
+                                 C=numpy.array([0, 1, 0, 0]),
+                                 G=numpy.array([0, 0, 1, 0]),
+                                 T=numpy.array([0, 0, 0, 1]),
+                                 N=numpy.array([.25, .25, .25, .25]))
+
+        self.chrom_to_lens = {"seq0": 10,
+                              "seq1": 25,
+                              "seq2": 24}
+        self.in_memory = False
+
+    def _get_small_genome(self):
+        return EncodedHDF5Genome("amber/utils/tests/files/small_genome.encoded.h5", in_memory=self.in_memory)
+
+    def test_load_mixed_case_sequence(self):
+        expected = [
+                [1, 0, 0, 0],
+                [1, 0, 0, 0],
+                [.25, .25, .25, .25],
+                [.25, .25, .25, .25],
+                [0, 1, 0, 0],
+                [0, 1, 0, 0],
+                [0, 0, 0, 1],
+                [0, 0, 0, 1],
+                [0, 0, 1, 0],
+                [0, 0, 1, 0]
+            ]
+        expected = numpy.array(expected)
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq0", 0, 10)
+        self.assertSequenceEqual(expected.tolist(), observed.tolist())
+
+
+    def test_load_sequence0(self):
+        expected = [[0, 1, 0, 0]]
+        expected = numpy.array(expected)
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq0", 5, 6)
+        self.assertSequenceEqual(expected.tolist(), observed.tolist())
+
+    def test_load_sequence1(self):
+        expected = [[0, 0, 0, 1],
+                    [0, 0, 1, 0]]
+        expected = numpy.array(expected)
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq2", 22, 24)
+        self.assertSequenceEqual(expected.tolist(), observed.tolist())
+
+    def test_load_rc(self):
+        expected = "CCAAGGNNTT"
+        expected = [[0, 1, 0, 0],
+                    [0, 1, 0, 0],
+                    [1, 0, 0, 0],
+                    [1, 0, 0, 0],
+                    [0, 0, 1, 0],
+                    [0, 0, 1, 0],
+                    [.25, .25, .25, .25],
+                    [.25, .25, .25, .25],
+                    [0, 0, 0, 1],
+                    [0, 0, 0, 1]]
+        expected = numpy.array(expected)
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq0", 0, 10, "-")
+        self.assertSequenceEqual(expected.tolist(), observed.tolist())
+
+    def test_coords_flipped(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 2, 1)
+        self.assertFalse(observed)
+
+    def test_end_coord_too_large(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 9, 11)
+        self.assertFalse(observed)
+
+    def test_both_coords_too_large(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 11, 14)
+        self.assertFalse(observed)
+
+    def test_coords_too_large_negative(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", -100, -99)
+        self.assertFalse(observed)
+
+    def test_coords_negative_flipped(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", -1, -2)
+        self.assertFalse(observed)
+
+    def test_coords_negative(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", -2, -1)
+        self.assertFalse(observed)
+
+    def test_bad_chrom(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq-1", 1, 2)
+        self.assertFalse(observed)
+
+    def test_too_small(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 1, 1)
+        self.assertFalse(observed)
+
+    def test_first_coord(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 0, 1)
+        self.assertTrue(observed)
+
+    def test_last_coord(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 9, 10)
+        self.assertTrue(observed)
+
+    def test_in_memory(self):
+        self.assertFalse(self.in_memory)
+
+    def test_genome_in_memory(self):
+        g = self._get_small_genome()
+        self.assertFalse(g.in_memory)
+
+
+class TestEncodedHDF5GenomeInMemory(TestEncodedHDF5Genome):
+    def setUp(self):
+        super(TestEncodedHDF5GenomeInMemory, self).setUp()
+        self.in_memory = True
+
+    def test_in_memory(self):
+        self.assertTrue(self.in_memory)
+
+    def test_genome_in_memory(self):
+        g = self._get_small_genome()
+        self.assertTrue(g.in_memory)
+
diff --git a/amber/utils/tests/test_genome.py b/amber/utils/tests/test_genome.py
new file mode 100644
index 0000000..787e7ad
--- /dev/null
+++ b/amber/utils/tests/test_genome.py
@@ -0,0 +1,141 @@
+import numpy
+import unittest
+
+from amber.utils.sequences import Genome
+
+class TestGenome(unittest.TestCase):
+    def setUp(self):
+        self.bases = ['A', 'C', 'G', 'T']
+        self.bases_to_arr = dict(A=numpy.array([1, 0, 0, 0]),
+                                 C=numpy.array([0, 1, 0, 0]),
+                                 G=numpy.array([0, 0, 1, 0]),
+                                 T=numpy.array([0, 0, 0, 1]),
+                                 N=numpy.array([.25, .25, .25, .25]))
+
+        self.chrom_to_lens = {"seq0": 10,
+                              "seq1": 25,
+                              "seq2": 24}
+        self.in_memory = False
+
+    def _get_small_genome(self):
+        return Genome("amber/utils/tests/files/small_genome.fa", in_memory=self.in_memory)
+
+    def test_load_mixed_case_sequence(self):
+        expected = "AANNCCTTGG"
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq0", 0, 10)
+        self.assertEqual(expected, observed)
+
+    def test_load_rc(self):
+        expected = "CCAAGGNNTT"
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq0", 0, 10, "-")
+        self.assertEqual(expected, observed)
+
+    def test_load_sequence0(self):
+        expected = "C"
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq1", 5, 6)
+        self.assertEqual(expected, observed)
+
+    def test_load_sequence1(self):
+        expected = "A"
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq2", 0, 1)
+        self.assertEqual(expected, observed)
+
+    def test_load_sequence2(self):
+        expected = ("A" * 5) + ("C" * 5) + ("T" * 5 ) + ("G" * 5) + ("N" * 5)
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq1", 0, 25)
+        self.assertEqual(expected, observed)
+
+    def test_load_sequence3(self):
+        expected = "ACTG" * 6
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq2", 0, 24)
+        self.assertEqual(expected, observed)
+
+    def test_load_sequence4(self):
+        expected = ("T" * 5 ) + ("G" * 5) + ("N" * 5)
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq1", 10, 25)
+        self.assertEqual(expected, observed)
+
+    def test_length(self):
+        expected = sum(self.chrom_to_lens.values())
+        g = self._get_small_genome()
+        observed = len(g)
+        self.assertEqual(expected, observed)
+
+    def test_coords_flipped(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 2, 1)
+        self.assertFalse(observed)
+
+    def test_end_coord_too_large(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 9, 11)
+        self.assertFalse(observed)
+
+    def test_both_coords_too_large(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 11, 14)
+        self.assertFalse(observed)
+
+    def test_coords_too_large_negative(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", -100, -99)
+        self.assertFalse(observed)
+
+    def test_coords_negative_flipped(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", -1, -2)
+        self.assertFalse(observed)
+
+    def test_coords_negative(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", -2, -1)
+        self.assertFalse(observed)
+
+    def test_bad_chrom(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq-1", 1, 2)
+        self.assertFalse(observed)
+
+    def test_too_small(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 1, 1)
+        self.assertFalse(observed)
+
+    def test_first_coord(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 0, 1)
+        self.assertTrue(observed)
+
+    def test_last_coord(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 9, 10)
+        self.assertTrue(observed)
+
+    def test_in_memory(self):
+        self.assertFalse(self.in_memory)
+
+    def test_genome_in_memory(self):
+        g = self._get_small_genome()
+        self.assertFalse(g.in_memory)
+
+
+class TestGenomeInMemory(TestGenome):
+    def setUp(self):
+        super(TestGenomeInMemory, self).setUp()
+        self.in_memory = True
+
+    def test_in_memory(self):
+        self.assertTrue(self.in_memory)
+
+    def test_genome_in_memory(self):
+        g = self._get_small_genome()
+        self.assertTrue(g.in_memory)
+
+
diff --git a/amber/utils/tests/test_hdf5_genome.py b/amber/utils/tests/test_hdf5_genome.py
new file mode 100644
index 0000000..e3850ea
--- /dev/null
+++ b/amber/utils/tests/test_hdf5_genome.py
@@ -0,0 +1,141 @@
+import numpy
+import unittest
+
+from amber.utils.sequences import HDF5Genome
+
+class TestHDF5Genome(unittest.TestCase):
+    def setUp(self):
+        self.bases = ['A', 'C', 'G', 'T']
+        self.bases_to_arr = dict(A=numpy.array([1, 0, 0, 0]),
+                                 C=numpy.array([0, 1, 0, 0]),
+                                 G=numpy.array([0, 0, 1, 0]),
+                                 T=numpy.array([0, 0, 0, 1]),
+                                 N=numpy.array([.25, .25, .25, .25]))
+
+        self.chrom_to_lens = {"seq0": 10,
+                              "seq1": 25,
+                              "seq2": 24}
+        self.in_memory = False
+
+    def _get_small_genome(self):
+        return HDF5Genome("amber/utils/tests/files/small_genome.h5", in_memory=self.in_memory)
+
+    def test_load_mixed_case_sequence(self):
+        expected = "AANNCCTTGG"
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq0", 0, 10)
+        self.assertEqual(expected, observed)
+
+    def test_load_rc(self):
+        expected = "CCAAGGNNTT"
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq0", 0, 10, "-")
+        self.assertEqual(expected, observed)
+
+    def test_load_sequence0(self):
+        expected = "C"
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq1", 5, 6)
+        self.assertEqual(expected, observed)
+
+    def test_load_sequence1(self):
+        expected = "A"
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq2", 0, 1)
+        self.assertEqual(expected, observed)
+
+    def test_load_sequence2(self):
+        expected = ("A" * 5) + ("C" * 5) + ("T" * 5 ) + ("G" * 5) + ("N" * 5)
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq1", 0, 25)
+        self.assertEqual(expected, observed)
+
+    def test_load_sequence3(self):
+        expected = "ACTG" * 6
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq2", 0, 24)
+        self.assertEqual(expected, observed)
+
+    def test_load_sequence4(self):
+        expected = ("T" * 5 ) + ("G" * 5) + ("N" * 5)
+        g = self._get_small_genome()
+        observed = g.get_sequence_from_coords("seq1", 10, 25)
+        self.assertEqual(expected, observed)
+
+    def test_length(self):
+        expected = sum(self.chrom_to_lens.values())
+        g = self._get_small_genome()
+        observed = len(g)
+        self.assertEqual(expected, observed)
+
+    def test_coords_flipped(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 2, 1)
+        self.assertFalse(observed)
+
+    def test_end_coord_too_large(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 9, 11)
+        self.assertFalse(observed)
+
+    def test_both_coords_too_large(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 11, 14)
+        self.assertFalse(observed)
+
+    def test_coords_too_large_negative(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", -100, -99)
+        self.assertFalse(observed)
+
+    def test_coords_negative_flipped(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", -1, -2)
+        self.assertFalse(observed)
+
+    def test_coords_negative(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", -2, -1)
+        self.assertFalse(observed)
+
+    def test_bad_chrom(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq-1", 1, 2)
+        self.assertFalse(observed)
+
+    def test_too_small(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 1, 1)
+        self.assertFalse(observed)
+
+    def test_first_coord(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 0, 1)
+        self.assertTrue(observed)
+
+    def test_last_coord(self):
+        g = self._get_small_genome()
+        observed = g.coords_are_valid("seq0", 9, 10)
+        self.assertTrue(observed)
+
+    def test_in_memory(self):
+        self.assertFalse(self.in_memory)
+
+    def test_genome_in_memory(self):
+        g = self._get_small_genome()
+        self.assertFalse(g.in_memory)
+
+
+class TestHDF5GenomeInMemory(TestHDF5Genome):
+    def setUp(self):
+        super(TestHDF5GenomeInMemory, self).setUp()
+        self.in_memory = True
+
+    def test_in_memory(self):
+        self.assertTrue(self.in_memory)
+
+    def test_genome_in_memory(self):
+        g = self._get_small_genome()
+        self.assertTrue(g.in_memory)
+
+