code: def main(config):
      log_dir = "./log/slimevolley"
      if not os.path.exists(log_dir):
          os.makedirs(log_dir, exist_ok=True)
      logger = util.create_logger(
          name="SlimeVolley", log_dir=log_dir, debug=config.debug
      )
      logger.info("EvoJAX SlimeVolley")
      logger.info("=" * 30)
  
      max_steps = 3000
      train_task = SlimeVolley(test=False, max_steps=max_steps)
      test_task = SlimeVolley(test=True, max_steps=max_steps)
      policy = MLPPolicy(
          input_dim=train_task.obs_shape[0],
          hidden_dims=[
              config.hidden_size,
          ],
          output_dim=train_task.act_shape[0],
          output_act_fn="tanh",
      )
      solver = CMA(
          pop_size=config.pop_size,
          param_size=policy.num_params,
          init_stdev=config.init_std,
          seed=config.seed,
          logger=logger,
      )
      # 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=config.n_repeats,
          n_evaluations=config.num_tests,
          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)
  
      # Visualize the policy.
      task_reset_fn = jax.jit(test_task.reset)
      policy_reset_fn = jax.jit(policy.reset)
      step_fn = jax.jit(test_task.step)
      action_fn = jax.jit(policy.get_actions)
      best_params = trainer.solver.best_params[None, :]
      key = jax.random.PRNGKey(0)[None, :]
  
      task_state = task_reset_fn(key)
      policy_state = policy_reset_fn(task_state)
      screens = []
      for _ in range(max_steps):
          action, policy_state = action_fn(
              task_state, best_params, policy_state
          )
          task_state, reward, done = step_fn(task_state, action)
          screens.append(SlimeVolley.render(task_state))
  
      gif_file = os.path.join(log_dir, "slimevolley.gif")
      screens[0].save(
          gif_file,
          save_all=True,
          append_images=screens[1:],
          duration=40,
          loop=0,
      )
      logger.info("GIF saved to {}.".format(gif_file))
  
file_path: evojax/examples/train_slimevolley.py
imports: import argparse
  import os
  import shutil
  import jax
  from evojax.task.slimevolley import SlimeVolley
  from evojax.policy.mlp import MLPPolicy
  from evojax.algo import CMA
  from evojax import Trainer
  from evojax import util
label_desc: Use image node_id 2 for calling a global function with example usage: main(config) without return types