text: """ This file defines Meta Imitation Learning (MIL). """
  from __future__ import division
  
  import numpy as np
  import random
  import tensorflow as tf
  from tensorflow.python.platform import flags
  from tf_utils import *
  from utils import Timer
  from natsort import natsorted
  
  FLAGS = flags.FLAGS
  
  class MIL(object):
      """ Initialize MIL. Need to call init_network to contruct the architecture after init. """
      def __init__(self, dU, state_idx=None, img_idx=None, network_config=None):
          # MIL hyperparams
          self.num_updates = FLAGS.num_updates
          self.update_batch_size = FLAGS.update_batch_size
          self.meta_batch_size = FLAGS.meta_batch_size
          self.meta_lr = FLAGS.meta_lr
          self.activation_fn = tf.nn.relu # by default, we use relu
          self.T = FLAGS.T
          self.network_params = network_config
          self.norm_type = FLAGS.norm
          # List of indices for state (vector) data and image (tensor) data in observation.
          self.state_idx, self.img_idx = state_idx, img_idx
          # Dimension of input and output of the model
          self._dO = len(img_idx) + len(state_idx)
          self._dU = dU
  
      def init_network(self, graph, input_tensors=None, restore_iter=0, prefix='Training_'):
          """ Helper method to initialize the tf networks used """
          with graph.as_default():
              with Timer('building TF network'):
                  result = self.construct_model(input_tensors=input_tensors, prefix=prefix, dim_input=self._dO, dim_output=self._dU,
                                            network_config=self.network_params)
              outputas, outputbs, test_output, lossesa, lossesb, final_eept_lossesb, flat_img_inputb, gradients = result
              if 'Testing' in prefix:
                  self.obs_tensor = self.obsa
                  self.state_tensor = self.statea
                  self.test_act_op = test_output
                  self.image_op = flat_img_inputb
  
              trainable_vars = tf.trainable_variables()
              total_loss1 = tf.reduce_sum(lossesa) / tf.to_float(self.meta_batch_size)
              total_losses2 = [tf.reduce_sum(lossesb[j]) / tf.to_float(self.meta_batch_size) for j in range(self.num_updates)]
              total_final_eept_losses2 = [tf.reduce_sum(final_eept_lossesb[j]) / tf.to_float(self.meta_batch_size) for j in range(self.num_updates)]
  
              if 'Training' in prefix:
                  self.total_loss1 = total_loss1
                  self.total_losses2 = total_losses2
                  self.total_final_eept_losses2 = total_final_eept_losses2
              elif 'Validation' in prefix:
                  self.val_total_loss1 = total_loss1
                  self.val_total_losses2 = total_losses2
                  self.val_total_final_eept_losses2 = total_final_eept_losses2
  
              if 'Training' in prefix:
                  self.train_op = tf.train.AdamOptimizer(self.meta_lr).minimize(self.total_losses2[self.num_updates - 1])
                  # Add summaries
                  summ = [tf.summary.scalar(prefix + 'Pre-update_loss', self.total_loss1)]
                  for j in xrange(self.num_updates):
                      summ.append(tf.summary.scalar(prefix + 'Post-update_loss_step_%d' % j, self.total_losses2[j]))
                      summ.append(tf.summary.scalar(prefix + 'Post-update_final_eept_loss_step_%d' % j, self.total_final_eept_losses2[j]))
                      for k in xrange(len(self.sorted_weight_keys)):
                          summ.append(tf.summary.histogram('Gradient_of_%s_step_%d' % (self.sorted_weight_keys[k], j), gradients[j][k]))
                  self.train_summ_op = tf.summary.merge(summ)
              elif 'Validation' in prefix:
                  # Add summaries
                  summ = [tf.summary.scalar(prefix + 'Pre-update_loss', self.val_total_loss1)]
                  for j in xrange(self.num_updates):
                      summ.append(tf.summary.scalar(prefix + 'Post-update_loss_step_%d' % j, self.val_total_losses2[j]))
                      summ.append(tf.summary.scalar(prefix + 'Post-update_final_eept_loss_step_%d' % j, self.val_total_final_eept_losses2[j]))
                  self.val_summ_op = tf.summary.merge(summ)
  
      def construct_image_input(self, nn_input, state_idx, img_idx, network_config=None):
          """ Preprocess images. """
          state_input = nn_input[:, 0:state_idx[-1]+1]
          flat_image_input = nn_input[:, state_idx[-1]+1:img_idx[-1]+1]
  
          # image goes through 3 convnet layers
          num_filters = network_config['num_filters']
  
          im_height = network_config['image_height']
          im_width = network_config['image_width']
          num_channels = network_config['image_channels']
          image_input = tf.reshape(flat_image_input, [-1, num_channels, im_width, im_height])
          image_input = tf.transpose(image_input, perm=[0,3,2,1])
          if FLAGS.pretrain_weight_path != 'N/A':
              image_input = image_input * 255.0 - tf.convert_to_tensor(np.array([103.939, 116.779, 123.68], np.float32))
              # 'RGB'->'BGR'
              image_input = image_input[:, :, :, ::-1]
          return image_input, flat_image_input, state_input
  
      def construct_weights(self, dim_input=27, dim_output=7, network_config=None):
          """ Construct weights for the network. """
          weights = {}
          num_filters = network_config['num_filters']
          strides = network_config.get('strides', [[1, 2, 2, 1], [1, 2, 2, 1], [1, 2, 2, 1]])
          filter_sizes = network_config.get('filter_size', [3]*len(strides)) # used to be 2
          if type(filter_sizes) is not list:
              filter_sizes = len(strides)*[filter_sizes]
          im_height = network_config['image_height']
          im_width = network_config['image_width']
          num_channels = network_config['image_channels']
          is_dilated = network_config.get('is_dilated', False)
          use_fp = FLAGS.fp
          pretrain = FLAGS.pretrain_weight_path != 'N/A'
          train_pretrain_conv1 = FLAGS.train_pretrain_conv1
          initialization = network_config.get('initialization', 'random')
          if pretrain:
              num_filters[0] = 64
          pretrain_weight_path = FLAGS.pretrain_weight_path
          n_conv_layers = len(num_filters)
          downsample_factor = 1
          for stride in strides:
              downsample_factor *= stride[1]
          if use_fp:
              self.conv_out_size = int(num_filters[-1]*2)
          else:
              self.conv_out_size = int(np.ceil(im_width/(downsample_factor)))*int(np.ceil(im_height/(downsample_factor)))*num_filters[-1]
  
          # conv weights
          fan_in = num_channels
          if FLAGS.conv_bt:
              fan_in += num_channels
          if FLAGS.conv_bt:
              weights['img_context'] = safe_get('img_context', initializer=tf.zeros([im_height, im_width, num_channels], dtype=tf.float32))
              weights['img_context'] = tf.clip_by_value(weights['img_context'], 0., 1.)
          for i in xrange(n_conv_layers):
              if not pretrain or i != 0:
                  if self.norm_type == 'selu':
                      weights['wc%d' % (i+1)] = init_conv_weights_snn([filter_sizes[i], filter_sizes[i], fan_in, num_filters[i]], name='wc%d' % (i+1)) # 5x5 conv, 1 input, 32 outputs
                  elif initialization == 'xavier':
                      weights['wc%d' % (i+1)] = init_conv_weights_xavier([filter_sizes[i], filter_sizes[i], fan_in, num_filters[i]], name='wc%d' % (i+1)) # 5x5 conv, 1 input, 32 outputs
                  elif initialization == 'random':
                      weights['wc%d' % (i+1)] = init_weights([filter_sizes[i], filter_sizes[i], fan_in, num_filters[i]], name='wc%d' % (i+1)) # 5x5 conv, 1 input, 32 outputs
                  else:
                      raise NotImplementedError
                  weights['bc%d' % (i+1)] = init_bias([num_filters[i]], name='bc%d' % (i+1))
                  fan_in = num_filters[i]
              else:
                  import h5py
  
                  assert num_filters[i] == 64
                  vgg_filter_size = 3
                  weights['wc%d' % (i+1)] = safe_get('wc%d' % (i+1), [vgg_filter_size, vgg_filter_size, fan_in, num_filters[i]], dtype=tf.float32, trainable=train_pretrain_conv1)
                  weights['bc%d' % (i+1)] = safe_get('bc%d' % (i+1), [num_filters[i]], dtype=tf.float32, trainable=train_pretrain_conv1)
                  pretrain_weight = h5py.File(pretrain_weight_path, 'r')
                  conv_weight = pretrain_weight['block1_conv%d' % (i+1)]['block1_conv%d_W_1:0' % (i+1)][...]
                  conv_bias = pretrain_weight['block1_conv%d' % (i+1)]['block1_conv%d_b_1:0' % (i+1)][...]
                  weights['wc%d' % (i+1)].assign(conv_weight)
                  weights['bc%d' % (i+1)].assign(conv_bias)
                  fan_in = conv_weight.shape[-1]
  
          # fc weights
          in_shape = self.conv_out_size
          if not FLAGS.no_state:
              in_shape += len(self.state_idx)
          if FLAGS.learn_final_eept:
              final_eept_range = range(FLAGS.final_eept_min, FLAGS.final_eept_max)
              final_eept_in_shape = self.conv_out_size
              if FLAGS.fc_bt:
                  weights['context_final_eept'] = safe_get('context_final_eept', initializer=tf.zeros([FLAGS.bt_dim], dtype=tf.float32))
                  final_eept_in_shape += FLAGS.bt_dim
              weights['w_ee'] = init_weights([final_eept_in_shape, len(final_eept_range)], name='w_ee')
              weights['b_ee'] = init_bias([len(final_eept_range)], name='b_ee')
              if FLAGS.two_head and FLAGS.no_final_eept:
                  weights['w_ee_two_heads'] = init_weights([final_eept_in_shape, len(final_eept_range)], name='w_ee_two_heads')
                  weights['b_ee_two_heads'] = init_bias([len(final_eept_range)], name='b_ee_two_heads')
              in_shape += (len(final_eept_range))
          if FLAGS.fc_bt:
              in_shape += FLAGS.bt_dim
          if FLAGS.fc_bt:
              weights['context'] = safe_get('context', initializer=tf.zeros([FLAGS.bt_dim], dtype=tf.float32))
          fc_weights = self.construct_fc_weights(in_shape, dim_output, network_config=network_config)
          self.conv_out_size_final = in_shape
          weights.update(fc_weights)
          return weights
  
      def construct_fc_weights(self, dim_input=27, dim_output=7, network_config=None):
          n_layers = network_config.get('n_layers', 4)
          dim_hidden = network_config.get('layer_size', [100]*(n_layers-1))
          if type(dim_hidden) is not list:
              dim_hidden = (n_layers - 1)*[dim_hidden]
          dim_hidden.append(dim_output)
          weights = {}
          in_shape = dim_input
          for i in xrange(n_layers):
              if FLAGS.two_arms and i == 0:
                  if self.norm_type == 'selu':
                      weights['w_%d_img' % i] = init_fc_weights_snn([in_shape-len(self.state_idx), dim_hidden[i]], name='w_%d_img' % i)
                      weights['w_%d_state' % i] = init_fc_weights_snn([len(self.state_idx), dim_hidden[i]], name='w_%d_state' % i)
                  else:
                      weights['w_%d_img' % i] = init_weights([in_shape-len(self.state_idx), dim_hidden[i]], name='w_%d_img' % i)
                      weights['w_%d_state' % i] = init_weights([len(self.state_idx), dim_hidden[i]], name='w_%d_state' % i)
                      weights['b_%d_state_two_arms' % i] = init_bias([dim_hidden[i]], name='b_%d_state_two_arms' % i)
                  weights['b_%d_img' % i] = init_bias([dim_hidden[i]], name='b_%d_img' % i)
                  weights['b_%d_state' % i] = init_bias([dim_hidden[i]], name='b_%d_state' % i)
                  in_shape = dim_hidden[i]
                  continue
              if i > 0 and FLAGS.all_fc_bt:
                  in_shape += FLAGS.bt_dim
                  weights['context_%d' % i] = init_bias([FLAGS.bt_dim], name='context_%d' % i)
              if self.norm_type == 'selu':
                  weights['w_%d' % i] = init_fc_weights_snn([in_shape, dim_hidden[i]], name='w_%d' % i)
              else:
                  weights['w_%d' % i] = init_weights([in_shape, dim_hidden[i]], name='w_%d' % i)
              weights['b_%d' % i] = init_bias([dim_hidden[i]], name='b_%d' % i)
              if (i == n_layers - 1 or (i == 0 and FLAGS.zero_state and not FLAGS.two_arms)) and FLAGS.two_head:
                  weights['w_%d_two_heads' % i] = init_weights([in_shape, dim_hidden[i]], name='w_%d_two_heads' % i)
                  weights['b_%d_two_heads' % i] = init_bias([dim_hidden[i]], name='b_%d_two_heads' % i)
              in_shape = dim_hidden[i]
          return weights
  
      def forward(self, image_input, state_input, weights, meta_testing=False, is_training=True, testing=False, network_config=None):
          """ Perform the forward pass. """
          if FLAGS.fc_bt:
              im_height = network_config['image_height']
              im_width = network_config['image_width']
              num_channels = network_config['image_channels']
              flatten_image = tf.reshape(image_input, [-1, im_height*im_width*num_channels])
              context = tf.transpose(tf.gather(tf.transpose(tf.zeros_like(flatten_image)), range(FLAGS.bt_dim)))
              context += weights['context']
              if FLAGS.learn_final_eept:
                  context_final_eept = tf.transpose(tf.gather(tf.transpose(tf.zeros_like(flatten_image)), range(FLAGS.bt_dim)))
                  context_final_eept += weights['context_final_eept']
          norm_type = self.norm_type
          decay = network_config.get('decay', 0.9)
          strides = network_config.get('strides', [[1, 2, 2, 1], [1, 2, 2, 1], [1, 2, 2, 1]])
          downsample_factor = strides[0][1]
          n_strides = len(strides)
          n_conv_layers = len(strides)
          use_dropout = FLAGS.dropout
          prob = FLAGS.keep_prob
          is_dilated = network_config.get('is_dilated', False)
          im_height = network_config['image_height']
          im_width = network_config['image_width']
          num_channels = network_config['image_channels']
          conv_layer = image_input
          if FLAGS.conv_bt:
              img_context = tf.zeros_like(conv_layer)
              img_context += weights['img_context']
              conv_layer = tf.concat(axis=3, values=[conv_layer, img_context])
          for i in xrange(n_conv_layers):
              if not use_dropout:
                  conv_layer = norm(conv2d(img=conv_layer, w=weights['wc%d' % (i+1)], b=weights['bc%d' % (i+1)], strides=strides[i], is_dilated=is_dilated), \
                                  norm_type=norm_type, decay=decay, id=i, is_training=is_training, activation_fn=self.activation_fn)
              else:
                  conv_layer = dropout(norm(conv2d(img=conv_layer, w=weights['wc%d' % (i+1)], b=weights['bc%d' % (i+1)], strides=strides[i], is_dilated=is_dilated), \
                                  norm_type=norm_type, decay=decay, id=i, is_training=is_training, activation_fn=self.activation_fn), keep_prob=prob, is_training=is_training, name='dropout_%d' % (i+1))
          if FLAGS.fp:
              _, num_rows, num_cols, num_fp = conv_layer.get_shape()
              if is_dilated:
                  num_rows = int(np.ceil(im_width/(downsample_factor**n_strides)))
                  num_cols = int(np.ceil(im_height/(downsample_factor**n_strides)))
              num_rows, num_cols, num_fp = [int(x) for x in [num_rows, num_cols, num_fp]]
              x_map = np.empty([num_rows, num_cols], np.float32)
              y_map = np.empty([num_rows, num_cols], np.float32)
  
              for i in range(num_rows):
                  for j in range(num_cols):
                      x_map[i, j] = (i - num_rows / 2.0) / num_rows
                      y_map[i, j] = (j - num_cols / 2.0) / num_cols
  
              x_map = tf.convert_to_tensor(x_map)
              y_map = tf.convert_to_tensor(y_map)
  
              x_map = tf.reshape(x_map, [num_rows * num_cols])
              y_map = tf.reshape(y_map, [num_rows * num_cols])
  
              # rearrange features to be [batch_size, num_fp, num_rows, num_cols]
              features = tf.reshape(tf.transpose(conv_layer, [0,3,1,2]),
                                    [-1, num_rows*num_cols])
              softmax = tf.nn.softmax(features)
  
              fp_x = tf.reduce_sum(tf.multiply(x_map, softmax), [1], keep_dims=True)
              fp_y = tf.reduce_sum(tf.multiply(y_map, softmax), [1], keep_dims=True)
  
              conv_out_flat = tf.reshape(tf.concat(axis=1, values=[fp_x, fp_y]), [-1, num_fp*2])
          else:
              conv_out_flat = tf.reshape(conv_layer, [-1, self.conv_out_size])
          fc_input = tf.add(conv_out_flat, 0)
          if FLAGS.learn_final_eept:
              final_eept_range = range(FLAGS.final_eept_min, FLAGS.final_eept_max)
              if testing:
                  T = 1
              else:
                  T = self.T
              conv_out_flat = tf.reshape(conv_out_flat, [-1, T, self.conv_out_size])
              conv_size = self.conv_out_size
              if FLAGS.fc_bt:
                  context_dim = FLAGS.bt_dim
                  conv_out_flat = tf.concat(axis=2, values=[conv_out_flat, tf.reshape(context_final_eept, [-1, T, context_dim])])
                  conv_size += context_dim
              # only predict the final eept using the initial image
              final_ee_inp = tf.reshape(conv_out_flat, [-1, conv_size])
              # use video for preupdate only if no_final_eept
              if (not FLAGS.learn_final_eept_whole_traj) or meta_testing:
                  final_ee_inp = conv_out_flat[:, 0, :]
              if FLAGS.two_head and not meta_testing and FLAGS.no_final_eept:
                  final_eept_pred = tf.matmul(final_ee_inp, weights['w_ee_two_heads']) + weights['b_ee_two_heads']
              else:
                  final_eept_pred = tf.matmul(final_ee_inp, weights['w_ee']) + weights['b_ee']
              if (not FLAGS.learn_final_eept_whole_traj) or meta_testing:
                  final_eept_pred = tf.reshape(tf.tile(tf.reshape(final_eept_pred, [-1]), [T]), [-1, len(final_eept_range)])
                  final_eept_concat = tf.identity(final_eept_pred)
              else:
                  # Assume tbs == 1
                  # Only provide the FC layers with final_eept_pred at first time step
                  final_eept_concat = final_eept_pred[0]
                  final_eept_concat = tf.reshape(tf.tile(tf.reshape(final_eept_concat, [-1]), [T]), [-1, len(final_eept_range)])
              fc_input = tf.concat(axis=1, values=[fc_input, final_eept_concat])
          else:
              final_eept_pred = None
          if FLAGS.fc_bt:
              fc_input = tf.concat(axis=1, values=[fc_input, context])
          return self.fc_forward(fc_input, weights, state_input=state_input, meta_testing=meta_testing, is_training=is_training, testing=testing, network_config=network_config), final_eept_pred
  
      def fc_forward(self, fc_input, weights, state_input=None, meta_testing=False, is_training=True, testing=False, network_config=None):
          n_layers = network_config.get('n_layers', 4)
          use_dropout = FLAGS.dropout
          prob = FLAGS.keep_prob
          fc_output = tf.add(fc_input, 0)
          use_selu = self.norm_type == 'selu'
          norm_type = self.norm_type
          if state_input is not None and not FLAGS.two_arms:
              fc_output = tf.concat(axis=1, values=[fc_output, state_input])
          for i in xrange(n_layers):
              if i > 0 and FLAGS.all_fc_bt:
                  context = tf.transpose(tf.gather(tf.transpose(tf.zeros_like(fc_output)), range(FLAGS.bt_dim)))
                  context += weights['context_%d' % i]
                  fc_output = tf.concat(axis=1, values=[fc_output, context])
              if (i == n_layers - 1 or (i == 0 and FLAGS.zero_state and not FLAGS.two_arms)) and FLAGS.two_head and not meta_testing:
                  fc_output = tf.matmul(fc_output, weights['w_%d_two_heads' % i]) + weights['b_%d_two_heads' % i]
              elif i == 0 and FLAGS.two_arms:
                  assert state_input is not None
                  if FLAGS.two_arms:
                      state_part = weights['b_%d_state_two_arms' % i]
                  else:
                      state_part = tf.matmul(state_input, weights['w_%d_state' % i]) + weights['b_%d_state' % i]
                  if not meta_testing:
                      fc_output = tf.matmul(fc_output, weights['w_%d_img' % i]) + weights['b_%d_img' % i] + state_part
                  else:
                      fc_output = tf.matmul(fc_output, weights['w_%d_img' % i]) + weights['b_%d_img' % i] + \
                                  tf.matmul(state_input, weights['w_%d_state' % i]) + weights['b_%d_state' % i]
              else:
                  fc_output = tf.matmul(fc_output, weights['w_%d' % i]) + weights['b_%d' % i]
              if i != n_layers - 1:
                  if use_selu:
                      fc_output = selu(fc_output)
                  else:
                      fc_output = self.activation_fn(fc_output)
                  # only use dropout for post-update
                  if use_dropout:
                      fc_output = dropout(fc_output, keep_prob=prob, is_training=is_training, name='dropout_fc_%d' % i, selu=use_selu)
          return fc_output
  
      def construct_model(self, input_tensors=None, prefix='Training_', dim_input=27, dim_output=7, network_config=None):
          """
          Construct the meta-learning graph.
          Args:
              input_tensors: tensors of input videos, if available
              prefix: indicate whether we are building training, validation or testing graph.
              dim_input: Dimensionality of input.
              dim_output: Dimensionality of the output.
              network_config: dictionary of network structure parameters
          Returns:
              a tuple of output tensors.
          """
          if input_tensors is None:
              self.obsa = obsa = tf.placeholder(tf.float32, name='obsa') # meta_batch_size x update_batch_size x dim_input
              self.obsb = obsb = tf.placeholder(tf.float32, name='obsb')
          else:
              self.obsa = obsa = input_tensors['inputa'] # meta_batch_size x update_batch_size x dim_input
              self.obsb = obsb = input_tensors['inputb']
  
          if not hasattr(self, 'statea'):
              self.statea = statea = tf.placeholder(tf.float32, name='statea')
              self.stateb = stateb = tf.placeholder(tf.float32, name='stateb')
              self.actiona = actiona = tf.placeholder(tf.float32, name='actiona')
              self.actionb = actionb = tf.placeholder(tf.float32, name='actionb')
          else:
              statea = self.statea
              stateb = self.stateb
              actiona = self.actiona
              actionb = self.actionb
  
          inputa = tf.concat(axis=2, values=[statea, obsa])
          inputb = tf.concat(axis=2, values=[stateb, obsb])
  
          with tf.variable_scope('model', reuse=None) as training_scope:
              # Construct layers weight & bias
              if 'weights' not in dir(self):
                  if FLAGS.learn_final_eept:
                      final_eept_range = range(FLAGS.final_eept_min, FLAGS.final_eept_max)
                      self.weights = weights = self.construct_weights(dim_input, dim_output-len(final_eept_range), network_config=network_config)
                  else:
                      self.weights = weights = self.construct_weights(dim_input, dim_output, network_config=network_config)
                  self.sorted_weight_keys = natsorted(self.weights.keys())
              else:
                  training_scope.reuse_variables()
                  weights = self.weights
  
              self.step_size = FLAGS.train_update_lr
              loss_multiplier = FLAGS.loss_multiplier
              final_eept_loss_eps = FLAGS.final_eept_loss_eps
              act_loss_eps = FLAGS.act_loss_eps
              use_whole_traj = FLAGS.learn_final_eept_whole_traj
  
              num_updates = self.num_updates
              lossesa, outputsa = [], []
              lossesb = [[] for _ in xrange(num_updates)]
              outputsb = [[] for _ in xrange(num_updates)]
  
              def batch_metalearn(inp):
                  inputa, inputb, actiona, actionb = inp
                  inputa = tf.reshape(inputa, [-1, dim_input])
                  inputb = tf.reshape(inputb, [-1, dim_input])
                  actiona = tf.reshape(actiona, [-1, dim_output])
                  actionb = tf.reshape(actionb, [-1, dim_output])
                  gradients_summ = []
                  testing = 'Testing' in prefix
  
                  final_eepta, final_eeptb = None, None
                  if FLAGS.learn_final_eept:
                      final_eept_range = range(FLAGS.final_eept_min, FLAGS.final_eept_max)
                      final_eepta = actiona[:, final_eept_range[0]:final_eept_range[-1]+1]
                      final_eeptb = actionb[:, final_eept_range[0]:final_eept_range[-1]+1]
                      actiona = actiona[:, :final_eept_range[0]]
                      actionb = actionb[:, :final_eept_range[0]]
                      if FLAGS.no_final_eept:
                          final_eepta = tf.zeros_like(final_eepta)
  
                  if FLAGS.no_action:
                      actiona = tf.zeros_like(actiona)
  
                  local_outputbs, local_lossesb, final_eept_lossesb = [], [], []
                  # Assume fixed data for each update
                  actionas = [actiona]*num_updates
  
                  # Convert to image dims
                  inputa, _, state_inputa = self.construct_image_input(inputa, self.state_idx, self.img_idx, network_config=network_config)
                  inputb, flat_img_inputb, state_inputb = self.construct_image_input(inputb, self.state_idx, self.img_idx, network_config=network_config)
                  inputas = [inputa]*num_updates
                  inputbs = [inputb]*num_updates
                  if FLAGS.zero_state:
                      state_inputa = tf.zeros_like(state_inputa)
                  state_inputas = [state_inputa]*num_updates
                  if FLAGS.no_state:
                      state_inputa = None
  
                  if FLAGS.learn_final_eept:
                      final_eeptas = [final_eepta]*num_updates
  
                  # Pre-update
                  if 'Training' in prefix:
                      local_outputa, final_eept_preda = self.forward(inputa, state_inputa, weights, network_config=network_config)
                  else:
                      local_outputa, final_eept_preda = self.forward(inputa, state_inputa, weights, is_training=False, network_config=network_config)
                  if FLAGS.learn_final_eept:
                      final_eept_lossa = euclidean_loss_layer(final_eept_preda, final_eepta, multiplier=loss_multiplier, use_l1=FLAGS.use_l1_l2_loss)
                  else:
                      final_eept_lossa = tf.constant(0.0)
                  local_lossa = act_loss_eps * euclidean_loss_layer(local_outputa, actiona, multiplier=loss_multiplier, use_l1=FLAGS.use_l1_l2_loss)
                  if FLAGS.learn_final_eept:
                      local_lossa += final_eept_loss_eps * final_eept_lossa
  
                  # Compute fast gradients
                  grads = tf.gradients(local_lossa, weights.values())
                  gradients = dict(zip(weights.keys(), grads))
                  # make fast gradient zero for weights with gradient None
                  for key in gradients.keys():
                      if gradients[key] is None:
                          gradients[key] = tf.zeros_like(weights[key])
                  if FLAGS.stop_grad:
                      gradients = {key:tf.stop_gradient(gradients[key]) for key in gradients.keys()}
                  if FLAGS.clip:
                      clip_min = FLAGS.clip_min
                      clip_max = FLAGS.clip_max
                      for key in gradients.keys():
                          gradients[key] = tf.clip_by_value(gradients[key], clip_min, clip_max)
                  if FLAGS.pretrain_weight_path != 'N/A':
                      gradients['wc1'] = tf.zeros_like(gradients['wc1'])
                      gradients['bc1'] = tf.zeros_like(gradients['bc1'])
                  gradients_summ.append([gradients[key] for key in self.sorted_weight_keys])
                  fast_weights = dict(zip(weights.keys(), [weights[key] - self.step_size*gradients[key] for key in weights.keys()]))
  
                  # Post-update
                  if FLAGS.no_state:
                      state_inputb = None
                  if 'Training' in prefix:
                      outputb, final_eept_predb = self.forward(inputb, state_inputb, fast_weights, meta_testing=True, network_config=network_config)
                  else:
                      outputb, final_eept_predb = self.forward(inputb, state_inputb, fast_weights, meta_testing=True, is_training=False, testing=testing, network_config=network_config)
                  local_outputbs.append(outputb)
                  if FLAGS.learn_final_eept:
                      final_eept_lossb = euclidean_loss_layer(final_eept_predb, final_eeptb, multiplier=loss_multiplier, use_l1=FLAGS.use_l1_l2_loss)
                  else:
                      final_eept_lossb = tf.constant(0.0)
                  local_lossb = act_loss_eps * euclidean_loss_layer(outputb, actionb, multiplier=loss_multiplier, use_l1=FLAGS.use_l1_l2_loss)
                  if FLAGS.learn_final_eept:
                      local_lossb += final_eept_loss_eps * final_eept_lossb
                  if use_whole_traj:
                      # assume tbs == 1
                      final_eept_lossb = euclidean_loss_layer(final_eept_predb[0], final_eeptb[0], multiplier=loss_multiplier, use_l1=FLAGS.use_l1_l2_loss)
                  final_eept_lossesb.append(final_eept_lossb)
                  local_lossesb.append(local_lossb)
  
                  for j in range(num_updates - 1):
                      # Pre-update
                      state_inputa_new = state_inputas[j+1]
                      if FLAGS.no_state:
                          state_inputa_new = None
                      if 'Training' in prefix:
                          outputa, final_eept_preda = self.forward(inputas[j+1], state_inputa_new, fast_weights, network_config=network_config)
                      else:
                          outputa, final_eept_preda = self.forward(inputas[j+1], state_inputa_new, fast_weights, is_training=False, testing=testing, network_config=network_config)
                      if FLAGS.learn_final_eept:
                          final_eept_lossa = euclidean_loss_layer(final_eept_preda, final_eeptas[j+1], multiplier=loss_multiplier, use_l1=FLAGS.use_l1_l2_loss)
                      else:
                          final_eept_lossa = tf.constant(0.0)
                      loss = act_loss_eps * euclidean_loss_layer(outputa, actionas[j+1], multiplier=loss_multiplier, use_l1=FLAGS.use_l1_l2_loss)
                      if FLAGS.learn_final_eept:
                          loss += final_eept_loss_eps * final_eept_lossa
  
                      # Compute fast gradients
                      grads = tf.gradients(loss, fast_weights.values())
                      gradients = dict(zip(fast_weights.keys(), grads))
                      # make fast gradient zero for weights with gradient None
                      for key in gradients.keys():
                          if gradients[key] is None:
                              gradients[key] = tf.zeros_like(fast_weights[key])
                      if FLAGS.stop_grad:
                          gradients = {key:tf.stop_gradient(gradients[key]) for key in gradients.keys()}
                      if FLAGS.clip:
                          clip_min = FLAGS.clip_min
                          clip_max = FLAGS.clip_max
                          for key in gradients.keys():
                              gradients[key] = tf.clip_by_value(gradients[key], clip_min, clip_max)
                      if FLAGS.pretrain_weight_path != 'N/A':
                          gradients['wc1'] = tf.zeros_like(gradients['wc1'])
                          gradients['bc1'] = tf.zeros_like(gradients['bc1'])
                      gradients_summ.append([gradients[key] for key in self.sorted_weight_keys])
                      fast_weights = dict(zip(fast_weights.keys(), [fast_weights[key] - self.step_size*gradients[key] for key in fast_weights.keys()]))
  
                      # Post-update
                      if FLAGS.no_state:
                          state_inputb = None
                      if 'Training' in prefix:
                          output, final_eept_predb = self.forward(inputbs[j+1], state_inputb, fast_weights, meta_testing=True, network_config=network_config)
                      else:
                          output, final_eept_predb = self.forward(inputbs[j+1], state_inputb, fast_weights, meta_testing=True, is_training=False, testing=testing, network_config=network_config)
                      local_outputbs.append(output)
                      if FLAGS.learn_final_eept:
                          final_eept_lossb = euclidean_loss_layer(final_eept_predb, final_eeptb, multiplier=loss_multiplier, use_l1=FLAGS.use_l1_l2_loss)
                      else:
                          final_eept_lossb = tf.constant(0.0)
                      lossb = act_loss_eps * euclidean_loss_layer(output, actionb, multiplier=loss_multiplier, use_l1=FLAGS.use_l1_l2_loss)
                      if FLAGS.learn_final_eept:
                          lossb += final_eept_loss_eps * final_eept_lossb
                      if use_whole_traj:
                          # assume tbs == 1
                          final_eept_lossb = euclidean_loss_layer(final_eept_predb[0], final_eeptb[0], multiplier=loss_multiplier, use_l1=FLAGS.use_l1_l2_loss)
                      final_eept_lossesb.append(final_eept_lossb)
                      local_lossesb.append(lossb)
                  local_fn_output = [local_outputa, local_outputbs, local_outputbs[-1], local_lossa, local_lossesb, final_eept_lossesb, flat_img_inputb, gradients_summ]
                  return local_fn_output
  
          if self.norm_type:
              # initialize batch norm vars.
              unused = batch_metalearn((inputa[0], inputb[0], actiona[0], actionb[0]))
  
          out_dtype = [tf.float32, [tf.float32]*num_updates, tf.float32, tf.float32, [tf.float32]*num_updates, [tf.float32]*num_updates, tf.float32, [[tf.float32]*len(self.weights.keys())]*num_updates]
          result = tf.map_fn(batch_metalearn, elems=(inputa, inputb, actiona, actionb), dtype=out_dtype)
          print 'Done with map.'
          return result
  