code: def setup_mnist(config, logger):
      from evojax.task.mnist import MNIST
  
      policy = ConvNetPolicy(logger=logger)
      train_task = MNIST(batch_size=config["batch_size"], test=False)
      test_task = MNIST(batch_size=config["batch_size"], test=True)
      return train_task, test_task, policy
  
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 4 for calling a global function with example usage: setup_mnist(config, logger) and returns: train_task,  test_task,  policy