jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _process_experience_global(self, experience):
    """Processes the experience in case the agent accepts only global features.

    In the experience coming from the replay buffer, the reward (and all other
    elements) have two batch dimensions `batch_size` and `time_steps`, where
    `time_steps` is the number of driver steps executed in each training loop.
    We flatten the tensors in order to reflect the effective batch size. Then,
    all the necessary processing on the observation is done, including splitting
    the action mask if it is present.

    Args:
      experience: An instance of trajectory. Every element in the trajectory has
        two batch dimensions.

    Returns:
      A tuple of reward, action, observation, and batch_size. All the outputs
        (except `batch_size`) have a single batch dimension of value
        `batch_size`.
    """
    reward, _ = nest_utils.flatten_multi_batched_nested_tensors(
        experience.reward, self._time_step_spec.reward
    )
    observation, _ = nest_utils.flatten_multi_batched_nested_tensors(
        experience.observation, self.training_data_spec.observation
    )
    action, _ = nest_utils.flatten_multi_batched_nested_tensors(
        experience.action, self._action_spec
    )
    batch_size = tf.cast(
        tf.compat.dimension_value(tf.shape(reward)[0]), dtype=tf.int64
    )

    if self._observation_and_action_constraint_splitter is not None:
      observation, _ = self._observation_and_action_constraint_splitter(
          observation
      )
    if self._add_bias:
      # The bias is added via a constant 1 feature.
      observation = tf.concat(
          [observation, tf.ones([batch_size, 1], dtype=observation.dtype)],
          axis=1,
      )

    observation = tf.reshape(
        tf.cast(observation, self._dtype), [batch_size, -1]
    )
    reward = tf.cast(reward, self._dtype)

    return reward, action, observation, batch_size
//  [python]
def errors(logits, labels, name=None):
    """Compute error mean and whether each unlabeled example is erroneous

    Assume unlabeled examples have label == -1.
    Compute the mean error over unlabeled examples.
    Mean error is NaN if there are no unlabeled examples.
    Note that unlabeled examples are treated differently in cost calculation.
    """
    with tf.name_scope(name, "errors") as scope:
        applicable = tf.not_equal(labels, -1)
        labels = tf.boolean_mask(labels, applicable)
        logits = tf.boolean_mask(logits, applicable)
        predictions = tf.argmax(logits, -1)
        labels = tf.to_int64(labels)
        per_sample = tf.to_float(tf.not_equal(predictions, labels))
        mean = tf.reduce_mean(per_sample, name=scope)
        return mean, per_sample
//  [python]
def errors(logits, labels, name=None):
    """Compute error mean and whether each unlabeled example is erroneous

    Assume unlabeled examples have label == -1.
    Compute the mean error over unlabeled examples.
    Mean error is NaN if there are no unlabeled examples.
    Note that unlabeled examples are treated differently in cost calculation.
    """
    with tf.name_scope(name, "errors") as scope:
        applicable = tf.not_equal(labels, -1)
        labels = tf.boolean_mask(labels, applicable)
        logits = tf.boolean_mask(logits, applicable)
        predictions = tf.argmax(logits, -1)
        labels = tf.cast(labels, tf.int64)
        per_sample = tf.to_float(tf.not_equal(predictions, labels))
        mean = tf.reduce_mean(per_sample, name=scope)
        return mean, per_sample
//  [python]
def crop_proposals(proposals, bbox, crop_size=64):
    n_proposals = tf.reduce_max(proposals) + 1
    pred_imgs = tf.one_hot(proposals, n_proposals, axis=0)
    pred_imgs = tf.expand_dims(pred_imgs, -1)

    return tf.image.crop_and_resize(
        tf.to_float(pred_imgs),
        bbox,
        tf.range(tf.shape(bbox)[0]),
        crop_size=(crop_size,crop_size),
    )
//  [python]
def crop_proposals(proposals, bbox, crop_size=64):
    n_proposals = tf.reduce_max(proposals) + 1
    pred_imgs = tf.one_hot(proposals, n_proposals, axis=0)
    pred_imgs = tf.expand_dims(pred_imgs, -1)

    return tf.image.crop_and_resize(
        tf.cast(pred_imgs, tf.float32),
        bbox,
        tf.range(tf.shape(bbox)[0]),
        crop_size=(crop_size,crop_size),
    )
//  [python]
def _process_experience_per_arm(self, experience):
    """Processes the experience in case the agent accepts per-arm features.

    In the experience coming from the replay buffer, the reward (and all other
    elements) have two batch dimensions `batch_size` and `time_steps`, where
    `time_steps` is the number of driver steps executed in each training loop.
    We flatten the tensors in order to reflect the effective batch size. Then,
    all the necessary processing on the observation is done, including splitting
    the action mask if it is present.

    After the preprocessing, the per-arm part of the observation is copied over
    from the respective policy info field and concatenated with the global
    observation. The action tensor will be replaced by zeros, since in the
    per-arm case, there is only one reward model to update.

    Args:
      experience: An instance of trajectory. Every element in the trajectory has
        two batch dimensions.

    Returns:
      A tuple of reward, action, observation, and batch_size. All the outputs
        (except `batch_size`) have a single batch dimension of value
        `batch_size`.
    """
    reward, _ = nest_utils.flatten_multi_batched_nested_tensors(
        experience.reward, self._time_step_spec.reward
    )
    observation, _ = nest_utils.flatten_multi_batched_nested_tensors(
        experience.observation, self.training_data_spec.observation
    )

    if self._observation_and_action_constraint_splitter is not None:
      observation, _ = self._observation_and_action_constraint_splitter(
          observation
      )
    batch_size = tf.cast(
        tf.compat.dimension_value(tf.shape(reward)[0]), dtype=tf.int64
    )
    global_observation = observation[bandit_spec_utils.GLOBAL_FEATURE_KEY]
    if self._add_bias:
      # The bias is added via a constant 1 feature.
      global_observation = tf.concat(
          [
              global_observation,
              tf.ones([batch_size, 1], dtype=global_observation.dtype),
          ],
          axis=1,
      )

    # The arm observation we train on needs to be copied from the respective
    # policy info field to the per arm observation field. Pretending there was
    # only one action, we fill the action field with zeros.
    action = tf.zeros(shape=[batch_size], dtype=tf.int64)
    chosen_action, _ = nest_utils.flatten_multi_batched_nested_tensors(
        experience.policy_info.chosen_arm_features,
        self.policy.info_spec.chosen_arm_features,
    )
    arm_observation = chosen_action
    overall_observation = tf.concat(
        [global_observation, arm_observation], axis=1
    )
    overall_observation = tf.reshape(
        tf.cast(overall_observation, self._dtype), [batch_size, -1]
    )
    reward = tf.cast(reward, self._dtype)

    return reward, action, overall_observation, batch_size
