class_docstr: Impala Network which also outputs penultimate representation layers.
code: def setup(self):
      logging.info("\t Creating %s ...", self.__class__.__name__)
      logging.info("\t Creating ImpalaDQNNetwork ...")
      logging.info("\t num_blocks: %s", self.num_blocks)
      logging.info("\t nn_scale: %s", self.nn_scale)
      logging.info("\t stack_sizes: %s", self.stack_sizes)
      self._stacks = [
          Stack(
              num_ch=stack_size * self.nn_scale,
              num_blocks=self.num_blocks,
          )
          for stack_size in self.stack_sizes
      ]
  
imports: import itertools
  import time
  from typing import Optional, Sequence, Tuple, Union
  from absl import logging
  from dopamine.discrete_domains import atari_lib
  from flax import linen
  import gin
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling the ImpalaEncoder obj's underlying member method code with example usage: obj.setup() without return types