class_docstr: JAX DQN Network for Minatar environments.
class_docstr_tok: ["JAX","DQN","Network","for","Minatar","environments","."]
code: def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
      initializer = nn.initializers.xavier_uniform()
      x = x.astype(jnp.float32)
      x = nn.Conv(
          features=16,
          kernel_size=(3, 3),
          strides=(1, 1),
          kernel_init=initializer,
      )(x)
      x = nn.relu(x)
      x = x.reshape(-1)  # flatten
      q_values = nn.Dense(
          features=self.num_actions, kernel_init=initializer
      )(x)
      return atari_lib.DQNNetworkType(q_values)
  
imports: import time
  from dopamine.discrete_domains import atari_lib
  from flax import linen
  import gin
  import jax
  import jax.numpy
  import minatar
label_desc: Use image node_id 1 for calling the MinatarDQNNetwork obj's underlying member method code with example usage: obj.__call__(x) and returns: atari_lib