//  [python]
def adjust_bbox(bbox, min_size=64, h=544, w=704):
    '''
    adjust bbox size to be 2x larger, but no smaller than min_size x min_size
    '''
    cy = (bbox[:,0] + bbox[:,2])/2
    cx = (bbox[:,1] + bbox[:,3])/2
    hh = (bbox[:,2] - bbox[:,0])*2
    ww = (bbox[:,3] - bbox[:,1])*2
    hh = tf.clip_by_value(hh, min_size, 9999)
    ww = tf.clip_by_value(ww, min_size, 9999)
    h = t.cast(h, t.float64)
    w = t.cast(w, t.float64)
    bbox_out = tf.stack([
          (cy - hh / 2) / h,
          (cx - ww / 2) / w,
          (cy + hh / 2) / h,
          (cx + ww / 2) / w,
      ], axis = -1)
    return t.cast(bbox_out, t.float32)
//  [python]
def adjust_bbox(bbox, min_size=64, h=544, w=704):
    '''
    adjust bbox size to be 2x larger, but no smaller than min_size x min_size
    '''
    cy = (bbox[:,0] + bbox[:,2])/2
    cx = (bbox[:,1] + bbox[:,3])/2
    hh = (bbox[:,2] - bbox[:,0])*2
    ww = (bbox[:,3] - bbox[:,1])*2
    hh = tf.clip_by_value(hh, min_size, 9999)
    ww = tf.clip_by_value(ww, min_size, 9999)
    h = tf.cast(h, tf.float64)
    w = tf.cast(w, tf.float64)
    bbox_out = tf.stack([
          (cy - hh / 2) / h,
          (cx - ww / 2) / w,
          (cy + hh / 2) / h,
          (cx + ww / 2) / w,
      ], axis = -1)
    return tf.cast(bbox_out, tf.float32)
//  [python]
def _distributed_train_step(self, experience, weights=None):
    """Distributed train fn to be passed as input to run()."""
    experience_reward, action, experience_observation, batch_size = (
        self._process_experience(experience)
    )
    self._train_step_counter.assign_add(batch_size)
    observation, reward = self._maybe_apply_per_example_weight(
        experience_observation, experience_reward, weights
    )
    loss = self._loss(experience, weights, training=True)

    for k in range(self._num_models):
      diag_mask = tf.linalg.tensor_diag(
          tf.cast(tf.equal(action, k), self._dtype)
      )
      observations_for_arm = tf.matmul(diag_mask, observation)
      rewards_for_arm = tf.matmul(diag_mask, tf.reshape(reward, [-1, 1]))

      # Compute local updates for the matrix A and b of this arm.
      cov_matrix_local_update = tf.matmul(
          observations_for_arm, observations_for_arm, transpose_a=True
      )
      data_vector_local_update = bandit_utils.sum_reward_weighted_observations(
          rewards_for_arm, observations_for_arm
      )

      def _merge_fn(
          strategy,
          per_replica_cov_matrix_update,
          per_replica_data_vector_update,
      ):
        """Merge the per-replica-updates."""
        # Reduce the per-replica-updates using SUM.
        # pylint: disable=cell-var-from-loop
        updates_and_vars = [
            (per_replica_cov_matrix_update, self._cov_matrix_list[k]),
            (per_replica_data_vector_update, self._data_vector_list[k]),
        ]

        reduced_updates = strategy.extended.batch_reduce_to(
            tf.distribute.ReduceOp.SUM, updates_and_vars
        )

        # Update the model variables.
        self._cov_matrix_list[k].assign_add(reduced_updates[0])
        self._data_vector_list[k].assign_add(reduced_updates[1])

        # Compute the eigendecomposition, if needed.
        if self._use_eigendecomp:
          eig_vals, eig_matrix = tf.linalg.eigh(self._cov_matrix_list[k])
          self._eig_vals_list[k].assign(eig_vals)
          self._eig_matrix_list[k].assign(eig_matrix)

      # Passes the local_updates to the _merge_fn() above that performs custom
      # computation on the per-replica values.
      # All replicas pause their execution until merge_call() is done and then,
      # execution is resumed.
      replica_context = tf.distribute.get_replica_context()
      replica_context.merge_call(
          _merge_fn, args=(cov_matrix_local_update, data_vector_local_update)
      )

    with tf.control_dependencies([loss.loss]):
      self._theta.assign(self.theta)
    return loss
//  [python]
def __init__(self, initial_learning_rate, hidden_size, warmup_steps):
    """Initialize configuration of the learning rate schedule.

    Args:
      initial_learning_rate: A float, the initial learning rate.
      hidden_size: An integer, the model dimension in the hidden layers.
      warmup_steps: An integer, the number of steps required for linear warmup.
    """
    super(LearningRateSchedule, self).__init__()
    self.initial_learning_rate = initial_learning_rate
    self.hidden_size = hidden_size
    self.warmup_steps = warmup_steps
    self.warmup_steps_tensor = tf.to_float(warmup_steps)
//  [python]
def __init__(self, initial_learning_rate, hidden_size, warmup_steps):
    """Initialize configuration of the learning rate schedule.

    Args:
      initial_learning_rate: A float, the initial learning rate.
      hidden_size: An integer, the model dimension in the hidden layers.
      warmup_steps: An integer, the number of steps required for linear warmup.
    """
    super(LearningRateSchedule, self).__init__()
    self.initial_learning_rate = initial_learning_rate
    self.hidden_size = hidden_size
    self.warmup_steps = warmup_steps
    self.warmup_steps_tensor = tf.cast(warmup_steps, tf.float32)
//  [python]
def __call__(self, global_step):
    """Calculate learning rate with linear warmup and rsqrt decay.

    Args:
      global_step: An integer, the current global step used for learning rate
        calculation.

    Returns:
      A float, the learning rate needs to be used for current global step.
    """
    with tf.name_scope('learning_rate_schedule'):
      global_step = tf.cast(global_step, tf.float32)
      learning_rate = self.initial_learning_rate
      learning_rate *= (self.hidden_size**-0.5)
      # Apply linear warmup
      learning_rate *= tf.minimum(1.0, global_step / self.warmup_steps_tensor)
      # Apply rsqrt decay
      learning_rate /= tf.sqrt(
          tf.maximum(global_step, self.warmup_steps_tensor))
      return learning_rate
