fix(parametric_action_distribution): sum kl divergence over event_ndi…

…ms in parametric action distribution (#142)

jumanji/training/networks/parametric_distribution.py

@@ -49,8 +49,11 @@ def __init__(
         """
         self._param_size = param_size
         self._postprocessor = postprocessor
+        if event_ndims not in [0, 1]:
+            raise ValueError(
+                f"Event ndims {event_ndims} is not supported, expected value in [0, 1]."
+            )
         self._event_ndims = event_ndims  # rank of events
-        assert event_ndims in [0, 1]
 
     @abc.abstractmethod
     def create_dist(self, parameters: chex.Array) -> Distribution:
@@ -91,8 +94,6 @@ def log_prob(self, parameters: chex.Array, raw_actions: chex.Array) -> chex.Arra
         log_probs -= self._postprocessor.forward_log_det_jacobian(raw_actions)
         if self._event_ndims == 1:
             log_probs = jnp.sum(log_probs, axis=-1)  # sum over action dimension
-        else:
-            assert self._event_ndims == 0
         return log_probs
 
     def entropy(self, parameters: chex.Array, seed: chex.PRNGKey) -> chex.Array:
@@ -102,17 +103,23 @@ def entropy(self, parameters: chex.Array, seed: chex.PRNGKey) -> chex.Array:
         entropy += self._postprocessor.forward_log_det_jacobian(dist.sample(seed=seed))
         if self._event_ndims == 1:
             entropy = jnp.sum(entropy, axis=-1)
-        else:
-            assert self._event_ndims == 0
         return entropy
 
     def kl_divergence(
         self, parameters: chex.Array, other_parameters: chex.Array
     ) -> chex.Array:
         """KL divergence is invariant with respect to transformation by the same bijector."""
+        if not isinstance(self._postprocessor, IdentityBijector):
+            raise ValueError(
+                f"The current post_processor used ({self._postprocessor}) is a non-identity"
+                "bijector which does not implement kl_divergence."
+            )
         dist = self.create_dist(parameters)
         other_dist = self.create_dist(other_parameters)
-        return dist.kl_divergence(other_dist)
+        kl_divergence = dist.kl_divergence(other_dist)
+        if self._event_ndims == 1:
+            kl_divergence = jnp.sum(kl_divergence, axis=-1)
+        return kl_divergence
 
 
 class CategoricalParametricDistribution(ParametricDistribution):