diff --git a/OpenELM_Paper.pdf b/OpenELM_Paper.pdf
index 0f9b9963..3adb3e9b 100644
Binary files a/OpenELM_Paper.pdf and b/OpenELM_Paper.pdf differ
diff --git a/README.md b/README.md
index ee237b64..723970db 100644
--- a/README.md
+++ b/README.md
@@ -10,6 +10,9 @@ We want to support users with many different compute profiles!
 3. Provide a simple interface to a range of example environments for evolutionary search, to let users adapt these easily for their domain.
 4. Demonstrate the potential of evolution with LLMs.
 
+# Install
+`pip install openelm`
+
 # Features
 
 ### LLM integration with evolutionary algorithms
@@ -43,7 +46,7 @@ Roughly, ELM consists of a pipeline of different components:
 All options for these classes are defined in `configs.py`, via dataclasses which are registered as a `hydra` config, and can be overriden via the command line when running one of the example scripts such as `run_elm.py`.
 
 ## Running ELM
-`python run_elm.py` will start an ELM evolutionary search using the defaults listed in `configs.py`. These can be overriden via the command line.
+`python run_elm.py` will start an ELM evolutionary search using the defaults listed in `configs.py`. These can be overriden via the command line. For example, you can use `run_elm.py env=image_evolution` to run the Image Evolution environment.
 
 ## Sandbox
 To use the code execution sandbox, see the [sandboxing readme](https://github.com/CarperAI/OpenELM/blob/main/src/openelm/sandbox/README.md) for instructions to set it up in a Docker container with the gVisor runtime.
diff --git a/run_p3.py b/run_p3.py
index eb86b524..a0a17fae 100644
--- a/run_p3.py
+++ b/run_p3.py
@@ -3,7 +3,6 @@
 import pathlib
 import time
 from collections import Counter
-from typing import List
 
 import hydra
 import requests
@@ -11,11 +10,9 @@
 from omegaconf import OmegaConf
 
 from openelm.codegen.codegen_utilities import set_seed
-from openelm.environments import P3Problem, P3ProbSol
-from openelm.mutation_model import DiffModel, MutationModel, PromptModel
 from openelm.configs import P3Config
-from openelm.environments import P3Problem, p3_long_init_args, p3_med_init_args
-from openelm.mutation_model import DiffModel, MutationModel, PromptModel
+from openelm.environments.p3.p3 import P3Problem, P3ProbSol
+from openelm.mutation_model import MutationModel, PromptModel
 from openelm.sandbox.server.sandbox_codex_execute import ExecResult
 from openelm.utils.code_eval import pass_at_k
 
@@ -30,6 +27,7 @@
 python run_p3.py probsol=True model.model_path=Salesforce/codegen-2B-mono env.batch_size=8 iterations_per_puzzle=16
 """
 
+
 class P3:
     def __init__(self, config: P3Config) -> None:
         """
@@ -38,40 +36,47 @@ def __init__(self, config: P3Config) -> None:
         self.config: P3Config = config
 
         # Model
-        if self.config.model.model_name == 'prompt':
+        if self.config.model.model_name == "prompt":
             self.mutation_model: MutationModel = PromptModel(self.config.model)
         # elif self.config.model.model_name == 'diff':
         #     self.mutation_model: MutationModel = DiffModel(self.config.model)
 
         self.log_dir = self.cfg.output_dir
 
-
     def run(self):
         """
         Query PromptModel to generate
             self.config.probsol=False: solutions to given programming puzzle problems
             self.config.probsol=True:  new problem+solution pairs
         """
-        puzzles = requests.get("https://raw.githubusercontent.com/microsoft/PythonProgrammingPuzzles/v0.2/puzzles/puzzles.json").json()
+        puzzles = requests.get(
+            "https://raw.githubusercontent.com/microsoft/PythonProgrammingPuzzles/v0.2/puzzles/puzzles.json"
+        ).json()
         run_start_time = time.time()
         for puzzle_id in self.config.starting_seeds:
             self.config.env.starting_seed = puzzle_id
 
             puzzle = puzzles[puzzle_id]
             puzzle_start_time = time.time()
-            puzzle_dict = {'name': puzzle['name']}
-            logging.info(puzzle['name'])
+            puzzle_dict = {"name": puzzle["name"]}
+            logging.info(puzzle["name"])
 
             if self.config.probsol:
-                env = P3ProbSol(config=self.config.env, mutation_model=self.mutation_model)
+                env = P3ProbSol(
+                    config=self.config.env, mutation_model=self.mutation_model
+                )
             else:
-                env = P3Problem(config=self.config.env, mutation_model=self.mutation_model)
+                env = P3Problem(
+                    config=self.config.env, mutation_model=self.mutation_model
+                )
 
             # Run
             solutions = []
             assert self.config.iterations_per_puzzle >= self.config.env.batch_size
-            for i in range(self.config.iterations_per_puzzle // self.config.env.batch_size):
-                set_seed(i) # Change seed for each query
+            for i in range(
+                self.config.iterations_per_puzzle // self.config.env.batch_size
+            ):
+                set_seed(i)  # Change seed for each query
 
                 solutions += env.random()
 
@@ -79,37 +84,38 @@ def run(self):
             res_sols_list = []
             solved = False
             for sol in solutions:
-                res_sol_dict = {'program_str': sol.program_str}
+                res_sol_dict = {"program_str": sol.program_str}
                 if self.config.save_result_obj is not None:
                     if isinstance(sol.result_obj, ExecResult):
-                        res_sol_dict['result_obj'] = sol.result_obj.name
+                        res_sol_dict["result_obj"] = sol.result_obj.name
                     else:
-                        res_sol_dict['result_obj'] = sol.result_obj
+                        res_sol_dict["result_obj"] = sol.result_obj
 
                 fitness = env.fitness(sol)
 
                 res_sol_dict["fitness"] = fitness
                 res_sols_list.append(res_sol_dict)
                 if fitness == 1.0:
-                    solved = True # just want to save if the current problem is solved by any attempt
+                    solved = True  # just want to save if the current problem is solved by any attempt
 
-            puzzle_dict['config'] = OmegaConf.to_container(self.config)
-            puzzle_dict['solutions'] = res_sols_list
-            puzzle_dict['solved'] = solved
-            puzzle_dict['time_elapsed'] = time.time() - puzzle_start_time
+            puzzle_dict["config"] = OmegaConf.to_container(self.config)
+            puzzle_dict["solutions"] = res_sols_list
+            puzzle_dict["solved"] = solved
+            puzzle_dict["time_elapsed"] = time.time() - puzzle_start_time
 
             # Save results
             if self.config.save_results:
                 dir = f'{self.log_dir}/{puzzle_dict["name"]}/{run_start_time}'
                 pathlib.Path(dir).mkdir(parents=True, exist_ok=True)
 
-                with open(f'{dir}/results.json', 'w') as file:
+                with open(f"{dir}/results.json", "w") as file:
                     file.write(json.dumps(puzzle_dict))
 
-        logging.info(f'Successfully ran on {len(self.config.starting_seeds)}' +
-                        f'/{len(self.config.starting_seeds)}' +
-                        f' puzzles and saved any results to {self.log_dir}')
-
+        logging.info(
+            f"Successfully ran on {len(self.config.starting_seeds)}"
+            + f"/{len(self.config.starting_seeds)}"
+            + f" puzzles and saved any results to {self.log_dir}"
+        )
 
     def eval_pass_at_k(self, timestamp: str, k: int):
         """
@@ -122,7 +128,7 @@ def eval_pass_at_k(self, timestamp: str, k: int):
         """
 
         path = pathlib.Path(self.log_dir)
-        puzzle_paths = sorted(list(path.iterdir())) # Get all logged puzzles
+        puzzle_paths = sorted(list(path.iterdir()))  # Get all logged puzzles
         paks = []
         for p in puzzle_paths:
             n = 0
@@ -131,7 +137,7 @@ def eval_pass_at_k(self, timestamp: str, k: int):
             if len(timestamp) == 0:
                 # Get latest run
                 path = pathlib.Path(p)
-                run_paths = sorted(list(path.iterdir())) # Get all the runs per puzzle
+                run_paths = sorted(list(path.iterdir()))  # Get all the runs per puzzle
                 run_path = run_paths[-1]
             else:
                 # Get 'timestamp' run
diff --git a/src/openelm/algorithms/genetic.py b/src/openelm/algorithms/genetic.py
new file mode 100644
index 00000000..d2d3193b
--- /dev/null
+++ b/src/openelm/algorithms/genetic.py
@@ -0,0 +1,159 @@
+import os
+import pickle
+import random
+from pathlib import Path
+from typing import Optional, Tuple
+
+import numpy as np
+
+from openelm.configs import QDConfig
+from openelm.environments import BaseEnvironment, Genotype
+
+Phenotype = Optional[np.ndarray]
+MapIndex = Optional[tuple]
+Individual = Tuple[np.ndarray, float]
+
+
+class Pool:
+    """The pool stores a set of solutions or individuals."""
+
+    def __init__(self, pool_size: int):
+        """Initializes an empty pool.
+
+        Args:
+            pool_size (int): The number of solutions to store in the pool.
+            history_length (int): The number of historical solutions
+                to maintain in the pool.
+        """
+        self.pool_size = pool_size
+        self.pool = []
+
+    def add(self, solution, fitness):
+        """Adds a solution to the pool.
+
+        If the pool is full, the oldest solution is removed. The solution
+        is also added to the history.
+
+        Args:
+            solution: The solution to add to the pool.
+        """
+        # if new fitness is better than the worst, add it to the pool
+        if fitness > self.pool[-1][1]:
+            if len(self.pool) >= self.pool_size:
+                self.pool.pop(0)
+            self.pool.append((solution, fitness))
+            # sort the pool by fitness
+            self.pool.sort(key=lambda x: x[1], reverse=True)
+
+
+class MAPElitesBase:
+    """
+    Base class for a genetic algorithm
+    """
+
+    def __init__(
+        self,
+        env,
+        config: QDConfig,
+        init_pool: Optional[Pool] = None,
+    ):
+        """
+        The base class for a genetic algorithm, implementing common functions and search.
+
+        Args:
+            env (BaseEnvironment): The environment to evaluate solutions in. This
+            should be a subclass of `BaseEnvironment`, and should implement
+            methods to generate random solutions, mutate existing solutions,
+            and evaluate solutions for their fitness in the environment.
+            config (QDConfig): The configuration for the algorithm.
+            init_pool (Pool, optional): A pool to use for the algorithm. If not passed,
+            a new pool will be created. Defaults to None.
+        """
+        self.env: BaseEnvironment = env
+        self.config: QDConfig = config
+        self.save_history = self.config.save_history
+        self.save_snapshot_interval = self.config.save_snapshot_interval
+        self.start_step = 0
+        self.save_np_rng_state = self.config.save_np_rng_state
+        self.load_np_rng_state = self.config.load_np_rng_state
+        self.rng = np.random.default_rng(self.config.seed)
+        self.rng_generators = None
+
+        self._init_pool(init_pool, self.config.log_snapshot_dir)
+
+    def to_mapindex(self, b: Phenotype) -> MapIndex:
+        """Converts a phenotype (position in behaviour space) to a map index."""
+        raise NotImplementedError
+
+    def _init_pool(
+        self, init_map: Optional[Pool] = None, log_snapshot_dir: Optional[str] = None
+    ):
+        if init_map is None and log_snapshot_dir is None:
+            self.pool = Pool(self.config.pool_size)
+        elif init_map is not None and log_snapshot_dir is None:
+            self.pool = init_map
+        elif init_map is None and log_snapshot_dir is not None:
+            self.pool = Pool(self.config.pool_size)
+            log_path = Path(log_snapshot_dir)
+            if log_snapshot_dir and os.path.isdir(log_path):
+                stem_dir = log_path.stem
+
+                assert (
+                    "step_" in stem_dir
+                ), f"loading directory ({stem_dir}) doesn't contain 'step_' in name"
+                self.start_step = (
+                    int(stem_dir.replace("step_", "")) + 1
+                )  # add 1 to correct the iteration steps to run
+
+                snapshot_path = log_path / "pool.pkl"
+                assert os.path.isfile(
+                    snapshot_path
+                ), f'{log_path} does not contain map snapshot "pool.pkl"'
+                # first, load arrays and set them in Maps
+                # Load maps from pickle file
+                with open(snapshot_path, "rb") as f:
+                    self.pool = pickle.load(f)
+
+        print("Loading finished")
+
+    def random_selection(self) -> MapIndex:
+        """Randomly select a niche (cell) in the map that has been explored."""
+        return random.choice(self.pool.pool)
+
+    def search(self, init_steps: int, total_steps: int, atol: float = 0.0) -> str:
+        """
+        Run the genetic algorithm.
+
+        Args:
+            initsteps (int): Number of initial random solutions to generate.
+            totalsteps (int): Total number of steps to run the algorithm for,
+                including initial steps.
+            atol (float, optional): Tolerance for how close the best performing
+                solution has to be to the maximum possible fitness before the
+                search stops early. Defaults to 1.
+
+        Returns:
+            str: A string representation of the best perfoming solution. The
+                best performing solution object can be accessed via the
+                `current_max_genome` class attribute.
+        """
+        total_steps = int(total_steps)
+        for n_steps in range(total_steps):
+            if n_steps < init_steps:
+                # Initialise by generating initsteps random solutions
+                new_individuals: list[Genotype] = self.env.random()
+            else:
+                # Randomly select a batch of individuals
+                batch: list[Genotype] = []
+                for _ in range(self.env.batch_size):
+                    item = self.random_selection()
+                    batch.append(item)
+                # Mutate
+                new_individuals = self.env.mutate(batch)
+
+            for individual in new_individuals:
+                # Evaluate fitness
+                fitness = self.env.fitness(individual)
+                if np.isinf(fitness):
+                    continue
+                self.pool.add(individual, fitness)
diff --git a/src/openelm/algorithms/map_elites.py b/src/openelm/algorithms/map_elites.py
index aef21547..c90a4fd6 100644
--- a/src/openelm/algorithms/map_elites.py
+++ b/src/openelm/algorithms/map_elites.py
@@ -77,6 +77,32 @@ def __setitem__(self, map_ix, value):
             self.top[map_ix] = top_val
             self.array[(self.top[map_ix], *map_ix)] = value
 
+    def assign_fitness_in_depth(self, map_ix, value: float) -> int:
+        indices_at_bin = (slice(None),) + map_ix
+        # expecting a non-empty index, only calling this method when we know
+        # current fitness can be placed somewhere
+        insert_idx = np.where(self.array[indices_at_bin] < value)[0][-1]
+        new_bin_fitnesses = np.concatenate(
+            (
+                self.array[indices_at_bin][1 : insert_idx + 1],
+                np.array([value]),
+                self.array[indices_at_bin][insert_idx + 1 :],
+            )
+        )
+        self.array[indices_at_bin] = new_bin_fitnesses
+        return insert_idx
+
+    def insert_individual_at_depth(self, map_ix, depth, individual):
+        indices_at_bin = (slice(None),) + map_ix
+        new_bin_individuals = np.concatenate(
+            (
+                self.array[indices_at_bin][1 : depth + 1],
+                np.array([individual]),
+                self.array[indices_at_bin][depth + 1 :],
+            )
+        )
+        self.array[indices_at_bin] = new_bin_individuals
+
     @property
     def latest(self) -> np.ndarray:
         """Returns the latest values in the history buffer."""
@@ -351,54 +377,31 @@ def search(self, init_steps: int, total_steps: int, atol: float = 0.0) -> str:
             else:
                 # Randomly select a batch of elites from the map.
                 batch: list[Genotype] = []
-                for _ in range(self.env.batch_size):
-                    map_ix = self.random_selection()
-                    batch.append(self.genomes[map_ix])
+                if self.config.crossover:
+                    crossover_parents = []
+                    previous_ix = None
+                    for i in range(self.config.crossover_parents):
+                        map_ix = self.random_selection()
+                        if map_ix != previous_ix:
+                            crossover_parents.append(self.genomes[map_ix])
+                            previous_ix = map_ix
+                    batch.append(crossover_parents)
+                else:
+                    for _ in range(self.env.batch_size):
+                        map_ix = self.random_selection()
+                        batch.append(self.genomes[map_ix])
                 # Mutate the elite.
                 new_individuals = self.env.mutate(batch)
 
-            # `new_individuals` is a list of generation/mutation. We put them
-            # into the behavior space one-by-one.
-            # TODO: account for the case where multiple new individuals are
-            # placed in the same niche, for saving histories.
-            for individual in new_individuals:
-                fitness = self.env.fitness(individual)
-                if np.isinf(fitness):
-                    continue
-                map_ix = self.to_mapindex(individual.to_phenotype())
-                # if the return is None, the individual is invalid and is thrown
-                # into the recycle bin.
-                if map_ix is None:
-                    self.recycled[self.recycled_count % len(self.recycled)] = individual
-                    self.recycled_count += 1
-                    continue
-
-                if self.save_history:
-                    # TODO: thresholding
-                    self.history[map_ix].append(individual)
-                self.nonzero[map_ix] = True
-
-                # If new fitness greater than old fitness in niche, replace.
-                if fitness > self.fitnesses[map_ix]:
-                    self.fitnesses[map_ix] = fitness
-                    self.genomes[map_ix] = individual
-                # If new fitness is the highest so far, update the tracker.
-                if fitness > max_fitness:
-                    max_fitness = fitness
-                    max_genome = individual
-
-                    tbar.set_description(f"{max_fitness=:.4f}")
-                # Stop if best fitness is within atol of maximum possible fitness.
-                if np.isclose(max_fitness, self.env.max_fitness, atol=atol):
-                    break
+            max_genome, max_fitness = self.update_map(
+                new_individuals, max_genome, max_fitness
+            )
+            tbar.set_description(f"{max_fitness=:.4f}")
 
             self.fitness_history["max"].append(self.max_fitness())
             self.fitness_history["min"].append(self.min_fitness())
             self.fitness_history["mean"].append(self.mean_fitness())
             self.fitness_history["qd_score"].append(self.qd_score())
-            self.fitness_history["niches_filled"].append(self.niches_filled())
-            self.fitness_history["qd_score"].append(self.qd_score())
-            self.fitness_history["niches_filled"].append(self.niches_filled())
 
             if (
                 self.save_snapshot_interval is not None
@@ -412,6 +415,53 @@ def search(self, init_steps: int, total_steps: int, atol: float = 0.0) -> str:
         self.visualize()
         return str(max_genome)
 
+    def update_map(self, new_individuals, max_genome, max_fitness):
+        """
+        Update the map if new individuals achieve better fitness scores.
+
+        Args:
+            new_individuals (list[Genotype]) : List of new solutions
+            max_fitness : current maximum fitness
+
+        Returns:
+            max_genome : updated maximum genome
+            max_fitness : updated maximum fitness
+
+        """
+        # `new_individuals` is a list of generation/mutation. We put them
+        # into the behavior space one-by-one.
+        for individual in new_individuals:
+            fitness = self.env.fitness(individual)
+            if np.isinf(fitness):
+                continue
+            phenotype = individual.to_phenotype()
+            map_ix = self.to_mapindex(phenotype)
+
+            # if the return is None, the individual is invalid and is thrown
+            # into the recycle bin.
+            if map_ix is None:
+                self.recycled[self.recycled_count % len(self.recycled)] = individual
+                self.recycled_count += 1
+                continue
+
+            if self.save_history:
+                # TODO: thresholding
+                self.history[map_ix].append(individual)
+
+            self.nonzero[map_ix] = True
+
+            # If new fitness greater than old fitness in niche, replace.
+            if fitness > self.fitnesses[map_ix]:
+                self.fitnesses[map_ix] = fitness
+                self.genomes[map_ix] = individual
+
+            # update if new fitness is the highest so far.
+            if fitness > max_fitness:
+                max_fitness = fitness
+                max_genome = individual
+
+        return max_genome, max_fitness
+
     def niches_filled(self):
         """Get the number of niches that have been explored in the map."""
         return self.fitnesses.niches_filled
@@ -447,8 +497,11 @@ def save_results(self, step: int):
             "nonzero": self.nonzero.array,
         }
         # Save maps as pickle file
-        with open((output_folder / "maps.pkl"), "wb") as f:
-            pickle.dump(maps, f)
+        try:
+            with open((output_folder / "maps.pkl"), "wb") as f:
+                pickle.dump(maps, f)
+        except Exception:
+            pass
         if self.save_history:
             with open((output_folder / "history.pkl"), "wb") as f:
                 pickle.dump(self.history, f)
@@ -483,18 +536,19 @@ def plot_fitness(self):
         plt.plot(self.fitness_history["min"], label="Min fitness")
         plt.legend()
         plt.savefig(f"{save_path}/MAPElites_fitness_history.png")
+        plt.close("all")
 
         plt.figure()
         plt.plot(self.fitness_history["qd_score"], label="QD score")
         plt.legend()
         plt.savefig(f"{save_path}/MAPElites_qd_score.png")
+        plt.close("all")
 
         plt.figure()
         plt.plot(self.fitness_history["niches_filled"], label="Niches filled")
         plt.legend()
         plt.savefig(f"{save_path}/MAPElites_niches_filled.png")
-
-        # self.visualize_individuals()
+        plt.close("all")
 
         if len(self.map_dims) > 1:
             if len(self.fitnesses.dims) == 2:
@@ -514,8 +568,8 @@ def plot_fitness(self):
 
             plt.figure()
             plt.pcolor(map2d, cmap="inferno")
-            plt.colorbar()
             plt.savefig(f"{save_path}/MAPElites_vis.png")
+        plt.close("all")
 
     def visualize_individuals(self):
         """Visualize the genes of the best performing solution."""
diff --git a/src/openelm/benchmarks/benchmark_bugs.py b/src/openelm/benchmarks/benchmark_bugs.py
index 95d93ebc..e79064a3 100644
--- a/src/openelm/benchmarks/benchmark_bugs.py
+++ b/src/openelm/benchmarks/benchmark_bugs.py
@@ -26,7 +26,7 @@ class BenchmarkBugsConfig(BaseConfig):
             "run": {"dir": "logs/benchmarks/bugs/${hydra.job.override_dirname}"}
         }
     )
