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

    # Configure.
    self._cell = 2
    obs_size = (5, 2, 4, 4, 3, 2)
    y = obs_size[2] // self._cell
    x = obs_size[3] // self._cell
    channels = np.prod(obs_size[4:])
    rew_size = (obs_size[0]-1, obs_size[1], x, y)

    # Input data.
    self._obs_np = np.random.uniform(size=obs_size)
    self._obs_tf = tf.placeholder(tf.float32, obs_size)

    # Expected pseudo-rewards.
    abs_diff = np.absolute(self._obs_np[1:] - self._obs_np[:-1])
    abs_diff = abs_diff.reshape((-1,) + obs_size[2:4] + (channels,))
    abs_diff = abs_diff.reshape((-1, y, self._cell, x, self._cell, channels))
    avg_abs_diff = abs_diff.mean(axis=(2, 4, 5))
    self._expected_pseudo_rewards = avg_abs_diff.reshape(rew_size)
Defines example data and expected result for the op.
trfl/pixel_control_ops_test.py