8ff49a4b2dbbaf9d76ee5cadd074a955036b3e19,tf_agents/bandits/agents/greedy_reward_prediction_agent_test.py,,_get_initial_and_final_steps_with_action_mask,#Any#Any#,93

Before Change

          name="step_type"),
      tf.constant(rewards, dtype=tf.float32, name="reward"),
      tf.constant(1.0, dtype=tf.float32, shape=[batch_size], name="discount"),
      (tf.constant(observations[0] + 100.0), tf.constant(observations[1])))
  return initial_step, final_step

After Change

def _get_initial_and_final_steps_with_action_mask(observations, rewards):
  batch_size = tf.nest.flatten(observations)[0].shape[0]
  initial_step = ts.TimeStep(
      tf.constant(
          ts.StepType.FIRST,
          dtype=tf.int32,
          shape=[batch_size],
          name="step_type"),
      tf.constant(0.0, dtype=tf.float32, shape=[batch_size], name="reward"),
      tf.constant(1.0, dtype=tf.float32, shape=[batch_size], name="discount"),
      (observations[0], observations[1]))
  final_step = ts.TimeStep(
      tf.constant(
          ts.StepType.LAST,
          dtype=tf.int32,
          shape=[batch_size],
          name="step_type"),
      tf.constant(rewards, dtype=tf.float32, name="reward"),
      tf.constant(1.0, dtype=tf.float32, shape=[batch_size],
                  name="discount"), (tf.nest.map_structure(
                      lambda x: x + 100., observations[0]), observations[1]))
  return initial_step, final_step

Instances