fix: score think_format_reward when <think> is prefilled in the prompt

[Ruinique]

What does this PR do?

think_format_reward currently requires the completion to start with <think>. GRPO only decodes generated tokens (grpo_trainer.py), so models whose chat template already prefills that tag into the prompt — DeepSeek-R1 distill, and Qwen3.5/3.6 with thinking enabled — score 0.0 on every well-formed completion. A constant reward gives zero advantage, so the function contributes nothing for the whole run.

This PR makes the opening <think> tag optional, keeps the existing no-nested-<think> lookahead, and still requires </think>. Completions that already start with <think> (Qwen3) keep scoring 1.0.

Fixes #6966

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

AI writing disclosure

We welcome the use of AI tools to help with contributions. For transparency and to help us improve our review process, please indicate the level of AI involvement in this PR.

• No AI usage: the PR was written entirely by a human.
• AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
• AI-generated: the PR was mostly or fully generated by an AI tool.

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag members/contributors who may be interested in your PR.

Made with Cursor

[behroozazarkhalili]

Verdict: correct, with two documentation changes requested. It resolves #6966 and lowers no score that was 1.0 before.

What I checked

/scratch/ermia/venvs/hf_trl/bin/python 3.11.5, on the PR worktree and on a copy of it with only trl/rewards/format_rewards.py reverted to d9e880c1 (PR tests kept). RED/GREEN on the exact changed ids:

# PR head: pytest tests/test_rewards.py::TestThinkFormatReward::{test_valid_format,test_invalid_format} -q
2 passed        PROBE_RC_HEAD=0
# pre-fix source, PR tests
E  assert [1.0, 1.0, 1....1.0, 0.0, ...] == [1.0, 1.0, 1....1.0, 1.0, ...]
E    At index 5 diff: 0.0 != 1.0
1 failed, 1 passed        PROBE_RC_BASE=1

Issue #6966 minimal case, scored on the completion string GRPO decodes (no download needed): head [1.0], pre-fix [0.0]. The mechanism holds in this tree: trl/trainer/grpo_trainer.py:1914 slices prompt_completion_ids[:, prompt_length:] and :2266 decodes only that, and trl/chat_templates/deepseek_r1_distill.jinja:1 ends with {{'<｜Assistant｜><think>\n'}}. The same prefill is in qwen3_5_think.jinja:152, qwen3_6.jinja:152, nemotron_3_nano.jinja:200 and deepseekv3.jinja, so the affected set is wider than the issue.

No regression. Scoring both patterns over every string of up to 4 atoms from <think>, </think>, a, \n, <t, hink>, >:

strings_tested 2801
REGRESSIONS_1_to_0 0 []
NEW_1_count 739
new_1_containing_open_think 0 []

That follows from the structure: (?:<think>)? is greedy, so the engine tries exactly the old pattern first and the relaxation can only add matches. Over 4681 strings the new pattern is exactly "starts with <think>, no second <think>, </think> follows" or "no <think> anywhere and a </think> present", 0 mismatches.

Mutants scored against the PR's own valid and invalid lists:

PR (as merged)                    PASSES_PR_TESTS=True
M1 drop-lookahead                 False  n_wrong=3
M2 any-closing ^.*?</think>.*$    False  n_wrong=3
M3 hardcoded "</think>" in text   False  n_wrong=3
M4 hardcoded 1.0                  False  n_wrong=7
M5 opening tag forbidden          False  n_wrong=5

All killed, so the test is a real oracle. ruff check and ruff format --check at the CI-pinned 0.13.3 pass on both files. think_format_reward has one definition, so there is no sibling copy in trl/trainer, trl/experimental, trl/scripts or trl/cli. The direction matches reasoning_accuracy_reward, which already keys only on the closing delimiter (trl/rewards/accuracy_rewards.py:283-285, pre-existing case at tests/test_rewards.py:250).

Findings

1. trl/rewards/format_rewards.py:45: the relaxation is unconditional, so it also changes scores under templates that do not prefill. 37 of the 63 bundled templates contain no <think> tag at all (qwen2_5, llama3, gemma3 and phi3 among them), so there the model must emit the tag itself, and a completion omitting it now scores 1.0, so the reward no longer teaches the model to open the tag. Newly accepted, all previously 0.0, all executed: "reasoning</think>answer" (intended), "</think>" alone, " \n</think>\n\nanswer" with empty reasoning, and "The closing tag </think> is written like that.", which has no reasoning at all. I think the tradeoff is acceptable, since keying on the closing tag is what reasoning_accuracy_reward already does, but it is a user-visible scoring change and belongs in the docstring rather than left implicit.
2. trl/rewards/format_rewards.py:20-21: the summary line still says reasoning must be "enclosed within <think> and </think> tags", which is not what is enforced now. Suggestion: say the completion must contain </think> with at most one optional leading <think>, keeping the added paragraph as the rationale.
3. trl/rewards/format_rewards.py:53-54: the comment names only DeepSeek-R1 distill and Qwen3.5/3.6, but deepseekv3.jinja, nemotron_3_nano.jinja:200, nemotron_3_super, nemotron_3_ultra and nemotron_3_5_lightning prefill the same tag. Suggestion: drop the list or add these, since a partial list reads as a limit on where the fix applies.
4. tests/test_rewards.py:40-43: the added cases all look like real reasoning, and nothing pins the boundary of the newly accepted set. Suggestion: add "</think>" and "An answer that merely mentions </think>." to the valid list with a comment saying they are accepted deliberately, so a later tightening fails loudly.

Not reproduced

Nothing. Every claim above was executed.

Unevaluable

The end-to-end GRPO run: I scored the string the trainer decodes rather than generating it, so the reward path is verified but not the tokenizer round trip on a real DeepSeek checkpoint.

Out of scope: with enable_thinking=False the prompt prefills <think>\n\n</think>\n\n (qwen3_5_think.jinja:150, nemotron_3_nano.jinja:202), so no </think> reaches the completion and the reward stays a constant 0.0. This PR neither claims nor fixes that.

[A-ryanVAT-S]

@behroozazarkhalili
Thanks for checking this so thoroughly. In the issue (#6966 ), I only mentioned a few model families in the issue to demonstrate the bug with some concrete examples, not to enumerate all the affected templates. Good catch on the wider set of models affected.

[Ruinique]

@behroozazarkhalili Thx for review. I updated the docstring and added the two boundary test cases you suggested.

[behroozazarkhalili]

@Ruinique after talking with @qgallouedec I reopened this work as #7141 from my fork so it can be reviewed. Your commits are cherry-picked there with your authorship intact and the PR body credits you; if you would rather carry it yourself, say so there and I will close mine.