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

Fixes NovaSky-AI#2075. The VLM generator's deferred obs-token extraction requires
each re-render to be a token-prefix extension of the previous one (the
NOTE in agent_loop). Thinking-model vendor templates strip reasoning
from history and violate that, silently corrupting trajectories --
loud only when images make the token/feature pairing fail inside the
model forward.

* _render_conversation forwards chat_template / chat_template_kwargs
  (both already exist on the base generator via generator.chat_template,
  including source=file; the VLM path never sent them). A thinking-
  preserving template restores the prefix property and makes training
  on-policy. Servers need trust_request_chat_template for per-request
  templates.
* Trajectory-end integrity guard via the render response's own
  mm_placeholders: actionable per-trajectory error instead of the
  opaque torch._check deep into training.

Validated end to end: Qwen3.5-9B multi-turn computer-use GRPO (64-turn
episodes with screenshots), 6 GRPO steps, one epoch, completed cleanly
with a thinking-preserving template supplied via source=file; the same
setup crashed at the first policy update twice without these changes
(traces in NovaSky-AI#2075).

Author: deniz-fleet

skyrl/train/generators/skyrl_vlm_generator.py

@@ -59,9 +59,20 @@ def _validate_cfg(self, generator_cfg: GeneratorConfig):
             )
 
     async def _render_conversation(self, conversation: ConversationType) -> RenderedConversation:
-        rendered = await self.inference_engine_client.render_chat_completion(
-            {"json": {"model": self.inference_engine_client.model_name, "messages": conversation}}
-        )
+        body: Dict[str, Any] = {"model": self.inference_engine_client.model_name, "messages": conversation}
+        # Honor generator.chat_template / chat_template_kwargs (already supported
+        # by the base generator) on the render path. The deferred obs-token
+        # extraction below requires each re-render to be a token-prefix
+        # extension of the previous one (see NOTE in agent_loop); thinking-model
+        # vendor templates strip reasoning from history and violate that, so a
+        # thinking-preserving custom template is the supported way to train
+        # them (issue #2075). Servers need trust_request_chat_template for
+        # per-request templates.
+        if self.custom_chat_template:
+            body["chat_template"] = self.custom_chat_template
+        if self.generator_cfg.chat_template_kwargs:
+            body["chat_template_kwargs"] = dict(self.generator_cfg.chat_template_kwargs)
+        rendered = await self.inference_engine_client.render_chat_completion({"json": body})
         return RenderedConversation(prompt_ids=rendered["token_ids"], features=rendered.get("features", None))
 
     async def agent_loop(
@@ -129,12 +140,14 @@ async def agent_loop(
             # sequences which are a prefix of subsequent turns. In general, this
             # will not hold for thinking models, e.g., Qwen3-Thinking.
             pending_obs_offset: Optional[int] = None
+            last_render_ids: Optional[List[int]] = None
 
             while not done:
                 # 1. Render full conversation for this turn's generation input
                 rendered_conversation = await self._render_conversation(conversation)
                 input_ids = rendered_conversation["prompt_ids"]
                 latest_features = rendered_conversation["features"]
+                last_render_ids = input_ids
 
                 # 1b. Flush pending obs tokens from the previous turn
                 if pending_obs_offset is not None:
@@ -209,6 +222,28 @@ async def agent_loop(
             pixel_values = mm_kwargs["pixel_values"]
             image_grid_thw = mm_kwargs["image_grid_thw"]
 
+            # ── Integrity guard: image tokens vs features ─────────────────
+            # The model forward pairs each image token with exactly one
+            # feature row; a mismatch fails there as an opaque torch._check
+            # deep into training. Detect it here, per trajectory, with an
+            # actionable message. Uses the render response's mm_placeholders,
+            # so no model constants are needed.
+            placeholders = (latest_features or {}).get("mm_placeholders", {}).get("image", [])
+            if placeholders:
+                render_ids = last_render_ids if last_render_ids is not None else prompt_ids
+                ph0 = placeholders[0]
+                image_token = render_ids[ph0["offset"] + ph0["length"] // 2]
+                expected = sum(ph["length"] for ph in placeholders)
+                actual = (list(prompt_ids) + list(response_ids)).count(image_token)
+                if actual != expected:
+                    raise RuntimeError(
+                        f"trajectory image-token mismatch: sequence has {actual} image tokens, "
+                        f"render declares {expected} across {len(placeholders)} images "
+                        f"(stop_reason={stop_reason}). The trajectory token stream diverged from "
+                        f"the render -- typically a history-editing chat template breaking the "
+                        f"prefix-extension assumption (issue #2075)."
+                    )
+
             # ── Cleanup ───────────────────────────────────────────────────
             env_metrics = env.get_metrics()
             await self._run_in_executor_if_available(env.close)

tests/train/generators/test_vlm_render_template.py

@@ -0,0 +1,41 @@
+"""The VLM render path honors generator.chat_template (issue #2075).
+
+Unit seam: _render_conversation built without __init__ so the test needs no
+engines. Asserts the request body carries the custom template exactly when
+configured, and stays untouched otherwise.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from skyrl.train.generators.skyrl_vlm_generator import SkyRLVLMGymGenerator
+
+
+def _bare_generator(custom_template, kwargs=None):
+    gen = SkyRLVLMGymGenerator.__new__(SkyRLVLMGymGenerator)
+    gen.custom_chat_template = custom_template
+    gen.generator_cfg = SimpleNamespace(chat_template_kwargs=kwargs or {})
+    gen.inference_engine_client = SimpleNamespace(
+        model_name="test-model",
+        render_chat_completion=AsyncMock(return_value={"token_ids": [1, 2], "features": None}),
+    )
+    return gen
+
+
+@pytest.mark.asyncio
+async def test_custom_template_reaches_render_body():
+    gen = _bare_generator("{{ messages }}", {"enable_thinking": True})
+    await gen._render_conversation([{"role": "user", "content": "hi"}])
+    body = gen.inference_engine_client.render_chat_completion.call_args.args[0]["json"]
+    assert body["chat_template"] == "{{ messages }}"
+    assert body["chat_template_kwargs"] == {"enable_thinking": True}
+
+
+@pytest.mark.asyncio
+async def test_no_template_leaves_body_untouched():
+    gen = _bare_generator(None)
+    await gen._render_conversation([{"role": "user", "content": "hi"}])
+    body = gen.inference_engine_client.render_chat_completion.call_args.args[0]["json"]
+    assert "chat_template" not in body and "chat_template_kwargs" not in body