diff --git a/docs/src/visualizer.md b/docs/src/visualizer.md index c6b73c90e..b241f0134 100644 --- a/docs/src/visualizer.md +++ b/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`. diff --git a/pufferlib/config/ocean/drive.ini b/pufferlib/config/ocean/drive.ini index f68434856..4d9e70231 100644 --- a/pufferlib/config/ocean/drive.ini +++ b/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 diff --git a/pufferlib/ocean/drive/drive.h b/pufferlib/ocean/drive/drive.h index e53c88b8d..94248cb15 100644 --- a/pufferlib/ocean/drive/drive.h +++ b/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 diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py index 908bd4537..1c80de553 100644 --- a/pufferlib/pufferl.py +++ b/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)