def METHOD_NAME(
    environment_spec: specs.EnvironmentSpec) -> networks_lib.FeedForwardNetwork:
  """Returns DQN networks used by the agent in the multigrid environment."""
  # Check that multigrid environment is defined with discrete actions, 0-indexed
  assert np.issubdtype(environment_spec.actions.dtype, np.integer), (
      'Expected multigrid environment to have discrete actions with int dtype'
      f' but environment_spec.actions.dtype == {environment_spec.actions.dtype}'
  )
  assert environment_spec.actions.minimum == 0, (
      'Expected multigrid environment to have 0-indexed action indices, but'
      f' environment_spec.actions.minimum == {environment_spec.actions.minimum}'
  )
  num_actions = environment_spec.actions.maximum + 1
  def network(inputs):
    model = hk.Sequential([
        hk.Flatten(),
        hk.nets.MLP([50, 50, num_actions]),
    ])
    processed_inputs = multigrid_obs_preproc(inputs)
    return model(processed_inputs)
  network_hk = hk.without_apply_rng(hk.transform(network))
  dummy_obs = acme_jax_utils.add_batch_dim(
      acme_jax_utils.zeros_like(environment_spec.observations))
  return networks_lib.FeedForwardNetwork(
      init=lambda rng: network_hk.init(rng, dummy_obs), apply=network_hk.apply)