class_docstr: The MAP-ELITES algorithm.
class_docstr_tok: ["The","MAP-ELITES","algorithm","."]
docstr: Initialization function.
  
  Args:
      pop_size - Population size.
      param_size - Parameter size.
      bd_extractors - A list of behavior descriptor extractors.
      init_params - Initial parameters, all zeros if not given.
      iso_sigma - Standard deviation for Gaussian sampling for mutation.
      line_sigma - Parameter for cross over.
      seed - Random seed for parameters sampling.
      logger - Logging utility.
docstr_tok: ["Initialization","function",".","Args",":","pop_size","-","Population","size",".","param_size","-","Parameter","size",".","bd_extractors","-","A","list","of","behavior","descriptor","extractors",".","init_params","-","Initial","parameters",",","all","zeros","if","not","given",".","iso_sigma","-","Standard","deviation","for","Gaussian","sampling","for","mutation",".","line_sigma","-","Parameter","for","cross","over",".","seed","-","Random","seed","for","parameters","sampling",".","logger","-","Logging","utility","."]
code: def __init__(
      self,
      pop_size: int,
      param_size: int,
      bd_extractor: BDExtractor,
      init_params: Optional[Union[jnp.ndarray, np.ndarray]] = None,
      iso_sigma: float = 0.05,
      line_sigma: float = 0.5,
      seed: int = 0,
      logger: logging.Logger = None,
  ):
      """Initialization function.
  
      Args:
          pop_size - Population size.
          param_size - Parameter size.
          bd_extractors - A list of behavior descriptor extractors.
          init_params - Initial parameters, all zeros if not given.
          iso_sigma - Standard deviation for Gaussian sampling for mutation.
          line_sigma - Parameter for cross over.
          seed - Random seed for parameters sampling.
          logger - Logging utility.
      """
  
      if logger is None:
          self._logger = create_logger(name="MAP_ELITES")
      else:
          self._logger = logger
  
      self.pop_size = abs(pop_size)
      self.param_size = param_size
      self.bd_names = [x[0] for x in bd_extractor.bd_spec]
      self.bd_n_bins = [x[1] for x in bd_extractor.bd_spec]
      self.params_lattice = jnp.zeros(
          (np.prod(self.bd_n_bins), param_size)
      )
      self.fitness_lattice = -float("Inf") * jnp.ones(
          np.prod(self.bd_n_bins)
      )
      self.occupancy_lattice = jnp.zeros(
          np.prod(self.bd_n_bins), dtype=jnp.int32
      )
      self.population = None
      self.bin_idx = jnp.zeros(self.pop_size, dtype=jnp.int32)
      self.key = jax.random.PRNGKey(seed)
  
      init_param = (
          jnp.array(init_params)
          if init_params is not None
          else jnp.zeros(param_size)
      )
  
      def sample_parents(key, occupancy, params):
          new_key, sample_key, mutate_key = jax.random.split(key, 3)
          parents = jax.lax.cond(
              occupancy.sum() > 0,
              lambda: jnp.take(
                  params,
                  axis=0,
                  indices=jax.random.choice(
                      key=sample_key,
                      a=jnp.arange(occupancy.size),
                      replace=True,
                      p=occupancy / occupancy.sum(),
                      shape=(2 * self.pop_size,),
                  ),
              ),
              lambda: init_param[None, :]
              + iso_sigma
              * jax.random.normal(
                  sample_key, shape=(2 * self.pop_size, self.param_size)
              ),
          )
          return (
              new_key,
              mutate_key,
              parents.reshape((2, self.pop_size, -1)),
          )
  
      self._sample_parents = jax.jit(sample_parents)
  
      def generate_population(parents, key):
          k1, k2 = jax.random.split(key, 2)
          return (
              parents[0]
              + iso_sigma
              * jax.random.normal(key=k1, shape=parents[0].shape)
              +
              # Uniform sampling instead of Gaussian.
              jax.random.uniform(
                  key=k2, minval=0, maxval=line_sigma, shape=()
              )
              * (parents[1] - parents[0])
          )
  
      self._gen_pop = jax.jit(generate_population)
  
      def get_bin_idx(task_state):
          bd_idx = [
              task_state.__dict__[name].astype(int)
              for name in self.bd_names
          ]
          return jnp.ravel_multi_index(
              bd_idx, self.bd_n_bins, mode="clip"
          )
  
      self._get_bin_idx = jax.jit(jax.vmap(get_bin_idx))
  
      def update_fitness_and_param(
          target_bin,
          bin_idx,
          fitness,
          fitness_lattice,
          param,
          param_lattice,
      ):
          best_ix = jnp.where(
              bin_idx == target_bin, fitness, fitness_lattice.min()
          ).argmax()
          best_fitness = fitness[best_ix]
          new_fitness_lattice = jnp.where(
              best_fitness > fitness_lattice[target_bin],
              best_fitness,
              fitness_lattice[target_bin],
          )
          new_param_lattice = jnp.where(
              best_fitness > fitness_lattice[target_bin],
              param[best_ix],
              param_lattice[target_bin],
          )
          return new_fitness_lattice, new_param_lattice
  
      self._update_lattices = jax.jit(
          jax.vmap(
              update_fitness_and_param,
              in_axes=(0, None, None, None, None, None),
          )
      )
  
file_path: evojax/evojax/algo/map_elites.py
imports: import logging
  import numpy
  from typing import Optional
  from typing import Union
  import jax
  import jax.numpy
  from evojax.algo.base import QualityDiversityMethod
  from evojax.task.base import TaskState
  from evojax.task.base import BDExtractor
  from evojax.util import create_logger
label_desc: Use image node_id 1 to create a new MAPElites object from inherited base classes: QualityDiversityMethod with example: obj = MAPElites(pop_size, param_size, bd_extractor, init_params, iso_sigma, line_sigma, seed, logger)