'Sample an action from a distribution.'

def sample_action(self, logits, sampling_dim, act_dim, act_type, greedy=False):

    if self.env_spec.is_discrete(act_type):
        if greedy:
            act = tf.argmax(logits, 1)
        else:
            act = tf.reshape(tf.multinomial(logits, 1), [(-1)])
    elif self.env_spec.is_box(act_type):
        means = logits[:, :(sampling_dim / 2)]
        std = logits[:, (sampling_dim / 2):]
        if greedy:
            act = means
        else:
            batch_size = tf.shape(logits)[0]
            act = (means + (std * tf.random_normal([batch_size, act_dim])))
    else:
        assert False
    return act
