diff --git a/.agents/exploration/interleavethinker-fastvideo-integration.md b/.agents/exploration/interleavethinker-fastvideo-integration.md
index 9ed03e5e9..ed093f20e 100644
--- a/.agents/exploration/interleavethinker-fastvideo-integration.md
+++ b/.agents/exploration/interleavethinker-fastvideo-integration.md
@@ -2712,3 +2712,154 @@ Conclusion:
- Next full-integration stage: add a planner-policy RL path so the
InterleaveThinker planner can be optimized with the same FastVideo GRPO loop,
using planner-specific rollouts and rewards instead of only critic RL/SFT.
+
+## Stage 10 Execution: Planner GRPO Path
+
+Status: validation complete; ready to commit and push.
+
+Goal:
+
+- Let `InterleaveThinkerPlannerModel` participate in the same FastVideo
+ managed GRPO loop as the critic actor.
+- Keep Qwen response-token policy optimization shared between planner and
+ critic wrappers.
+- Add a planner-specific reward scorer for valid
+ `...{"execution_plan": [...]}` outputs, with
+ optional external/fixture plan scores.
+- Add a public planner GRPO YAML using a trainable planner student and frozen
+ planner reference model.
+
+Planned implementation:
+
+- Move or duplicate the Qwen GRPO actor update onto `Qwen3VLActorBase` so both
+ planner and critic actors can train over response-token logprobs.
+- Add `InterleaveThinkerPlannerModel.generate_interleave_responses(...)`
+ returning rollouts compatible with `InterleaveThinkerRLMethod`:
+ `response`, parsed plan metadata, `old_logprobs`, `response_mask`,
+ `sample_index`, `generation_index`, and `group_key`.
+- Add `planner_rl` dataset normalization that can consume prompt-only planner
+ RL files without requiring an assistant response.
+- Add `InterleavePlannerRewardScorer` under
+ `fastvideo/train/methods/rl/rewards/`.
+- Add `examples/train/configs/rl/interleave_thinker/planner_grpo.yaml`.
+- Add focused unit tests for:
+ - planner RL data normalization;
+ - planner reward parsing/scoring;
+ - planner rollout generation with old logprobs;
+ - planner public YAML parsing;
+ - shared Qwen actor policy update behavior remains intact.
+
+Validation plan:
+
+- Local syntax checks and `git diff --check`.
+- Modal focused pytest for planner model/data/reward/method/GRPO tests.
+- Modal pre-commit on changed files.
+- If feasible after unit validation, run a one-step Modal smoke with real
+ `InterleaveThinker/InterleaveThinker-Planner-8B` student and frozen
+ reference planner using format-only planner rewards.
+
+Implemented changes:
+
+- Shared Qwen actor update:
+ - Moved the response-token GRPO actor update into
+ `Qwen3VLActorBase.train_interleave_rollouts(...)`.
+ - Moved `_grpo_logprob_batch(...)` and weighted actor-metric aggregation into
+ `qwen_actor.py`.
+ - Removed the duplicate critic-local GRPO update; the critic now inherits the
+ shared Qwen actor path.
+- Planner rollouts:
+ - Added `InterleaveThinkerPlannerModel.generate_interleave_responses(...)`.
+ - Planner rollouts now include `response`, parsed plan metadata, `steps`,
+ `sample_index`, `generation_index`, `group_key`, `old_logprobs`, and
+ `response_mask`.
+ - Planner image prompt construction now accepts the upstream `images` list as
+ an image-source alias.
+- Planner RL data:
+ - Added `planner_rl` to `InterleaveDatasetKind`.
+ - Added `planner_rl.jsonl` default filename.
+ - Added `load_planner_rl_records(...)` and
+ `normalize_planner_rl_record(...)`.
+ - Planner RL records can be prompt-only and do not require an assistant SFT
+ response.
+- Planner rewards:
+ - Added `InterleavePlannerRewardScorer`.
+ - Added `extract_interleave_plan_payload(...)`,
+ `interleave_planner_format_reward(...)`, and
+ `score_interleave_planner_rewards(...)`.
+ - Exported planner reward utilities from
+ `fastvideo.train.methods.rl.rewards`.
+ - `InterleaveThinkerRLMethod` now accepts optional
+ `method.reward_scorer` for scorer injection; when absent, it keeps the
+ existing critic reward scorer behavior.
+- Public config:
+ - Added
+ `examples/train/configs/rl/interleave_thinker/planner_grpo.yaml` with a
+ trainable LoRA planner student, frozen planner reference, planner reward
+ scorer, and GRPO/reference-KL knobs.
+- Tests:
+ - Added planner RL data normalization coverage.
+ - Added planner reward parsing/scoring coverage.
+ - Added planner RL rollout generation coverage with old logprobs.
+ - Added planner GRPO public YAML parse coverage.
+ - Reused existing critic GRPO tests to verify the shared Qwen refactor.
+
+Validation completed:
+
+- Local lightweight checks:
+ - `python -m py_compile` passed for the changed InterleaveThinker model,
+ reward, method, and test modules.
+ - `git diff --check` passed.
+ - Local pre-commit applied yapf/ruff formatting; local mypy still cannot run
+ in this hyphenated worktree path and reports
+ `fastvideo-interleavethinker is not a valid Python package name`.
+ Modal mypy is authoritative.
+- Modal focused pytest:
+ - First app URL:
+ `https://modal.com/apps/hao-ai-lab/main/ap-aHGkQGhbijME5oLBFIdZOH`
+ - Result: `36 passed`, `1 failed`.
+ - Failure was a test fake issue: the fake planner processor returned one
+ decoded response even when generation returned two sequences.
+ - Fixed by making the fake `batch_decode(...)` return one response per
+ sequence.
+- Modal focused pytest after fake fix:
+ - App URL:
+ `https://modal.com/apps/hao-ai-lab/main/ap-CUch3nJsv7MWZfzik5y1wP`
+ - Result: `37 passed, 14 warnings in 19.09s`.
+- Modal pre-commit before local formatter application:
+ - App URL:
+ `https://modal.com/apps/hao-ai-lab/main/ap-TM7tfdNJ9ZfGReagJGpw8P`
+ - Result: yapf and ruff modified files; codespell and mypy passed.
+ - Applied the same formatter changes locally.
+- Modal focused pytest after formatter changes:
+ - App URL:
+ `https://modal.com/apps/hao-ai-lab/main/ap-wccYJm9hNWdmEqyLZZfcnv`
+ - Result: `37 passed, 14 warnings in 15.25s`.
+- Final Modal pre-commit:
+ - App URL:
+ `https://modal.com/apps/hao-ai-lab/main/ap-xdYf7BFSVVyZqKx3Qdd0Yt`
+ - Result: yapf, ruff, codespell, mypy, filename, and suggestion hooks
+ passed; PyMarkdown/actionlint skipped with no files to check.
+- Real one-step planner GRPO smoke:
+ - App URL:
+ `https://modal.com/apps/hao-ai-lab/main/ap-PDBijC8opxsMiMU0Uc064A`
+ - Result:
+ - Loaded the trainable LoRA student
+ `InterleaveThinker/InterleaveThinker-Planner-8B`.
+ - Enabled PEFT LoRA with rank 4 and alpha 8 on
+ `['q_proj', 'k_proj', 'v_proj', 'o_proj']`.
+ - Loaded the frozen planner reference
+ `InterleaveThinker/InterleaveThinker-Planner-8B`.
+ - Generated two planner rollouts, computed old and reference response-token
+ logprobs, scored with `InterleavePlannerRewardScorer`, and completed one
+ GRPO update step with `kl_coef=0.01`.
+ - The job printed `INTERLEAVE_PLANNER_RL_SMOKE_OK`.
+
+Conclusion:
+
+- FastVideo now has planner SFT, critic SFT, critic GRPO, and planner GRPO
+ paths for the InterleaveThinker Qwen actors.
+- The planner and critic share the same Qwen response-token GRPO actor update.
+- The planner and critic GRPO paths both have real one-step Modal smoke
+ validation against actual InterleaveThinker 8B checkpoints.
+- Next action: commit this Stage 10 planner GRPO integration and push it to
+ `origin/interleavethinker-fastvideo`.
diff --git a/examples/train/configs/rl/interleave_thinker/planner_grpo.yaml b/examples/train/configs/rl/interleave_thinker/planner_grpo.yaml
new file mode 100644
index 000000000..e47719c20
--- /dev/null
+++ b/examples/train/configs/rl/interleave_thinker/planner_grpo.yaml
@@ -0,0 +1,101 @@
+# InterleaveThinker planner GRPO bridge.
+#
+# This config trains the Qwen3-VL planner policy with FastVideo's managed GRPO
+# loop. It uses a lightweight planner reward by default: valid
+# ...{"execution_plan": [...]} format, plus an
+# optional scalar plan score when present in the dataset.
+
+models:
+ student:
+ _target_: fastvideo.train.models.interleave_thinker.InterleaveThinkerPlannerModel
+ init_from: InterleaveThinker/InterleaveThinker-Planner-8B
+ processor_from: Qwen/Qwen3-VL-8B-Instruct
+ load_backend: true
+ trainable: true
+ dataset_kind: planner_rl
+ image_dir: data/InterleaveThinker/Train-Data
+ torch_dtype: auto
+ attn_implementation: flash_attention_2
+ freeze_vision_tower: true
+ freeze_multi_modal_projector: true
+ enable_gradient_checkpointing: true
+ max_prompt_length: 16384
+ max_response_length: 4096
+ lora:
+ enable: true
+ rank: 16
+ alpha: 32
+ target_modules:
+ - q_proj
+ - k_proj
+ - v_proj
+ - o_proj
+ reference:
+ _target_: fastvideo.train.models.interleave_thinker.InterleaveThinkerPlannerModel
+ init_from: InterleaveThinker/InterleaveThinker-Planner-8B
+ processor_from: Qwen/Qwen3-VL-8B-Instruct
+ load_backend: true
+ trainable: false
+ dataset_kind: planner_rl
+ image_dir: data/InterleaveThinker/Train-Data
+ torch_dtype: auto
+ attn_implementation: flash_attention_2
+ freeze_vision_tower: true
+ freeze_multi_modal_projector: true
+ enable_gradient_checkpointing: false
+ max_prompt_length: 16384
+ max_response_length: 4096
+
+method:
+ _target_: fastvideo.train.methods.rl.interleave_thinker.InterleaveThinkerRLMethod
+ reward_scorer:
+ _target_: fastvideo.train.methods.rl.rewards.InterleavePlannerRewardScorer
+ format_weight: 1.0
+ fallback_plan_reward: 0.0
+ num_generations: 8
+ num_batches_per_step: 1
+ temperature: 1.0
+ top_p: 1.0
+ max_new_tokens: 2048
+ clip_range: 0.2
+ kl_coef: 0.01
+ micro_batch_size_per_device_for_update: 1
+ advantage_eps: 0.0001
+ advantage_clip: 5.0
+ max_grad_norm: 1.0
+ terminal_progress: true
+
+training:
+ distributed:
+ num_gpus: 8
+ sp_size: 1
+ tp_size: 4
+ hsdp_replicate_dim: 1
+ hsdp_shard_dim: 8
+
+ data:
+ data_path: data/InterleaveThinker/Train-Data/planner_rl.jsonl
+ preprocessed_data_type: text_only
+ train_batch_size: 16
+ dataloader_num_workers: 0
+ seed: 42
+
+ optimizer:
+ learning_rate: 2.0e-6
+ betas: [0.9, 0.999]
+ weight_decay: 0.0
+ lr_scheduler: constant
+ lr_warmup_steps: 0
+
+ loop:
+ max_train_steps: 1000
+ gradient_accumulation_steps: 1
+
+ checkpoint:
+ output_dir: outputs/interleave_thinker_planner_grpo
+ training_state_checkpointing_steps: 50
+ checkpoints_total_limit: 3
+
+ tracker:
+ project_name: InterleaveThinker-GRPO
+ run_name: interleave_thinker_planner_grpo_fastvideo
diff --git a/fastvideo/train/methods/rl/interleave_thinker.py b/fastvideo/train/methods/rl/interleave_thinker.py
index 3e95e6e1e..3884e5d1f 100644
--- a/fastvideo/train/methods/rl/interleave_thinker.py
+++ b/fastvideo/train/methods/rl/interleave_thinker.py
@@ -86,14 +86,7 @@ class InterleaveThinkerRLMethod(TrainingMethod):
)
self._max_grad_norm = self._read_float("max_grad_norm", 0.0)
self._terminal_progress = bool(self.method_config.get("terminal_progress", True))
- self._reward_scorer = InterleaveThinkerRewardScorer(
- format_weight=self._read_float("format_weight", 0.5),
- judge_accuracy_weight=self._read_float("judge_accuracy_weight", 0.2),
- semantic_weight=self._read_float("semantic_weight", 0.6),
- quality_weight=self._read_float("quality_weight", 0.2),
- fallback_edit_reward=self._read_float("fallback_edit_reward", 0.5),
- edit_scorer=self._build_edit_scorer(self.method_config.get("edit_scorer")),
- )
+ self._reward_scorer = self._build_reward_scorer()
self._student_optimizer: torch.optim.Optimizer | None = None
self._student_lr_scheduler: Any | None = None
self._init_optimizer_and_scheduler()
@@ -222,6 +215,23 @@ class InterleaveThinkerRLMethod(TrainingMethod):
return instantiate(dict(raw))
raise TypeError("method.edit_scorer must be a callable or a mapping with _target_")
+ def _build_reward_scorer(self) -> Any:
+ raw = self.method_config.get("reward_scorer")
+ if raw is not None:
+ if callable(raw):
+ return raw
+ if isinstance(raw, Mapping):
+ return instantiate(dict(raw))
+ raise TypeError("method.reward_scorer must be a callable or a mapping with _target_")
+ return InterleaveThinkerRewardScorer(
+ format_weight=self._read_float("format_weight", 0.5),
+ judge_accuracy_weight=self._read_float("judge_accuracy_weight", 0.2),
+ semantic_weight=self._read_float("semantic_weight", 0.6),
+ quality_weight=self._read_float("quality_weight", 0.2),
+ fallback_edit_reward=self._read_float("fallback_edit_reward", 0.5),
+ edit_scorer=self._build_edit_scorer(self.method_config.get("edit_scorer")),
+ )
+
def _freeze_reference_model(self) -> None:
if self.reference is None:
return
diff --git a/fastvideo/train/methods/rl/rewards/__init__.py b/fastvideo/train/methods/rl/rewards/__init__.py
index 8ae3afb00..238146d52 100644
--- a/fastvideo/train/methods/rl/rewards/__init__.py
+++ b/fastvideo/train/methods/rl/rewards/__init__.py
@@ -7,15 +7,20 @@ from typing import TYPE_CHECKING
from fastvideo.train.methods.rl.rewards.interleave_thinker import (
EditScoreProvider,
+ InterleavePlannerRewardResult,
+ InterleavePlannerRewardScorer,
InterleaveThinkerAnswer,
InterleaveThinkerEditRequest,
InterleaveThinkerEditScore,
InterleaveThinkerRewardResult,
InterleaveThinkerRewardScorer,
extract_interleave_answer,
+ extract_interleave_plan_payload,
interleave_format_reward,
+ interleave_planner_format_reward,
interleave_judge_accuracy_reward,
normalize_interleave_response,
+ score_interleave_planner_rewards,
score_interleave_thinker_rewards,
)
@@ -104,6 +109,8 @@ __all__ = [
"EditScoreProvider",
"GeminiInterleaveImageScorer",
"GeminiNanoBananaEditScorer",
+ "InterleavePlannerRewardResult",
+ "InterleavePlannerRewardScorer",
"InterleaveThinkerAnswer",
"InterleaveThinkerEditRequest",
"InterleaveThinkerEditScore",
@@ -114,9 +121,12 @@ __all__ = [
"RewardScorer",
"build_multi_reward_scorer",
"extract_interleave_answer",
+ "extract_interleave_plan_payload",
"interleave_format_reward",
+ "interleave_planner_format_reward",
"interleave_judge_accuracy_reward",
"normalize_interleave_response",
+ "score_interleave_planner_rewards",
"score_interleave_thinker_rewards",
"select_first_frame",
]
diff --git a/fastvideo/train/methods/rl/rewards/interleave_thinker.py b/fastvideo/train/methods/rl/rewards/interleave_thinker.py
index 6bcc03453..f5466d92f 100644
--- a/fastvideo/train/methods/rl/rewards/interleave_thinker.py
+++ b/fastvideo/train/methods/rl/rewards/interleave_thinker.py
@@ -87,6 +87,24 @@ class InterleaveThinkerRewardResult:
}
+@dataclass(frozen=True, slots=True)
+class InterleavePlannerRewardResult:
+ """One scored InterleaveThinker planner rollout."""
+
+ overall: float
+ format_reward: float
+ planner_score: float
+ index: int
+
+ def as_dict(self) -> dict[str, float]:
+ return {
+ "overall": float(self.overall),
+ "format_reward": float(self.format_reward),
+ "planner_score": float(self.planner_score),
+ "idx": float(self.index),
+ }
+
+
EditScoreProvider = Callable[[InterleaveThinkerEditRequest], InterleaveThinkerEditScore | Mapping[str, Any] | None]
@@ -141,6 +159,40 @@ def interleave_format_reward(response: str) -> float:
return 1.0 if extract_interleave_answer(normalized) is not None else 0.0
+def extract_interleave_plan_payload(response: str) -> dict[str, Any] | None:
+ """Extract the planner ``execution_plan`` answer payload."""
+ match = _ANSWER_PATTERN.search(normalize_interleave_response(response))
+ if not match:
+ return None
+ payload = _jsonish_loads(match.group(1))
+ if payload is None:
+ return None
+ raw_steps = payload.get("execution_plan")
+ if not isinstance(raw_steps, Sequence) or isinstance(raw_steps, str | bytes) or not raw_steps:
+ return None
+ for step in raw_steps:
+ if not isinstance(step, Mapping):
+ return None
+ has_prompt = _non_empty_optional_text(step.get("prompt")) is not None
+ has_instruction = _non_empty_optional_text(step.get("instruction")) is not None
+ has_auxiliary = _non_empty_optional_text(step.get("auxiliary_text")) is not None
+ if not (has_prompt or has_instruction or has_auxiliary):
+ return None
+ return payload
+
+
+def interleave_planner_format_reward(response: str) -> float:
+ """Return 1 when a planner response has valid reasoning and plan JSON."""
+ normalized = normalize_interleave_response(response)
+ think_match = _THINK_PATTERN.search(normalized)
+ answer_match = _ANSWER_PATTERN.search(normalized)
+ if think_match is None or answer_match is None:
+ return 0.0
+ if think_match.start() >= answer_match.start():
+ return 0.0
+ return 1.0 if extract_interleave_plan_payload(normalized) is not None else 0.0
+
+
def interleave_judge_accuracy_reward(
predicted_previous_step_success: bool | None,
ground_truth_previous_step_success: bool | None,
@@ -161,6 +213,15 @@ def _coerce_float(value: Any, *, default: float = 0.0) -> float:
return float(value)
+def _non_empty_optional_text(value: Any) -> str | None:
+ if value is None:
+ return None
+ text = str(value).strip()
+ if not text or text.lower() in {"none", "null"}:
+ return None
+ return text
+
+
def _mapping_get_bool(mapping: Mapping[str, Any], *keys: str) -> bool | None:
for key in keys:
if key in mapping:
@@ -244,6 +305,24 @@ def _normalized_edit_rewards(
)
+def _planner_score_from_item(
+ item: Mapping[str, Any],
+ *,
+ fallback: float,
+) -> float:
+ for key in ("planner_score", "plan_score", "reward", "score"):
+ if key in item:
+ return _coerce_float(item.get(key), default=fallback)
+ raw = item.get("ground_truth", item.get("evaluation", {}))
+ if isinstance(raw, str):
+ raw = _jsonish_loads(raw) or {}
+ if isinstance(raw, Mapping):
+ for key in ("planner_score", "plan_score", "reward", "score"):
+ if key in raw:
+ return _coerce_float(raw.get(key), default=fallback)
+ return fallback
+
+
class InterleaveThinkerRewardScorer:
"""Batch scorer for InterleaveThinker critic outputs.
@@ -303,10 +382,8 @@ class InterleaveThinkerRewardScorer:
previous_quality=previous_quality,
fallback=self.fallback_edit_reward,
)
- non_format_reward = (
- self.judge_accuracy_weight * judge_reward + self.semantic_weight * semantic_reward +
- self.quality_weight * quality_reward
- ) / inner_total
+ non_format_reward = (self.judge_accuracy_weight * judge_reward + self.semantic_weight * semantic_reward +
+ self.quality_weight * quality_reward) / inner_total
overall = self.format_weight * format_reward + (1.0 - self.format_weight) * non_format_reward
results.append(
InterleaveThinkerRewardResult(
@@ -344,6 +421,64 @@ class InterleaveThinkerRewardScorer:
return tensors
+class InterleavePlannerRewardScorer:
+ """Batch scorer for InterleaveThinker planner outputs.
+
+ This scorer is intentionally lightweight: by default it rewards valid
+ planner response format and can blend in an externally supplied scalar plan
+ score from each rollout for stronger supervision.
+ """
+
+ def __init__(
+ self,
+ *,
+ format_weight: float = 1.0,
+ fallback_plan_reward: float = 0.0,
+ ) -> None:
+ self.format_weight = float(format_weight)
+ self.fallback_plan_reward = float(fallback_plan_reward)
+ if not 0.0 <= self.format_weight <= 1.0:
+ raise ValueError("format_weight must be in [0, 1]")
+
+ def __call__(
+ self,
+ reward_inputs: Sequence[Mapping[str, Any]],
+ ) -> list[InterleavePlannerRewardResult]:
+ results: list[InterleavePlannerRewardResult] = []
+ for index, item in enumerate(reward_inputs):
+ response = str(item.get("response", "") or "")
+ format_reward = interleave_planner_format_reward(response)
+ planner_score = _planner_score_from_item(
+ item,
+ fallback=self.fallback_plan_reward,
+ )
+ overall = self.format_weight * format_reward + (1.0 - self.format_weight) * planner_score
+ results.append(
+ InterleavePlannerRewardResult(
+ overall=float(overall),
+ format_reward=float(format_reward),
+ planner_score=float(planner_score),
+ index=index,
+ ))
+ return results
+
+ def as_tensors(
+ self,
+ reward_inputs: Sequence[Mapping[str, Any]],
+ *,
+ device: torch.device | str = "cpu",
+ ) -> dict[str, torch.Tensor]:
+ results = self(reward_inputs)
+ names = ["overall", "format_reward", "planner_score"]
+ tensors: dict[str, torch.Tensor] = {}
+ for name in names:
+ tensors[name] = torch.tensor([getattr(result, name) for result in results],
+ device=device,
+ dtype=torch.float32)
+ tensors["avg"] = tensors["overall"]
+ return tensors
+
+
def score_interleave_thinker_rewards(
reward_inputs: Sequence[Mapping[str, Any]],
**kwargs: Any,
@@ -353,16 +488,30 @@ def score_interleave_thinker_rewards(
return [result.as_dict() for result in scorer(reward_inputs)]
+def score_interleave_planner_rewards(
+ reward_inputs: Sequence[Mapping[str, Any]],
+ **kwargs: Any,
+) -> list[dict[str, float]]:
+ """Convenience wrapper for planner rollout rewards."""
+ scorer = InterleavePlannerRewardScorer(**kwargs)
+ return [result.as_dict() for result in scorer(reward_inputs)]
+
+
__all__ = [
"EditScoreProvider",
+ "InterleavePlannerRewardResult",
"InterleaveThinkerAnswer",
"InterleaveThinkerEditRequest",
"InterleaveThinkerEditScore",
"InterleaveThinkerRewardResult",
+ "InterleavePlannerRewardScorer",
"InterleaveThinkerRewardScorer",
"extract_interleave_answer",
+ "extract_interleave_plan_payload",
"interleave_format_reward",
+ "interleave_planner_format_reward",
"interleave_judge_accuracy_reward",
"normalize_interleave_response",
+ "score_interleave_planner_rewards",
"score_interleave_thinker_rewards",
]
diff --git a/fastvideo/train/models/interleave_thinker/__init__.py b/fastvideo/train/models/interleave_thinker/__init__.py
index 06ec7da2e..f57e7b97f 100644
--- a/fastvideo/train/models/interleave_thinker/__init__.py
+++ b/fastvideo/train/models/interleave_thinker/__init__.py
@@ -14,11 +14,13 @@ from fastvideo.train.models.interleave_thinker.data import (
load_critic_rl_records,
load_critic_sft_records,
load_interleave_dataset,
+ load_planner_rl_records,
load_planner_sft_records,
normalize_critic_rl_record,
normalize_critic_sft_record,
normalize_ground_truth,
normalize_interleave_dataset_record,
+ normalize_planner_rl_record,
normalize_planner_sft_record,
resolve_interleave_image_path,
validate_image_path,
@@ -49,11 +51,13 @@ __all__ = [
"load_critic_rl_records",
"load_critic_sft_records",
"load_interleave_dataset",
+ "load_planner_rl_records",
"load_planner_sft_records",
"normalize_critic_rl_record",
"normalize_critic_sft_record",
"normalize_ground_truth",
"normalize_interleave_dataset_record",
+ "normalize_planner_rl_record",
"normalize_planner_sft_record",
"resolve_interleave_image_path",
"validate_image_path",
diff --git a/fastvideo/train/models/interleave_thinker/critic.py b/fastvideo/train/models/interleave_thinker/critic.py
index d5612aa18..689e861dc 100644
--- a/fastvideo/train/models/interleave_thinker/critic.py
+++ b/fastvideo/train/models/interleave_thinker/critic.py
@@ -3,7 +3,7 @@
from __future__ import annotations
-from collections.abc import Mapping, Sequence
+from collections.abc import Mapping
from typing import Any, TYPE_CHECKING
import torch
@@ -12,14 +12,10 @@ from fastvideo.train.models.interleave_thinker.qwen_actor import (
Qwen3VLActorBase,
_PlaceholderActorModule as _PlaceholderActorModule,
batch_to_items,
- coerce_logprob_vector,
- coerce_response_mask,
first_string,
- pad_1d_tensors,
rollout_group_key,
)
from fastvideo.train.models.interleave_thinker.data import InterleaveDatasetKind
-from fastvideo.train.methods.rl.common.grpo import compute_grpo_loss
if TYPE_CHECKING:
from fastvideo.train.utils.lora import LoraConfig
@@ -143,89 +139,6 @@ class InterleaveThinkerCriticModel(Qwen3VLActorBase):
rollouts.append(rollout)
return rollouts
- def train_interleave_rollouts(
- self,
- **kwargs: Any,
- ) -> tuple[dict[str, torch.Tensor], dict[str, float]]:
- self._require_backend()
- rollouts = kwargs["rollouts"]
- advantages = kwargs["advantages"].detach().to(self.device).float()
- optimizer = kwargs.get("optimizer")
- lr_scheduler = kwargs.get("lr_scheduler")
- clip_range = float(kwargs.get("clip_range", 0.2) or 0.0)
- kl_coef = float(kwargs.get("kl_coef", 0.0) or 0.0)
- update_micro_batch_size = kwargs.get("update_micro_batch_size")
- if update_micro_batch_size is None:
- update_micro_batch_size = len(rollouts)
- update_micro_batch_size = max(1, int(update_micro_batch_size or 1))
- max_grad_norm = float(kwargs.get("max_grad_norm", 0.0) or 0.0)
- grad_accum = max(1, int(kwargs.get("gradient_accumulation_steps", 1) or 1))
-
- if optimizer is None:
- raise RuntimeError("InterleaveThinkerCriticModel.train_interleave_rollouts() requires an optimizer")
- optimizer.zero_grad(set_to_none=True)
- self.transformer.train()
-
- if int(advantages.shape[0]) != len(rollouts):
- raise ValueError("advantages count must match rollout count")
-
- loss_results = []
- token_counts: list[float] = []
- loss_scale = max(1, grad_accum)
- for start in range(0, len(rollouts), update_micro_batch_size):
- end = min(start + update_micro_batch_size, len(rollouts))
- current_logprobs, old_logprobs, response_masks, reference_logprobs = self._grpo_logprob_batch(
- rollouts[start:end])
- result = compute_grpo_loss(
- current_logprobs=current_logprobs,
- old_logprobs=old_logprobs,
- advantages=advantages[start:end],
- response_mask=response_masks,
- clip_range=clip_range,
- reference_logprobs=reference_logprobs,
- kl_coef=kl_coef,
- )
- (result.total_loss / loss_scale).backward()
- loss_results.append(result)
- token_counts.append(float(result.token_count.detach().cpu()))
-
- if max_grad_norm > 0.0:
- torch.nn.utils.clip_grad_norm_(
- [p for p in self.transformer.parameters() if p.requires_grad],
- max_grad_norm,
- )
- optimizer.step()
- if lr_scheduler is not None:
- lr_scheduler.step()
-
- total_tokens = max(1.0, sum(token_counts))
- total_loss = _weighted_result_mean(loss_results, "total_loss", token_counts, self.device, total_tokens)
- policy_loss = _weighted_result_mean(loss_results, "policy_loss", token_counts, self.device, total_tokens)
- kl_loss = _weighted_result_mean(loss_results, "kl_loss", token_counts, self.device, total_tokens)
- approx_kl = _weighted_result_mean(loss_results, "approx_kl", token_counts, self.device, total_tokens)
- clipped_fraction = _weighted_result_mean(
- loss_results,
- "clipped_fraction",
- token_counts,
- self.device,
- total_tokens,
- )
- mean_ratio = _weighted_result_mean(loss_results, "mean_ratio", token_counts, self.device, total_tokens)
- return (
- {
- "total_loss": total_loss.detach()
- },
- {
- "actor/policy_loss": float(policy_loss.detach().cpu()),
- "actor/kl_loss": float(kl_loss.detach().cpu()),
- "actor/approx_kl": float(approx_kl.detach().cpu()),
- "actor/clipped_fraction": float(clipped_fraction.detach().cpu()),
- "actor/mean_ratio": float(mean_ratio.detach().cpu()),
- "actor/mean_advantage": float(advantages.mean().detach().cpu()),
- "actor/response_tokens": float(total_tokens / max(1, len(rollouts))),
- },
- )
-
def build_messages(
self,
item: Mapping[str, Any],
@@ -236,67 +149,6 @@ class InterleaveThinkerCriticModel(Qwen3VLActorBase):
)
return self.build_text_image_messages(prompt, _item_image_paths(item))
- def _response_nll(
- self,
- rollout: Mapping[str, Any],
- ) -> tuple[torch.Tensor, int]:
- return self.response_nll_from_messages(
- self.build_messages(rollout),
- str(rollout.get("response", "") or ""),
- )
-
- def _grpo_logprob_batch(
- self,
- rollouts: Sequence[Mapping[str, Any]],
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
- current_logprob_rows: list[torch.Tensor] = []
- old_logprob_rows: list[torch.Tensor] = []
- mask_rows: list[torch.Tensor] = []
- reference_rows: list[torch.Tensor] = []
- has_reference = False
- for rollout in rollouts:
- response = str(rollout.get("response", "") or "")
- current_logprobs, default_mask = self.response_logprobs_from_messages(
- self.build_messages(rollout),
- response,
- )
- expected_len = int(current_logprobs.numel())
- old_logprobs = coerce_logprob_vector(
- rollout.get("old_logprobs", rollout.get("old_logprob")),
- expected_len=expected_len,
- device=self.device,
- name="old_logprobs",
- )
- if old_logprobs is None:
- old_logprobs = current_logprobs.detach()
- response_mask = coerce_response_mask(
- rollout.get("response_mask"),
- expected_len=expected_len,
- device=self.device,
- )
- if response_mask is None:
- response_mask = default_mask
- reference_logprobs = coerce_logprob_vector(
- rollout.get("reference_logprobs", rollout.get("ref_logprobs")),
- expected_len=expected_len,
- device=self.device,
- name="reference_logprobs",
- )
- current_logprob_rows.append(current_logprobs)
- old_logprob_rows.append(old_logprobs)
- mask_rows.append(response_mask)
- if reference_logprobs is None:
- reference_rows.append(current_logprobs.detach())
- else:
- has_reference = True
- reference_rows.append(reference_logprobs)
- return (
- pad_1d_tensors(current_logprob_rows),
- pad_1d_tensors(old_logprob_rows),
- pad_1d_tensors(mask_rows),
- pad_1d_tensors(reference_rows) if has_reference else None,
- )
-
def backward(
self,
loss: torch.Tensor,
@@ -317,19 +169,4 @@ def _item_image_paths(item: Mapping[str, Any], ) -> list[str]:
return image_paths[:2]
-def _weighted_result_mean(
- results: Sequence[Any],
- field: str,
- weights: Sequence[float],
- device: torch.device,
- total_weight: float,
-) -> torch.Tensor:
- if not results:
- return torch.zeros((), device=device)
- values = []
- for result, weight in zip(results, weights, strict=True):
- values.append(getattr(result, field).detach() * float(weight))
- return torch.stack(values).sum() / float(total_weight)
-
-
__all__ = ["INTERLEAVE_CRITIC_PROMPT", "InterleaveThinkerCriticModel", "_PlaceholderActorModule"]
diff --git a/fastvideo/train/models/interleave_thinker/data.py b/fastvideo/train/models/interleave_thinker/data.py
index b839c5521..5dcc83e10 100644
--- a/fastvideo/train/models/interleave_thinker/data.py
+++ b/fastvideo/train/models/interleave_thinker/data.py
@@ -9,10 +9,11 @@ import os
from pathlib import Path
from typing import Any, Literal
-InterleaveDatasetKind = Literal["planner_sft", "critic_sft", "critic_rl"]
+InterleaveDatasetKind = Literal["planner_sft", "planner_rl", "critic_sft", "critic_rl"]
DEFAULT_FILENAMES: dict[InterleaveDatasetKind, str] = {
"planner_sft": "planner_sft.json",
+ "planner_rl": "planner_rl.jsonl",
"critic_sft": "critic_sft.json",
"critic_rl": "critic_rl.jsonl",
}
@@ -80,6 +81,20 @@ def load_planner_sft_records(
)
+def load_planner_rl_records(
+ data_path: str | os.PathLike[str],
+ *,
+ image_dir: str | os.PathLike[str] = "",
+ validate_image_files: bool = False,
+) -> list[dict[str, Any]]:
+ return load_interleave_dataset(
+ data_path,
+ kind="planner_rl",
+ image_dir=image_dir,
+ validate_image_files=validate_image_files,
+ )
+
+
def load_critic_sft_records(
data_path: str | os.PathLike[str],
*,
@@ -122,6 +137,8 @@ def normalize_interleave_dataset_record(
)
if kind == "planner_sft":
return normalize_planner_sft_record(normalized)
+ if kind == "planner_rl":
+ return normalize_planner_rl_record(normalized)
if kind == "critic_sft":
return normalize_critic_sft_record(normalized)
if kind == "critic_rl":
@@ -140,6 +157,29 @@ def normalize_planner_sft_record(record: Mapping[str, Any]) -> dict[str, Any]:
raise ValueError("planner_sft record requires an assistant message")
normalized.setdefault("instruction", user_message)
normalized.setdefault("response", assistant_message)
+ images = _string_sequence(normalized.get("images"))
+ if images:
+ normalized.setdefault("input_image_paths", images)
+ return normalized
+
+
+def normalize_planner_rl_record(record: Mapping[str, Any]) -> dict[str, Any]:
+ normalized = dict(record)
+ if "messages" in normalized:
+ messages = _require_messages(normalized, where="planner_rl")
+ normalized.setdefault("instruction", _first_message_content(messages, "user"))
+ instruction = _first_text(
+ normalized.get("instruction"),
+ normalized.get("text_input"),
+ normalized.get("origin_prompt"),
+ normalized.get("prompt"),
+ )
+ if not instruction:
+ raise ValueError("planner_rl record requires instruction, text_input, origin_prompt, or prompt")
+ normalized["instruction"] = instruction
+ images = _string_sequence(normalized.get("images"))
+ if images:
+ normalized.setdefault("input_image_paths", images)
return normalized
@@ -199,8 +239,7 @@ def normalize_ground_truth(record: Mapping[str, Any]) -> dict[str, Any]:
if not isinstance(raw, Mapping):
raw = {}
- success = _optional_bool(
- raw.get("success", raw.get("previous_step_success", record.get("previous_step_success"))))
+ success = _optional_bool(raw.get("success", raw.get("previous_step_success", record.get("previous_step_success"))))
if success is None:
raise ValueError("critic_rl ground_truth requires boolean success or previous_step_success")
return {
@@ -380,12 +419,14 @@ __all__ = [
"load_critic_rl_records",
"load_critic_sft_records",
"load_interleave_dataset",
+ "load_planner_rl_records",
"load_planner_sft_records",
"looks_like_uri",
"normalize_critic_rl_record",
"normalize_critic_sft_record",
"normalize_ground_truth",
"normalize_interleave_dataset_record",
+ "normalize_planner_rl_record",
"normalize_planner_sft_record",
"resolve_interleave_image_path",
"validate_image_path",
diff --git a/fastvideo/train/models/interleave_thinker/planner.py b/fastvideo/train/models/interleave_thinker/planner.py
index 4347d492a..371cc2c23 100644
--- a/fastvideo/train/models/interleave_thinker/planner.py
+++ b/fastvideo/train/models/interleave_thinker/planner.py
@@ -15,6 +15,7 @@ import torch
from fastvideo.train.models.interleave_thinker.qwen_actor import (
Qwen3VLActorBase,
batch_to_items,
+ rollout_group_key,
)
from fastvideo.train.models.interleave_thinker.data import InterleaveDatasetKind
@@ -193,6 +194,45 @@ class InterleaveThinkerPlannerModel(Qwen3VLActorBase):
})
return outputs
+ @torch.no_grad()
+ def generate_interleave_responses(
+ self,
+ batch: dict[str, Any],
+ **kwargs: Any,
+ ) -> list[dict[str, Any]]:
+ num_generations = max(1, int(kwargs.get("num_generations", 1) or 1))
+ temperature = float(kwargs.get("temperature", 1.0) or 1.0)
+ top_p = float(kwargs.get("top_p", 1.0) or 1.0)
+ max_new_tokens = int(kwargs.get("max_new_tokens") or self.max_response_length)
+
+ rollouts: list[dict[str, Any]] = []
+ for item_idx, item in enumerate(batch_to_items(batch)):
+ messages = self.build_messages(item)
+ decoded = self.generate_qwen_responses(
+ messages,
+ num_generations=num_generations,
+ temperature=temperature,
+ top_p=top_p,
+ max_new_tokens=max_new_tokens,
+ )
+ for generation_idx, response in enumerate(decoded):
+ parsed = extract_interleave_plan(response)
+ rollout = dict(item)
+ rollout["response"] = response
+ rollout["plan"] = parsed
+ rollout["steps"] = list(parsed.steps) if parsed is not None else []
+ rollout.setdefault("sample_index", item_idx)
+ rollout.setdefault("generation_index", generation_idx)
+ rollout.setdefault("group_key", _planner_group_key(rollout, item_idx))
+ old_logprobs, response_mask = self.response_logprobs_from_messages(
+ messages,
+ response,
+ )
+ rollout["old_logprobs"] = old_logprobs.detach().cpu().tolist()
+ rollout["response_mask"] = response_mask.detach().cpu().tolist()
+ rollouts.append(rollout)
+ return rollouts
+
def generate_interleave_plan(
self,
instruction: str,
@@ -298,7 +338,7 @@ def _planner_instruction(item: Mapping[str, Any]) -> str:
def _planner_image_paths(item: Mapping[str, Any]) -> list[str]:
paths: list[str] = []
- for key in ("input_image_paths", "image_paths"):
+ for key in ("input_image_paths", "image_paths", "images"):
value = item.get(key)
if isinstance(value, Sequence) and not isinstance(value, str | bytes):
paths.extend(str(path) for path in value if path)
@@ -309,6 +349,17 @@ def _planner_image_paths(item: Mapping[str, Any]) -> list[str]:
return paths
+def _planner_group_key(
+ rollout: Mapping[str, Any],
+ index: int,
+) -> str:
+ for key in ("group_key", "problem_id", "instruction", "text_input", "origin_prompt", "prompt"):
+ value = rollout.get(key)
+ if value is not None:
+ return str(value)
+ return rollout_group_key(rollout, index)
+
+
__all__ = [
"INTERLEAVE_GUIDANCE_PLANNER_PROMPT",
"INTERLEAVE_PLANNER_PROMPT",
diff --git a/fastvideo/train/models/interleave_thinker/qwen_actor.py b/fastvideo/train/models/interleave_thinker/qwen_actor.py
index 158e762bd..8e0d838dd 100644
--- a/fastvideo/train/models/interleave_thinker/qwen_actor.py
+++ b/fastvideo/train/models/interleave_thinker/qwen_actor.py
@@ -16,6 +16,7 @@ from torch.utils.data import DataLoader, Dataset
from fastvideo.logger import init_logger
from fastvideo.pipelines import TrainingBatch
+from fastvideo.train.methods.rl.common.grpo import compute_grpo_loss
from fastvideo.train.models.base import ModelBase
from fastvideo.train.models.interleave_thinker.data import (
InterleaveDatasetKind,
@@ -374,6 +375,141 @@ class Qwen3VLActorBase(ModelBase):
if was_training:
self.transformer.train()
+ def train_interleave_rollouts(
+ self,
+ **kwargs: Any,
+ ) -> tuple[dict[str, torch.Tensor], dict[str, float]]:
+ self._require_backend()
+ rollouts = kwargs["rollouts"]
+ advantages = kwargs["advantages"].detach().to(self.device).float()
+ optimizer = kwargs.get("optimizer")
+ lr_scheduler = kwargs.get("lr_scheduler")
+ clip_range = float(kwargs.get("clip_range", 0.2) or 0.0)
+ kl_coef = float(kwargs.get("kl_coef", 0.0) or 0.0)
+ update_micro_batch_size = kwargs.get("update_micro_batch_size")
+ if update_micro_batch_size is None:
+ update_micro_batch_size = len(rollouts)
+ update_micro_batch_size = max(1, int(update_micro_batch_size or 1))
+ max_grad_norm = float(kwargs.get("max_grad_norm", 0.0) or 0.0)
+ grad_accum = max(1, int(kwargs.get("gradient_accumulation_steps", 1) or 1))
+
+ if optimizer is None:
+ raise RuntimeError(f"{type(self).__name__}.train_interleave_rollouts() requires an optimizer")
+ optimizer.zero_grad(set_to_none=True)
+ self.transformer.train()
+
+ if int(advantages.shape[0]) != len(rollouts):
+ raise ValueError("advantages count must match rollout count")
+
+ loss_results = []
+ token_counts: list[float] = []
+ loss_scale = max(1, grad_accum)
+ for start in range(0, len(rollouts), update_micro_batch_size):
+ end = min(start + update_micro_batch_size, len(rollouts))
+ current_logprobs, old_logprobs, response_masks, reference_logprobs = self._grpo_logprob_batch(
+ rollouts[start:end])
+ result = compute_grpo_loss(
+ current_logprobs=current_logprobs,
+ old_logprobs=old_logprobs,
+ advantages=advantages[start:end],
+ response_mask=response_masks,
+ clip_range=clip_range,
+ reference_logprobs=reference_logprobs,
+ kl_coef=kl_coef,
+ )
+ (result.total_loss / loss_scale).backward()
+ loss_results.append(result)
+ token_counts.append(float(result.token_count.detach().cpu()))
+
+ if max_grad_norm > 0.0:
+ torch.nn.utils.clip_grad_norm_(
+ [p for p in self.transformer.parameters() if p.requires_grad],
+ max_grad_norm,
+ )
+ optimizer.step()
+ if lr_scheduler is not None:
+ lr_scheduler.step()
+
+ total_tokens = max(1.0, sum(token_counts))
+ total_loss = _weighted_result_mean(loss_results, "total_loss", token_counts, self.device, total_tokens)
+ policy_loss = _weighted_result_mean(loss_results, "policy_loss", token_counts, self.device, total_tokens)
+ kl_loss = _weighted_result_mean(loss_results, "kl_loss", token_counts, self.device, total_tokens)
+ approx_kl = _weighted_result_mean(loss_results, "approx_kl", token_counts, self.device, total_tokens)
+ clipped_fraction = _weighted_result_mean(
+ loss_results,
+ "clipped_fraction",
+ token_counts,
+ self.device,
+ total_tokens,
+ )
+ mean_ratio = _weighted_result_mean(loss_results, "mean_ratio", token_counts, self.device, total_tokens)
+ return (
+ {
+ "total_loss": total_loss.detach()
+ },
+ {
+ "actor/policy_loss": float(policy_loss.detach().cpu()),
+ "actor/kl_loss": float(kl_loss.detach().cpu()),
+ "actor/approx_kl": float(approx_kl.detach().cpu()),
+ "actor/clipped_fraction": float(clipped_fraction.detach().cpu()),
+ "actor/mean_ratio": float(mean_ratio.detach().cpu()),
+ "actor/mean_advantage": float(advantages.mean().detach().cpu()),
+ "actor/response_tokens": float(total_tokens / max(1, len(rollouts))),
+ },
+ )
+
+ def _grpo_logprob_batch(
+ self,
+ rollouts: Sequence[Mapping[str, Any]],
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
+ current_logprob_rows: list[torch.Tensor] = []
+ old_logprob_rows: list[torch.Tensor] = []
+ mask_rows: list[torch.Tensor] = []
+ reference_rows: list[torch.Tensor] = []
+ has_reference = False
+ for rollout in rollouts:
+ response = str(rollout.get("response", "") or "")
+ current_logprobs, default_mask = self.response_logprobs_from_messages(
+ self.build_messages(rollout),
+ response,
+ )
+ expected_len = int(current_logprobs.numel())
+ old_logprobs = coerce_logprob_vector(
+ rollout.get("old_logprobs", rollout.get("old_logprob")),
+ expected_len=expected_len,
+ device=self.device,
+ name="old_logprobs",
+ )
+ if old_logprobs is None:
+ old_logprobs = current_logprobs.detach()
+ response_mask = coerce_response_mask(
+ rollout.get("response_mask"),
+ expected_len=expected_len,
+ device=self.device,
+ )
+ if response_mask is None:
+ response_mask = default_mask
+ reference_logprobs = coerce_logprob_vector(
+ rollout.get("reference_logprobs", rollout.get("ref_logprobs")),
+ expected_len=expected_len,
+ device=self.device,
+ name="reference_logprobs",
+ )
+ current_logprob_rows.append(current_logprobs)
+ old_logprob_rows.append(old_logprobs)
+ mask_rows.append(response_mask)
+ if reference_logprobs is None:
+ reference_rows.append(current_logprobs.detach())
+ else:
+ has_reference = True
+ reference_rows.append(reference_logprobs)
+ return (
+ pad_1d_tensors(current_logprob_rows),
+ pad_1d_tensors(old_logprob_rows),
+ pad_1d_tensors(mask_rows),
+ pad_1d_tensors(reference_rows) if has_reference else None,
+ )
+
def compute_interleave_sft_loss(
self,
batch: Mapping[str, Any],
@@ -687,6 +823,21 @@ def pad_1d_tensors(
return padded
+def _weighted_result_mean(
+ results: Sequence[Any],
+ field: str,
+ weights: Sequence[float],
+ device: torch.device,
+ total_weight: float,
+) -> torch.Tensor:
+ if not results:
+ return torch.zeros((), device=device)
+ values = []
+ for result, weight in zip(results, weights, strict=True):
+ values.append(getattr(result, field).detach() * float(weight))
+ return torch.stack(values).sum() / float(total_weight)
+
+
def first_string(
item: Mapping[str, Any],
*keys: str,
diff --git a/tests/local_tests/test_interleave_thinker_data.py b/tests/local_tests/test_interleave_thinker_data.py
index 6ef66abe7..aafd1cc93 100644
--- a/tests/local_tests/test_interleave_thinker_data.py
+++ b/tests/local_tests/test_interleave_thinker_data.py
@@ -11,6 +11,7 @@ from fastvideo.train.models.interleave_thinker import (
load_critic_rl_records,
load_critic_sft_records,
load_interleave_dataset,
+ load_planner_rl_records,
load_planner_sft_records,
normalize_critic_rl_record,
resolve_interleave_image_path,
@@ -44,6 +45,27 @@ def test_load_planner_sft_records_normalizes_sharegpt_images(tmp_path: Path) ->
assert records[0]["instruction"] == "draw a cat step by step"
assert records[0]["response"] == '{"execution_plan": []}'
assert records[0]["images"] == [str(image_dir / "planner/cat_step.png")]
+ assert records[0]["input_image_paths"] == [str(image_dir / "planner/cat_step.png")]
+
+
+def test_load_planner_rl_records_accepts_prompt_only_rows(tmp_path: Path) -> None:
+ image_dir = tmp_path / "images"
+ data_path = tmp_path / "planner_rl.jsonl"
+ data_path.write_text(
+ json.dumps({
+ "text_input": "draw a cat in three clear steps",
+ "images": ["planner/start.png"],
+ "plan_score": 0.75,
+ }) + "\n",
+ encoding="utf-8",
+ )
+
+ records = load_planner_rl_records(data_path, image_dir=image_dir)
+
+ assert records[0]["instruction"] == "draw a cat in three clear steps"
+ assert records[0]["images"] == [str(image_dir / "planner/start.png")]
+ assert records[0]["input_image_paths"] == [str(image_dir / "planner/start.png")]
+ assert records[0]["plan_score"] == 0.75
def test_load_critic_sft_records_adds_image_pair_aliases(tmp_path: Path) -> None:
diff --git a/tests/local_tests/test_interleave_thinker_planner_model.py b/tests/local_tests/test_interleave_thinker_planner_model.py
index 01b4a41e6..5461844ed 100644
--- a/tests/local_tests/test_interleave_thinker_planner_model.py
+++ b/tests/local_tests/test_interleave_thinker_planner_model.py
@@ -38,13 +38,15 @@ class _FakeProcessor:
):
del tokenize, add_generation_prompt, return_dict, return_tensors
assert messages[0]["role"] == "user"
+ has_assistant = any(message["role"] == "assistant" for message in messages)
+ length = 5 if has_assistant else 3
return {
- "input_ids": torch.arange(1, 4).unsqueeze(0),
- "attention_mask": torch.ones(1, 3, dtype=torch.long),
+ "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 sequences, kwargs
+ del kwargs
return [
"""
plan
@@ -54,6 +56,7 @@ class _FakeProcessor:
]}
"""
+ for _ in sequences
]
@@ -67,8 +70,23 @@ class _FakeQwen(torch.nn.Module):
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, suffix], dim=1)
+ return torch.cat([input_ids.repeat(num_return_sequences, 1), suffix.repeat(num_return_sequences, 1)], dim=1)
+
+ def forward(self, **kwargs):
+ input_ids = kwargs["input_ids"]
+ vocab_size = max(16, int(input_ids.max().detach().cpu()) + 1)
+ logits = torch.zeros(
+ *input_ids.shape,
+ vocab_size,
+ dtype=self.weight.dtype,
+ device=input_ids.device,
+ )
+ for idx in range(input_ids.shape[1] - 1):
+ next_token = input_ids[:, idx + 1]
+ logits[:, idx].scatter_(1, next_token[:, None], self.weight.expand(input_ids.shape[0], 1))
+ return SimpleNamespace(logits=logits)
def test_extract_interleave_plan_accepts_json_answer_block():
@@ -168,6 +186,32 @@ def test_interleave_thinker_planner_fake_backend_generates_parseable_plan():
assert model.transformer.generate_kwargs["num_return_sequences"] == 1
+def test_interleave_thinker_planner_fake_backend_generates_rl_rollouts():
+ model = _FakeBackendPlanner(load_backend=False)
+ model.processor = _FakeProcessor()
+ model.transformer = _FakeQwen()
+
+ rollouts = model.generate_interleave_responses(
+ {
+ "items": [{
+ "instruction": "Show how to draw a cat step by step."
+ }]
+ },
+ num_generations=2,
+ temperature=0.7,
+ top_p=0.9,
+ max_new_tokens=12,
+ )
+
+ assert len(rollouts) == 2
+ assert all(rollout["plan"] is not None for rollout in rollouts)
+ assert rollouts[0]["group_key"] == rollouts[1]["group_key"]
+ assert rollouts[0]["group_key"] == "Show how to draw a cat step by step."
+ assert len(rollouts[0]["old_logprobs"]) == 2
+ assert rollouts[0]["response_mask"] == [1.0, 1.0]
+ assert model.transformer.generate_kwargs["num_return_sequences"] == 2
+
+
def test_interleave_thinker_planner_config_parses_public_yaml():
cfg = load_run_config("examples/train/configs/interleave_thinker/planner_smoke.yaml")
@@ -176,3 +220,20 @@ def test_interleave_thinker_planner_config_parses_public_yaml():
assert cfg.models["student"]["init_from"] == "InterleaveThinker/InterleaveThinker-Planner-8B"
assert cfg.models["student"]["processor_from"] == "Qwen/Qwen3-VL-8B-Instruct"
assert cfg.models["student"]["trainable"] is False
+
+
+def test_interleave_thinker_planner_grpo_config_parses_public_yaml():
+ cfg = load_run_config("examples/train/configs/rl/interleave_thinker/planner_grpo.yaml")
+
+ assert cfg.models["student"]["_target_"] == (
+ "fastvideo.train.models.interleave_thinker.InterleaveThinkerPlannerModel")
+ assert cfg.models["student"]["dataset_kind"] == "planner_rl"
+ assert cfg.models["student"]["lora"]["enable"] is True
+ assert cfg.models["reference"]["_target_"] == (
+ "fastvideo.train.models.interleave_thinker.InterleaveThinkerPlannerModel")
+ assert cfg.models["reference"]["trainable"] is False
+ assert cfg.method["_target_"] == "fastvideo.train.methods.rl.interleave_thinker.InterleaveThinkerRLMethod"
+ assert cfg.method["reward_scorer"]["_target_"] == (
+ "fastvideo.train.methods.rl.rewards.InterleavePlannerRewardScorer")
+ assert cfg.method["kl_coef"] == 0.01
+ assert cfg.training.data.data_path.endswith("planner_rl.jsonl")
diff --git a/tests/local_tests/test_interleave_thinker_reward.py b/tests/local_tests/test_interleave_thinker_reward.py
index 92a412cf3..81cafe4dc 100644
--- a/tests/local_tests/test_interleave_thinker_reward.py
+++ b/tests/local_tests/test_interleave_thinker_reward.py
@@ -1,10 +1,14 @@
import pytest
from fastvideo.train.methods.rl.rewards import (
+ InterleavePlannerRewardScorer,
InterleaveThinkerEditScore,
InterleaveThinkerRewardScorer,
extract_interleave_answer,
+ extract_interleave_plan_payload,
interleave_format_reward,
+ interleave_planner_format_reward,
+ score_interleave_planner_rewards,
score_interleave_thinker_rewards,
)
@@ -96,3 +100,43 @@ def test_reward_scorer_uses_injected_edit_scorer_and_json_string_ground_truth():
"success": True
},
}])[0].overall >= 0.0
+
+
+def _planner_response(prompt="a clean cat sketch"):
+ prompt = prompt.replace('"', '\\"')
+ return f"""
+ Plan the sequence.
+
+ {{"execution_plan": [
+ {{"step_number": 1, "step_name": "Sketch", "instruction": "Draw a cat", "prompt": "{prompt}", "auxiliary_text": null}}
+ ]}}
+
+ """
+
+
+def test_planner_reward_accepts_execution_plan_answer_block():
+ payload = extract_interleave_plan_payload(_planner_response())
+
+ assert payload is not None
+ assert interleave_planner_format_reward(_planner_response()) == 1.0
+ assert interleave_planner_format_reward("{}") == 0.0
+
+
+def test_planner_reward_scorer_blends_format_and_scalar_plan_score():
+ scorer = InterleavePlannerRewardScorer(format_weight=0.25, fallback_plan_reward=0.1)
+
+ result = scorer([{
+ "response": _planner_response(),
+ "plan_score": 0.9,
+ }])[0]
+ wrapped = score_interleave_planner_rewards([{
+ "response": _planner_response(),
+ "ground_truth": {
+ "score": 0.5
+ },
+ }], format_weight=0.5)[0]
+
+ assert result.format_reward == 1.0
+ assert result.planner_score == pytest.approx(0.9)
+ assert result.overall == pytest.approx(0.25 * 1.0 + 0.75 * 0.9)
+ assert wrapped["overall"] == pytest.approx(0.5 * 1.0 + 0.5 * 0.5)