docstr: Keras network for tests.
docstr_tok: ["Keras","network","for","tests","."]
code: def _create_test_agent(self):
      """Keras network for tests."""
      # This dummy network allows us to deterministically anticipate that
      # action 0 will be selected by an argmax.
  
      # In Quantile we are dealing with a distribution over Q-values,
      # which are represented as num_atoms quantiles.
      # The output layer will have num_actions * num_atoms elements,
      # so each group of num_atoms weights represent the value quantiles for
      # a particular action. By setting 1s everywhere, except for the first
      # num_atoms (representing the quantiles for the first action), which
      # are set to onp.arange(num_atoms), we are ensuring that the first action
      # has a higher expected Q-value; this results in the first
      # action being chosen.
      class MockQuantileNetwork(linen.Module):
          """Custom Jax network used in tests."""
  
          num_actions: int
          num_atoms: int
          inputs_preprocessed: bool = False
  
          @linen.compact
          def __call__(self, x):
              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 = linen.Dense(
                  features=self.num_actions * self.num_atoms,
                  kernel_init=custom_init,
                  bias_init=linen.initializers.ones,
              )(x)
              logits = x.reshape((self.num_actions, self.num_atoms))
              probabilities = linen.softmax(logits)
              qs = jnp.mean(logits, axis=1)
              return atari_lib.RainbowNetworkType(
                  qs, logits, probabilities
              )
  
      agent = quantile_agent.JaxQuantileAgent(
          network=MockQuantileNetwork,
          num_actions=self.num_actions,
          num_atoms=self._num_atoms,
          min_replay_history=self._min_replay_history,
          epsilon_fn=lambda w, x, y, z: 0.0,  # No exploration.
          epsilon_eval=0.0,
          epsilon_decay_period=self._epsilon_decay_period,
      )
      # This ensures non-random action choices (since epsilon_eval = 0.0) and
      # skips the train_step.
      agent.eval_mode = True
      return agent
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import absltest
  from dopamine.discrete_domains import atari_lib
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.jax.agents.quantile import quantile_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 JaxQuantileAgentTest obj's underlying member method code with example usage: obj._create_test_agent() and returns: agent, atari_lib, to_pick_first_action