docstr: Check that embeddings are added appropriately to a custom memory.
docstr_tok: ["Check","that","embeddings","are","added","appropriately","to","a","custom","memory","."]
code: def test_custom_memory(self):
      """Check that embeddings are added appropriately to a custom memory."""
  
      @self.variant
      def episodic_memory_intrinsic_rewards(
          embeddings, memory, reward_scale
      ):
          return exploration.episodic_memory_intrinsic_rewards(
              embeddings,
              self.num_neighbors,
              reward_scale,
              exploration.IntrinsicRewardState(
                  memory=memory, next_memory_index=2
              ),
              max_memory_size=4,
          )
  
      embeddings = np.array([[2.0, 2.0], [3.0, 3.0], [4.0, 4.0]])
      memory = np.array(
          [
              [
                  -1.0,
                  -1.0,
              ],
              [1.0, 1.0],
              [0.0, 0.0],
              [0.0, 0.0],
          ]
      )
      _, intrinsic_reward_state = episodic_memory_intrinsic_rewards(
          embeddings, memory, self.reward_scale
      )
  
      np.testing.assert_array_equal(
          intrinsic_reward_state.memory,
          # Embeddings should have been added in a ring buffer way.
          np.array(
              [
                  [
                      4.0,
                      4.0,
                  ],
                  [1.0, 1.0],
                  [2.0, 2.0],
                  [3.0, 3.0],
              ]
          ),
      )
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import chex
  import jax
  import numpy
  from rlax._src import exploration
label_desc: Use image node_id 3 for calling the EMIntrinsicRewardTest obj's underlying member method code with example usage: obj.test_custom_memory() and returns: exploration