-    model_path: str = "/fsx/diff_models/diff_16b_prelim"
+    model_path: str = "CarperAI/diff-codegen-2b-v2"
     bugs_data_path: str = "/fsx/shared/diff_benchmark.json"
     mode: str = "diff"
     seed: Optional[int] = None
@@ -40,12 +40,12 @@ class BenchmarkBugsConfig(BaseConfig):
     top_p: float = 0.95
     gen_max_len: int = 256
     batch_size: int = 32
-    n_trials: int = 3200
+    n_trials: int = 500
     n_bugs_trials: int = 100
     timeout: float = 5.0
     verbose: bool = True
-    tasks: list[str] = field(default_factory=lambda: ["bugs"])
-    n_bugs: list[int] = field(default_factory=lambda: [1, 2, 3, 4, 5])
+    tasks: list[str] = field(default_factory=lambda: ["parity"])
+    n_bugs: list[int] = field(default_factory=lambda: [1])
     temp_samples: list[float] = field(default_factory=lambda: [0.8])
     sweep: bool = False
 
@@ -95,6 +95,7 @@ def benchmark_parity(self, n_bugs, **kwargs):
             elif self.cfg.mode == "diff":
                 eval_results = []
                 for text in completions:
+                    # print(text)
                     # split the diff text according to <NME>, <BEF>, <MSG>, <DFF>.
                     parsed: dict = split_diff(text)
                     # truncate the diff hunk at the first line not starting with