//  [python]
def rescale_tf(x, scale, method=tf.image.ResizeMethod.BILINEAR, is_3d=False):
    if is_3d:
        shp = tf.shape(x)
        b, d, h, w, c = shp[0], shp[1], shp[2], shp[3], shp[4]
        hw = tf.reshape(tf.transpose(x, [0,2,3,1,4]), [b,h,w,d*c])
        h = tf.cast(tf.cast(h, tf.float32)*scale, tf.int32)
        w = tf.cast(tf.cast(w, tf.float32)*scale, tf.int32)

        hw = tf.compat.v1.image.resize(hw, (h,w), method=method)
        hw = tf.reshape(hw, [b,h,w,d,c])
        dh = tf.reshape(tf.transpose(hw, [0,3,1,2,4]), [b,d,h,w*c])
        d = tf.cast(tf.cast(d, tf.float32)*scale, tf.int32)
        dh = tf.compat.v1.image.resize(dh, (d,h), method=method)
        x = tf.reshape(dh, [b,d,h,w,c])
    else:
        shp = tf.shape(x)
        b, h, w, c = shp[0], shp[1], shp[2], shp[3]
        h = tf.cast(tf.cast(h, tf.float32)*scale, tf.int32)
        w = tf.cast(tf.cast(w, tf.float32)*scale, tf.int32)
        x = tf.compat.v1.image.resize(x, (h,w), method=method)
    return x
//  [python]
def rescale_tf(x, scale, method=tf.image.ResizeMethod.BILINEAR, is_3d=False):
    if is_3d:
        # b, d, h, w, c = int_shape(x)
        shp = tf.shape(x)
        b, d, h, w, c = shp[0], shp[1], shp[2], shp[3], shp[4]
        hw = tf.reshape(tf.transpose(x, [0,2,3,1,4]), [b,h,w,d*c])
        h = tf.cast(tf.cast(h, tf.float32)*scale, tf.int32)
        w = tf.cast(tf.cast(w, tf.float32)*scale, tf.int32)

        hw = tf.compat.v1.image.resize(hw, (h,w), method=method)
        hw = tf.reshape(hw, [b,h,w,d,c])
        dh = tf.reshape(tf.transpose(hw, [0,3,1,2,4]), [b,d,h,w*c])
        d = tf.cast(tf.cast(d, tf.float32)*scale, tf.int32)
        dh = tf.compat.v1.image.resize(dh, (d,h), method=method)
        x = tf.reshape(dh, [b,d,h,w,c])
    else:
        # b, h, w, c = int_shape(x)
        shp = tf.shape(x)
        b, h, w, c = shp[0], shp[1], shp[2], shp[3]
        h = tf.cast(tf.cast(h, tf.float32)*scale, tf.int32)
        w = tf.cast(tf.cast(w, tf.float32)*scale, tf.int32)
        x = tf.compat.v1.image.resize(x, (h,w), method=method)
    return x
//  [python]
def hsv2rgb(h,s,v):
    c = s * v
    m = v - c
    dh = h * 6
    h_category = t.cast(dh, t.int32)
    fmodu = t.mod(dh, 2)
    x = c * (1 - t.abs(fmodu - 1))
    component_shape = t.shape(h)
    dtype = h.dtype
    rr = t.zeros(component_shape, dtype=dtype)
    gg = t.zeros(component_shape, dtype=dtype)
    bb = t.zeros(component_shape, dtype=dtype)
    h0 = t.equal(h_category, 0)
    rr = t.where(h0, c, rr)
    gg = t.where(h0, x, gg)
    h1 = t.equal(h_category, 1)
    rr = t.where(h1, x, rr)
    gg = t.where(h1, c, gg)
    h2 = t.equal(h_category, 2)
    gg = t.where(h2, c, gg)
    bb = t.where(h2, x, bb)
    h3 = t.equal(h_category, 3)
    gg = t.where(h3, x, gg)
    bb = t.where(h3, c, bb)
    h4 = t.equal(h_category, 4)
    rr = t.where(h4, x, rr)
    bb = t.where(h4, c, bb)
    h5 = t.equal(h_category, 5)
    rr = t.where(h5, c, rr)
    bb = t.where(h5, x, bb)
    r = rr + m
    g = gg + m
    b = bb + m
    return r,g,b
//  [python]
def hsv2rgb(h,s,v):
    c = s * v
    m = v - c
    dh = h * 6
    h_category = tf.cast(dh, tf.int32)
    fmodu = tf.mod(dh, 2)
    x = c * (1 - tf.abs(fmodu - 1))
    component_shape = tf.shape(h)
    dtype = h.dtype
    rr = tf.zeros(component_shape, dtype=dtype)
    gg = tf.zeros(component_shape, dtype=dtype)
    bb = tf.zeros(component_shape, dtype=dtype)
    h0 = tf.equal(h_category, 0)
    rr = tf.where(h0, c, rr)
    gg = tf.where(h0, x, gg)
    h1 = tf.equal(h_category, 1)
    rr = tf.where(h1, x, rr)
    gg = tf.where(h1, c, gg)
    h2 = tf.equal(h_category, 2)
    gg = tf.where(h2, c, gg)
    bb = tf.where(h2, x, bb)
    h3 = tf.equal(h_category, 3)
    gg = tf.where(h3, x, gg)
    bb = tf.where(h3, c, bb)
    h4 = tf.equal(h_category, 4)
    rr = tf.where(h4, x, rr)
    bb = tf.where(h4, c, bb)
    h5 = tf.equal(h_category, 5)
    rr = tf.where(h5, c, rr)
    bb = tf.where(h5, x, bb)
    r = rr + m
    g = gg + m
    b = bb + m
    return r,g,b
