import tensorflow as tf

        epoch = tf.where(rand < 0.1, tf.zeros_like(epoch), epoch)
        # Embed the epoch number.
        emb_epoch = common_layers.embedding(epoch, 32, 32)  # [batch, 32]
        flat_x = tf.concat([flat_x, emb_epoch], axis=1)
      flat_x = tf.layers.dropout(flat_x, rate=dropout)
      x = tf.layers.dense(flat_x, 128, activation=tf.nn.relu)

      logits = tf.layers.dense(
          x, self.hparams.problem.num_actions, name="dense2"
      )
