code: def setup_problem(config, logger):
      if config["problem_type"] == "cartpole_easy":
          return setup_cartpole(config, False)
      elif config["problem_type"] == "cartpole_hard":
          return setup_cartpole(config, True)
      elif config["problem_type"] == "brax":
          return setup_brax(config)
      elif config["problem_type"] == "mnist":
          return setup_mnist(config, logger)
      elif config["problem_type"] == "waterworld":
          return setup_waterworld(config)
      elif config["problem_type"] == "waterworld_ma":
          return setup_waterworld_ma(config)
  
file_path: evojax/scripts/benchmarks/problems.py
imports: import ast
  import re
  import yaml
  import json
  import numpy
  from evojax.policy import MLPPolicy
  from evojax.policy.convnet import ConvNetPolicy
label_desc: Use image node_id 1 for calling a global function with example usage: setup_problem(config, logger) and returns: setup_cartpole, setup_cartpole, setup_brax, setup_mnist, setup_waterworld, setup_waterworld_ma