//  [python]
def fast_parser(inputs, det_model, min_iou = 0.1, crop_size=64, min_size=64):
    ''' this parser is based on bbox instead of mask roi '''
    def process_one_img(k):
        m_boxes, matched, m_ious, ious = box_matching(
            t.cast(bboxes[k][None,...], t.float32),
            t.cast(gt_bboxes[k][None,...], t.float32),
            )
        bb = adjust_bbox(bboxes[k], min_size=min_size, h=height, w=width)
        matched = t.clip_by_value(matched[0,:], 0, 9999) # remove batch dimension, and -1 values(unmatched)
        proposal_crops = crop_proposals(proposals[k], bb, crop_size=crop_size)
        img_crops = crop_img(imgs[k], bb, crop_size=crop_size)
        mask_crops = crop_masks(labels['mask_indices'][k], bb, matched, crop_size=crop_size)
        ind = t.cast(m_ious > min_iou, t.int32) * (labels['class'][k] + 1)
        return proposal_crops, img_crops, mask_crops, ind[0,:], bb

    with t.device("/gpu:0"):
        imgs, labels = inputs
        height = t.shape(imgs)[1]
        width = t.shape(imgs)[2]

        model_out=det_model((imgs, labels))
        proposals = pred_labels(model_out['offsets'], model_out['weights'])
        bboxes = bbox_of_proposals(proposals)
        mi = labels['mask_indices']
        if not type(mi) is t.RaggedTensor:
            mi = t.RaggedTensor.from_row_starts(mi[0], [0])
        #gt_bboxes = bbox_of_masks(mi)
        gt_bboxes = labels['bboxes']

        proposal_crops, img_crops, mask_crops, indices, boxes = t.map_fn(
            process_one_img,
            t.range(t.shape(imgs)[0]),
            fn_output_signature = (
                t.RaggedTensorSpec((None,64,64,1), t.float32, 0),
                t.RaggedTensorSpec((None,64,64,1), t.float32, 0),
                t.RaggedTensorSpec((None,64,64,1), t.float32, 0),
                t.RaggedTensorSpec((None,), t.int32, 0),
                t.RaggedTensorSpec((None, 4), t.float32, 0),
            ),
        )
    return {
          'source_image': img_crops.merge_dims(0,1),
          'proposal': proposal_crops.merge_dims(0,1),
          'gt_cell_type': indices.merge_dims(0,1),
          'gt_mask': mask_crops.merge_dims(0,1),
          'bbox': boxes.merge_dims(0,1),
          }
//  [python]
def fast_parser(inputs, det_model, min_iou = 0.1, crop_size=64, min_size=64):
    ''' this parser is based on bbox instead of mask roi '''
    def process_one_img(k):
        m_boxes, matched, m_ious, ious = box_matching(
            tf.cast(bboxes[k][None,...], tf.float32),
            tf.cast(gt_bboxes[k][None,...], tf.float32),
            )
        bb = adjust_bbox(bboxes[k], min_size=min_size, h=height, w=width)
        matched = tf.clip_by_value(matched[0,:], 0, 9999) # remove batch dimension, and -1 values(unmatched)
        proposal_crops = crop_proposals(proposals[k], bb, crop_size=crop_size)
        img_crops = crop_img(imgs[k], bb, crop_size=crop_size)
        mask_crops = crop_masks(labels['mask_indices'][k], bb, matched, crop_size=crop_size)
        ind = tf.cast(m_ious > min_iou, tf.int32) * (labels['class'][k] + 1)
        return proposal_crops, img_crops, mask_crops, ind[0,:], bb

    with tf.device("/gpu:0"):
        imgs, labels = inputs
        height = tf.shape(imgs)[1]
        width = tf.shape(imgs)[2]

        model_out=det_model((imgs, labels))
        proposals = pred_labels(model_out['offsets'], model_out['weights'])
        bboxes = bbox_of_proposals(proposals)
        mi = labels['mask_indices']
        if not type(mi) is tf.RaggedTensor:
            mi = tf.RaggedTensor.from_row_starts(mi[0], [0])
        #gt_bboxes = bbox_of_masks(mi)
        gt_bboxes = labels['bboxes']

        proposal_crops, img_crops, mask_crops, indices, boxes = tf.map_fn(
            process_one_img,
            tf.range(tf.shape(imgs)[0]),
            fn_output_signature = (
                tf.RaggedTensorSpec((None,64,64,1), tf.float32, 0),
                tf.RaggedTensorSpec((None,64,64,1), tf.float32, 0),
                tf.RaggedTensorSpec((None,64,64,1), tf.float32, 0),
                tf.RaggedTensorSpec((None,), tf.int32, 0),
                tf.RaggedTensorSpec((None, 4), tf.float32, 0),
            ),
        )
    return {
          'source_image': img_crops.merge_dims(0,1),
          'proposal': proposal_crops.merge_dims(0,1),
          'gt_cell_type': indices.merge_dims(0,1),
          'gt_mask': mask_crops.merge_dims(0,1),
          'bbox': boxes.merge_dims(0,1),
          }
//  [python]
def step_decay(step):
  curr = t.cast(step, t.float32)
  warmup_steps =4000
  d_model = 512
  d_model = t.cast(d_model, t.float32)
  arg1=  t.math.rsqrt(curr)
  arg2 = t.scalar_mul((warmup_steps ** -1.5),curr)
  return t.math.rsqrt(d_model) * t.math.minimum(arg1, arg2)
//  [python]
def step_decay(step):
  curr = tf.cast(step, tf.float32)
  warmup_steps =4000
  d_model = 512
  d_model = tf.cast(d_model, tf.float32)
  arg1=  tf.math.rsqrt(curr)
  arg2 = tf.scalar_mul((warmup_steps ** -1.5),curr)
  return tf.math.rsqrt(d_model) * tf.math.minimum(arg1, arg2)
//  [python]
def create_padding_mask(self, seq , seq_length  ): 
    return tf.to_int32(tf.sequence_mask(seq.shape[0], maxlen=seq_length))
//  [python]
def create_padding_mask(self, seq , seq_length  ): 
    return tf.cast(tf.sequence_mask(seq.shape[0], maxlen=seq_length),  dtype = tf.int32 )
//  [python]
def act(self, states, data_aug=False):
    """Act with batch of states.

    Args:
      states: tf.tensor n_batch x 64 x 64 x 3
      data_aug: bool, whether to use stochastic data aug (else deterministic)

    Returns:
      action: tf.tensor
    """
    if data_aug and self.num_augmentations > 0:
      states = states[0]
    if self.num_augmentations > 0:
      # use pad of 2 to bump 64 to 68 with 2 + 64 + 2 on each side
      img_pad = 2
      paddings = tf.constant(
          [[0, 0], [img_pad, img_pad], [img_pad, img_pad], [0, 0]],
          dtype=tf.int32)
      states = tf.pad(tf.cast(states * 255., tf.int32), paddings, 'SYMMETRIC')
      states = tf.to_float(states) / 255.
    return self.policy(states, sample=False)
