class_docstr: Deep torso for Atari, from the IMPALA paper.
code: def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
      output = self.resnet(x)
      output = jax.nn.relu(output)
      output = hk.Flatten(preserve_dims=-3)(output)
      output = self.mlp_head(output)
      return output
  
imports: from typing import Any, Optional, Sequence, Tuple
  from acme.jax.networks import base
  from acme.jax.networks import duelling
  from acme.jax.networks import embedding
  from acme.jax.networks import policy_value
  from acme.jax.networks import resnet
  from acme.wrappers import observation_action_reward
  import haiku
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling the DeepAtariTorso obj's underlying member method code with example usage: obj.__call__(x) and returns: output