code: def plot_figure(lattice, log_dir, title):
      grid = lattice.reshape((10, 10, 10, 10))
      fig, axes = plt.subplots(10, 10, figsize=(8, 8))
      for i in range(10):
          for j in range(10):
              ax = axes[i][j]
              ax.imshow(grid[i, j])
              ax.set_axis_off()
      fig.suptitle(title, fontsize=20, fontweight="bold")
      plt.savefig(os.path.join(log_dir, "{}.png".format(title)))
  
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 2 for calling a global function with example usage: plot_figure(lattice, log_dir, title) without return types