ABLATION_NOTES.md

# Ground-R1 `rl-projects` — Ablation Branch Notes

Notes from a code read of the experimental forks of `r1-v/` on the `rl-projects` branch
(`Irisicy4/Ground-R1-project`). Each branch keeps a sibling copy of the trainer so the
baseline `r1-v/` stays untouched. Below: what each variant changes vs `r1-v`, what
hypothesis it isolates, and any bugs/risks found.

---

## At a glance

| Branch | Bbox / crop | Round structure | Image(s) carried across rounds | History across rounds |
|---|---|---|---|---|
| `r1-v` (baseline) | yes (`<box>` + crop + 2nd image) | adaptive: stop on `<answer>` or iter==4 | original + each crop | full conversation |
| `r1-v-no-hist` | yes (still crops) | adaptive (same cap) | **only the latest crop** | **reset each round** |
| `r1-v-no123-match-round` | **no** | adaptive (same cap) | original only | full conversation |
| `r1-v-no123` | **no** | **fixed: exactly 2 rounds** | original only | full conversation |
| `r1-v-no23-1round` | `<box>`+`<answer>` in one turn | adaptive (same cap), retry only on missing `<answer>` | original re-sent every retry | full conversation |

The naming reads as: **no-hist** = remove cross-round memory; **no123** = remove the three
grounding pieces (bbox-step, crop, second image); **match-round** = keep r1-v's adaptive
iteration cap; **1round** = intended single fused turn.

---

## `r1-v-no-hist` — drop conversation history between grounding rounds

Keeps the full Ground-R1 bbox/crop machinery, but at the start of each new round throws
away prior turns and shows the model only the cropped image + the next-round prompt.

- **Trainer `_prepare_for_stage2`** is the only behavior change:
  - before: `origin_prompt.extend([{assistant: bbox_str}, next_user_with_crop])`,
    `combined_images = [original, crop]`
  - after: `origin_prompt = [next_stage_entry]`, `combined_images = [crop]`
- **`grpo.py`** resume hardening (not algorithmic): replaces the dead hardcoded
  `/home/meng/GRPO/...` resume path with `get_last_checkpoint(output_dir)`, adds a
  `configure_torch_checkpoint_resume()` helper (sets `TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1`,
  registers `ZeroStageEnum` as a `torch.serialization` safe global for PyTorch 2.6+),
  and fixes `trainer.save_state(output_dir)` → `trainer.save_state()`.

**Hypothesis:** does iterative grounding still work when the model can't see what it
grounded before?

**Risk / note (by design, easy to overlook):** because the prompt is reset every round,
the final `prompt_completion_ids` sequence contains only the **last** round's
user+assistant turns. The loss-mask scan (`[151644, 77091, 198]` … `151645`) therefore
catches **only the final round's** assistant tokens. Intermediate grounding rounds
produce **no gradient** — they only decide which crop the final round sees. So a 4-round
rollout yields the same per-rollout RL signal as a 1-round one. Intended for the
"no-history" ablation, but it sharply reduces training signal per rollout; expect slower
learning than baseline at equal steps.

---

## `r1-v-no123-match-round` — pure CoT, no bbox/crop, keep r1-v's adaptive cap

Strips all bounding-box / crop logic; the model just thinks → maybe answers → if not, is
re-asked the same prompt. Up to 4 extra rounds (same cap as `r1-v`).

- **Prompt** (`grpo.py` + trainer `STAGE_PROMPT_TEMPLATE`): no `<box>` language; "if no
  further thinking is needed, provide `<answer>`." Format: `<think>…</think>` or
  `<think>…</think><answer>…</answer>`.
- **Format reward:** `pattern_stage1` = `^<think>(.+?)</think>$` (no `<box>`, no
  `<answer>`); `pattern_stage2` keeps `<think>…</think><answer>…</answer>`.
- **Trainer surgery:** deletes `bbox_adjust`, `cal_bbox_for_iou`,
  `_crop_image_for_next_stage`, `_get_bbox_for_last_stage`; renames `_prepare_for_stage2`
  → `_prepare_for_next_round` (appends `{assistant: previous_response}` + a fresh
  text-only user turn; no image added — original image stays in history once);
  `_generate_for_stage2_batch` passes `images=None` when there are none.
- **`grpo.py`:** deletes all bbox/IoU reward+score fns (`compute_iou`, `compute_giou`,
  `bbox_reward_stage2`, `bbox_score_stage{1,2}`, `bbox_iou_stage{1,2,3}`); registries
  trimmed to `{accuracy, format}` / `{refine_times}`. Same resume + `save_state` fixes.
- **`prepare_data.py`:** `item.pop('bboxs', None)`; checked-in jsonl already has bboxs
  stripped.

**Hypothesis:** is the "think → maybe answer, else re-ask" outer loop alone (no grounding)
competitive with the full pipeline? Clean; no bug found.

