code: def build_network(
      num_hidden_units: int, num_actions: int
  ) -> hk.Transformed:
      """Factory for a simple MLP network for approximating Q-values."""
  
      def q(obs):
          flatten = lambda x: jnp.reshape(x, (-1,))
          network = hk.Sequential(
              [flatten, nets.MLP([num_hidden_units, num_actions])]
          )
          return network(obs)
  
      return hk.without_apply_rng(hk.transform(q))
  
imports: import collections
  from absl import app
  from absl import flags
  from bsuite.environments import catch
  import dm_env
  import haiku
  from haiku import nets
  import jax
  import jax.numpy
  import numpy
  import optax
  import rlax
  from rlax.examples import experiment
label_desc: Use image node_id 1 for calling a global function with example usage: build_network(num_hidden_units, num_actions) and returns: hk, network