@@ -108,6 +109,7 @@ def benchmark_parity(self, n_bugs, **kwargs):
                         #   2. it ignores invalid lines (not starting with " ",
                         #   "+" or "-" and not being "@@ ... @@").
                         eval_results.append(apply_diff(function_str, diff_hunk))
+                        # print("Patched!")
                     else:
                         # Invalid format. No patching.
                         eval_results.append(function_str)
diff --git a/src/openelm/configs.py b/src/openelm/configs.py
index 3b75a549..932e8426 100644
--- a/src/openelm/configs.py
+++ b/src/openelm/configs.py
@@ -52,6 +52,8 @@ class QDConfig(BaseConfig):
     seed: Optional[int] = 42
     save_np_rng_state: bool = False
     load_np_rng_state: bool = False
+    crossover: bool = False
+    crossover_parents: int = 2
 
 
 @dataclass
@@ -155,7 +157,7 @@ class PromptEnvConfig(EnvConfig):
 defaults_elm = [
     {"model": "prompt"},
     {"qd": "mapelites"},
-    {"env": "qdaif"},
+    {"env": "sodarace"},
     "_self_",
 ]
 
diff --git a/src/openelm/elm.py b/src/openelm/elm.py
index d670ab42..ca23eb0e 100644
--- a/src/openelm/elm.py
+++ b/src/openelm/elm.py
@@ -1,16 +1,61 @@
-from typing import Optional
+from typing import Any, Optional
 
 from hydra.core.hydra_config import HydraConfig
 
 from openelm.configs import DiffModelConfig, ELMConfig, PromptModelConfig
