diff --git a/.agents/exploration/interleavethinker-fastvideo-integration.md b/.agents/exploration/interleavethinker-fastvideo-integration.md
index 01ff9c88d..f5817fb30 100644
--- a/.agents/exploration/interleavethinker-fastvideo-integration.md
+++ b/.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
diff --git a/examples/train/configs/rl/interleave_thinker/critic_grpo.yaml b/examples/train/configs/rl/interleave_thinker/critic_grpo.yaml
index 5c924c252..ebb84ff3d 100644
--- a/examples/train/configs/rl/interleave_thinker/critic_grpo.yaml
+++ b/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
diff --git a/fastvideo/train/models/interleave_thinker/critic.py b/fastvideo/train/models/interleave_thinker/critic.py
index 1356dff36..05b236fa9 100644
--- a/fastvideo/train/models/interleave_thinker/critic.py
+++ b/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,
diff --git a/tests/local_tests/test_interleave_thinker_critic_model.py b/tests/local_tests/test_interleave_thinker_critic_model.py
index bec6a091b..0f818b632 100644
--- a/tests/local_tests/test_interleave_thinker_critic_model.py
+++ b/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 [
+ 'ok{"previous_step_success": true, "refine_prompt": "better"}'
+ 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": 'ok{"previous_step_success": true, "refine_prompt": "better"}',
+ }],
+ 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:])
diff --git a/tests/local_tests/test_interleave_thinker_method.py b/tests/local_tests/test_interleave_thinker_method.py
index c46b24c1c..992f70bdb 100644
--- a/tests/local_tests/test_interleave_thinker_method.py
+++ b/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")