[feat] harden InterleaveThinker critic backend

.agents/exploration/interleavethinker-fastvideo-integration.md

@@ -435,6 +435,98 @@ Pending:
 - Commit this model/API integration slice.
 - Push the commit immediately after committing, per user instruction.
 
+Commit completed:
+
+- `fc04d019` — `[feat] integrate InterleaveThinker model backends`
+- Pushed to `origin/interleavethinker-fastvideo`.
+
+## Critic Backend Hardening Extension
+
+Status: implemented and Modal-validated; pending commit/push.
+
+Goal for this slice:
+
+- Tighten the Qwen3-VL critic wrapper from config/import-ready toward
+  executable-backend-ready without pulling real Qwen weights into unit tests.
+- Preserve user constraints:
+  - work only in `/tmp/fastvideo-interleavethinker`,
+  - do not touch `/home/toolbox/FastVideo`,
+  - validate via Modal,
+  - commit and push immediately after commit.
+
+Planned changes:
+
+1. Add `image_dir` handling to `InterleaveThinkerCriticModel`, matching
+   upstream EasyR1's convention that relative `origin_image_path` and
+   `edited_image_path` are resolved under a configured image root.
+2. Normalize image paths at dataset load/collate time, and expose an explicit
+   before/after image pair:
+   - before: `previous_image_path`, then `origin_image_path`, then
+     `input_image_path`;
+   - after: `edited_image_path`, then `generated_image_path`, then
+     `output_image_path`.
+3. Improve actor-update scaling so per-rollout backward computes the mean batch
+   policy gradient instead of summing all rollout losses.
+4. Add fake Qwen backend tests that exercise:
+   - `generate_interleave_responses(...)`,
+   - `train_interleave_rollouts(...)`,
+   - image path resolution,
+   - public YAML parse for the new config knobs.
+
+Validation plan:
+
+- Modal pytest:
+  `pytest tests/local_tests/test_interleave_thinker_reward.py
+  tests/local_tests/test_interleave_thinker_method.py
+  tests/local_tests/test_interleave_thinker_api_models.py
+  tests/local_tests/test_interleave_thinker_critic_model.py
+  tests/local_tests/test_train_rl_sampling.py -q`
+- Modal pre-commit on changed files.
+
+Implemented changes:
+
+- `InterleaveThinkerCriticModel` accepts `image_dir` and passes it into the
+  JSON/JSONL dataset loader.
+- Dataset loading normalizes relative image paths under `image_dir` while
+  leaving absolute paths, URI-style paths, and `data:` URLs unchanged.
+- Qwen3-VL message assembly now selects an explicit before/after image pair.
+  This avoids sending `origin`, `previous`, and `edited` together when the
+  critic task expects two images.
+- Actor update scaling now averages over `len(rollouts) *
+  gradient_accumulation_steps`, preventing the loss from growing with rollout
+  count.
+- Added fake-backend tests for generation and response-token-only actor
+  training without downloading real Qwen weights.
+- Added `image_dir: data` to
+  `examples/train/configs/rl/interleave_thinker/critic_grpo.yaml`.
+
+Validation completed:
+
+- Modal pytest:
+  - App URL: `https://modal.com/apps/hao-ai-lab/main/ap-fRlUUdFx19OWpSvDSrWoAe`
+  - Command:
+    `pytest tests/local_tests/test_interleave_thinker_reward.py
+    tests/local_tests/test_interleave_thinker_method.py
+    tests/local_tests/test_interleave_thinker_api_models.py
+    tests/local_tests/test_interleave_thinker_critic_model.py
+    tests/local_tests/test_train_rl_sampling.py -q`
+  - Result: `30 passed, 14 warnings in 19.47s`.
+- Modal pre-commit:
+  - App URL: `https://modal.com/apps/hao-ai-lab/main/ap-DplMFq23YYfBx34e6TcsRc`
+  - Command:
+    `pre-commit run --files
+    examples/train/configs/rl/interleave_thinker/critic_grpo.yaml
+    fastvideo/train/models/interleave_thinker/critic.py
+    tests/local_tests/test_interleave_thinker_critic_model.py
+    tests/local_tests/test_interleave_thinker_method.py`
+  - Result: yapf, ruff, codespell, mypy, filename check, and suggestion hooks
+    passed; PyMarkdown/actionlint skipped with no files to check.
+
+Pending:
+
+- Commit this critic-backend hardening slice.
+- Push the commit immediately after committing, per user instruction.
+
 ## Training Loop Extension
 
 Status: in progress as of the user request to "Add a training loop" for

