import tensorflow as tf


            U = tf.reshape(flat_logits, [-1, num_acts, self.k])

            # Calculate w
            cut_g = tf.stop_gradient(self.g)
            cut_g = tf.expand_dims(cut_g, [1])
            gstack = tf.concat([self.prev_g, cut_g], axis=1)

