[WIP] Trajviz: Vulkan offline trajectory renderer

[eugenevinitsky]

Summary

A new headless Vulkan-backed renderer that turns saved Drive trajectories into MP4 videos. Independent of the existing raylib visualizer (scripts/build_ocean.sh visualize) — opt-in via TRAJVIZ=1 python setup.py build_ext --inplace. Drive sim code (drive.{h,c,py}, binding.c) is untouched.

Status: WIP. Builds, runs, renders correctly on real Drive sims. There is more to do (see below) before this is merge-ready.

• Public Python API: Renderer.render_episode(...) and Renderer.render_batch([...]) (up to 16 episodes per batch)
• CLI: python -m pufferlib.ocean.drive.trajviz <inputs> --maps-dir ... --out ...
• Random-rollout smoke test: python -m pufferlib.ocean.drive.trajviz.tools.random_rollout
• Standalone C harness: pufferlib/ocean/drive/trajviz/tests/test_main.c (no Python required)

Performance

On RTX 4080 + 16-core CPU, 1280×720, 90-frame episodes, both views, libx264 -preset veryfast:

batch_size	wall time	per-episode	ep/s
1	345 ms	345 ms	2.9
4	1094 ms	274 ms	3.7
8	2136 ms	267 ms	3.7

Pure Vulkan + readback floor (encoder bypassed) is ~30 ms / ep ≈ 32 ep/s. The remaining gap is libx264 — NVENC integration (deferred) would close it.

Headline optimizations applied during development:

• HOST_CACHED readback memory — single biggest win (~6-7×). Default HOST_VISIBLE | HOST_COHERENT on NVIDIA picks write-combined PCIe BAR memory, fast for the GPU but ~250 MB/s for CPU reads. HOST_CACHED puts it in regular RAM (>5 GB/s).
• LINE_STRIP polylines — one draw call per polyline instead of per segment.
• Vertically-tiled atlas for batched rendering — N episodes' tiles stacked vertically so each tile's bytes are row-contiguous in the readback buffer (one fwrite per pipe per frame, no row stitching).
• Threaded fwrite fan-out — each ffmpeg pipe owns a writer thread; per-frame write phase costs max(single fwrite) instead of sum(fwrites).
• F_SETPIPE_SZ — bumps the kernel pipe buffer up to whatever the per-process limit allows so fwrites don't ping-pong on a 64 KB pipe.

Architecture

__init__.py / _native.c / trajviz.{h,c}
                ↓
   vk_renderer (single)   vk_batch_renderer (tiled atlas)
                ↓                  ↓
   vk_pipeline / vk_context / ffmpeg_pipe (+ writer thread)
                ↓                  ↓
        Vulkan 1.3             ffmpeg subprocess

Two views, matching the existing live raylib path:

• Top-down = RenderView.FULL_SIM_STATE — orthographic full-map
• BEV = RenderView.BEV_AGENT_OBS — agent-centric ~100m × 178m window, ego at center facing up

Documentation in docs/src/trajviz.md covers prerequisites (apt packages), build, Python and CLI usage, performance tuning (sysctl knobs, HOST_CACHED explanation, batch size guidance), debugging env vars, troubleshooting, and architecture overview.

Why Vulkan, not raylib

• Headless on Linux clusters with no X server (raylib needs xvfb)
• Throughput-oriented batching (impossible without command-buffer control)
• Independent build path so the optional Vulkan dep doesn't pollute the live drive sim build

Why WIP — known gaps before merge

• No NPC / expert-replay agents — currently only renders the controlled agents from get_sim_trajectories. The other ~18 vehicles in a typical Waymo scenario (the WOSAC "context" tracks) aren't shown. Adding them needs a separate Drive API to expose expert trajectories.
• No 3D follow-cam — RenderView.AGENT_PERSP (3D car meshes from .glb) is not implemented in trajviz. Top-down + BEV only.
• CPU-bound by libx264 once batched — NVENC integration (Vulkan video encode or libnvidia-encode + CUDA-Vulkan interop) would unlock the remaining ~12% gap to the pure-GPU ceiling.
• batch_size capped at 16 — atlas image height grows linearly. Past ~22 we'd need multiple atlas passes or a 2-D tile grid.
• Uniform num_steps per batch — short episodes get padded with zeros. Wastes a tiny bit of GPU work on the trailing zeros.
• Buffer/image helpers duplicated between vk_renderer.c and vk_batch_renderer.c (~50 lines each). Should be consolidated into a shared helper before merge.

