code: def sampled_policy_distillation_loss(
      distribution: distrax.DistributionLike,
      policy_targets: PolicyTarget,
      stop_target_gradients: bool = True,
  ) -> chex.Numeric:
      """Compute a sampled cross-entropy-like loss.
  
      The loss corresponds to taking the mean of `-weights * log_prob(actions)`,
      where the weights and actions come from a `PolicyTarget` object, and the mean
      is computed over the N samples in the `policy_target` as well as any batch
      dimension. The loss is suitable for both discrete and continuous actions.
  
      Args:
        distribution: a predicted `distrax` or `tfp` distribution.
        policy_targets: a policy target to learn from.
        stop_target_gradients: bool indicating whether or not to apply a stop
          gradient to the policy_targets, default True.
  
      Returns:
        a scalar loss.
      """
      # Optionally, stop gradients from propagating into the targets and into
      # the actions; the latter is mostly relevant in continuous control.
      weights = jax.lax.select(
          stop_target_gradients,
          jax.lax.stop_gradient(policy_targets.weights),
          policy_targets.weights,
      )
      actions = jax.lax.select(
          stop_target_gradients,
          jax.lax.stop_gradient(policy_targets.actions),
          policy_targets.actions,
      )
      # Compute log-probabilities.
      log_probs = distribution.log_prob(actions)
      # Assert shapes are compatible.
      chex.assert_equal_shape([weights, log_probs])
      # We avoid NaNs from `0 * (-inf)` by using `0 * min_logp` in that case.
      min_logp = jnp.finfo(log_probs.dtype).min
      # We average over the samples, over time and batch, and if the actions are
      # a continuous vector also over the actions.
      return -jnp.mean(weights * jnp.maximum(log_probs, min_logp))
  
file_path: rlax/rlax/_src/policy_targets.py
imports: import functools
  import chex
  import distrax
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling a global function with example usage: sampled_policy_distillation_loss(distribution, policy_targets, stop_target_gradients) and returns: unknown