examples/train/configs/rl/interleave_thinker/critic_grpo.yaml

@@ -11,6 +11,7 @@ models:
     init_from: InterleaveThinker/Critic-SFT-8B
     processor_from: Qwen/Qwen3-VL-8B-Instruct
     load_backend: true
+    image_dir: data
     torch_dtype: auto
     attn_implementation: flash_attention_2
     freeze_vision_tower: true

fastvideo/train/models/interleave_thinker/critic.py

@@ -60,8 +60,13 @@ class _InterleaveJSONLDataset(Dataset):
     def __init__(
         self,
         data_path: str,
+        *,
+        image_dir: str = "",
     ) -> None:
-        self.records = list(_load_interleave_records(data_path))
+        self.records = list(_load_interleave_records(
+            data_path,
+            image_dir=image_dir,
+        ))
         if not self.records:
             raise ValueError(f"No InterleaveThinker records found at {data_path!r}")
 
@@ -101,6 +106,7 @@ class InterleaveThinkerCriticModel(ModelBase):
         training_config: TrainingConfig | None = None,
         trainable: bool = True,
         load_backend: bool = True,
+        image_dir: str = "",
         torch_dtype: str = "auto",
         device_map: str | dict[str, Any] | None = None,
         attn_implementation: str | None = None,
@@ -119,6 +125,7 @@ class InterleaveThinkerCriticModel(ModelBase):
         super().__init__(trainable=trainable, lora=lora)
         self.init_from = str(init_from)
         self.processor_from = str(processor_from or init_from)
+        self.image_dir = str(image_dir or "")
         self.training_config = training_config
         self.max_prompt_length = int(max_prompt_length)
         self.max_response_length = int(max_response_length)
@@ -214,7 +221,10 @@ class InterleaveThinkerCriticModel(ModelBase):
             return
         if not os.path.exists(os.path.expanduser(data_path)):
             raise FileNotFoundError(f"InterleaveThinker data_path not found: {data_path}")
