[fix][generators] VLM generator: honor generator.chat_template in renders + image-token/feature integrity guard

[dzorlu]

Fixes #2075.

SkyRLVLMGymGenerator extracts observation tokens by slicing each turn's re-render at a predicted offset (pending_obs_offset = len(prev_render) + len(gen_ids)), which requires renders to be token-prefix extensions of one another — the assumption named in the file's own NOTE. Thinking-model vendor templates (Qwen3/3.5, DeepSeek-R1, Kimi-K2-Thinking, GLM-4.6 — templates quoted in the issue) strip reasoning from historical assistant turns, violating it: the stale offset then swallows the head of the next observation. With images in observations this crashes the first policy update (Image features and image tokens do not match); with text observations it corrupts training sequences silently.

Two minimal changes; no behavior change for existing users:

1. _render_conversation forwards chat_template / chat_template_kwargs — both already exist on the base generator via generator.chat_template (including source=file, so users supply their own template file); the VLM path just never sent them. A thinking-preserving template (vendor template rendering assistant history verbatim inside the <think> scaffold) restores the prefix property and makes training on-policy for thinking models. Note: vLLM requires trust_request_chat_template server-side for per-request templates (reachable via engine_init_kwargs).
2. Trajectory-end integrity guard using the render response's own mm_placeholders: image tokens present in the assembled sequence vs declared placeholder lengths, failing with an actionable per-trajectory message instead of the opaque torch._check deep inside training. No model constants needed.

Tests: a unit test asserting the render body carries the configured template (and is untouched without one).

Validated end to end: Qwen3.5-9B multi-turn computer-use GRPO — 64-turn episodes with screenshots in observations, 6 GRPO steps, one full epoch, completed cleanly with a thinking-preserving template supplied via source=file; the identical setup crashed at the first policy update twice without these changes (traces in #2075). Property tests for such a template (token-prefix extension against the real tokenizer; offset arithmetic landing exactly on observation boundaries) are available if maintainers want a registered template shipped — kept out of this PR to stay minimal.

Follow-up offered: a template-agnostic default fix — measuring the observation boundary (one extra render before appending the observation) or independent observation tokenization as the text generator does — so arbitrary templates are safe by default. Happy to send as a second PR.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces a custom chat template, qwen35_vl_with_thinking, designed to preserve thinking tokens in assistant history, along with a fail-fast contract check (_check_render_contract) and an integrity guard to prevent image-token mismatches during training. Feedback on these changes highlights a potential race condition and thundering herd issue on startup, as multiple concurrent workers may execute the asynchronous contract check simultaneously. To resolve this, it is recommended to serialize the check using an asyncio.Lock.

[gemini-code-assist]

Race Condition / Thundering Herd on Startup

Since _check_render_contract is an asynchronous method that performs network requests (await self._render_conversation(...)) and is called concurrently for all parallel trajectories at the start of agent_loop, multiple concurrent tasks will check _render_contract_ok simultaneously while it is still None.

This creates a race condition where every concurrent generation worker (which can be hundreds of workers in parallel) will execute the probe and send redundant render requests to the inference server at the exact same time. This thundering herd can easily overload the inference server, cause connection timeouts, or significantly delay startup.

Using an asyncio.Lock to serialize the check ensures that only the first task performs the probe, while subsequent tasks wait and then safely skip it once the contract is verified.

    _render_contract_ok: Optional[bool] = None
    _render_contract_lock: Any = None

    async def _check_render_contract(self) -> None:
        """One-time fail-fast probe: the render endpoint must honor our template.

        Renders a synthetic thinking conversation and its one-message extension
        through the REAL endpoint and asserts token-prefix extension. Catches
        (a) history-editing templates (thinking models under vendor templates)
        and (b) the endpoint silently ignoring the chat_template field -- both
        otherwise corrupt trajectories silently or crash 30 minutes later in
        the model forward with an image token/feature mismatch.
        """
        if SkyRLVLMGymGenerator._render_contract_ok:
            return
        if SkyRLVLMGymGenerator._render_contract_lock is None:
            import asyncio
            SkyRLVLMGymGenerator._render_contract_lock = asyncio.Lock()
        async with SkyRLVLMGymGenerator._render_contract_lock:
            if SkyRLVLMGymGenerator._render_contract_ok:
                return
            base = [
                {"role": "user", "content": "probe"},
                {"role": "assistant", "content": "thinking about it\n</think>\n\nanswer"},
            ]
            extended = base + [{"role": "user", "content": "next"}]
            r1 = (await self._render_conversation(base))["prompt_ids"]
            r2 = (await self._render_conversation(extended))["prompt_ids"]
            if r2[: len(r1) - self._gen_prompt_tail_len(r1)] != r1[: len(r1) - self._gen_prompt_tail_len(r1)]:
                raise RuntimeError(
                    "render contract violated: re-renders are not token-prefix extensions. "
                    "The chat template edits history (thinking model vendor template?) or the "
                    "render endpoint ignored the chat_template field. Set generator.chat_template "
                    "to a thinking-preserving template; refusing to train on corrupt offsets."
                )
            SkyRLVLMGymGenerator._render_contract_ok = True