class_docstr: Jax Rainbow network for Full Rainbow.
  
  Attributes:
    num_actions: int, number of actions the agent can take at any state.
    num_atoms: int, the number of buckets of the value function distribution.
    noisy: bool, Whether to use noisy networks.
    dueling: bool, Whether to use dueling network architecture.
    distributional: bool, whether to use distributional RL.
class_docstr_tok: ["Jax","Rainbow","network","for","Full","Rainbow",".","Attributes",":","num_actions",":","int",",","number","of","actions","the","agent","can","take","at","any","state",".","num_atoms",":","int",",","the","number","of","buckets","of","the","value","function","distribution",".","noisy",":","bool",",","Whether","to","use","noisy","networks",".","dueling",":","bool",",","Whether","to","use","dueling","network","architecture",".","distributional",":","bool",",","whether","to","use","distributional","RL","."]
code: def setup(self):
      if self.use_impala_encoder:
          self.encoder = ImpalaEncoder(nn_scale=self.nn_scale)
      else:
          self.encoder = CNNEncoder(nn_scale=self.nn_scale)
  
imports: import collections
  from typing import Tuple
  from absl import logging
  from dopamine.discrete_domains import atari_lib
  from flax import linen
  import gin
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling the ParameterizedRainbowNetwork obj's underlying member method code with example usage: obj.setup() without return types