Lunar Lander Relanded: Using Dask to Distribute Evaluations

This example extends the Lunar Lander tutorial by using Dask to distribute evaluations and thus speed things up. It also adds in more logging capabilities and a CLI.

  1"""Uses CMA-ME to train agents with linear policies in Lunar Lander.
  2
  3Install the following dependencies before running this example -- swig must be
  4installed before box2d can be installed, hence it is a separate command:
  5    pip install swig
  6    pip install ribs[visualize] tqdm fire gymnasium[box2d]==0.29.1 moviepy>=1.0.0 dask>=2.0.0 distributed>=2.0.0 bokeh>=2.0.0
  7
  8This script uses the same setup as the tutorial, but it also uses Dask instead
  9of Python's multiprocessing to parallelize evaluations on a single machine and
 10adds in a CLI. Refer to the tutorial here:
 11https://docs.pyribs.org/en/stable/tutorials/lunar_lander.html for more info.
 12
 13You should not need much familiarity with Dask to read this example. However, if
 14you would like to know more about Dask, we recommend referring to the quickstart
 15for Dask distributed: https://distributed.dask.org/en/latest/quickstart.html.
 16
 17This script creates an output directory (defaults to `lunar_lander_output/`, see
 18the --outdir flag) with the following files:
 19
 20    - archive.csv: The CSV representation of the final archive, obtained with
 21      as_pandas().
 22    - archive_ccdf.png: A plot showing the (unnormalized) complementary
 23      cumulative distribution function of objectives in the archive. For
 24      each objective p on the x-axis, this plot shows the number of
 25      solutions that had an objective of at least p.
 26    - heatmap.png: A heatmap showing the performance of solutions in the
 27      archive.
 28    - metrics.json: Metrics about the run, saved as a mapping from the metric
 29      name to a list of x values (iteration number) and a list of y values
 30      (metric value) for that metric.
 31    - {metric_name}.png: Plots of the metrics, currently just `archive_size` and
 32      `max_score`.
 33
 34In evaluation mode (--run-eval flag), the script will read in the archive from
 35the output directory and simulate 10 random solutions from the archive. It will
 36write videos of these simulations to a `videos/` subdirectory in the output
 37directory.
 38
 39Usage:
 40    # Basic usage - should take ~1 hour with 4 cores.
 41    python lunar_lander.py NUM_WORKERS
 42    # Now open the Dask dashboard at http://localhost:8787 to view worker
 43    # status.
 44
 45    # Evaluation mode. If you passed a different outdir and/or env_seed when
 46    # running the algorithm with the command above, you must pass the same
 47    # outdir and/or env_seed here.
 48    python lunar_lander.py --run-eval
 49Help:
 50    python lunar_lander.py --help
 51"""
 52import json
 53import time
 54from pathlib import Path
 55
 56import fire
 57import gymnasium as gym
 58import matplotlib.pyplot as plt
 59import numpy as np
 60import pandas as pd
 61import tqdm
 62from dask.distributed import Client, LocalCluster
 63
 64from ribs.archives import GridArchive
 65from ribs.emitters import EvolutionStrategyEmitter
 66from ribs.schedulers import Scheduler
 67from ribs.visualize import grid_archive_heatmap
 68
 69
 70def simulate(model, seed=None, video_env=None):
 71    """Simulates the lunar lander model.
 72
 73    Args:
 74        model (np.ndarray): The array of weights for the linear policy.
 75        seed (int): The seed for the environment.
 76        video_env (gym.Env): If passed in, this will be used instead of creating
 77            a new env. This is used primarily for recording video during
 78            evaluation.
 79    Returns:
 80        total_reward (float): The reward accrued by the lander throughout its
 81            trajectory.
 82        impact_x_pos (float): The x position of the lander when it touches the
 83            ground for the first time.
 84        impact_y_vel (float): The y velocity of the lander when it touches the
 85            ground for the first time.
 86    """
 87    if video_env is None:
 88        # Since we are using multiple processes, it is simpler if each worker
 89        # just creates their own copy of the environment instead of trying to
 90        # share the environment. This also makes the function "pure." However,
 91        # we should use the video_env if it is passed in.
 92        env = gym.make("LunarLander-v2")
 93    else:
 94        env = video_env
 95
 96    action_dim = env.action_space.n
 97    obs_dim = env.observation_space.shape[0]
 98    model = model.reshape((action_dim, obs_dim))
 99