Test plan

• TRAJVIZ=1 python setup.py build_ext --inplace — clean build, no warnings beyond the standard PyCFunction cast
• Single-episode render via random_rollout.py — 90-frame MP4 with valid h264, 1280×720, both views
• Batched render of 4 episodes — 4 valid MP4s, ~270 ms/ep
• Visual sanity check on extracted frames (ego centered in BEV, Waymo road geometry slides correctly)
• Render an actual saved trajectories_*.npz from a real training run end-to-end via render_npz
• Try on a non-NVIDIA GPU (AMD radv, Intel) to confirm HOST_CACHED fallback works
• CI build with TRAJVIZ=1 so the extension keeps compiling
• Decide whether to delete notebooks/visualize_trajectories.py once map_io.py is the canonical parser

🤖 Generated with Claude Code

[Copilot]

Pull request overview

Adds an opt-in, headless Vulkan-backed “trajviz” renderer for offline Drive trajectory visualization (MP4 output), and wires up trajectory capture/saving during training checkpoints to feed the renderer—without changing the live raylib visualizer path.

Changes:

• Introduces pufferlib.ocean.drive.trajviz (C/Vulkan renderer + CPython extension + CLI/tools) built only when TRAJVIZ=1.
• Adds C-side sim-trajectory recording and Python/C bindings to retrieve + save per-checkpoint trajectories (including map context) for offline rendering.
• Adds docs and a shared .bin map parser (map_io.py) to support loading and rendering saved trajectories.

Reviewed changes

Copilot reviewed 35 out of 37 changed files in this pull request and generated 11 comments.

Show a summary per file

