Cot ablation script - #55
Open
jaagat-prashar-ai wants to merge 15 commits into
Open
Conversation
Module docstring lays out the CoT ablation design (adapted from masking/training/run.py's experiments a-d, reimplemented at the text/token level since AutoVLA has no equivalent of MaskedAlpamayo1_5's attention-masking primitives), plus load_config/parse_args and the injected_mistake word-swap table.
Drops words matching (by prefix, so plurals count) any concept in the --concepts list from the reasoning text, mirroring masking's per-word concept-set ablation but as a plain text edit instead of attention masking.
Keeps only the first n words of the reasoning text; n=0 fully removes it, letting the model reconstruct reasoning from scratch on continuation.
Keeps only the last n words of the reasoning text; n<=0 yields an empty suffix (full removal), completing the prefix/suffix threshold pair.
Single-pass regex substitution over MISTAKE_SWAPS (stop<->accelerate, red<->green, left<->right, pedestrian->no pedestrian, ...) so a matched word is only ever rewritten once, avoiding double-flip bugs from chained substitutions. In the spirit of masking experiment D's clause-reversal probe: does the decoded trajectory follow an injected wrong claim?
Per-waypoint L2 xy distance between a condition's trajectory and the baseline, plus its mean (ade_m) and last value (endpoint_m) -- mirrors masking's ade_m/endpoint_m/delta_xy_per_waypoint fields.
heading-rate (dtheta/dt, wrapped to [-pi, pi]) and speed are derived via finite differences over the decoded (x, y, heading) trajectory, then diffed against the baseline the same way masking compares curvature/accel -- AutoVLA's action_tokenizer has no `controls` dict like Alpamayo's, so these are recomputed here from the raw pose sequence.
Same dict-comprehension idiom AutoVLA.predict()/generate_sample() already use to move only the tensor entries of a processor output onto the target device, factored out so both generate_full and continue_from_text share it.
Same get_prompt -> vlm.generate -> trim-trailing-eos flow as AutoVLA.predict(), but additionally splits the completion at the first action-tagged token so the reasoning-text span can be extracted and edited by later conditions; predict() itself only returns the final decoded trajectory + full text.
…ion) Core mechanism for concept_mask/prefix/suffix/injected_mistake: tokenizes the edited reasoning text, appends it to the prompt's input_ids/ attention_mask (video tensors carried over unchanged since the visual prompt doesn't change), and lets vlm.generate continue from there. This is the AutoVLA-native stand-in for masking's attention-column masking, which has no equivalent here.
Generates one baseline rollout, the no_cot comparison, then builds and runs concept_mask/prefix_*/suffix_*/injected_mistake off the same baseline reasoning text and prompt inputs so every condition is isolated to a single text edit. Tracks when no mistake-swap rule matched a scene so silent gaps in that condition are visible in the output rather than just missing.
Follows tools/eval/nusc_eval.py's exact config/checkpoint/dataset-loading sequence (SFTDataset + SFTAutoVLA + state_dict load), then runs run_scene per val scene and appends each result to a JSONL, warning up front if the loaded config wasn't trained with use_cot=true.
The module-level import of WaymoE2ECoTAnnotationDataset pulls in tensorflow, which unconditionally breaks NuPlan-only preprocessing even though the nuplan path never touches Waymo code. Move the import inside the waymo branch so it's only required when actually selected. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Adds a No-CoT preprocessing config and a training-style eval config pointed at the HF AutoVLA_PDMS_89 checkpoint and Qwen2.5-VL-3B base model, plus full-run and 30-scene smoke-test eval scripts wired to reuse an existing NFS-mirrored navtest dataset (maps/logs/sensor blobs) instead of re-downloading the full ~51GB navsim navtest_eval bundle. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The camera/BEV/trajectory visualization block was dead code (commented out) with a stray bug: it referenced `cot_results`/`trajectory` from the try block above but recomputed `scene` unconditionally even though AutoVLAAgent.requires_scene is False (so `scene` was never bound otherwise). Re-enable it gated behind an opt-in `save_visualization` config flag so normal eval runs don't pay the matplotlib rendering cost, and add a small eval script variant for generating a handful of example figures. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.