---

## `r1-v-no123` — pure CoT, forced exactly 2 rounds

Same bbox/crop removal as `match-round`, but **does not** mimic r1-v's adaptive loop.
Always two passes.

- **Two distinct prompts:** `STAGE_ONE_TEMPLATE` ("think only … **do not provide the
  final answer yet**"); `STAGE_TWO_TEMPLATE` ("rethink using image + history, then answer").
- **Loop:** round 1 generates `<think>…</think>` only; round 2 appends that as assistant +
  a `STAGE_TWO_TEMPLATE` user turn → one final `generate` for all. Exactly 2 forward
  passes, no `<answer>`/iter early-exit.
- Same deletions, registries, resume + `save_state` fixes as `match-round`.

**Hypothesis:** is a fixed think→rethink+answer schedule enough (vs adaptive)? Mechanically
the simplest of the set; hard to break. No bug found.

---

## `r1-v-no23-1round` — fused single turn (`<think>`+`<box>`+`<answer>`), retry on missing answer

Asks for thinking, one bbox, and the final answer in a **single** response. No image
cropping; r1-v's adaptive retry cap is reused only as a fallback when the answer is
missing.

- **Single fused prompt:** "...provide one bounding box [x1,y1,x2,y2] inside `<box>` …
  Then directly provide the final answer inside `<answer>`." Format example:
  `<think>…</think> <box>[…]</box> <answer>…</answer>`.
- **Format reward:** `pattern_stage1` (intermediate) accepts `<think>…</think><box>[…]</box>`
  with `(?!.*<answer>)`; `pattern_stage2` (final) requires the full
  `<think>…</think><box>[…]</box><answer>…</answer>`. Both add single-occurrence
  anti-duplication lookaheads.
- **Loop:** stage-1 generate; if `<answer>` present → finalize; else
  `_prepare_for_next_round(..., include_image=True)` appends the failed response + a fresh
  user turn that **re-includes the original image** and the same prompt; loop to
  `<answer>` or iter==4.
- All bbox/crop infra + IoU rewards deleted; registries trimmed; jsonl bboxs stripped;
  same resume + `save_state` fixes.

**Hypothesis:** does emitting the box and answer in one shot (no separate crop trajectory)
match the multi-round pipeline?

**Risk found — wasteful image re-injection:** every retry appends **another copy of the
same original image** (both an `{"type":"image"}` placeholder and the image in the per-
rollout image list). With `max_pixels=401408`, each image ≈ 250 tokens, so after 4 retries
the prompt carries ~5 identical images (~1.25k tokens of duplicate). Risks: hitting
`max_prompt_length=8192` (left-truncation could drop the leading image position) and
extra memory/time on long retries. **Suggested fix:** pass the image only on the first
turn (`include_image=False` on retries) — the chat already shows it once. Not a
correctness bug, but a real perf/robustness footgun.

---

## Cross-cutting: a real pre-existing bug in baseline `r1-v` (fixed in all four forks)

`r1-v/src/open_r1/grpo.py` reward/score fns (`accuracy_reward_stage2`,
`accuracy_score_stage1`, `accuracy_score_stage2`) iterate per-sample over
`zip(contents, solution)` but route the dataset handler via **`dataset[0]`** and log
**`problem_id[0]`** — the batch-wide kwargs collapsed to the first sample. For any batch
spanning >1 dataset (`per_device_train_batch_size>1`, or grad-accum across prompts),
every sample gets the **first sample's** handler → silently wrong rewards (no crash). The
shipped config (`per_device=1`, `accum=1`) masks it because a GRPO group is one prompt.

All four `no-*` forks fix it by iterating `cur_dataset`/`cur_problem_id` from the zip and
defaulting `reward, student_answer = 0.0, ""` in the missing-handler branch. Worth
upstreaming to baseline `r1-v`.

---

## Shared mechanical fixes in every fork

- Dead hardcoded resume path (`/home/meng/GRPO/...`) → `get_last_checkpoint(output_dir)`
  (baseline would never actually resume on a fresh box).
- `trainer.save_state(output_dir)` → `trainer.save_state()` (kwarg-less is the supported
  Trainer API).
- `r1-v-no-hist` additionally adds the PyTorch-2.6+ `torch.load` weights-only workaround
  for DeepSpeed resume — the other forks should copy it if they need to resume under
  torch ≥ 2.6.

---

## Bottom line

- **Clean:** `r1-v-no123`, `r1-v-no123-match-round`.
- **By-design caveat:** `r1-v-no-hist` — only the final round contributes gradient.
- **Perf footgun:** `r1-v-no23-1round` — duplicate original image re-sent on every retry.
- **Real bug (baseline):** per-batch reward mis-routing via `dataset[0]`/`problem_id[0]`;
  fixed in the forks, should go upstream.