code: def make_policy_network(
      action_spec: specs.DiscreteArray,
  ) -> snt.Module:
      return snt.Sequential(
          [
              snt.Flatten(),
              snt.nets.MLP([64, 64, action_spec.num_values]),
          ]
      )
  
imports: import functools
  import operator
  from absl import app
  from absl import flags
  import acme
  from acme import specs
  from acme import types
  from acme.agents.tf import actors
  from acme.agents.tf.bc import learning
  from acme.agents.tf.dqfd import bsuite_demonstrations
  from acme.tf import utils
  from acme.utils import counting
  from acme.utils import loggers
  from acme.wrappers import single_precision
  import bsuite
  import reverb
  import sonnet
  import tensorflow
  import tree
  import trfl
label_desc: Use image node_id 1 for calling a global function with example usage: make_policy_network(action_spec) and returns: snt