code: def constant_policy_targets(
      distribution: distrax.DistributionLike,
      rng_key: chex.PRNGKey,
      num_samples: int,
      weights_scale: float = 1.0,
  ) -> PolicyTarget:
      """Create policy targets with constant weights.
  
      The actions will be sampled from `distribution` and will all be associated to
      a constant `weights_scale`. If `distribution` is a (discrete or continuous)
      uniform probability distribution, distilling these targets will push the
      agent's policy towards a uniform distribution. The strength of the penalty
      associated with a non-uniform policy depends on `weights_scale` (e.g. in the
      extreme case `weights_scale==0`, distillation loss is 0 for any policy).
  
      Args:
        distribution: a `distrax` or `tfp` distribution for sampling actions.
        rng_key: a JAX pseudo random number generator.
        num_samples: number of action sampled in the PolicyTarget.
        weights_scale: constant weight associated with each action.
  
      Returns:
        a PolicyTarget to learn from.
      """
      random_actions = distribution.sample(
          seed=rng_key, sample_shape=(num_samples,)
      )
      return PolicyTarget(
          actions=random_actions,
          weights=weights_scale
          * jnp.ones((num_samples,) + distribution.batch_shape),
      )
  
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 1 for calling a global function with example usage: constant_policy_targets(distribution, rng_key, num_samples, weights_scale) and returns: PolicyTarget