def generate_train_config(model_spec, train_batch_spec, val_batch_spec,  epoch_size,
                        n_epochs=100, lr=0.001):
    """ Generates a spec fully specifying a run of the experiment

    Args:
      model_spec: a model spec
      train_batch_spec: a batch spec for train
      val_batch_spec: batch spec for validation
      epoch_size: int
      n_epochs: int
      lr: learning rate - float

    Returns:
      train_spec
    """

    train_spec = {
        'model_spec': model_spec,
        'train_batch_spec': train_batch_spec,
        'val_batch_spec':  val_batch_spec,
        'train_spec': {
            'epoch_size': epoch_size,
            'n_epochs': n_epochs,
            'lr': lr
        }
    }
    return train_spec