code: def _create_test_agent(self):
      # This dummy network allows us to deterministically anticipate that
      # action 0 will be selected by an argmax.
  
      # In Rainbow we are dealing with a distribution over Q-values,
      # which are represented as num_atoms bins, ranging from -vmax to vmax.
      # The output layer will have num_actions * num_atoms elements,
      # so each group of num_atoms weights represent the logits for a
      # particular action. By setting 1s everywhere, except for the first
      # num_atoms (representing the logits for the first action), which are
      # set to onp.arange(num_atoms), we are ensuring that the first action
      # places higher weight on higher Q-values; this results in the first
      # action being chosen.
      class MockFullRainbowNetwork(nn.Module):
          """Custom Jax network used in tests."""
  
          num_actions: int
          num_atoms: int
          noisy: bool
          dueling: bool
          distributional: bool
          inputs_preprocessed: bool = False
  
          @nn.compact
          def __call__(self, x, support, eval_mode=False, key=None):
              def custom_init(key, shape, dtype=jnp.float32):
                  del key
                  to_pick_first_action = onp.ones(shape, dtype)
                  to_pick_first_action[
                      :, : self.num_atoms
                  ] = onp.arange(1, self.num_atoms + 1)
                  return to_pick_first_action
  
              x = x.astype(jnp.float32)
              x = x.reshape((-1))  # flatten
              x = nn.Dense(
                  features=self.num_actions * self.num_atoms,
                  kernel_init=custom_init,
                  bias_init=nn.initializers.ones,
              )(x)
              logits = x.reshape((self.num_actions, self.num_atoms))
              if not self.distributional:
                  qs = jnp.sum(
                      logits, axis=-1
                  )  # Sum over all the num_atoms
                  return atari_lib.DQNNetworkType(qs)
              probabilities = nn.softmax(logits)
              qs = jnp.sum(support * probabilities, axis=1)
              return atari_lib.RainbowNetworkType(
                  qs, logits, probabilities
              )
  
      agent = full_rainbow_agent.JaxFullRainbowAgent(
          network=MockFullRainbowNetwork,
          num_actions=self._num_actions,
          num_atoms=self._num_atoms,
          vmax=self._vmax,
          distributional=True,
          epsilon_fn=lambda w, x, y, z: 0.0,  # No exploration.
      )
      # This ensures non-random action choices (since epsilon_eval = 0.0) and
      # skips the train_step.
      agent.eval_mode = True
      return agent
  
file_path: dopamine/tests/dopamine/jax/agents/full_rainbow/full_rainbow_agent_test.py
imports: from absl.testing import absltest
  from dopamine.discrete_domains import atari_lib
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.jax.agents.full_rainbow import full_rainbow_agent
  from dopamine.utils import test_utils
  from flax import linen
  import gin
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling the FullRainbowAgentTest obj's underlying member method code with example usage: obj._create_test_agent() and returns: agent, atari_lib, to_pick_first_action, atari_lib