File	Description
setup.py	Adds opt-in build path for the trajviz Vulkan CPython extension and shader build step.
pufferlib/vector.py	Tags worker envs with _worker_idx; adds save_worker_trajectories() fan-out helper.
pufferlib/pufferl.py	Saves trajectories at checkpoint time; adds reproducibility snapshot; threads traj temp dir into env kwargs.
pufferlib/ocean/env_binding.h	Adds vectorized bindings to pull sim trajectories and world_mean from C envs.
pufferlib/ocean/drive/drive.py	Caches world_mean; adds get_sim_trajectories() and notify() to write per-worker .npz.
pufferlib/ocean/drive/drive.h	Allocates/records per-step sim trajectories in C; exposes c_get_sim_trajectories().
pufferlib/ocean/drive/datatypes.h	Extends Agent with sim-trajectory buffers and frees them on teardown.
pufferlib/ocean/drive/map_io.py	Adds canonical parser/transform helpers for Drive .bin maps for offline tooling.
pufferlib/ocean/drive/trajviz/init.py	Python Renderer wrapper + render_npz() convenience loader.
pufferlib/ocean/drive/trajviz/main.py	CLI entry point to render one or more trajectories_*.npz inputs.
pufferlib/ocean/drive/trajviz/_native.c	CPython extension: numpy validation + GIL release around render calls.
pufferlib/ocean/drive/trajviz/trajviz.h	Public C API for single-episode and batched rendering.
pufferlib/ocean/drive/trajviz/trajviz.c	Orchestrates Vulkan renderers + ffmpeg pipes; implements batch tiling path.
pufferlib/ocean/drive/trajviz/vk_context.h	Declares Vulkan instance/device/queue lifecycle and error helpers.
pufferlib/ocean/drive/trajviz/vk_context.c	Implements Vulkan init (1.3 + dynamic rendering + sync2) and teardown.
pufferlib/ocean/drive/trajviz/vk_pipeline.h	Declares shared pipeline/push-constant and instance formats.
pufferlib/ocean/drive/trajviz/vk_pipeline.c	Creates Vulkan graphics pipelines for polylines and agent boxes.
pufferlib/ocean/drive/trajviz/vk_renderer.h	Declares single-episode renderer (frame slots, readback, ffmpeg drain).
pufferlib/ocean/drive/trajviz/vk_renderer.c	Implements per-frame rendering + readback + ffmpeg write for one episode.
pufferlib/ocean/drive/trajviz/vk_batch_renderer.h	Declares batched atlas renderer API.
pufferlib/ocean/drive/trajviz/vk_batch_renderer.c	Implements batched atlas rendering + threaded pipe write fan-out.
pufferlib/ocean/drive/trajviz/vk_math.h	Adds small mat4 helpers for ortho fit and BEV camera.
pufferlib/ocean/drive/trajviz/ffmpeg_pipe.h	Declares ffmpeg subprocess pipe + writer-thread API.
pufferlib/ocean/drive/trajviz/ffmpeg_pipe.c	Implements popen-based ffmpeg piping and async writer thread.
pufferlib/ocean/drive/trajviz/shaders.h	Declares externs for generated SPIR-V blobs.
pufferlib/ocean/drive/trajviz/shaders/build_shaders.sh	Builds GLSL → SPIR-V and generates shaders.c.
pufferlib/ocean/drive/trajviz/shaders/polyline.vert	Adds GLSL for road polyline vertex stage.
pufferlib/ocean/drive/trajviz/shaders/polyline.frag	Adds GLSL for road polyline fragment stage.
pufferlib/ocean/drive/trajviz/shaders/agent_box.vert	Adds GLSL for instanced agent quad expansion.
pufferlib/ocean/drive/trajviz/shaders/agent_box.frag	Adds GLSL for flat-colored agent box fragment stage.
pufferlib/ocean/drive/trajviz/tools/random_rollout.py	End-to-end smoke test that rolls out Drive and renders via trajviz.
pufferlib/ocean/drive/trajviz/tools/init.py	Marks tools as a package (module discovery).
pufferlib/ocean/drive/trajviz/tests/test_main.c	Standalone C harness to validate Vulkan+ffmpeg path without Python.
docs/src/trajviz.md	Adds end-user documentation: build/run/tuning/architecture/troubleshooting.
docs/src/SUMMARY.md	Links trajviz documentation into the docs sidebar.
notebooks/visualize_trajectories.py	Adds/updates notebook for analyzing and plotting saved trajectories.
.gitignore	Ignores generated pufferlib/ocean/drive/trajviz/shaders.c.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The multiprocessing stitching path doesn’t include agent_offsets, which render_npz() requires to slice per-env episodes. Simply concatenating per-worker (agent-local) offsets wouldn’t be correct anyway; you likely need to (a) collect each worker’s agent_offsets, (b) shift them by the cumulative agent count, and (c) concatenate to produce a global agent_offsets aligned with the concatenated traj_* arrays.

[Copilot]

traj_tmp is a shared directory (.../traj_tmp) that is never cleaned, and glob(traj_worker_*.npz) will pick up stale worker files (e.g., if a previous run used more workers, or a crashed worker left an old file). Consider writing into an epoch-scoped subdir, or deleting existing traj_worker_*.npz files before triggering save_worker_trajectories(), and/or validating the expected worker count before stitching.

[Copilot]

as_array() leaks a reference: PyObject_Str(value) creates a new object that isn’t DECREF’d (it’s passed directly into PyUnicode_AsUTF8(...)). Store the PyObject_Str result in a temporary, use it for formatting, then Py_DECREF it to avoid per-call leaks on shape/type errors.

Suggested change

	PyErr_Fetch(&type, &value, &tb);
	PyErr_Format(PyExc_TypeError, "%s: %s", name,
	value ? PyUnicode_AsUTF8(PyObject_Str(value)) : "type/shape mismatch");
	PyObject *value_str = NULL;
	const char *message = "type/shape mismatch";
	PyErr_Fetch(&type, &value, &tb);
	if (value) {
	value_str = PyObject_Str(value);
	if (value_str) {
	const char *utf8 = PyUnicode_AsUTF8(value_str);
	if (utf8)
	message = utf8;
	}
	}
	PyErr_Format(PyExc_TypeError, "%s: %s", name, message);
	Py_XDECREF(value_str);

