import tensorflow as tf

      self._reward = tf.Variable(
          lambda: tf.zeros(batch_dims, tf.float32),
          name='reward', trainable=False)
      self._done = tf.Variable(
          lambda: tf.cast(tf.ones(batch_dims), tf.bool),
          name='done', trainable=False)

  def __getattr__(self, name):
