code: def main(config):
      log_dir = "./log/mnist"
      if not os.path.exists(log_dir):
          os.makedirs(log_dir, exist_ok=True)
      logger = util.create_logger(
          name="MNIST", log_dir=log_dir, debug=config.debug
      )
      logger.info("EvoJAX MNIST Demo")
      logger.info("=" * 30)
  
      policy = ConvNetPolicy(logger=logger)
      train_task = MNIST(batch_size=config.batch_size, test=False)
      test_task = MNIST(batch_size=config.batch_size, test=True)
      solver = PGPE(
          pop_size=config.pop_size,
          param_size=policy.num_params,
          optimizer="adam",
          center_learning_rate=config.center_lr,
          stdev_learning_rate=config.std_lr,
          init_stdev=config.init_std,
          logger=logger,
          seed=config.seed,
      )
  
      # Train.
      trainer = Trainer(
          policy=policy,
          solver=solver,
          train_task=train_task,
          test_task=test_task,
          max_iter=config.max_iter,
          log_interval=config.log_interval,
          test_interval=config.test_interval,
          n_repeats=1,
          n_evaluations=1,
          seed=config.seed,
          log_dir=log_dir,
          logger=logger,
      )
      trainer.run(demo_mode=False)
  
      # Test the final model.
      src_file = os.path.join(log_dir, "best.npz")
      tar_file = os.path.join(log_dir, "model.npz")
      shutil.copy(src_file, tar_file)
      trainer.model_dir = log_dir
      trainer.run(demo_mode=True)
  
file_path: evojax/examples/train_mnist.py
imports: import argparse
  import os
  import shutil
  from evojax import Trainer
  from evojax.task.mnist import MNIST
  from evojax.policy.convnet import ConvNetPolicy
  from evojax.algo import PGPE
  from evojax import util
label_desc: Use image node_id 2 for calling a global function with example usage: main(config) without return types