Reduce the number of reshapes in `pixel_control_rewards` using `tf.nn.avg_pool3d`.
Reduce the number of reshapes in `pixel_control_rewards` using `tf.nn.avg_pool3d`.

PiperOrigin-RevId: 242452517

deepmind/trfl