Add support for rendering during training in benchmark configurations

README.md

@@ -145,7 +145,8 @@ puffer eval puffer_drive carla \
 ```
 
 Scenario renders are captured during the benchmark pass and written as
-interactive HTML with retained `.replay.zlib` files. To render only episodes
+self-contained interactive HTML; the intermediate `.replay.zlib` bundles are
+deleted afterwards unless `eval.keep_zlib_replays=true`. To render only episodes
 where `offroad_rate > 0` instead:
 
 ```bash
@@ -158,6 +159,8 @@ puffer eval puffer_drive carla \
 
 Set `env.eval_training_render=true` to evaluate and render the Gigaflow
 environment distribution saved in the checkpoint's adjacent `config.yaml`.
+The selected benchmark's `map_dir` and `num_maps` override the checkpoint values,
+so checkpoints remain portable across machines with different dataset paths.
 
 Use `eval.num_agents`, not `env.num_agents`, to configure evaluation capacity.
 Evaluation outputs are written under

docs/evaluation.md

@@ -105,8 +105,10 @@ puffer eval puffer_drive carla_fast \
 
 `eval.render_scenarios=true` records each of the benchmark's configured
 `num_scenarios` during the metrics rollout. It writes the completed
-`.replay.zlib` files incrementally, then renders one interactive HTML page per
-scenario and builds a navigable `index.html`.
+`.replay.zlib` files incrementally in `replays_zlib/`, then renders one
+interactive HTML page per scenario and builds a navigable `index.html`. The
+compressed files are removed after rendering by default; use
+`eval.keep_zlib_replays=true` to keep them.
 
 `eval.capture_observations=true` also stores policy observations.
 
@@ -115,8 +117,7 @@ To evaluate and render the environment distribution used during training, run:
 ```bash
 puffer eval puffer_drive carla_fast \
   load_model_path=path/to/model.pt \
-  env.eval_training_render=true \
-  env.map_dir=pufferlib/resources/drive/binaries/carla
+  env.eval_training_render=true
 ```
 
 ## Filtered replay and rendering
@@ -168,7 +169,7 @@ eval/<benchmark>[_<output_name>]/<timestamp>/
 ├── resolved_benchmark.yaml
 ├── episode_metrics.csv
 ├── evaluation_summary.json
-├── replays/                        # only when render_scenarios=true
+├── replays_zlib/                   # only when render_scenarios=true and keep_zlib_replays=true
 │   └── *.replay.zlib
 ├── rendered_replays/               # only when render_scenarios=true
 │   ├── *.html
@@ -177,7 +178,7 @@ eval/<benchmark>[_<output_name>]/<timestamp>/
     ├── selected_failures.csv
     ├── episode_metrics.csv
     ├── evaluation_summary.json
-    ├── replays/
+    ├── replays_zlib/                # only when keep_zlib_replays=true
     │   └── *.replay.zlib
     └── rendered_replays/
         ├── *.html

pufferlib/config/evaluation/benchmark.yaml

@@ -58,6 +58,15 @@ benchmarks:
       control_mode: control_vehicles
       map_dir: pufferlib/resources/drive/binaries/carla
 
+  - name: carla_render
+    seed: 42
+    num_scenarios: 8
+    env:
+      simulation_mode: gigaflow
+      num_maps: 8
+      map_dir: pufferlib/resources/drive/binaries/carla
+      eval_training_render: true
+
   - name: womd_multi
     seed: 42
     num_scenarios: 1000

pufferlib/config/puffer_drive.yaml

@@ -303,6 +303,8 @@ eval:
   output_dir_name: eval
   # Capture and render every scenario during the standard benchmark pass.
   render_scenarios: false
+  # Keep the intermediate .replay.zlib bundles after rendering HTML.
+  keep_zlib_replays: false
   # Render scenarios where any configured metric is positive.
   # Use all_infractions for every infraction type; null disables it.
   render_filter: null

pufferlib/config_schema.py

@@ -418,6 +418,7 @@ class EvaluationConfig:
     output_name: str | None = MISSING
     output_dir_name: str = _constrained_field(NONEMPTY_STRING_CONSTRAINT)
     render_scenarios: bool = MISSING
