class_docstr: Network for Atari.
code: def _head(self, core_output):
      policy_logits = hk.Linear(self._num_actions)(core_output)
      value = hk.Linear(1)(core_output)
      value = jnp.squeeze(value, axis=-1)
      return NetOutput(policy_logits=policy_logits, value=value)
  
imports: import collections
  import dm_env
  import haiku
  import jax.nn
  import jax.numpy
label_desc: Use image node_id 5 for calling the AtariNet obj's underlying member method code with example usage: obj._head(core_output) and returns: NetOutput