//  [python]
def act(self, states, data_aug=False):
    """Act with batch of states.

    Args:
      states: tf.tensor n_batch x 64 x 64 x 3
      data_aug: bool, whether to use stochastic data aug (else deterministic)

    Returns:
      action: tf.tensor
    """
    if data_aug and self.num_augmentations > 0:
      states = states[0]
    if self.num_augmentations > 0:
      # use pad of 2 to bump 64 to 68 with 2 + 64 + 2 on each side
      img_pad = 2
      paddings = tf.constant(
          [[0, 0], [img_pad, img_pad], [img_pad, img_pad], [0, 0]],
          dtype=tf.int32)
      states = tf.cast(
          tf.pad(tf.cast(states * 255., tf.int32), paddings, 'SYMMETRIC'),
          tf.float32) / 255.
    return self.policy(states, sample=False)
//  [python]
def mean(self):
        return tf.to_float(self._dist.mean())
//  [python]
def mean(self):
        return tf.cast(self._dist.mean(), self._dtype)
//  [python]
def mode(self):
        return t.cast(self._dist.mode(), self._dtype)
//  [python]
def mode(self):
        return tf.cast(self._dist.mode(), self._dtype)
//  [python]
def sparsity_loss(self, post, prior_prob, free, scale):
    kld = tfd.kl_divergence
    u_post = tfd.Independent(
        BernoulliDist(probs=t.cast(post["u_logit"], tf.float32)), 1
    )
    u_prior = tfd.Independent(
        BernoulliDist(probs=tf.ones(post["u_logit"].shape) * prior_prob), 1
    )
    loss = kld(u_post, u_prior)
    loss = tf.maximum(tf.reduce_mean(loss), free)
    return loss * scale
//  [python]
def sparsity_loss(self, post, prior_prob, free, scale):
    kld = tfd.kl_divergence
    u_post = tfd.Independent(
        BernoulliDist(probs=tf.cast(post["u_logit"], tf.float32)), 1
    )
    u_prior = tfd.Independent(
        BernoulliDist(probs=tf.ones(post["u_logit"].shape) * prior_prob), 1
    )
    loss = kld(u_post, u_prior)
    loss = tf.maximum(tf.reduce_mean(loss), free)
    return loss * scale
//  [python]
def entropy(self):
        return tf.cast(self._dist.entropy(), self._dtype)
//  [python]
def __make_train_step(self):
        trainable_vars = self.sess.graph.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)

        learning_rate = self.params['learning_rate']

        lr_for_num_graphs_per_batch = self.params.get('lr_for_num_graphs_per_batch')
        if lr_for_num_graphs_per_batch is not None:
            # This ensures that the learning rate _per_ graph in the batch stays the same,
            # which can be important for tasks in which the loss is defined per-graph
            # (e.g., full graph regression tasks, or one-node-per-graph classification)
            lr_norm_factor = (tf.to_float(self.__placeholders['num_graphs'])
                              / tf.constant(lr_for_num_graphs_per_batch, dtype=tf.float32))
            learning_rate *= lr_norm_factor

        optimizer_name = self.params['optimizer'].lower()
        if optimizer_name == 'sgd':
            optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)
        elif optimizer_name == 'rmsprop':
            optimizer = tf.train.RMSPropOptimizer(learning_rate=learning_rate,
                                                  decay=self.params['learning_rate_decay'],
                                                  momentum=self.params['momentum'])
        elif optimizer_name == 'adam':
            optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
        else:
            raise Exception('Unknown optimizer "%s".' % (self.params['optimizer']))

        grads_and_vars = optimizer.compute_gradients(self.__ops['task_metrics']['loss'], var_list=trainable_vars)
        clipped_grads = []
        for grad, var in grads_and_vars:
            if grad is not None:
                clipped_grads.append((tf.clip_by_norm(grad, self.params['clamp_gradient_norm']), var))
            else:
                clipped_grads.append((grad, var))
        self.__ops['train_step'] = optimizer.apply_gradients(clipped_grads)
//  [python]
def __make_train_step(self):
        trainable_vars = self.sess.graph.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)

        learning_rate = self.params['learning_rate']

        lr_for_num_graphs_per_batch = self.params.get('lr_for_num_graphs_per_batch')
        if lr_for_num_graphs_per_batch is not None:
            # This ensures that the learning rate _per_ graph in the batch stays the same,
            # which can be important for tasks in which the loss is defined per-graph
            # (e.g., full graph regression tasks, or one-node-per-graph classification)
            lr_norm_factor = (tf.cast(self.__placeholders['num_graphs'], tf.float32)
                              / tf.constant(lr_for_num_graphs_per_batch, dtype=tf.float32))
            learning_rate *= lr_norm_factor

        optimizer_name = self.params['optimizer'].lower()
        if optimizer_name == 'sgd':
            optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)
        elif optimizer_name == 'rmsprop':
            optimizer = tf.train.RMSPropOptimizer(learning_rate=learning_rate,
                                                  decay=self.params['learning_rate_decay'],
                                                  momentum=self.params['momentum'])
        elif optimizer_name == 'adam':
            optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
        else:
            raise Exception('Unknown optimizer "%s".' % (self.params['optimizer']))

        grads_and_vars = optimizer.compute_gradients(self.__ops['task_metrics']['loss'], var_list=trainable_vars)
        clipped_grads = []
        for grad, var in grads_and_vars:
            if grad is not None:
                clipped_grads.append((tf.clip_by_norm(grad, self.params['clamp_gradient_norm']), var))
            else:
                clipped_grads.append((grad, var))
        self.__ops['train_step'] = optimizer.apply_gradients(clipped_grads)
