Distributions

jaxnasium.algorithms.core.EpsilonGreedy

Like distrax.EpsilonGreedy but wrapped in a Joint and Independent for PyTree support. Assumes independent actions. Also supports action masking.

__init__(preference, *, epsilon: float = 0.0, action_mask: PyTree | None = None, dtype: jnp.dtype | type[Any] = int)

jaxnasium.algorithms.core._distributions.TanhNormal

Normal squashed through tanh, affinely rescaled to [low, high].

NOTE: obtain log-probs with sample_and_log_prob, instead of sample followed by log_prob.

https://github.com/google-deepmind/distrax/issues/7 https://github.com/google-deepmind/distrax/issues/216

mode()

jaxnasium.algorithms.core._distributions.TanhNormalFactory(low, high) -> Callable[..., TanhNormal]

jaxnasium.algorithms.core._distributions.make_independent(dist: distrax.Distribution) -> distrax.Distribution

Wraps a distrax distribution in an Independent distribution if the output space is multi-dimensional and sets the event shape accordingly.