+    keep_zlib_replays: bool = MISSING
     render_filter: Any = MISSING
     max_rendered_failures: int | None = _constrained_field(POSITIVE_INT_CONSTRAINT)
     failure_replay_csv: str | None = MISSING
@@ -788,9 +789,24 @@ def normalize_puffer_drive_benchmarks(environment_config, benchmarks, context, v
         simulation_mode = benchmark_environment.get("simulation_mode")
         if simulation_mode not in ("gigaflow", "replay"):
             _raise_config_error(context, f"{benchmark_path}.env.simulation_mode", "must be 'gigaflow' or 'replay'")
+        eval_training_render = benchmark_environment.get("eval_training_render", False)
+        if not isinstance(eval_training_render, bool):
+            _raise_config_error(
+                context,
+                f"{benchmark_path}.env.eval_training_render",
+                "must be a boolean",
+            )
         control_mode = benchmark_environment.get("control_mode")
-        if not isinstance(control_mode, str) or not control_mode:
+        if control_mode is not None and (not isinstance(control_mode, str) or not control_mode):
+            _raise_config_error(context, f"{benchmark_path}.env.control_mode", "must be a non-empty string")
+        if control_mode is None and not eval_training_render:
             _raise_config_error(context, f"{benchmark_path}.env.control_mode", "must be a non-empty string")
+        if eval_training_render and simulation_mode != "gigaflow":
+            _raise_config_error(
+                context,
+                f"{benchmark_path}.env.eval_training_render",
+                "is only supported in gigaflow mode",
+            )
 
         seed = benchmark.get("seed")
         if seed is None:
@@ -814,7 +830,7 @@ def normalize_puffer_drive_benchmarks(environment_config, benchmarks, context, v
 
         max_agents_per_env = benchmark_environment.get("max_agents_per_env")
         single_agent_replay = simulation_mode == "replay" and control_mode == "control_sdc_only"
-        if max_agents_per_env is None and not single_agent_replay:
+        if max_agents_per_env is None and not single_agent_replay and not eval_training_render:
             _raise_config_error(context, f"{benchmark_path}.env.max_agents_per_env", "must be a positive integer")
         if max_agents_per_env is not None:
             _validate_value_constraint(

pufferlib/ocean/evaluation_utils/eval_replay.py

@@ -1,6 +1,7 @@
 import numbers
 import os
 import pickle
+import shutil
 import zlib
 from concurrent.futures import ThreadPoolExecutor
 
@@ -12,6 +13,10 @@
 import pufferlib.viz
 
 
+ZLIB_REPLAY_DIR_NAME = "replays_zlib"
+ZLIB_REPLAY_SUFFIX = ".replay.zlib"
+
+
 def _eval_replay_stem(summary, episode_id):
     map_stem = os.path.splitext(os.path.basename(summary["map_name"]))[0]
     return f"{map_stem}__seed_{summary['seed']}__episode_{episode_id:06d}"
@@ -140,7 +145,7 @@ def queue_replay(self, summary, episode_id):
         for replay_key, history_values in self.policy_history.items():
             replay[replay_key] = history_values[:episode_length, global_agent_start:global_agent_end]
         replay_stem = _eval_replay_stem(summary, self.episode_id_offset + episode_id)
-        replay_path = os.path.abspath(os.path.join(self.replay_output_dir, f"{replay_stem}.replay.zlib"))
+        replay_path = os.path.abspath(os.path.join(self.replay_output_dir, f"{replay_stem}{ZLIB_REPLAY_SUFFIX}"))
         self.pending_replays.append((replay_environment["scenario"], replay, replay_path))
         summary["has_replay"] = 1
         summary["replay_path"] = replay_path
@@ -159,7 +164,7 @@ def write_pending(self):
         self.pending_replays = []
 
 
-def _render_eval_replays(episode_summaries, out_dir):
+def _render_eval_replays(episode_summaries, out_dir, keep_zlib_replays):
     """Render captured eval replays as navigable HTML pages plus an index."""
     render_dir = os.path.join(out_dir, "rendered_replays")
     os.makedirs(render_dir, exist_ok=True)
@@ -172,10 +177,9 @@ def _render_eval_replays(episode_summaries, out_dir):
         if not replay_path or not os.path.isfile(replay_path):
             raise RuntimeError(f"Cannot render episode {episode_id}: replay file is missing: {replay_path}")
         replay_filename = os.path.basename(replay_path)
-        replay_suffix = ".replay.zlib"
-        if not replay_filename.endswith(replay_suffix):
+        if not replay_filename.endswith(ZLIB_REPLAY_SUFFIX):
             raise RuntimeError(f"Cannot render episode {episode_id}: unexpected replay filename: {replay_filename}")
-        html_filename = f"{replay_filename[: -len(replay_suffix)]}.html"
+        html_filename = f"{replay_filename[: -len(ZLIB_REPLAY_SUFFIX)]}.html"
         output_path = os.path.join(render_dir, html_filename)
         replay_paths.append(replay_path)
         output_paths.append(output_path)
@@ -197,4 +201,6 @@ def _render_eval_replays(episode_summaries, out_dir):
     pufferlib.viz.build_gallery_index(render_dir, file_metrics=file_metrics)
     print(f"Rendered {len(episode_summaries)} replay pages into {render_dir}")
     print(f"Wrote replay index to {os.path.join(render_dir, 'index.html')}")
+    if not keep_zlib_replays:
+        shutil.rmtree(os.path.join(out_dir, ZLIB_REPLAY_DIR_NAME))
     return render_dir

pufferlib/ocean/evaluation_utils/evaluation_utils.py

@@ -201,6 +201,8 @@ def _build_benchmark_args(base_args, benchmark, environment_config):
     args["vec"]["seed"] = seed
     if eval_training_render:
         args["eval"]["action_selection"] = pufferlib.pytorch.ACTION_SELECT_SAMPLE
+        args["env"]["map_dir"] = benchmark_environment_config["map_dir"]
+        args["env"]["num_maps"] = benchmark_environment_config["num_maps"]
     else:
         args["env"].update(copy.deepcopy(environment_config))
         args["env"].update(copy.deepcopy(benchmark_environment_config))
@@ -239,6 +241,9 @@ def _finalize_benchmark_args(args, cli_overrides, eval_training_render, validati
 
 def build_benchmark_args(base_args, benchmark, environment_config, cli_overrides=()):
     """Compose and validate final arguments for one benchmark evaluation."""
+    if benchmark["env"].get("eval_training_render"):
+        base_args = copy.deepcopy(base_args)
+        base_args["env"]["eval_training_render"] = True
     eval_training_render = base_args["env"]["eval_training_render"]
     validation_context = f"evaluation.{benchmark['name']}"
     if eval_training_render:

pufferlib/pufferl.py

@@ -1681,12 +1681,10 @@ def eval(
     selected_benchmarks = benchmark_names if benchmark_names is not None else eval_config["benchmarks"]
     eval_config["benchmarks"] = selected_benchmarks
     output_name = eval_config["output_name"]
-    render_scenarios = eval_config["render_scenarios"]
     render_filter = eval_config["render_filter"]
     max_rendered_failures = eval_config["max_rendered_failures"]
     failure_replay_csv = eval_config["failure_replay_csv"]
     eval_training_render = args["env"]["eval_training_render"]
-    render_scenarios = render_scenarios or eval_training_render
 
     report_to_wandb = bool(args["wandb"]) and not use_training_config
     environment_config, benchmarks = drive_benchmark.load_benchmark_config(benchmark_config_path, selected_benchmarks)
@@ -1722,6 +1720,7 @@ def eval(
             environment_config,
             cli_overrides,
         )
+        render_scenarios = eval_config["render_scenarios"] or run_args["env"]["eval_training_render"]
         output_directory_name = benchmark["name"]
         if output_name is not None:
             output_directory_name = f"{output_directory_name}_{output_name}"
@@ -1767,7 +1766,9 @@ def eval(
             capture_replay=render_scenarios,
         )
         print(f"Evaluation {benchmark['name']}: {num_scenarios} scenarios across {num_workers} workers")
-        replay_output_dir = os.path.join(benchmark_output_dir, "replays") if render_scenarios else None
+        replay_output_dir = (
+            os.path.join(benchmark_output_dir, drive_eval_replay.ZLIB_REPLAY_DIR_NAME) if render_scenarios else None
+        )
         summaries = _run_eval_rollout(
             run_args,
             env_name,
@@ -1787,7 +1788,7 @@ def eval(
         }
 
         if render_scenarios:
-            drive_eval_replay._render_eval_replays(summaries, benchmark_output_dir)
+            drive_eval_replay._render_eval_replays(summaries, benchmark_output_dir, eval_config["keep_zlib_replays"])
         elif render_filter is not None:
             _render_eval_failures(
                 env_name,
@@ -2272,7 +2273,7 @@ def _render_eval_failures(
         )
         replay_agent_capacity = failure_args["env"]["max_agents_per_env"]
         failure_args["env"]["num_agents"] = replay_agent_capacity
-    replay_output_dir = os.path.join(failures_dir, "replays")
+    replay_output_dir = os.path.join(failures_dir, drive_eval_replay.ZLIB_REPLAY_DIR_NAME)
     os.makedirs(replay_output_dir, exist_ok=True)
     agents_per_batch_values = selected_rows["agents_per_batch"].unique()
     if len(agents_per_batch_values) != 1:
@@ -2309,7 +2310,7 @@ def _render_eval_failures(
         )
         summaries.extend(wave_summaries)
     summary = drive_benchmark._write_eval_reports(summaries, failures_dir, len(pairs))
-    drive_eval_replay._render_eval_replays(summaries, failures_dir)
+    drive_eval_replay._render_eval_replays(summaries, failures_dir, run_args["eval"]["keep_zlib_replays"])
     return {
         "episodes": summaries,
         "summary": summary,

tests/eval/test_eval.py

@@ -25,6 +25,7 @@
 CARLA_WORKER_COUNT = 2
 CARLA_MAP_COUNT = 4
 CARLA_SCENARIO_LENGTH = 32
+BENCHMARK_SCENARIO_LENGTH = 48
 TRAIN_EPOCH_COUNT = 7
 TRAIN_EVAL_INTERVAL = 3
 TRAIN_HORIZON = 16
@@ -410,7 +411,7 @@ def record_vector_make(*make_args, **make_kwargs):
         return original_vector_make(*make_args, **make_kwargs)
 
     monkeypatch.setattr(pufferlib.vector, "make", record_vector_make)
-    replay_output_dir = tmp_path / "replays"
+    replay_output_dir = tmp_path / drive_eval_replay.ZLIB_REPLAY_DIR_NAME
     summaries = pufferl._run_eval_rollout(
         args,
         "puffer_drive",
@@ -446,11 +447,17 @@ def record_vector_make(*make_args, **make_kwargs):
         assert header["obs_dim"] > 0
         assert required_chunks <= set(header["chunks"])
 
-    render_dir = Path(drive_eval_replay._render_eval_replays(summaries, str(tmp_path)))
+    render_dir = Path(drive_eval_replay._render_eval_replays(summaries, str(tmp_path), keep_zlib_replays=True))
     rendered_pages = sorted(path for path in render_dir.glob("*.html") if path.name != "index.html")
     assert len(rendered_pages) == 2
     assert (render_dir / "index.html").is_file()
     assert all('class="payload-chunk"' in page.read_text() for page in rendered_pages)
+    assert all(replay_path.is_file() for replay_path in replay_paths)
+
+    drive_eval_replay._render_eval_replays(summaries, str(tmp_path), keep_zlib_replays=False)
+    assert not replay_output_dir.exists()
+    assert (render_dir / "index.html").is_file()
+    assert len(sorted(path for path in render_dir.glob("*.html") if path.name != "index.html")) == 2
 
 
 def _write_training_benchmark(tmp_path):