100    total_reward = 0.0
101    impact_x_pos = None
102    impact_y_vel = None
103    all_y_vels = []
104    obs, _ = env.reset(seed=seed)
105    done = False
106
107    while not done:
108        action = np.argmax(model @ obs)  # Linear policy.
109        obs, reward, terminated, truncated, _ = env.step(action)
110        done = terminated or truncated
111        total_reward += reward
112
113        # Refer to the definition of state here:
114        # https://gymnasium.farama.org/environments/box2d/lunar_lander/
115        x_pos = obs[0]
116        y_vel = obs[3]
117        leg0_touch = bool(obs[6])
118        leg1_touch = bool(obs[7])
119        all_y_vels.append(y_vel)
120
121        # Check if the lunar lander is impacting for the first time.
122        if impact_x_pos is None and (leg0_touch or leg1_touch):
123            impact_x_pos = x_pos
124            impact_y_vel = y_vel
125
126    # If the lunar lander did not land, set the x-pos to the one from the final
127    # timestep, and set the y-vel to the max y-vel (we use min since the lander
128    # goes down).
129    if impact_x_pos is None:
130        impact_x_pos = x_pos
131        impact_y_vel = min(all_y_vels)
132
133    # Only close the env if it was not a video env.
134    if video_env is None:
135        env.close()
136
137    return total_reward, impact_x_pos, impact_y_vel
138
139
140def create_scheduler(seed, n_emitters, sigma0, batch_size):
141    """Creates the Scheduler based on given configurations.
142
143    See lunar_lander_main() for description of args.
144
145    Returns:
146        A pyribs scheduler set up for CMA-ME (i.e. it has
147        EvolutionStrategyEmitter's and a GridArchive).
148    """
149    env = gym.make("LunarLander-v2")
150    action_dim = env.action_space.n
151    obs_dim = env.observation_space.shape[0]
152    initial_model = np.zeros((action_dim, obs_dim))
153    archive = GridArchive(
154        solution_dim=initial_model.size,
155        dims=[50, 50],  # 50 cells in each dimension.
156        # (-1, 1) for x-pos and (-3, 0) for y-vel.
157        ranges=[(-1.0, 1.0), (-3.0, 0.0)],
158        seed=seed)
159
160    # If we create the emitters with identical seeds, they will all output the
161    # same initial solutions. The algorithm should still work -- eventually, the
162    # emitters will produce different solutions because they get different
163    # responses when inserting into the archive. However, using different seeds
164    # avoids this problem altogether.
165    seeds = ([None] * n_emitters
166             if seed is None else [seed + i for i in range(n_emitters)])
167
168    # We use the EvolutionStrategyEmitter to create an ImprovementEmitter.
169    emitters = [
170        EvolutionStrategyEmitter(
171            archive,
172            x0=initial_model.flatten(),
173            sigma0=sigma0,
174            ranker="2imp",
175            batch_size=batch_size,
176            seed=s,
177        ) for s in seeds
178    ]
179
180    scheduler = Scheduler(archive, emitters)
181    return scheduler
182
183
184def run_search(client, scheduler, env_seed, iterations, log_freq):
185    """Runs the QD algorithm for the given number of iterations.
186
187    Args:
188        client (Client): A Dask client providing access to workers.
189        scheduler (Scheduler): pyribs scheduler.
190        env_seed (int): Seed for the environment.
191        iterations (int): Iterations to run.
192        log_freq (int): Number of iterations to wait before recording metrics.
193    Returns:
194        dict: A mapping from various metric names to a list of "x" and "y"
195        values where x is the iteration and y is the value of the metric. Think
196        of each entry as the x's and y's for a matplotlib plot.
197    """
198    print(
199        "> Starting search.\n"
200        "  - Open Dask's dashboard at http://localhost:8787 to monitor workers."
201    )
202
203    metrics = {
204        "Max Score": {
205            "x": [],
206            "y": [],
207        },
208        "Archive Size": {
209            "x": [0],
210            "y": [0],
211        },
212    }
213
214    start_time = time.time()
215    for itr in tqdm.trange(1, iterations + 1):
216        # Request models from the scheduler.
217        sols = scheduler.ask()
218
219        # Evaluate the models and record the objectives and measures.
220        objs, meas = [], []
221
222        # Ask the Dask client to distribute the simulations among the Dask
223        # workers, then gather the results of the simulations.
224        futures = client.map(lambda model: simulate(model, env_seed), sols)
225        results = client.gather(futures)
226
227        # Process the results.
228        for obj, impact_x_pos, impact_y_vel in results:
229            objs.append(obj)
230            meas.append([impact_x_pos, impact_y_vel])
231
232        # Send the results back to the scheduler.
233        scheduler.tell(objs, meas)
234
235        # Logging.
236        if itr % log_freq == 0 or itr == iterations:
237            elapsed_time = time.time() - start_time
238            metrics["Max Score"]["x"].append(itr)
239            metrics["Max Score"]["y"].append(scheduler.archive.stats.obj_max)
240            metrics["Archive Size"]["x"].append(itr)
241            metrics["Archive Size"]["y"].append(len(scheduler.archive))
242            tqdm.tqdm.write(
243                f"> {itr} itrs completed after {elapsed_time:.2f} s\n"
244                f"  - Max Score: {metrics['Max Score']['y'][-1]}\n"
245                f"  - Archive Size: {metrics['Archive Size']['y'][-1]}")
246
247    return metrics
248
249
250def save_heatmap(archive, filename):
251    """Saves a heatmap of the scheduler's archive to the filename.
252
253    Args:
254        archive (GridArchive): Archive with results from an experiment.
255        filename (str): Path to an image file.
256    """
257    fig, ax = plt.subplots(figsize=(8, 6))
258    grid_archive_heatmap(archive, vmin=-300, vmax=300, ax=ax)
259    ax.invert_yaxis()  # Makes more sense if larger velocities are on top.
260    ax.set_ylabel("Impact y-velocity")
261    ax.set_xlabel("Impact x-position")
262    fig.savefig(filename)
263
264
265def save_metrics(outdir, metrics):
266    """Saves metrics to png plots and a JSON file.
267
268    Args:
269        outdir (Path): output directory for saving files.
270        metrics (dict): Metrics as output by run_search.
271    """
272    # Plots.
273    for metric in metrics:
274        fig, ax = plt.subplots()
275        ax.plot(metrics[metric]["x"], metrics[metric]["y"])
276        ax.set_title(metric)
277        ax.set_xlabel("Iteration")
278        fig.savefig(str(outdir / f"{metric.lower().replace(' ', '_')}.png"))
279
280    # JSON file.
281    with (outdir / "metrics.json").open("w") as file:
282        json.dump(metrics, file, indent=2)
283
284
285def save_ccdf(archive, filename):
286    """Saves a CCDF showing the distribution of the archive's objectives.
287
288    CCDF = Complementary Cumulative Distribution Function (see
289    https://en.wikipedia.org/wiki/Cumulative_distribution_function#Complementary_cumulative_distribution_function_(tail_distribution)).
290    The CCDF plotted here is not normalized to the range (0,1). This may help
291    when comparing CCDF's among archives with different amounts of coverage
292    (i.e. when one archive has more cells filled).
293
294    Args:
295        archive (GridArchive): Archive with results from an experiment.
296        filename (str): Path to an image file.
297    """
298    fig, ax = plt.subplots()
299    ax.hist(
300        archive.as_pandas(include_solutions=False)["objective"],
301        50,  # Number of cells.
302        histtype="step",
303        density=False,
304        cumulative=-1)  # CCDF rather than CDF.
305    ax.set_xlabel("Objectives")
306    ax.set_ylabel("Num. Entries")
307    ax.set_title("Distribution of Archive Objectives")
308    fig.savefig(filename)
309
310
311def run_evaluation(outdir, env_seed):
312    """Simulates 10 random archive solutions and saves videos of them.
313
314    Videos are saved to outdir / videos.
315
316    Args:
317        outdir (Path): Path object for the output directory from which to
318            retrieve the archive and save videos.
319        env_seed (int): Seed for the environment.
320    """
321    df = pd.read_csv(outdir / "archive.csv")
322    indices = np.random.permutation(len(df))[:10]
323
324    # Use a single env so that all the videos go to the same directory.
325    video_env = gym.wrappers.RecordVideo(
326        gym.make("LunarLander-v2", render_mode="rgb_array"),
327        video_folder=str(outdir / "videos"),
328        # This will ensure all episodes are recorded as videos.
329        episode_trigger=lambda idx: True,
330    )
331
332    for idx in indices:
333        model = np.array(df.loc[idx, "solution_0":])
334        reward, impact_x_pos, impact_y_vel = simulate(model, env_seed,
335                                                      video_env)
336        print(f"=== Index {idx} ===\n"
337              "Model:\n"
338              f"{model}\n"
339              f"Reward: {reward}\n"
340              f"Impact x-pos: {impact_x_pos}\n"
341              f"Impact y-vel: {impact_y_vel}\n")
342
343    video_env.close()
344
345
346def lunar_lander_main(workers=4,
347                      env_seed=52,
348                      iterations=500,
349                      log_freq=25,
350                      n_emitters=5,
351                      batch_size=30,
352                      sigma0=1.0,
353                      seed=None,
354                      outdir="lunar_lander_output",
355                      run_eval=False):
356    """Uses CMA-ME to train linear agents in Lunar Lander.
357
358    Args:
359        workers (int): Number of workers to use for simulations.
360        env_seed (int): Environment seed. The default gives the flat terrain
361            from the tutorial.
362        iterations (int): Number of iterations to run the algorithm.
363        log_freq (int): Number of iterations to wait before recording metrics
364            and saving heatmap.
365        n_emitters (int): Number of emitters.
366        batch_size (int): Batch size of each emitter.
367        sigma0 (float): Initial step size of each emitter.
368        seed (seed): Random seed for the pyribs components.
369        outdir (str): Directory for Lunar Lander output.
370        run_eval (bool): Pass this flag to run an evaluation of 10 random
371            solutions selected from the archive in the `outdir`.
372    """
373    outdir = Path(outdir)
374
375    if run_eval:
376        run_evaluation(outdir, env_seed)
377        return
378
379    # Make the directory here so that it is not made when running eval.
380    outdir.mkdir(exist_ok=True)
381
382    # Setup Dask. The client connects to a "cluster" running on this machine.
383    # The cluster simply manages several concurrent worker processes. If using
384    # Dask across many workers, we would set up a more complicated cluster and
385    # connect the client to it.
386    cluster = LocalCluster(
387        processes=True,  # Each worker is a process.
388        n_workers=workers,  # Create this many worker processes.
389        threads_per_worker=1,  # Each worker process is single-threaded.
390    )
391    client = Client(cluster)
392
393    # CMA-ME.
394    scheduler = create_scheduler(seed, n_emitters, sigma0, batch_size)
395    metrics = run_search(client, scheduler, env_seed, iterations, log_freq)
396
397    # Outputs.
398    scheduler.archive.as_pandas().to_csv(outdir / "archive.csv")
399    save_ccdf(scheduler.archive, str(outdir / "archive_ccdf.png"))
400    save_heatmap(scheduler.archive, str(outdir / "heatmap.png"))
401    save_metrics(outdir, metrics)
402
403
404if __name__ == "__main__":
405    fire.Fire(lunar_lander_main)