How to get entropy from normal distribution?

[MyNameIsArko]

Hi,
I'm in the process of writing PPO in Flax and Jax. I found out that I can get normal distribution using mean + std * jax.random.normal() and I can get log_prob from jax.scipy.stats.norm.logpdf(). But I don't know how to get entropy from it, because even though scipy has scipy.stats.norm.entropy(), it didn't get into jax implementation of scipy.

So now I'm wondering how could I get entropy of normal distribution? Anyone got an idea?

[jakevdp]

JAX doesn't have any built-in entropy function, but fortunately the formula for the normal distribution entropy is pretty straightforward to define yourself:

import jax.numpy as jnp

def normal_entropy(sigma):
  return 0.5 * (jnp.log(2 * jnp.pi * sigma ** 2) + 1)

[sharadmv]

You could also use TFP:

import tensorflow_probability.substrates.jax as tfp

tfp.distributions.Normal(loc=0., scale=1.).entropy()

[MyNameIsArko]

Thanks, this seems a lot easier