[Copilot]

rs.lineWidth is set to 1.5f, but Vulkan requires the wideLines device feature to be enabled for line widths != 1.0. Since vk_ctx_init() doesn’t enable wideLines, this can trigger validation errors or pipeline creation failure on some devices. Either set the line width back to 1.0 or explicitly query+enable VkPhysicalDeviceFeatures::wideLines when supported.

Suggested change

	d->rs.lineWidth = 1.5f; /* used for line topology only; ignored for tris */
	d->rs.lineWidth = 1.0f; /* Vulkan-safe default; wider lines require wideLines */

[Copilot]

The image layout barrier uses oldLayout = VK_IMAGE_LAYOUT_UNDEFINED every frame. After the first frame the image will actually be in VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL (set later in the same command buffer), so this will trip validation and is undefined behavior. Track the current layout per render target (or always transition from TRANSFER_SRC_OPTIMAL back to COLOR_ATTACHMENT_OPTIMAL after the copy) and use the real old layout in the next frame’s barrier.

[Copilot]

vec_get_sim_trajectories casts tuple items directly to PyArrayObject* and calls PyArray_DATA without validating they’re actually NumPy arrays / contiguous / expected dtype+shape. Other vec_* helpers in this file do PyArray_Check and shape-derived stride math. Add PyArray_Check (and ideally dtype/ndim checks) here too to prevent segfaults on misuse from Python.

[Copilot]

save_worker_trajectories() spins indefinitely using Python’s any(self.buf['notify']) over a NumPy array. This is both slower than self.buf['notify'].any() and can hang forever if a worker dies or never clears its flag. Consider using np.any(...)/.any() plus a timeout (and surfacing an error) to avoid deadlocking the training loop.

[Copilot]

This comment points users to docs/trajviz.md, but the documentation added in this PR lives under docs/src/trajviz.md. Update the path so the install instructions are discoverable from the repo layout.

[Copilot]

trajviz_render_episodes_batch allows all_road_offsets / all_road_types to be NULL (they’re not included in the required-pointer check), but later computes num_polys_s from poly_meta_offsets and can call vk_batch_renderer_set_episode with num_polys_s > 0 and off_s/typ_s == NULL, which is likely to crash. Either require these pointers when any episode has polylines, or validate per-episode and force num_polys_s=0 when offsets/types are absent.

Suggested change

	}
	}
	{
	int any_episode_has_polylines = 0;
	for (int s = 0; s < batch_size; ++s) {
	if (poly_meta_offsets[s + 1] > poly_meta_offsets[s]) {
	any_episode_has_polylines = 1;
	break;
	}
	}
	if (any_episode_has_polylines && (!all_road_offsets || !all_road_types)) {
	snprintf(ctx->last_error, sizeof(ctx->last_error),
	"road offset/type arrays are required when any episode has polylines");
	return TRAJVIZ_ERR_BAD_ARG;
	}
	}

[Copilot]

The per-episode road slicing can yield num_polys_s > 0 while off_s/typ_s are NULL (because all_road_offsets / all_road_types are treated as optional). Before calling vk_batch_renderer_set_episode, add a consistency check that offsets/types are present whenever num_polys_s > 0, and fail with a clear error if not.

Suggested change

	const uint32_t *typ_s = (all_road_types && num_polys_s > 0) ? (all_road_types + pt_start) : NULL;
	const uint32_t *typ_s = (all_road_types && num_polys_s > 0) ? (all_road_types + pt_start) : NULL;
	if (num_polys_s > 0 && (!off_s || !typ_s)) {
	snprintf(ctx->last_error, sizeof(ctx->last_error),
	"episode %d has %u road polygons but missing %s%s",
	s, num_polys_s,
	!off_s ? "road offsets" : "",
	(!off_s && !typ_s) ? " and road types" : (!typ_s ? "road types" : ""));
	err = TRAJVIZ_ERR_BAD_ARGS;
	goto cleanup;
	}