-from openelm.environments import BaseEnvironment, load_algorithm, load_env
 from openelm.mutation_model import DiffModel, MutationModel, PromptModel
 
 
+def load_env(env_name: str) -> Any:
+    if env_name == "sodarace":
+        from openelm.environments.sodaracer.sodarace import Sodarace
+
+        return Sodarace
+    elif env_name == "image_evolution":
+        from openelm.environments.base import ImageOptim
+
+        return ImageOptim
+    elif env_name == "match_string":
+        from openelm.environments.base import MatchString
+
+        return MatchString
+    elif env_name == "function_optim":
+        from openelm.environments.base import FunctionOptim
+
+        return FunctionOptim
+    elif env_name == "p3_probsol":
+        from openelm.environments.p3.p3 import P3ProbSol
+
+        return P3ProbSol
+    elif env_name == "p3_problem":
+        from openelm.environments.p3.p3 import P3Problem
+
+        return P3Problem
+    elif env_name == "prompt_evolution":
+        from openelm.environments.prompt.prompt import PromptEvolution
+
+        return PromptEvolution
+    elif env_name == "qdaif":
+        from openelm.environments.poetry import PoetryEvolution
+
+        return PoetryEvolution
+    else:
+        raise ValueError(f"Unknown environment {env_name}")
+
+
+def load_algorithm(algorithm_name: str) -> Any:
+    if algorithm_name == "mapelites":
+        from openelm.algorithms.map_elites import MAPElites
+
+        return MAPElites
+    elif algorithm_name == "cvtmapelites":
+        from openelm.algorithms.map_elites import CVTMAPElites
+
+        return CVTMAPElites
+
+
 class ELM:
