import tensorflow as tf

            w = tf.expand_dims(w, [2])
            # Calculate policy and sample
            logits = tf.reshape(tf.matmul(U, w), [-1, num_acts])
            self.pi = tf.nn.softmax(logits)
            self.log_pi = tf.nn.log_softmax(logits)
            self.sample = policy_utils.categorical_sample(
