import tensorflow as tf

  def body(self, features):
    observations = features["inputs"]
    flat_x = tf.layers.flatten(observations)
    with tf.variable_scope("dense_bitwise"):
      flat_x = discretization.int_to_bit_embed(flat_x, 8, 32)

