class_docstr: Convolutional network used to compute the agent's return distributions.
class_docstr_tok: ["Convolutional","network","used","to","compute","the","agent","'s","return","distributions","."]
code: def __call__(self, x, support):
      # We need to reset the list otherwise it would have the previous values each
      # time we create a new network.
      if self.is_initializing():
          for _ in range(len(self.layer_names)):
              self.layer_names.pop()
  
      def _scale_width(n):
          return int(math.ceil(n * self.width))
  
      initializer = nn.initializers.variance_scaling(
          scale=1.0 / jnp.sqrt(3.0),
          mode="fan_in",
          distribution="uniform",
      )
      if not self.inputs_preprocessed:
          x = x.astype(jnp.float32) / 255.0
      features = (32, 64, 64)
      kernel_sizes = (8, 4, 3)
      strides = (4, 2, 1)
  
      for n_feature, kernel_size, stride in zip(
          features, kernel_sizes, strides
      ):
          layer = nn.Conv(
              features=_scale_width(n_feature),
              kernel_size=(kernel_size, kernel_size),
              strides=(stride, stride),
              kernel_init=initializer,
          )
          x = layer(x)
          x = nn.relu(x)
          x = self._record_activations(x, layer)
  
      x = x.reshape((-1))  # flatten
      layer = nn.Dense(
          features=_scale_width(512), kernel_init=initializer
      )
      x = layer(x)
      x = nn.relu(x)
      x = self._record_activations(x, layer)
      layer = nn.Dense(
          features=self.num_actions * self.num_atoms,
          kernel_init=initializer,
          name="final_layer",
      )
      x = layer(x)
      x = self._record_activations(x, layer)
      logits = x.reshape((self.num_actions, self.num_atoms))
      probabilities = nn.softmax(logits)
      q_values = jnp.sum(support * probabilities, axis=1)
      return atari_lib.RainbowNetworkType(
          q_values, logits, probabilities
      )
  
imports: import math
  from dopamine.discrete_domains import atari_lib
  from flax import linen
  import gin
  import jax.numpy
label_desc: Use image node_id 2 for calling the ScalableRainbowNetwork obj's underlying member method code with example usage: obj.__call__(x, support) and returns: atari_lib, int