Fix/visualizer inherit config

docs/src/data.md

@@ -55,7 +55,7 @@ Refer to [Simulator](simulator.md) for how the binaries are consumed during rese
 ## Verifying data availability
 
 - After conversion, `ls resources/drive/binaries | head` should show numbered `.bin` files.
-- If you see `Required directory resources/drive/binaries/map_000.bin not found` during training, rerun the conversion or check paths.
+- If you see `No .bin map files found in <map_dir>` during training, rerun the conversion or check paths.
 - With binaries in place, run `puffer train puffer_drive` from [Getting Started](getting-started.md) as a smoke test that the build, data, and bindings are wired together.
 - To inspect the binary output, convert a single JSON file with `load_map(<json>, <id>, <output_path>)` inside `drive.py`.
 

pufferlib/ocean/drive/drive.py

@@ -129,15 +129,19 @@ def __init__(
 
         self._action_type_flag = 0 if action_type == "discrete" else 1
 
-        # Check if resources directory exists
-        binary_path = f"{map_dir}/map_000.bin"
-        if not os.path.exists(binary_path):
+        # Check if map directory contains any bin files
+        if not os.path.isdir(map_dir):
             raise FileNotFoundError(
-                f"Required directory {binary_path} not found. Please ensure the Drive maps are downloaded and installed correctly per docs."
+                f"Map directory {map_dir} not found. Please ensure the Drive maps are downloaded and installed correctly per docs."
+            )
+        bin_files = [name for name in os.listdir(map_dir) if name.endswith(".bin")]
+        if not bin_files:
+            raise FileNotFoundError(
+                f"No .bin map files found in {map_dir}. Please ensure the Drive maps are downloaded and installed correctly per docs."
             )
 
         # Check maps availability
-        available_maps = len([name for name in os.listdir(map_dir) if name.endswith(".bin")])
+        available_maps = len(bin_files)
         if num_maps > available_maps:
             if allow_map_resampling:
                 print("\n" + "=" * 80)

pufferlib/ocean/drive/visualize.c

@@ -189,13 +189,13 @@ static int make_gif_from_frames(const char *pattern, int fps, const char *palett
     return 0;
 }
 
-int eval_gif(const char *map_name, const char *policy_name, int show_grid, int obs_only, int lasers,
-             int show_human_logs, int frame_skip, const char *view_mode, const char *output_topdown,
+int eval_gif(const char *map_name, const char *policy_name, const char *config_file, int show_grid, int obs_only,
+             int lasers, int show_human_logs, int frame_skip, const char *view_mode, const char *output_topdown,
              const char *output_agent, int num_maps, int zoom_in) {
 
     // Parse configuration from INI file
     env_init_config conf = {0};
-    const char *ini_file = "pufferlib/config/ocean/drive.ini";
+    const char *ini_file = config_file ? config_file : "pufferlib/config/ocean/drive.ini";
     if (ini_parse(ini_file, handler, &conf) < 0) {
         fprintf(stderr, "Error: Could not load %s. Cannot determine environment configuration.\n", ini_file);
         return -1;
@@ -417,6 +417,7 @@ int main(int argc, char *argv[]) {
     // File paths and num_maps (not in [env] section)
     const char *map_name = NULL;
     const char *policy_name = "resources/drive/puffer_drive_weights.bin";
+    const char *config_file = NULL;
     const char *output_topdown = NULL;
     const char *output_agent = NULL;
     int num_maps = 1;
@@ -485,10 +486,18 @@ int main(int argc, char *argv[]) {
                 num_maps = atoi(argv[i + 1]);
                 i++;
             }
+        } else if (strcmp(argv[i], "--config") == 0) {
+            if (i + 1 < argc) {
+                config_file = argv[i + 1];
+                i++;
+            } else {
+                fprintf(stderr, "Error: --config option requires a config file path\n");
+                return 1;
+            }
         }
     }
 
-    eval_gif(map_name, policy_name, show_grid, obs_only, lasers, show_human_logs, frame_skip, view_mode, output_topdown,
-             output_agent, num_maps, zoom_in);
+    eval_gif(map_name, policy_name, config_file, show_grid, obs_only, lasers, show_human_logs, frame_skip, view_mode,
+             output_topdown, output_agent, num_maps, zoom_in);
     return 0;
 }

pufferlib/utils.py

@@ -167,6 +167,38 @@ def run_wosac_eval_in_subprocess(config, logger, global_step):
         print(f"Failed to run WOSAC evaluation: {type(e).__name__}: {e}")
 
 
+def _write_visualizer_config(env_cfg, output_path, base_ini="pufferlib/config/ocean/drive.ini"):
+    """Write a config file for the visualizer, inheriting from base INI but overriding with env_cfg."""
+    import configparser
+
+    cfg = configparser.ConfigParser()
+    cfg.read(base_ini)
+
+    if env_cfg is not None and "env" in cfg:
+        # Special attribute mappings for Drive class (attr stored differently than INI key)
+        special_mappings = {
+            "action_type": ("_action_type_flag", lambda v: "discrete" if v == 0 else "continuous"),
+        }
+
+        # Override [env] values with attributes from env_cfg
+        for key in cfg["env"]:
+            attr_name = key.replace("-", "_")
+
+            if attr_name in special_mappings:
+                # Handle special cases where attribute name/format differs
+                real_attr, transform = special_mappings[attr_name]
+                if hasattr(env_cfg, real_attr):
+                    cfg["env"][key] = transform(getattr(env_cfg, real_attr))
+            elif hasattr(env_cfg, attr_name):
+                cfg["env"][key] = str(getattr(env_cfg, attr_name))
+            elif hasattr(env_cfg, f"{attr_name}_str"):
+                # For attributes like control_mode that have a _str variant
+                cfg["env"][key] = str(getattr(env_cfg, f"{attr_name}_str"))
+
+    with open(output_path, "w") as f:
+        cfg.write(f)
+
+
 def render_videos(config, vecenv, logger, epoch, global_step, bin_path):
     """
     Generate and log training videos using C-based rendering.
@@ -205,8 +237,13 @@ def render_videos(config, vecenv, logger, epoch, global_step, bin_path):
         env_vars = os.environ.copy()
         env_vars["ASAN_OPTIONS"] = "exitcode=0"
 
-        # Base command with only visualization flags (env config comes from INI)
-        base_cmd = ["xvfb-run", "-a", "-s", "-screen 0 1280x720x24", "./visualize"]
+        # Write temp config inheriting base INI but with training env overrides
+        env_cfg = getattr(vecenv, "driver_env", None)
+        temp_config_path = os.path.join(model_dir, "visualizer_config.ini")
+        _write_visualizer_config(env_cfg, temp_config_path)
+
+        # Base command with config file
+        base_cmd = ["xvfb-run", "-a", "-s", "-screen 0 1280x720x24", "./visualize", "--config", temp_config_path]
 
         # Visualization config flags only
         if config.get("show_grid", False):
@@ -230,14 +267,26 @@ def render_videos(config, vecenv, logger, epoch, global_step, bin_path):
         base_cmd.extend(["--view", view_mode])
 
         # Get num_maps if available
-        env_cfg = getattr(vecenv, "driver_env", None)
         if env_cfg is not None and getattr(env_cfg, "num_maps", None):
             base_cmd.extend(["--num-maps", str(env_cfg.num_maps)])
 
         # Handle single or multiple map rendering
         render_maps = config.get("render_map", None)
-        if render_maps is None:
-            render_maps = [None]
+        if render_maps is None or render_maps == "none":
+            # Pick a random map from the training map_dir
+            map_dir = getattr(env_cfg, "map_dir", None) if env_cfg else None
+            if map_dir and os.path.isdir(map_dir):
+                import random
+
+                bin_files = [f for f in os.listdir(map_dir) if f.endswith(".bin")]
+                if bin_files:
+                    render_maps = [os.path.join(map_dir, random.choice(bin_files))]
+                else:
+                    print(f"Warning: No .bin files found in {map_dir}, skipping render")
+                    return
+            else:
+                print(f"Warning: map_dir not found or invalid ({map_dir}), skipping render")
+                return
         elif isinstance(render_maps, (str, os.PathLike)):
             render_maps = [render_maps]
         else: