sae: add counterfactual image-edit evaluation - #25
Conversation
The paired-intervention pipeline behind the paper's counterfactual figure:
hold ego history and route intent fixed, change only the front camera image,
and attribute the resulting SAE latent and planner shifts to the visual edit.
* analyze_sae_visual_gen.py - stage 1. YOLO detection provides scene
grounding to a VLM prompt generator, which
emits scene-consistent edit prompts; a
diffusion editor applies them and writes a
manifest of matched (original, edited) pairs.
* analyze_sae_visual_gen_pt2.py - stage 2. Re-encodes both halves of each
pair, and reports per-feature latent deltas
alongside PLANNER_DELTA_NAMES.
Stage 1 pulls in heavy, single-purpose dependencies (ultralytics, diffusers,
google-genai, openai), so they go in a `visual-edits` optional-dependency
group rather than the base install. Stage 2 and every other SAE script import
without them - verified.
Also drops pt2's private model_inputs_from_batch in favour of
sae_utils.planner_inputs_from_collated_batch: it was a duplicate that carried
the same IMAGES_JPEG-only assumption fixed in the foundation commit, and
would have raised KeyError on main's worker-decoded batches.
Import order in stage 1 tidied to stdlib / third-party / local. The cluster
recipe in the module docstring and the BASE_URL default are left as-is; the
latter is already environment-overridable.
There was a problem hiding this comment.
Pull request overview
Adds a two-stage counterfactual image-edit evaluation pipeline for SAE analysis: stage 1 generates scene-consistent visual edits and a manifest of (original, edited) pairs, and stage 2 replays those pairs through the planner/SAE to report latent + planner metric deltas.
Changes:
- Add stage 1 script to detect scene elements, generate an edit plan prompt, apply a diffusion edit, and emit a JSONL manifest.
- Add stage 2 script to reload edited images, replay planner inference, compute SAE latent deltas + planner deltas, and export CSV/PT artifacts.
- Add a
visual-editsoptional dependency group for stage 1’s heavier, single-purpose dependencies.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| src/camera-based-e2e/analyze_sae_visual_gen.py | Stage 1: YOLO-grounded prompt generation + diffusion editing; writes edited images and a manifest. |
| src/camera-based-e2e/analyze_sae_visual_gen_pt2.py | Stage 2: re-encodes original/edited pairs; computes SAE/planner deltas; exports summaries. |
| pyproject.toml | Introduces visual-edits optional dependencies to keep base installs lighter. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| def inference_yolo(model, images): | ||
| # Predict with the model | ||
| results = model(images, verbose=False, device="cuda:0") # predict on an image | ||
|
|
||
| # Access the results | ||
| for result in results: | ||
| xywh = result.boxes.xywh # center-x, center-y, width, height | ||
| xywhn = result.boxes.xywhn # normalized | ||
| xyxy = result.boxes.xyxy # top-left-x, top-left-y, bottom-right-x, bottom-right-y | ||
| xyxyn = result.boxes.xyxyn # normalized | ||
| names = [result.names[cls.item()] for cls in result.boxes.cls.int()] # class name of each box | ||
| confs = result.boxes.conf # confidence score of each box | ||
|
|
||
| return results |
| results = [] | ||
| for batch_images, batch_samples, batch_indices in batches: | ||
| batch_results = inference_yolo(yolo_model, batch_images) | ||
| results.extend(batch_results) |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0644654892
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| gen.add_argument("--n_items", type=int, default=5_000) | ||
| gen.add_argument("--start_idx", type=int, default=0) | ||
| gen.add_argument("--max_items", type=int, default=10) | ||
| gen.add_argument("--camera_idx", type=int, default=1) |
There was a problem hiding this comment.
When --camera_idx is anything other than 1, stage 1 edits that camera, but stage 2 calls collate_with_images(edited_samples) with its default cam_idxs=(1,), and DeepMonocularModel.prepare_visual_tokens exclusively consumes images[1]. The generated edit is therefore never passed to the planner and the reported deltas measure an unchanged front image; either constrain this argument to camera 1 or explicitly support the selected camera throughout replay.
Useful? React with 👍 / 👎.
| contents=[ | ||
| types.Content( | ||
| role="user", | ||
| parts=[types.Part.from_text(text=scene_description)], | ||
| ), |
There was a problem hiding this comment.
With --generator gemini, this request sends only YOLO class names, confidence values, and bounding boxes. That text contains neither traffic-light color nor enough visual context to determine whether a sidewalk, crosswalk, or intersection is present, so Gemini cannot reliably follow the edit-selection rules and can produce invalid counterfactual prompts. Include the image as a multimodal part, as the local generator already does.
Useful? React with 👍 / 👎.
| # diffusion image editor. Stage 2 (analyze_sae_visual_gen_pt2.py) and every | ||
| # other SAE script run without these. | ||
| visual-edits = [ | ||
| "diffusers>=0.31", |
There was a problem hiding this comment.
diffusers==0.31 satisfies this declared dependency but predates the QwenImageEditPipeline, QwenImageEditPlusPipeline, and Flux2KleinPipeline symbols imported unconditionally by stage 1. In an otherwise valid environment constrained to that allowed version, the script fails during import before any editor can be selected. Set the lower bound to the first release that exports all three pipelines, or import each optional pipeline lazily.
Useful? React with 👍 / 👎.
| std_delta = delta.std(dim=0, unbiased=False) | ||
| stderr_delta = std_delta / (n ** 0.5) | ||
| paired_t = mean_delta / stderr_delta.clamp_min(1e-6) | ||
| cohen_dz = mean_delta / std_delta.clamp_min(1e-6) |
There was a problem hiding this comment.
When an edit type or direction contains one sample, std_delta and stderr_delta are zero, so clamping the denominator to 1e-6 turns every nonzero singleton delta into an enormous finite paired_t and cohen_dz. With the default ten-item run, singleton groups are plausible and the summary CSV then reports statistically meaningful-looking values for undefined statistics; emit NaN or omit these metrics when n < 2 or variance is zero.
Useful? React with 👍 / 👎.
run_generate_edits built three full-length lists (decoded frames, samples, and YOLO Results) before editing anything. Each ultralytics Results keeps its own copy of the source frame, so the frames were held twice: ~6 GB per 1k items at 960x1088. A 5k-item run OOM-killed 122 seconds in under --mem=64g (slurm 2258791); the 1k run behind the current results fit only by luck of the constant. Decode, detect, and edit one batch of 16 at a time, and reduce each Results to its scene description before the next batch is read. Peak memory becomes O(batch) rather than O(items). The loop body is unchanged apart from indentation and reading NAME from the streamed batch. Verified output-identical: over a 200-frame stubbed run the manifest and the edited-image set are byte-for-byte the same as the previous implementation, with peak RSS 1422 MB -> 702 MB (understated, since the stub's Results carry calloc pages that are never faulted). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MRhP7oSUw8jucGB3afZHwW
summarize_group clamped both denominators to 1e-6. For a singleton group std_delta is 0 by construction, so any nonzero delta became a paired_t in the millions -- an undefined statistic rendered as a very significant-looking number. Default runs (--max_items 10) produce singleton groups readily. Report NaN when n < 2, or when a feature has zero variance but a nonzero mean. A feature whose delta is identically zero keeps its 0: it did not move, which is an answer rather than the absence of one. That distinction matters -- 42.6% of the (group, feature) cells in the 615-pair run are identically zero, and blanket-NaN would have churned all of them for no gain. Verified no effect on existing artifacts: replaying build_summary_rows over the 615-pair block-3 run reproduces all 5376 rows with max numeric drift 0 and no change in NaN-ness. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MRhP7oSUw8jucGB3afZHwW
run_generate_edits opened the manifest with "w", so any rerun destroyed the previous one. A 5k-item FireRed run is many hours of diffusion under a 36h wall clock, and the run that OOM-killed at 122 seconds took its manifest with it. Losing hours of GPU work to a timeout or a preemption is the expensive failure here, not the recompute. Read back any manifest already in --output_dir, keep the items that completed, and generate only the rest. Rows that cannot be trusted are dropped and their items redone: a truncated final line from a job killed mid-write, and rows marked "edited" whose image is no longer on disk. Resuming across incompatible settings (index_file, n_items, camera_idx, generator, editor) raises rather than silently merging two different experiments. The kept rows are staged to a temp file and moved into place, so an interruption during the rewrite cannot lose them either. On by default, since the common case is resubmitting the same sbatch after a timeout; --no-resume restores the old regenerate-everything behaviour. Verified against a clean reference run of 200 stubbed items: interrupting at 100 and resuming, truncating the last manifest line, and deleting an edited image all converge on byte-identical manifests and the same 133 edited images, completed rows are not regenerated, rerunning a finished manifest is a no-op, and a mismatched --editor raises. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MRhP7oSUw8jucGB3afZHwW
The paired-intervention pipeline behind the paper's counterfactual figure:
hold ego history and route intent fixed, change only the front camera image,
and attribute the resulting SAE latent and planner shifts to the visual edit.
grounding to a VLM prompt generator, which
emits scene-consistent edit prompts; a
diffusion editor applies them and writes a
manifest of matched (original, edited) pairs.
pair, and reports per-feature latent deltas
alongside PLANNER_DELTA_NAMES.
Stage 1 pulls in heavy, single-purpose dependencies (ultralytics, diffusers,
google-genai, openai), so they go in a
visual-editsoptional-dependencygroup rather than the base install. Stage 2 and every other SAE script import
without them - verified.
Also drops pt2's private model_inputs_from_batch in favour of
sae_utils.planner_inputs_from_collated_batch: it was a duplicate that carried
the same IMAGES_JPEG-only assumption fixed in the foundation commit, and
would have raised KeyError on main's worker-decoded batches.
Import order in stage 1 tidied to stdlib / third-party / local. The cluster
recipe in the module docstring and the BASE_URL default are left as-is; the
latter is already environment-overridable.