//  [python]
def update_step(self, replay_buffer_iter):
    """Performs a single training step for critic and embedding.

    Args:
      replay_buffer_iter: A tensorflow graph iteratable object.

    Returns:
      Dictionary with losses to track.
    """
    transition = next(replay_buffer_iter)
    numpy_dataset = isinstance(replay_buffer_iter, np.ndarray)
    # observation: n_batch x n_timesteps x 1 x H*W*3*n_frames x 1 ->
    # n_batch x H x W x 3*n_frames
    if not numpy_dataset:
      states = transition.observation[:, 0]
      next_states = transition.observation[:, 1]
      actions = transition.action[:, 0]

      if transition.observation.dtype == tf.uint8:
        states = tf.cast(states, tf.float32) / 255.
        next_states = tf.cast(next_states, tf.float32) / 255.
    else:
      states, actions, _, next_states, _ = transition

    if self.num_augmentations > 0:
      states, next_states = tf_utils.image_aug(
          states,
          next_states,
          img_pad=4,
          num_augmentations=self.num_augmentations,
          obs_dim=64,
          channels=3,
          cropped_shape=[self.batch_size, 68, 68, 3])
      states = states[0]
      next_states = next_states[0]

    # actions = tf.gather(self.PROCGEN_ACTION_MAP.astype(np.int32).argmax(1),
    #                     actions,axis=0)

    variables = self.policy.trainable_variables

    with tf.GradientTape(watch_accessed_variables=False) as tape:
      tape.watch(variables)

      log_probs, entropy = self.policy.log_probs(
          states, actions, with_entropy=True)

      loss = -tf.reduce_mean(log_probs)  # self.alpha * entropy +

    grads = tape.gradient(loss, variables)

    self.optimizer.apply_gradients(zip(grads, variables))

#     with tf.GradientTape(watch_accessed_variables=False) as tape:
#       tape.watch([self.log_alpha])
#       alpha_loss = tf.reduce_mean(self.alpha * (entropy - self.target_entropy)

#     alpha_grads = tape.gradient(alpha_loss, [self.log_alpha])
#     self.alpha_optimizer.apply_gradients(zip(alpha_grads, [self.log_alpha]))
    alpha_loss = tf.constant(0.)

    return {
        'bc_actor_loss': loss,
        'bc_alpha': self.alpha,
        'bc_alpha_loss': alpha_loss,
        'bc_log_probs': tf.reduce_mean(log_probs),
        'bc_entropy': tf.reduce_mean(entropy)
    }
//  [python]
def sample(self, *args, **kwargs):
        return t.cast(self._dist.sample(*args, **kwargs), self._dtype)
//  [python]
def sample(self, *args, **kwargs):
        return tf.cast(self._dist.sample(*args, **kwargs), self._dtype)
//  [python]
def batch_norm_template(inputs, is_training, scope, moments_dims, bn_decay):
  """ Batch normalization on convolutional maps and beyond...
  Ref.: http://stackoverflow.com/questions/33949786/how-could-i-use-batch-normalization-in-tensorflow
  
  Args:
      inputs:        Tensor, k-D input ... x C could be BC or BHWC or BDHWC
      is_training:   boolean tf.Varialbe, true indicates training phase
      scope:         string, variable scope
      moments_dims:  a list of ints, indicating dimensions for moments calculation
      bn_decay:      float or float tensor variable, controling moving average weight
  Return:
      normed:        batch-normalized maps
  """
  is_training = tf.cast(True, tf.bool)

  with tf.variable_scope(scope) as sc:
    num_channels = inputs.get_shape()[-1].value
    beta = tf.Variable(tf.constant(0.0, shape=[num_channels]),
                       name='beta', trainable=True)
    gamma = tf.Variable(tf.constant(1.0, shape=[num_channels]),
                        name='gamma', trainable=True)
    batch_mean, batch_var = tf.nn.moments(inputs, moments_dims, name='moments')
    decay = bn_decay if bn_decay is not None else 0.9
    ema = tf.train.ExponentialMovingAverage(decay=decay)
    # Operator that maintains moving averages of variables.

    ema_apply_op = tf.cond(is_training,
                           lambda: ema.apply([batch_mean, batch_var]),
                           lambda: tf.no_op())
    
    # Update moving average and return current batch's avg and var.
    def mean_var_with_update():
      with tf.control_dependencies([ema_apply_op]):
        return tf.identity(batch_mean), tf.identity(batch_var)
    
    # ema.average returns the Variable holding the average of var.
    mean, var = tf.cond(is_training,
                        mean_var_with_update,
                        lambda: (ema.average(batch_mean), ema.average(batch_var)))
    normed = tf.nn.batch_normalization(inputs, mean, var, beta, gamma, 1e-3)
  return normed
//  [python]
def sample(self, sample_shape=(), seed=None):
        return tf.to_float(super().sample(sample_shape, seed))
//  [python]
def sample(self, sample_shape=(), seed=None):
        return tf.cast(super().sample(sample_shape, seed), self._sample_dtype)
//  [python]
def kl_loss(self, post, prior, forward, balance, free, scale, mask_kl=False):
    kld = tfd.kl_divergence
    dist = lambda x: self.get_dist(x, tf.float32)
    sg = lambda x: tf.nest.map_structure(tf.stop_gradient, x)
    lhs, rhs = (prior, post) if forward else (post, prior)
    mix = balance if forward else (1 - balance)
    #FIXME: change the 0.5 case
    if balance == 0.5:
      value = kld(dist(lhs), dist(rhs))
      loss = tf.reduce_mean(tf.maximum(value, free))
    else:
      if mask_kl:
          value_lhs = value = tf.reduce_sum(
              tf.to_float(post["u_sample"])
              * kld(dist(lhs), dist(sg(rhs))),
              axis=-1,
          )
          value_rhs = tf.reduce_sum(
              tf.to_float(post["u_sample"])
              * kld(dist(sg(lhs)), dist(rhs)),
              axis=-1,
          )
      else:
          value_lhs = value = tf.reduce_sum(
              kld(dist(lhs), dist(sg(rhs))), axis=-1
          )
          value_rhs = tf.reduce_sum(
              kld(dist(sg(lhs)), dist(rhs)), axis=-1)

      loss_lhs = tf.maximum(tf.reduce_mean(value_lhs), free)
      loss_rhs = tf.maximum(tf.reduce_mean(value_rhs), free)
      loss = mix * loss_lhs + (1 - mix) * loss_rhs
    loss *= scale
    return loss, value
