class_docstr: Deep torso for Atari, from the IMPALA paper.
code: def __call__(self, x):
      torso_out = x / 255.0
      for i, (num_channels, num_blocks) in enumerate(
          [(16, 2), (32, 2), (32, 2)]
      ):
          conv = hk.Conv2D(
              num_channels,
              kernel_shape=[3, 3],
              stride=[1, 1],
              padding="SAME",
          )
          torso_out = conv(torso_out)
          torso_out = hk.max_pool(
              torso_out,
              window_shape=[1, 3, 3, 1],
              strides=[1, 2, 2, 1],
              padding="SAME",
          )
          for j in range(num_blocks):
              block = ResidualBlock(
                  num_channels, name=f"residual_{i}_{j}"
              )
              torso_out = block(torso_out)
  
      torso_out = jax.nn.relu(torso_out)
      torso_out = hk.Flatten()(torso_out)
      torso_out = hk.Linear(256)(torso_out)
      torso_out = jax.nn.relu(torso_out)
      return torso_out
  
imports: import collections
  import dm_env
  import haiku
  import jax.nn
  import jax.numpy
label_desc: Use image node_id 2 for calling the AtariDeepTorso obj's underlying member method code with example usage: obj.__call__(x) and returns: torso_out