[feat] add InterleaveThinker planner GRPO path

.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
+  `<think>...</think><answer>{"execution_plan": [...]}</answer>` 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`.

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
+# <think>...</think><answer>{"execution_plan": [...]}</answer> 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

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

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",
 ]

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",
 ]

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",

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"]

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",

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",

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,

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"] == '<answer>{"execution_plan": []}</answer>'
     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:

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 [
             """
             <think>plan</think>
@@ -54,6 +56,7 @@ class _FakeProcessor:
             ]}
             </answer>
             """
+            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")

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"""
+    <think>Plan the sequence.</think>
+    <answer>
+    {{"execution_plan": [
+      {{"step_number": 1, "step_name": "Sketch", "instruction": "Draw a cat", "prompt": "{prompt}", "auxiliary_text": null}}
+    ]}}
+    </answer>
+    """
+
+
+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("<answer>{}</answer>") == 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)