Render videos eval mode

docs/src/visualizer.md

@@ -33,3 +33,13 @@ xvfb-run -s "-screen 0 1280x720x24" ./visualize
 ```
 
 Adjust the screen size and color depth as needed. The `xvfb-run` wrapper allows Raylib to render without an attached display, which is convenient for servers and CI jobs.
+
+## Rendering Mode
+You can batch render videos for multiple maps using the evaluation mode. This will render the first `num_maps` maps (capped by the number of maps in the directory) from `map_dir` in parallel using the `visualize` binary and create these videos in an `output_dir`(All configs in [render] of `drive.ini`).
+
+After setting the configs run:
+```bash
+puffer render puffer_drive
+```
+
+This mode parallelizes rendering based on `vec.num_workers`.

pufferlib/config/ocean/drive.ini

@@ -132,6 +132,31 @@ human_replay_eval = False
 ; Control only the self-driving car
 human_replay_control_mode = "control_sdc_only"
 
+[render]
+; Mode to render a bunch of maps with a given policy
+; Path to dataset used for rendering
+map_dir = "resources/drive/binaries/training"
+; Directory to output rendered videos
+output_dir = "resources/drive/render_videos"
+; Evaluation will run on the first num_maps maps in the map_dir directory
+num_maps = 100
+; "both", "topdown", "agent"; Other args are passed from train confs
+view_mode = "both"
+; Policy bin file used for rendering videos
+policy_path = "resources/drive/puffer_drive_weights_resampling_300.bin"
+; Allows more than cpu cores workers for rendering
+overwork = True
+; If True, show exactly what the agent sees in agent observation
+obs_only = True
+; Show grid lines
+show_grid = True
+; Draws lines from ego agent observed ORUs and road elements to show detection range
+show_lasers = True
+; Display human xy logs in the background
+show_human_logs = False
+; If True, zoom in on a part of the map. Otherwise, show full map
+zoom_in = True
+
 [sweep.train.learning_rate]
 distribution = log_normal
 min = 0.001

pufferlib/ocean/drive/drive.h

@@ -332,6 +332,7 @@ float point_to_segment_distance_2d(float px, float py, float x1, float y1, float
 void init_goal_positions(Drive *env);
 float clipSpeed(float speed);
 void sample_new_goal(Drive *env, int agent_idx);
+int check_lane_aligned(Entity *car, Entity *lane, int geometry_idx);
 
 // ========================================
 // Utility Functions

pufferlib/pufferl.py

@@ -24,6 +24,7 @@
 
 import numpy as np
 import psutil
+from multiprocessing.pool import ThreadPool
 
 import torch
 import torch.distributed
@@ -1232,6 +1233,7 @@ def eval(env_name, args=None, vecenv=None, policy=None):
         print("HUMAN_REPLAY_METRICS_END")
 
         return results
+
     else:  # Standard evaluation: Render
         backend = args["vec"]["backend"]
         if backend != "PufferEnv":
@@ -1246,10 +1248,6 @@ def eval(env_name, args=None, vecenv=None, policy=None):
         num_agents = vecenv.observation_space.shape[0]
         device = args["train"]["device"]
 
-        # Rebuild visualize binary if saving frames (for C-based rendering)
-        if args["save_frames"] > 0:
-            ensure_drive_binary()
-
         state = {}
         if args["train"]["use_rnn"]:
             state = dict(
@@ -1643,8 +1641,90 @@ def puffer_type(value):
     return args
 
 
+def render(env_name, args=None):
+    args = args or load_config(env_name)
+    render_configs = args.get("render", {})
+
+    # Renders first num_maps from map_dir using visualize binary
+    try:
+        map_dir = render_configs["map_dir"]
+        num_maps = render_configs.get("num_maps", 1)
+        view_mode = render_configs["view_mode"]
+        render_policy_path = render_configs["policy_path"]
+        overwork = render_configs.get("overwork", False)
+        num_workers = args["vec"]["num_workers"]
+        output_dir = render_configs["output_dir"]
+    except KeyError as e:
+        raise pufferlib.APIUsageError(f"Missing render config: {e}")
+
+    cpu_cores = psutil.cpu_count(logical=False)
+    if num_workers > cpu_cores and not overwork:
+        raise pufferlib.APIUsageError(
+            " ".join(
+                [
+                    f"num_workers ({num_workers}) > hardware cores ({cpu_cores}) is disallowed by default.",
+                    "PufferLib multiprocessing is heavily optimized for 1 process per hardware core.",
+                    "If you really want to do this, set overwork=True (--vec-overwork in our demo.py).",
+                ]
+            )
+        )
+
+    if num_maps > len(os.listdir(map_dir)):
+        num_maps = len(os.listdir(map_dir))
+
+    render_maps = [os.path.join(map_dir, f) for f in sorted(os.listdir(map_dir)) if f.endswith(".bin")][:num_maps]
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Rebuild visualize binary
+    ensure_drive_binary()
+
+    def render_task(map_path):
+        base_cmd = (
+            ["./visualize"]
+            if sys.platform == "darwin"
+            else ["xvfb-run", "-a", "-s", "-screen 0 1280x720x24", "./visualize"]
+        )
+        cmd = base_cmd.copy()
+        cmd.extend(["--map-name", map_path])
+        if render_configs.get("show_grid", False):
+            cmd.append("--show-grid")
+        if render_configs.get("obs_only", False):
+            cmd.append("--obs-only")
+        if render_configs.get("show_lasers", False):
+            cmd.append("--lasers")
+        if render_configs.get("show_human_logs", False):
+            cmd.append("--show-human-logs")
+        if render_configs.get("zoom_in", False):
+            cmd.append("--zoom-in")
+        cmd.extend(["--view", view_mode])
+        if render_policy_path is not None:
+            cmd.extend(["--policy-name", render_policy_path])
+
+        map_name = os.path.basename(map_path).replace(".bin", "")
+
+        if view_mode == "topdown" or view_mode == "both":
+            cmd.extend(["--output-topdown", os.path.join(output_dir, f"topdown_{map_name}.mp4")])
+        if view_mode == "agent" or view_mode == "both":
+            cmd.extend(["--output-agent", os.path.join(output_dir, f"agent_{map_name}.mp4")])
+
+        env_vars = os.environ.copy()
+        env_vars["ASAN_OPTIONS"] = "exitcode=0"
+        try:
+            result = subprocess.run(cmd, cwd=os.getcwd(), capture_output=True, text=True, timeout=600, env=env_vars)
+            if result.returncode != 0:
+                print(f"Error rendering {map_name}: {result.stderr}")
+        except subprocess.TimeoutExpired:
+            print(f"Timeout rendering {map_name}: exceeded 600 seconds")
+
+    if render_maps:
+        print(f"Rendering {len(render_maps)} from {map_dir} with {num_workers} workers...")
+        with ThreadPool(num_workers) as pool:
+            pool.map(render_task, render_maps)
+        print(f"Finished rendering videos to {output_dir}")
+
+
 def main():
-    err = "Usage: puffer [train, eval, sweep, controlled_exp, autotune, profile, export, sanity] [env_name] [optional args]. --help for more info"
+    err = "Usage: puffer [train, eval, sweep, controlled_exp, autotune, profile, export, sanity, render] [env_name] [optional args]. --help for more info"
     if len(sys.argv) < 3:
         raise pufferlib.APIUsageError(err)
 
@@ -1666,6 +1746,8 @@ def main():
         export(env_name=env_name)
     elif mode == "sanity":
         sanity(env_name=env_name)
+    elif mode == "render":
+        render(env_name=env_name)
     else:
         raise pufferlib.APIUsageError(err)