-    def __init__(
-        self, config: ELMConfig, env: Optional[BaseEnvironment] = None
-    ) -> None:
+    def __init__(self, config: ELMConfig, env: Optional[Any] = None) -> None:
         """
         The main class of ELM.
 
diff --git a/src/openelm/environments/__init__.py b/src/openelm/environments/__init__.py
index e6f2801d..44b52b0b 100644
--- a/src/openelm/environments/__init__.py
+++ b/src/openelm/environments/__init__.py
@@ -1,56 +1,6 @@
-from typing import Any
-
-from openelm.algorithms.map_elites import CVTMAPElites, MAPElites
 from openelm.environments.base import BaseEnvironment, Genotype
 
-
-def load_env(env_name: str) -> Any:
-    if env_name == "sodarace":
-        from openelm.environments.sodaracer.sodarace import Sodarace
-
-        return Sodarace
-    elif env_name == "image_evolution":
-        from openelm.environments.base import ImageOptim
-
-        return ImageOptim
-    elif env_name == "match_string":
-        from openelm.environments.base import MatchString
-
-        return MatchString
-    elif env_name == "function_optim":
-        from openelm.environments.base import FunctionOptim
-
-        return FunctionOptim
-    elif env_name == "p3_probsol":
-        from openelm.environments.p3.p3 import P3ProbSol
-
-        return P3ProbSol
-    elif env_name == "p3_problem":
-        from openelm.environments.p3.p3 import P3Problem
-
-        return P3Problem
-    elif env_name == "prompt_evolution":
-        from openelm.environments.prompt.prompt import PromptEvolution
-
-        return PromptEvolution
-    elif env_name == "qdaif":
-        from openelm.environments.poetry import PoetryEvolution
-
-        return PoetryEvolution
-    else:
-        raise ValueError(f"Unknown environment {env_name}")
-
-
-def load_algorithm(algorithm_name: str) -> Any:
-    if algorithm_name == "mapelites":
-        return MAPElites
-    elif algorithm_name == "cvtmapelites":
-        return CVTMAPElites
-
-
 __all__ = [
     "Genotype",
     "BaseEnvironment",
-    "load_algorithm",
-    "load_env",
 ]
diff --git a/src/openelm/environments/poetry.py b/src/openelm/environments/poetry.py
index 7e5b1de1..2b2579f8 100644
--- a/src/openelm/environments/poetry.py
+++ b/src/openelm/environments/poetry.py
@@ -139,7 +139,7 @@ def random(self) -> list[PoetryGenotype]:
         results = []
         for prompt in prompt_list:
             results.append(
-                self.mutation_model([HumanMessage(content=prompt["prompt"])]).content
+                self.mutation_model(HumanMessage(content=prompt["prompt"]).content)
             )
         return [PoetryGenotype(poem=c) for c in results]
 
@@ -148,7 +148,7 @@ def mutate(self, genomes: list[PoetryGenotype]) -> list[PoetryGenotype]:
         results = []
         for prompt in prompt_list:
             results.append(
-                self.mutation_model([HumanMessage(content=prompt["prompt"])]).content
+                self.mutation_model(HumanMessage(content=prompt["prompt"]).content)
             )
         return [PoetryGenotype(poem=c) for c in results]
 
diff --git a/src/openelm/environments/sodaracer/simulator.py b/src/openelm/environments/sodaracer/simulator.py
index 29ffd471..1836994d 100644
--- a/src/openelm/environments/sodaracer/simulator.py
+++ b/src/openelm/environments/sodaracer/simulator.py
@@ -31,7 +31,7 @@ class IESoRWorld(Framework):
     name: "IESoRWorld"
 
     def __init__(self, canvas_size: tuple[int, int] = (150, 200)):
-        super(IESoRWorld, self).__init__()
+        # super(IESoRWorld, self).__init__()
         """
         Initialize the world.
 
