code: def setup_waterworld_ma(config, num_agents=16, max_steps=500):
      from evojax.task.ma_waterworld import MultiAgentWaterWorld
  
      train_task = MultiAgentWaterWorld(
          num_agents=num_agents, test=False, max_steps=max_steps
      )
      test_task = MultiAgentWaterWorld(
          num_agents=num_agents, test=True, max_steps=max_steps
      )
      policy = MLPPolicy(
          input_dim=train_task.obs_shape[-1],
          hidden_dims=[
              config["hidden_size"],
          ],
          output_dim=train_task.act_shape[-1],
          output_act_fn="softmax",
      )
      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 6 for calling a global function with example usage: setup_waterworld_ma(config, num_agents, max_steps) and returns: train_task,  test_task,  policy