class_docstr: Convolutional network used to compute the agent's return quantiles.
class_docstr_tok: ["Convolutional","network","used","to","compute","the","agent","'s","return","quantiles","."]
code: def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
      initializer = jax.nn.initializers.variance_scaling(
          scale=1.0 / jnp.sqrt(3.0),
          mode="fan_in",
          distribution="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
      x = nn.Dense(
          features=self.num_actions * self.num_atoms,
          kernel_init=initializer,
      )(x)
      logits = x.reshape((self.num_actions, self.num_atoms))
      probabilities = nn.softmax(logits)
      q_values = jnp.mean(logits, axis=1)
      return atari_lib.RainbowNetworkType(
          q_values, logits, probabilities
      )
  
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 MinatarQuantileNetwork obj's underlying member method code with example usage: obj.__call__(x) and returns: atari_lib