@@ -58,8 +58,8 @@ def __init__(self, canvas_size: tuple[int, int] = (150, 200)):
         self.gravity = b2.b2Vec2(0.0, -25.0)  # ??? Magic numbers.
         # Construct a world object, which will hold and simulate the rigid bodies.
 
-        self.world.gravity = self.gravity
-        # self.world: b2.b2World = b2.b2World(self.gravity)
+        # self.world.gravity = self.gravity
+        self.world: b2.b2World = b2.b2World(self.gravity)
         self.world.autoClearForces = False
 
         self.groundBodyDef: b2.b2BodyDef = b2.b2BodyDef()
@@ -632,8 +632,8 @@ def evaluate(self, time: float) -> float:
                 + [muscle.joint.bodyA.position[0] for muscle in self.world.muscle_list]
             )
             return abs(end + self.morphology["offsetX"])
-        except Exception as e:
-            print(e)
+        except Exception:
+            # print(e)
             # print(self.world.bone_list)
             # print(self.world.muscle_list)
             return -np.inf
diff --git a/src/openelm/utils/utils.py b/src/openelm/utils/utils.py
index 625e803f..fee9a0a6 100644
--- a/src/openelm/utils/utils.py
+++ b/src/openelm/utils/utils.py
@@ -1,3 +1,4 @@
+import os
 from dataclasses import is_dataclass
 from pathlib import Path
 
@@ -22,3 +23,9 @@ def validate_config(config):
             raise IOError(
                 "Invalid config type. Must be a path to a yaml, a dict, or dataclass."
             )
+
+
+def safe_open_w(path, *args, **kwargs):
+    """Open "path" for writing, creating any parent directories as needed."""
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    return open(path, *args, **kwargs)