-        dataset = _InterleaveJSONLDataset(data_path)
+        dataset = _InterleaveJSONLDataset(
+            data_path,
+            image_dir=self.image_dir,
+        )
         self.dataloader = DataLoader(
             dataset,
             batch_size=max(1, int(training_config.data.train_batch_size or 1)),
@@ -295,10 +305,11 @@ class InterleaveThinkerCriticModel(ModelBase):
         weighted_losses: list[torch.Tensor] = []
         nll_values: list[torch.Tensor] = []
         token_counts: list[int] = []
+        loss_scale = max(1, len(rollouts) * grad_accum)
         for idx, rollout in enumerate(rollouts):
             nll, token_count = self._response_nll(rollout)
             weighted = nll * advantages[idx]
-            (weighted / grad_accum).backward()
+            (weighted / loss_scale).backward()
             weighted_losses.append(weighted.detach())
             nll_values.append(nll.detach())
             token_counts.append(token_count)
@@ -446,7 +457,11 @@ class InterleaveThinkerCriticModel(ModelBase):
         raise NotImplementedError("InterleaveThinkerCriticModel uses train_interleave_rollouts().")
 
 
-def _load_interleave_records(data_path: str, ) -> list[dict[str, Any]]:
+def _load_interleave_records(
+    data_path: str,
+    *,
+    image_dir: str = "",
+) -> list[dict[str, Any]]:
     path = Path(os.path.expanduser(data_path))
     files = sorted(path.glob("*.json*")) if path.is_dir() else [path]
     records: list[dict[str, Any]] = []
@@ -456,18 +471,55 @@ def _load_interleave_records(data_path: str, ) -> list[dict[str, Any]]:
                 if line.strip():
                     raw = json.loads(line)
                     if isinstance(raw, Mapping):
-                        records.append(dict(raw))
+                        records.append(_normalize_interleave_record(raw, image_dir=image_dir))
         elif file_path.suffix == ".json":
             raw_json = json.loads(file_path.read_text())
             if isinstance(raw_json, list):
-                records.extend(dict(item) for item in raw_json if isinstance(item, Mapping))
+                records.extend(
+                    _normalize_interleave_record(item, image_dir=image_dir) for item in raw_json
+                    if isinstance(item, Mapping))
             elif isinstance(raw_json, Mapping):
-                records.append(dict(raw_json))
+                records.append(_normalize_interleave_record(raw_json, image_dir=image_dir))
         else:
             raise ValueError(f"Unsupported InterleaveThinker data file: {file_path}")
     return records
 
 
+def _normalize_interleave_record(
+    record: Mapping[str, Any],
+    *,
+    image_dir: str = "",
+) -> dict[str, Any]:
+    normalized = dict(record)
+    for key in (
+            "origin_image_path",
+            "previous_image_path",
+            "edited_image_path",
+            "generated_image_path",
+            "input_image_path",
+            "output_image_path",
+    ):
+        value = normalized.get(key)
+        if isinstance(value, str) and value:
+            normalized[key] = _resolve_image_path(value, image_dir=image_dir)
+    return normalized
+
+
+def _resolve_image_path(
+    value: str,
+    *,
+    image_dir: str = "",
+) -> str:
+    expanded = os.path.expanduser(value)
+    if not image_dir or os.path.isabs(expanded) or _looks_like_uri(expanded):
+        return expanded
+    return str(Path(os.path.expanduser(image_dir)) / expanded)
+
+
+def _looks_like_uri(value: str) -> bool:
+    return "://" in value or value.startswith("data:")
+
+
 def _collate_interleave_records(records: Sequence[Mapping[str, Any]], ) -> dict[str, Any]:
     return {"items": [dict(record) for record in records]}
 
@@ -493,16 +545,25 @@ def _batch_to_items(batch: Mapping[str, Any], ) -> list[dict[str, Any]]:
 
 
 def _item_image_paths(item: Mapping[str, Any], ) -> list[str]:
-    image_paths: list[str] = []
-    for key in ("origin_image_path", "previous_image_path", "edited_image_path"):
-        value = item.get(key)
-        if isinstance(value, str) and value:
-            image_paths.append(value)
+    before = _first_string(item, "previous_image_path", "origin_image_path", "input_image_path")
+    after = _first_string(item, "edited_image_path", "generated_image_path", "output_image_path")
+    image_paths = [value for value in (before, after) if value]
     if len(image_paths) == 1:
         image_paths.append(image_paths[0])
     return image_paths[:2]
 
 
+def _first_string(
+    item: Mapping[str, Any],
+    *keys: str,
+) -> str | None:
+    for key in keys:
+        value = item.get(key)
+        if isinstance(value, str) and value:
+            return value
+    return None
+
+
 def _move_to_device(
     value: Any,
     device: torch.device,

tests/local_tests/test_interleave_thinker_critic_model.py

@@ -1,5 +1,11 @@
 import json
+from types import SimpleNamespace
 
+import torch
+
+from fastvideo.train.models.interleave_thinker.critic import (
+    _PlaceholderActorModule,
+)
 from fastvideo.train.models.interleave_thinker import (
     INTERLEAVE_CRITIC_PROMPT,
     InterleaveThinkerCriticModel,
@@ -10,6 +16,65 @@ from fastvideo.train.utils.training_config import (
 )
 
 
+class _FakeBackendCritic(InterleaveThinkerCriticModel):
+
+    @property
+    def device(self):
+        return torch.device("cpu")
+
+
+class _FakeProcessor:
+
+    def __init__(self):
+        self.tokenizer = SimpleNamespace(pad_token_id=0, eos_token_id=2)
+
+    def apply_chat_template(
+        self,
+        messages,
+        *,
+        tokenize,
+        add_generation_prompt,
+        return_dict,
+        return_tensors,
+    ):
+        del tokenize, add_generation_prompt, return_dict, return_tensors
+        has_assistant = any(message["role"] == "assistant" for message in messages)
+        length = 5 if has_assistant else 3
+        return {
+            "input_ids": torch.arange(1, length + 1).unsqueeze(0),
+            "attention_mask": torch.ones(1, length, dtype=torch.long),
+        }
+
+    def batch_decode(self, sequences, **kwargs):
+        del kwargs
+        return [
+            '<think>ok</think><answer>{"previous_step_success": true, "refine_prompt": "better"}</answer>'
+            for _ in sequences
+        ]
+
+
+class _FakeQwen(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.tensor(1.0))
+        self.generate_kwargs = None
+        self.last_labels = None
+
+    def generate(self, **kwargs):
+        self.generate_kwargs = kwargs
+        input_ids = kwargs["input_ids"]
+        num_return_sequences = int(kwargs["num_return_sequences"])
+        suffix = torch.tensor([[9, 10]], dtype=input_ids.dtype)
+        return torch.cat([input_ids.repeat(num_return_sequences, 1), suffix.repeat(num_return_sequences, 1)], dim=1)
+
+    def forward(self, **kwargs):
+        labels = kwargs["labels"]
+        self.last_labels = labels.detach().clone()
+        trainable_fraction = (labels != -100).float().mean()
+        return SimpleNamespace(loss=self.weight.pow(2).sum() * trainable_fraction)
+
+
 def test_interleave_thinker_critic_builds_qwen_vl_messages_without_loading_backend():
     model = InterleaveThinkerCriticModel(load_backend=False)
 
@@ -35,25 +100,107 @@ def test_interleave_thinker_critic_builds_qwen_vl_messages_without_loading_backe
     assert "a vase on a table" in text
 
 
+def test_interleave_thinker_critic_prefers_previous_image_as_before_image():
+    model = InterleaveThinkerCriticModel(load_backend=False)
+
+    messages = model.build_messages({
+        "origin_prompt": "draw a vase",
+        "previous_prompt": "a vase on a table",
+        "origin_image_path": "original.png",
+        "previous_image_path": "before.png",
+        "edited_image_path": "after.png",
+    })
+
+    images = [part for part in messages[0]["content"] if part["type"] == "image"]
+    assert images == [{
+        "type": "image",
+        "image": "before.png",
+    }, {
+        "type": "image",
+        "image": "after.png",
+    }]
+
+
 def test_interleave_thinker_critic_initializes_jsonl_dataloader(tmp_path):
     data_path = tmp_path / "critic_rl.jsonl"
     data_path.write_text(
         json.dumps({
             "origin_prompt": "draw a chair",
             "rewritten_prompt": "a wooden chair",
-            "origin_image_path": "before.png",
-            "edited_image_path": "after.png",
+            "origin_image_path": "interleave/before.png",
+            "edited_image_path": "interleave/after.png",
             "evaluation": {
                 "success": True,
                 "semantics": 7.0,
                 "quality": 8.0,
             },
         }) + "\n")
-    model = InterleaveThinkerCriticModel(load_backend=False)
+    image_dir = tmp_path / "images"
+    model = InterleaveThinkerCriticModel(load_backend=False, image_dir=str(image_dir))
 
     model.init_preprocessors(TrainingConfig(data=DataConfig(data_path=str(data_path), train_batch_size=1)))
     batch = next(iter(model.dataloader))
 
     assert batch["items"][0]["origin_prompt"] == "draw a chair"
+    assert batch["items"][0]["origin_image_path"] == str(image_dir / "interleave/before.png")
+    assert batch["items"][0]["edited_image_path"] == str(image_dir / "interleave/after.png")
     assert batch["items"][0]["evaluation"]["success"] is True
     assert "{original_instruction}" in INTERLEAVE_CRITIC_PROMPT
+
+
+def test_interleave_thinker_critic_fake_backend_generates_rollouts():
+    model = _FakeBackendCritic(load_backend=False)
+    model.processor = _FakeProcessor()
+    model.transformer = _FakeQwen()
+
+    rollouts = model.generate_interleave_responses(
+        {
+            "items": [{
+                "origin_prompt": "draw a chair",
+                "previous_prompt": "a wooden chair",
+                "origin_image_path": "before.png",
+                "edited_image_path": "after.png",
+            }]
+        },
+        num_generations=2,
+        temperature=0.7,
+        top_p=0.9,
+        max_new_tokens=8,
+    )
+
+    assert len(rollouts) == 2
+    assert all("previous_step_success" in rollout["response"] for rollout in rollouts)
+    assert rollouts[0]["group_key"] == rollouts[1]["group_key"]
+    assert rollouts[0]["sample_index"] == 0
+    assert model.transformer.generate_kwargs["num_return_sequences"] == 2
+    assert model.transformer.generate_kwargs["max_new_tokens"] == 8
+
+
+def test_interleave_thinker_critic_fake_backend_trains_response_tokens_only():
+    model = _FakeBackendCritic(load_backend=False)
+    assert isinstance(model.transformer, _PlaceholderActorModule)
+    model.processor = _FakeProcessor()
+    model.transformer = _FakeQwen()
+    optimizer = torch.optim.SGD(model.transformer.parameters(), lr=0.1)
+    before = float(model.transformer.weight.detach())
+
+    loss_map, metrics = model.train_interleave_rollouts(
+        rollouts=[{
+            "origin_prompt": "draw a chair",
+            "previous_prompt": "a wooden chair",
+            "origin_image_path": "before.png",
+            "edited_image_path": "after.png",
+            "response": '<think>ok</think><answer>{"previous_step_success": true, "refine_prompt": "better"}</answer>',
+        }],
+        advantages=torch.tensor([1.0]),
+        rewards={},
+        optimizer=optimizer,
+        gradient_accumulation_steps=1,
+        max_grad_norm=0.0,
+    )
+
+    assert "total_loss" in loss_map
+    assert metrics["actor/response_tokens"] == 2.0
+    assert float(model.transformer.weight.detach()) < before
+    assert model.transformer.last_labels.tolist()[0][:3] == [-100, -100, -100]
+    assert all(label != -100 for label in model.transformer.last_labels.tolist()[0][3:])

tests/local_tests/test_interleave_thinker_method.py

@@ -218,6 +218,7 @@ def test_interleave_thinker_config_parses_public_yaml():
     assert cfg.models["student"]["_target_"] == (
         "fastvideo.train.models.interleave_thinker.InterleaveThinkerCriticModel")
     assert cfg.models["student"]["init_from"] == "InterleaveThinker/Critic-SFT-8B"
+    assert cfg.models["student"]["image_dir"] == "data"
     assert cfg.method["_target_"] == "fastvideo.train.methods.rl.interleave_thinker.InterleaveThinkerRLMethod"
     assert cfg.method["edit_scorer"]["_target_"] == (
         "fastvideo.train.methods.rl.rewards.GeminiNanoBananaEditScorer")