//  [python]
def kl_loss(self, post, prior, forward, balance, free, scale, mask_kl=False):
    kld = tfd.kl_divergence
    dist = lambda x: self.get_dist(x, tf.float32)
    sg = lambda x: tf.nest.map_structure(tf.stop_gradient, x)
    lhs, rhs = (prior, post) if forward else (post, prior)
    mix = balance if forward else (1 - balance)
    #FIXME: change the 0.5 case
    if balance == 0.5:
      value = kld(dist(lhs), dist(rhs))
      loss = tf.reduce_mean(tf.maximum(value, free))
    else:
      if mask_kl:
          value_lhs = value = tf.reduce_sum(
              tf.cast(post["u_sample"], tf.float32)
              * kld(dist(lhs), dist(sg(rhs))),
              axis=-1,
          )
          value_rhs = tf.reduce_sum(
              tf.cast(post["u_sample"], tf.float32)
              * kld(dist(sg(lhs)), dist(rhs)),
              axis=-1,
          )
      else:
          value_lhs = value = tf.reduce_sum(
              kld(dist(lhs), dist(sg(rhs))), axis=-1
          )
          value_rhs = tf.reduce_sum(
              kld(dist(sg(lhs)), dist(rhs)), axis=-1)

      loss_lhs = tf.maximum(tf.reduce_mean(value_lhs), free)
      loss_rhs = tf.maximum(tf.reduce_mean(value_rhs), free)
      loss = mix * loss_lhs + (1 - mix) * loss_rhs
    loss *= scale
    return loss, value
//  [python]
def __init__(self, loc, scale, threshold=1, **kwargs):
        self._threshold = t.cast(threshold, loc.dtype)
        super().__init__(loc, scale, **kwargs)
//  [python]
def __init__(self, loc, scale, threshold=1, **kwargs):
        self._threshold = tf.cast(threshold, loc.dtype)
        super().__init__(loc, scale, **kwargs)
//  [python]
def sample(self, sample_shape=(), seed=None):
        # Straight through biased gradient estimator.
        sample = t.cast(super().sample(sample_shape, seed), self._sample_dtype)
        probs = super().probs_parameter()
        while len(probs.shape) < len(sample.shape):
            probs = probs[None]
        sample += t.cast(probs - t.stop_gradient(probs), self._sample_dtype)
        return sample
//  [python]
def sample(self, sample_shape=(), seed=None):
        # Straight through biased gradient estimator.
        sample = tf.cast(super().sample(sample_shape, seed), self._sample_dtype)
        probs = super().probs_parameter()
        while len(probs.shape) < len(sample.shape):
            probs = probs[None]
        sample += tf.cast(probs - tf.stop_gradient(probs), self._sample_dtype)
        return sample
//  [python]
def _make_grid(nx=20, ny=20):
        # yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
        # return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
        xv, yv = tf.meshgrid(tf.range(nx), tf.range(ny))
        return tf.reshape(tf.stack([xv, yv], 2), [1, 1, ny * nx, 2]).astype(tf.float32)
//  [python]
def _make_grid(nx=20, ny=20):
        # yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
        # return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
        xv, yv = tf.meshgrid(tf.range(nx), tf.range(ny))
        return tf.cast(tf.reshape(tf.stack([xv, yv], 2), [1, 1, ny * nx, 2]), dtype=tf.float32)
//  [python]
def time_gate_fast(phase, r_on, leak_rate, training_phase, dtype=tf.float16):
    if not training_phase:
        leak_rate = 0.0
    cond_1 = t.cast(tf.less_equal(phase, 0.5 * r_on), dtype=dtype)
    cond_2 = t.cast(
        tf.logical_and(tf.less(0.5 * r_on, phase), tf.less(phase, r_on)), dtype=dtype
    )
    cond_3 = t.cast(tf.greater_equal(phase, r_on), dtype=dtype)

    term_1 = tf.multiply(cond_1, 2.0 * phase / r_on)
    term_2 = tf.multiply(cond_2, 2.0 - 2.0 * phase / r_on)
    term_3 = tf.multiply(cond_3, leak_rate * phase)
    return term_1 + term_2 + term_3
//  [python]
def time_gate_fast(phase, r_on, leak_rate, training_phase, dtype=tf.float16):
    if not training_phase:
        leak_rate = 0.0
    cond_1 = tf.cast(tf.less_equal(phase, 0.5 * r_on), dtype=dtype)
    cond_2 = tf.cast(
        tf.logical_and(tf.less(0.5 * r_on, phase), tf.less(phase, r_on)), dtype=dtype
    )
    cond_3 = tf.cast(tf.greater_equal(phase, r_on), dtype=dtype)

    term_1 = tf.multiply(cond_1, 2.0 * phase / r_on)
    term_2 = tf.multiply(cond_2, 2.0 - 2.0 * phase / r_on)
    term_3 = tf.multiply(cond_3, leak_rate * phase)
    return term_1 + term_2 + term_3
//  [python]
def schedule(string, step):
    try:
        return float(string)
    except ValueError:
        step = t.cast(step, t.float32)
        match = re.match(r"linear\((.+),(.+),(.+)\)", string)
        if match:
            initial, final, duration = [float(group) for group in match.groups()]
            mix = tf.clip_by_value(step / duration, 0, 1)
            return (1 - mix) * initial + mix * final
        match = re.match(r"warmup\((.+),(.+)\)", string)
        if match:
            warmup, value = [float(group) for group in match.groups()]
            scale = tf.clip_by_value(step / warmup, 0, 1)
            return scale * value
        match = re.match(r"exp\((.+),(.+),(.+)\)", string)
        if match:
            initial, final, halflife = [float(group) for group in match.groups()]
            return (initial - final) * 0.5 ** (step / halflife) + final
        match = re.match(r"horizon\((.+),(.+),(.+)\)", string)
        if match:
            initial, final, duration = [float(group) for group in match.groups()]
            mix = tf.clip_by_value(step / duration, 0, 1)
            horizon = (1 - mix) * initial + mix * final
            return 1 - 1 / horizon
        raise NotImplementedError(string)
//  [python]
def schedule(string, step):
    try:
        return float(string)
    except ValueError:
        step = tf.cast(step, tf.float32)
        match = re.match(r"linear\((.+),(.+),(.+)\)", string)
        if match:
            initial, final, duration = [float(group) for group in match.groups()]
            mix = tf.clip_by_value(step / duration, 0, 1)
            return (1 - mix) * initial + mix * final
        match = re.match(r"warmup\((.+),(.+)\)", string)
        if match:
            warmup, value = [float(group) for group in match.groups()]
            scale = tf.clip_by_value(step / warmup, 0, 1)
            return scale * value
        match = re.match(r"exp\((.+),(.+),(.+)\)", string)
        if match:
            initial, final, halflife = [float(group) for group in match.groups()]
            return (initial - final) * 0.5 ** (step / halflife) + final
        match = re.match(r"horizon\((.+),(.+),(.+)\)", string)
        if match:
            initial, final, duration = [float(group) for group in match.groups()]
            mix = tf.clip_by_value(step / duration, 0, 1)
            horizon = (1 - mix) * initial + mix * final
            return 1 - 1 / horizon
        raise NotImplementedError(string)
