class_docstr: A simple network.
class_docstr_tok: ["A","simple","network","."]
docstr: Process a batch of observations.
docstr_tok: ["Process","a","batch","of","observations","."]
code: def __call__(
      self,
      timestep: dm_env.TimeStep,
  ) -> tuple[jax.Array, jax.Array]:
      """Process a batch of observations."""
      torso = hk.Sequential(
          [
              hk.Flatten(),
              hk.Linear(128),
              jax.nn.relu,
              hk.Linear(64),
              jax.nn.relu,
          ]
      )
      hidden = torso(timestep.observation)
      policy_logits = hk.Linear(self._num_actions)(hidden)
      baseline = hk.Linear(1)(hidden)
      baseline = jnp.squeeze(baseline, axis=-1)
      return policy_logits, baseline
  
imports: import functools
  import queue
  import threading
  from typing import Any, Callable, NamedTuple
  from absl import app
  from absl import logging
  from bsuite.environments import catch
  import dm_env
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  import rlax
label_desc: Use image node_id 2 for calling the SimpleNet obj's underlying member method code with example usage: obj.__call__(timestep) and returns: policy_logits,  baseline