code: def main(config):
      log_dir = "./log/ant_map_elites"
      if not os.path.exists(log_dir):
          os.makedirs(log_dir, exist_ok=True)
      logger = util.create_logger(
          name="AntMapElites", log_dir=log_dir, debug=config.debug
      )
  
      logger.info("EvoJAX AntMapElites Demo")
      logger.info("=" * 30)
  
      bd_extractor = AntBDExtractor(logger=logger)
      train_task = BraxTask(
          env_name="ant",
          max_steps=500,
          bd_extractor=bd_extractor,
          test=False,
      )
      test_task = BraxTask(
          env_name="ant", bd_extractor=bd_extractor, test=True
      )
      policy = MLPPolicy(
          input_dim=train_task.obs_shape[0],
          hidden_dims=[32, 32, 32, 32],
          output_dim=train_task.act_shape[0],
      )
      solver = MAPElites(
          pop_size=config.pop_size,
          param_size=policy.num_params,
          bd_extractor=bd_extractor,
          iso_sigma=config.iso_sigma,
          line_sigma=config.line_sigma,
          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)
  
      # Visualize the results.
      qd_file = os.path.join(log_dir, "qd_lattices.npz")
      with open(qd_file, "rb") as f:
          data = np.load(f)
          params_lattice = data["params_lattice"]
          fitness_lattice = data["fitness_lattice"]
          occupancy_lattice = data["occupancy_lattice"]
      plot_figure(occupancy_lattice, log_dir, "occupancy")
      plot_figure(fitness_lattice, log_dir, "score")
  
      # Visualize the top policies.
      if config.save_gif:
          import jax
          import jax.numpy as jnp
          from brax import envs
          from brax.io import image
  
          @partial(jax.jit, static_argnums=(1,))
          def get_qp(state, ix):
              return jax.tree_map(lambda x: x[ix], state.qp)
  
          num_viz = 3
          idx = fitness_lattice.argsort()[-num_viz:]
          bins = [np.unravel_index(ix, (10, 10, 10, 10)) for ix in idx]
          logger.info(
              "Best {} policies: indices={}, bins={}".format(
                  num_viz, idx, bins
              )
          )
  
          policy_params = jnp.array(params_lattice[idx])
          task_reset_fn = jax.jit(test_task.reset)
          policy_reset_fn = jax.jit(policy.reset)
          step_fn = jax.jit(test_task.step)
          act_fn = jax.jit(policy.get_actions)
  
          total_reward = jnp.zeros(num_viz)
          valid_masks = jnp.ones(num_viz)
          rollouts = {i: [] for i in range(num_viz)}
          keys = jnp.repeat(
              jax.random.PRNGKey(seed=42)[None, :],
              repeats=num_viz,
              axis=0,
          )
          task_state = task_reset_fn(key=keys)
          policy_state = policy_reset_fn(task_state)
  
          for step in range(test_task.max_steps):
              for i in range(num_viz):
                  rollouts[i].append(get_qp(task_state.state, i))
              act, policy_state = act_fn(
                  task_state, policy_params, policy_state
              )
              task_state, reward, done = step_fn(task_state, act)
              total_reward = total_reward + reward * valid_masks
              valid_masks = valid_masks * (1 - done)
          logger.info("test_rewards={}".format(total_reward))
  
          logger.info("Start saving GIFs, this can take some time ...")
          env_fn = envs.create_fn(env_name="ant", legacy_spring=True)
          env = env_fn()
          for i in range(num_viz):
              qps = jax.tree_map(lambda x: np.array(x), rollouts[i])
              frames = [
                  Image.fromarray(
                      image.render_array(
                          env.sys, qp, 320, 240, None, None, 2
                      )
                  )
                  for qp in qps
              ]
              frames[0].save(
                  os.path.join(log_dir, "bin_{}.gif".format(bins[i])),
                  format="png",
                  append_images=frames[1:],
                  save_all=True,
                  duration=env.sys.config.dt * 1000,
                  loop=0,
              )
  
file_path: evojax/examples/train_ant_map_elites.py
imports: import argparse
  import os
  import numpy
  import matplotlib.pyplot
  from PIL import Image
  from functools import partial
  from evojax import Trainer
  from evojax.task.brax_task import BraxTask
  from evojax.task.brax_task import AntBDExtractor
  from evojax.policy import MLPPolicy
  from evojax.algo import MAPElites
  from evojax import util
label_desc: Use image node_id 3 for calling a global function with example usage: main(config) and returns: jax