//  [python]
def _nms(x, topk_all=100, iou_thres=0.45, conf_thres=0.25):  # agnostic NMS
        boxes, classes, scores = x
        class_inds = tf.to_float(tf.argmax(classes, axis=-1))
        scores_inp = tf.reduce_max(scores, -1)
        selected_inds = tf.image.non_max_suppression(boxes,
                                                     scores_inp,
                                                     max_output_size=topk_all,
                                                     iou_threshold=iou_thres,
                                                     score_threshold=conf_thres)
        selected_boxes = tf.gather(boxes, selected_inds)
        padded_boxes = tf.pad(selected_boxes,
                              paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]],
                              mode="CONSTANT",
                              constant_values=0.0)
        selected_scores = tf.gather(scores_inp, selected_inds)
        padded_scores = tf.pad(selected_scores,
                               paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
                               mode="CONSTANT",
                               constant_values=-1.0)
        selected_classes = tf.gather(class_inds, selected_inds)
        padded_classes = tf.pad(selected_classes,
                                paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
                                mode="CONSTANT",
                                constant_values=-1.0)
        valid_detections = tf.shape(selected_inds)[0]
        return padded_boxes, padded_scores, padded_classes, valid_detections
//  [python]
def _nms(x, topk_all=100, iou_thres=0.45, conf_thres=0.25):  # agnostic NMS
        boxes, classes, scores = x
        class_inds = tf.cast(tf.argmax(classes, axis=-1), tf.float32)
        scores_inp = tf.reduce_max(scores, -1)
        selected_inds = tf.image.non_max_suppression(boxes,
                                                     scores_inp,
                                                     max_output_size=topk_all,
                                                     iou_threshold=iou_thres,
                                                     score_threshold=conf_thres)
        selected_boxes = tf.gather(boxes, selected_inds)
        padded_boxes = tf.pad(selected_boxes,
                              paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]],
                              mode="CONSTANT",
                              constant_values=0.0)
        selected_scores = tf.gather(scores_inp, selected_inds)
        padded_scores = tf.pad(selected_scores,
                               paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
                               mode="CONSTANT",
                               constant_values=-1.0)
        selected_classes = tf.gather(class_inds, selected_inds)
        padded_classes = tf.pad(selected_classes,
                                paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]],
                                mode="CONSTANT",
                                constant_values=-1.0)
        valid_detections = tf.shape(selected_inds)[0]
        return padded_boxes, padded_scores, padded_classes, valid_detections
//  [python]
def positional_encoding(position, d_model_size):
    # create the sinusoidal pattern for the positional encoding
    angle_rads = angle_defn(np.arange(position)[:, np.newaxis], np.arange(d_model_size)[np.newaxis, :], d_model_size)

    sines = np.sin(angle_rads[:, 0::2])
    cosines = np.cos(angle_rads[:, 1::2])

    # pos_encoding = tf.cast(np.concatenate([sines, cosines], axis=-1)[np.newaxis, ...], dtype=tf.float32)
    pos_encoding = t.cast(np.concatenate([sines, cosines], axis=-1), dtype=t.float32)
    return pos_encoding
//  [python]
def positional_encoding(position, d_model_size):
    # create the sinusoidal pattern for the positional encoding
    angle_rads = angle_defn(np.arange(position)[:, np.newaxis], np.arange(d_model_size)[np.newaxis, :], d_model_size)

    sines = np.sin(angle_rads[:, 0::2])
    cosines = np.cos(angle_rads[:, 1::2])

    # pos_encoding = tf.cast(np.concatenate([sines, cosines], axis=-1)[np.newaxis, ...], dtype=tf.float32)
    pos_encoding = tf.cast(np.concatenate([sines, cosines], axis=-1), dtype=tf.float32)
    return pos_encoding
//  [python]
def grad_planner(
    state,
    num_actions,
    horizon,
    proposals,
    iterations,
    imagine,
    objective,
    kl_scale,
    step_size,
):
    dtype = prec.global_policy().compute_dtype
    B, P = list(state.values())[0].shape[0], proposals
    H, A = horizon, num_actions
    flat_state = {k: tf.repeat(v, P, 0) for k, v in state.items()}
    mean = tf.zeros((B, H, A), dtype)
    rawstd = 0.54 * tf.ones((B, H, A), dtype)
    for _ in range(iterations):
        proposals = tf.random.normal((B, P, H, A), dtype=dtype)
        with tf.GradientTape(watch_accessed_variables=False) as tape:
            tape.watch(mean)
            tape.watch(rawstd)
            std = tf.nn.softplus(rawstd)
            proposals = proposals * std[:, None] + mean[:, None]
            proposals = (
                tf.stop_gradient(tf.clip_by_value(proposals, -1, 1))
                + proposals
                - tf.stop_gradient(proposals)
            )
            flat_proposals = tf.reshape(proposals, (B * P, H, A))
            states = imagine(flat_proposals, flat_state)
            scores = objective(states)
            scores = tf.reshape(tf.reduce_sum(scores, -1), (B, P))
            div = tfd.kl_divergence(
                tfd.Normal(mean, std),
                tfd.Normal(tf.zeros_like(mean), tf.ones_like(std)),
            )
            elbo = tf.reduce_sum(scores) - kl_scale * div
            elbo /= t.cast(tf.reduce_prod(tf.shape(scores)), dtype)
        grad_mean, grad_rawstd = tape.gradient(elbo, [mean, rawstd])
        e, v = tf.nn.moments(grad_mean, [1, 2], keepdims=True)
        grad_mean /= tf.sqrt(e * e + v + 1e-4)
        e, v = tf.nn.moments(grad_rawstd, [1, 2], keepdims=True)
        grad_rawstd /= tf.sqrt(e * e + v + 1e-4)
        mean = tf.clip_by_value(mean + step_size * grad_mean, -1, 1)
        rawstd = rawstd + step_size * grad_rawstd
    return mean[:, 0, :]