code: def episodic_memory_intrinsic_rewards(
      embeddings: Array,
      num_neighbors: int,
      reward_scale: float,
      intrinsic_reward_state: Optional[IntrinsicRewardState] = None,
      constant: float = 1e-3,
      epsilon: float = 1e-4,
      cluster_distance: float = 8e-3,
      max_similarity: float = 8.0,
      max_memory_size: int = 30_000,
  ):
      """Compute intrinsic rewards for exploration via episodic memory.
  
      This method is adopted from the intrinsic reward computation used in "Never
      Give Up: Learning Directed Exploration Strategies" by Puigdomènech Badia et
      al., (2020) (https://arxiv.org/abs/2003.13350) and "Agent57: Outperforming the
      Atari Human Benchmark" by Puigdomènech Badia et al., (2020)
      (https://arxiv.org/abs/2002.06038).
  
      From an embedding, we compute the intra-episode intrinsic reward with respect
      to a pre-existing set of embeddings.
  
      NOTE: For this function to be jittable, static_argnums=[1,] must be passed, as
      the internal jax.lax.top_k(neg_distances, num_neighbors) computation in
      knn_query cannot be jitted with a dynamic num_neighbors that is passed as an
      argument.
  
      Args:
        embeddings: Array, shaped [M, D] for number of new state embeddings M and
          feature dim D.
        num_neighbors: int for K neighbors used in kNN query
        reward_scale: The β term used in the Agent57 paper to scale the reward.
        intrinsic_reward_state: An IntrinsicRewardState namedtuple, containing
          memory, next_memory_index, distance_sum, and distance_count.
          NOTE- On (only) the first call to episodic_memory_intrinsic_rewards, the
          intrinsic_reward_state is optional, if None is given, an
          IntrinsicRewardState will be initialized with default parameters,
          specifically, the memory will be initialized to an array of jnp.inf of
          shape [max_memory_size x feature dim D], and default values of 0 will be
          provided for next_memory_index, distance_sum, and distance_count.
        constant: float; small constant used for numerical stability used during
          normalizing distances.
        epsilon: float; small constant used for numerical stability when computing
          kernel output.
        cluster_distance: float; the ξ term used in the Agent57 paper to bound the
          distance rate used in the kernel computation.
        max_similarity: float; max limit of similarity; used to zero rewards when
          similarity between memories is too high to be considered 'useful' for an
          agent.
        max_memory_size: int; the maximum number of memories to store. Note that
          performance will be marginally faster if max_memory_size is an exact
          multiple of M (the number of embeddings to add to memory per call to
          episodic_memory_intrinsic_reward).
  
      Returns:
        Intrinsic reward for each embedding computed by using similarity measure to
        memories and next IntrinsicRewardState.
      """
  
      # Initialize IntrinsicRewardState if not provided to default values.
      if not intrinsic_reward_state:
          intrinsic_reward_state = IntrinsicRewardState(
              memory=jnp.inf
              * jnp.ones(shape=(max_memory_size, embeddings.shape[-1]))
          )
          # Pad the first num_neighbors entries with zeros.
          padding = jnp.zeros((num_neighbors, embeddings.shape[-1]))
          intrinsic_reward_state.memory = (
              intrinsic_reward_state.memory.at[:num_neighbors, :].set(
                  padding
              )
          )
      else:
          chex.assert_shape(
              intrinsic_reward_state.memory,
              (max_memory_size, embeddings.shape[-1]),
          )
  
      # Compute the KNN from the embeddings using the square distances from
      # the KNN d²(xₖ, x). Results are not guaranteed to be ordered.
      jit_knn_query = jax.jit(
          episodic_memory.knn_query,
          static_argnums=[
              2,
          ],
      )
      knn_query_result = jit_knn_query(
          intrinsic_reward_state.memory, embeddings, num_neighbors
      )
  
      # Insert embeddings into memory in a ring buffer fashion.
      memory = intrinsic_reward_state.memory
      start_index = (
          intrinsic_reward_state.next_memory_index % memory.shape[0]
      )
      indices = (
          jnp.arange(embeddings.shape[0]) + start_index
      ) % memory.shape[0]
      memory = jnp.asarray(memory).at[indices].set(embeddings)
  
      nn_distances_sq = knn_query_result.neighbor_neg_distances
  
      # Unpack running distance statistics, and update the running mean dₘ²
      distance_sum = intrinsic_reward_state.distance_sum
      distance_sum += jnp.sum(nn_distances_sq)
      distance_counts = intrinsic_reward_state.distance_count
      distance_counts += nn_distances_sq.size
  
      # We compute the sum of a kernel similarity with the KNN and set to zero
      # the reward when this similarity exceeds a given value (max_similarity)
      # Compute rate = d(xₖ, x)² / dₘ²
      mean_distance = distance_sum / distance_counts
      distance_rate = nn_distances_sq / (mean_distance + constant)
  
      # The distance rate becomes 0 if already small: r <- max(r-ξ, 0).
      distance_rate = jnp.maximum(
          distance_rate - cluster_distance,
          jnp.zeros_like(distance_rate),
      )
  
      # Compute the Kernel value K(xₖ, x) = ε/(rate + ε).
      kernel_output = epsilon / (distance_rate + epsilon)
  
      # Compute the similarity for the embedding x:
      # s = √(Σ_{xₖ ∈ Nₖ} K(xₖ, x)) + c
      similarity = jnp.sqrt(jnp.sum(kernel_output, axis=-1)) + constant
  
      # Compute the intrinsic reward:
      # r = 1 / s.
      reward_new = jnp.ones_like(embeddings[..., 0]) / similarity
  
      # Zero the reward if similarity is greater than max_similarity
      # r <- 0 if s > sₘₐₓ otherwise r.
      max_similarity_reached = similarity > max_similarity
      reward = jnp.where(max_similarity_reached, 0, reward_new)
  
      # r <- β * r
      reward *= reward_scale
  
      return reward, IntrinsicRewardState(
          memory=memory,
          next_memory_index=start_index
          + embeddings.shape[0] % max_memory_size,
          distance_sum=distance_sum,
          distance_count=distance_counts,
      )
  
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 4 for calling a global function with example usage: episodic_memory_intrinsic_rewards(embeddings, num_neighbors, reward_scale, intrinsic_reward_state, constant, epsilon, cluster_distance, max_similarity, max_memory_size) and returns: reward,  IntrinsicRewardState