def setUp(self):
    """Defines example data and expected result for the op."""
    super(PixelControlLossTest, self).setUp()

    # Observation shape is (2,2,3) (i.e., height 2, width 2, and 3 channels).
    # We will use no cropping, and a cell size of 1. We have num_actions = 3,
    # meaning our Q values should be (2,2,3). We will set the Q value equal to
    # the observation.
    self.seq_length = 3
    self.batch_size = 1
    num_actions = 3
    obs_shape = (2, 2, num_actions)
    self.discount = 0.9
    self.cell_size = 1
    self.scale = 1.0

    # Create ops to feed actions and rewards.
    self.observations_ph = tf.placeholder(
        shape=(self.seq_length+1, self.batch_size)+obs_shape, dtype=tf.float32)
    self.action_values_ph = tf.placeholder(
        shape=(self.seq_length+1, self.batch_size)+obs_shape, dtype=tf.float32)
    self.actions_ph = tf.placeholder(
        shape=(self.seq_length, self.batch_size), dtype=tf.int32)

    # Observations.
    obs1 = np.array([[[1, 2, 3], [3, 4, 5]], [[5, 6, 7], [7, 8, 9]]])
    obs2 = np.array([[[7, 8, 9], [1, 2, 3]], [[3, 4, 5], [5, 6, 7]]])
    obs3 = np.array([[[5, 6, 7], [7, 8, 9]], [[1, 2, 3], [3, 4, 5]]])
    obs4 = np.array([[[3, 4, 5], [5, 6, 7]], [[7, 8, 9], [1, 2, 3]]])

    # Actions.
    action1 = 0
    action2 = 1
    action3 = 2

    # Compute loss for constant discount.
    qa_tm1 = obs3[:, :, action3]
    reward3 = np.mean(np.abs(obs4 - obs3), axis=2)
    qmax_t = np.amax(obs4, axis=2)
    target = reward3 + self.discount * qmax_t
    error3 = target - qa_tm1

    qa_tm1 = obs2[:, :, action2]
    reward2 = np.mean(np.abs(obs3 - obs2), axis=2)
    target = reward2 + self.discount * target
    error2 = target - qa_tm1

    qa_tm1 = obs1[:, :, action1]
    reward1 = np.mean(np.abs(obs2 - obs1), axis=2)
    target = reward1 + self.discount * target
    error1 = target - qa_tm1

    # Compute loss for episode termination with discount 0.
    qa_tm1 = obs1[:, :, action1]
    reward1 = np.mean(np.abs(obs2 - obs1), axis=2)
    target = reward1 + 0. * target
    error1_term = target - qa_tm1

    self.error = np.sum(
        np.square(error1) + np.square(error2) + np.square(error3)) * 0.5
    self.error_term = np.sum(
        np.square(error1_term) + np.square(error2) + np.square(error3)) * 0.5

    # Placeholder data.
    self.observations = np.expand_dims(
        np.stack([obs1, obs2, obs3, obs4], axis=0), axis=1)
    self.action_values = self.observations
    self.actions = np.stack(
        [np.array([action1]), np.array([action2]), np.array([action3])], axis=0)
Defines example data and expected result for the op.
trfl/pixel_control_ops_test.py