class_docstr: Test the `pixel_control_loss` op.
class_docstr_tok: ["Test","the","`","pixel_control_loss","`","op","."]
docstr: Defines example data and expected result for the op.
docstr_tok: ["Defines","example","data","and","expected","result","for","the","op","."]
code: 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.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,
      )
  
file_path: trfl/trfl/pixel_control_ops_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  import tensorflow.compat.v1
  from trfl import pixel_control_ops
label_desc: Use image node_id 1 for calling the PixelControlLossTest obj's underlying member method code with example usage: obj.setUp() without return types