code: def add_dirichlet_noise(
      key: Array,
      prior: Array,
      dirichlet_alpha: float,
      dirichlet_fraction: float,
  ) -> Array:
      """Returns discrete actions with noise drawn from a Dirichlet distribution.
  
      See "Mastering the Game of Go without Human Knowledge" by Silver et. al. 2017
      (https://discovery.ucl.ac.uk/id/eprint/10045895/1/agz_unformatted_nature.pdf),
      "A General Reinforcement Learning Algorithm that Masters Chess, Shogi and
      Go Through Self-Play" by Silver et. al. 2018
      (http://airesearch.com/wp-content/uploads/2016/01/deepmind-mastering-go.pdf),
      and "Mastering Atari, Go, Chess and  Shogi by Planning with a Learned Model"
      by Schrittwieser et. al., 2019 (https://arxiv.org/abs/1911.08265).
  
      The AlphaZero family of algorithms adds noise sampled from a symmetric
      Dirichlet distribution to the prior policy generated by MCTS. Because the
      agent then samples from this new, noisy prior over actions, this encourages
      better exploration of the root node's children.
  
      Specifically, this computes:
  
              noise ~ Dirichlet(alpha)
              noisy_prior = (1 - fraction) * prior + fraction * noise
  
      Note that alpha is a single float to draw from a symmetric Dirichlet.
  
      For reference values, AlphaZero uses 0.3, 0.15, 0.03 for Chess, Shogi, and
      Go respectively, and MuZero uses 0.25 for Atari.
  
  
      Args:
        key: a key from `jax.random`.
        prior: 2-dim continuous prior policy vector of shapes [B, N], for B batch
          size and N num_actions.
        dirichlet_alpha: concentration parameter to parametrize Dirichlet
          distribution.
        dirichlet_fraction: float from 0 to 1 interpolating between using only the
          prior policy or just the noise.
  
      Returns:
        noisy action, of the same shape as input action.
      """
      chex.assert_rank(prior, 2)
      chex.assert_type([dirichlet_alpha, dirichlet_fraction], float)
  
      batch_size, num_actions = prior.shape
      noise = jax.random.dirichlet(
          key=key,
          alpha=jnp.full(
              shape=(num_actions,), fill_value=dirichlet_alpha
          ),
          shape=(batch_size,),
      )
  
      noisy_prior = (
          1 - dirichlet_fraction
      ) * prior + dirichlet_fraction * noise
  
      return noisy_prior
  
file_path: rlax/rlax/_src/exploration.py
imports: from typing import Optional, Union
  import chex
  import jax
  import jax.numpy
  from rlax._src import episodic_memory
label_desc: Use image node_id 3 for calling a global function with example usage: add_dirichlet_noise(key, prior, dirichlet_alpha, dirichlet_fraction) and returns: noisy_prior