code: def setup_brax(config):
      from evojax.task.brax_task import BraxTask
  
      train_task = BraxTask(env_name=config["env_name"], test=False)
      test_task = BraxTask(env_name=config["env_name"], test=True)
      policy = MLPPolicy(
          input_dim=train_task.obs_shape[0],
          output_dim=train_task.act_shape[0],
          hidden_dims=[32, 32, 32, 32],
      )
      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 3 for calling a global function with example usage: setup_brax(config) and returns: train_task,  test_task,  policy