Fix GRPO tool loop handling for image-returning tools

[vieveks]

What does this PR do?

This PR fixes multimodal tool responses in GRPOTrainer when a tool returns images for a VLM.

Previously, image blocks returned by tools were kept inside role="tool" messages. Some VLM chat templates render image placeholders only for role="user" messages, so the image tensors could be collected while the tokenized prompt contained no matching image tokens. That can lead to image feature/token mismatches.

This PR splits multimodal tool results so text remains in the tool message, while returned image blocks are surfaced as synthetic user messages for prompt rendering. The synthetic user image messages are kept out of completions so reward tokenizers do not receive structured image blocks. It also makes tool suffix token extraction robust to suffixes that include non-tool messages and templates with assistant-final scaffolding.

Fixes #5663.

Tests

• python3 -m py_compile trl/trainer/grpo_trainer.py tests/test_grpo_trainer.py
• env -u PYTHONPATH PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 /home/vieveks/miniconda3/envs/trl-dev/bin/pytest tests/test_grpo_trainer.py::TestGRPOTrainerVLM::test_tool_suffix_ids_include_user_image_from_tool_response tests/test_grpo_trainer.py::TestGRPOTrainerVLM::test_train_with_tools_multimodal_response -q
• env -u PYTHONPATH PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 /home/vieveks/miniconda3/envs/trl-dev/bin/pytest tests/test_grpo_trainer.py::TestGRPOTrainer::test_train_with_tools -q
• env -u PYTHONPATH PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 /home/vieveks/miniconda3/envs/trl-dev/bin/pytest tests/test_grpo_trainer.py::TestGRPOTrainer::test_train_with_environment_factory -q
• env -u PYTHONPATH PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 /home/vieveks/miniconda3/envs/trl-dev/bin/pytest tests/test_grpo_trainer.py::TestGRPOTrainer::test_train_with_malformed_tool_calls -q

Also ran tests/test_grpo_trainer.py -q -k "tool or vlm"; interrupted during a slow HF Hub download after 18 passed, 3 skipped, no failures observed.

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

• 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.

---

Note

Medium Risk
Changes GRPO tool-loop tokenization and message shaping on a critical VLM path; misalignment could still break training or image feature/token matching, though scoped to tool follow-ups with new tests.

Overview
Fixes GRPOTrainer multimodal tool handling when tools return images: many VLM chat templates only emit image placeholder tokens on user messages, so keeping image blocks inside tool messages could collect image tensors without matching tokens in the prompt.

The tool loop now splits list-shaped tool results—text stays in the tool message; image blocks are appended as a synthetic user message on the prompt transcript only (not in completions, so reward tokenizers avoid structured image blocks). _get_tool_suffix_ids picks the first tool role’s name (not tool_messages[0]), falls back to suffix alignment when prefix slicing fails (templates that drop assistant scaffolding after tool/user follow-ups), and suffix extraction walks the prompt transcript for trailing tool and user messages instead of only tool rows from completions.

Adds test_tool_suffix_ids_include_user_image_from_tool_response to assert image pad tokens appear in suffix IDs when a synthetic user image follows a tool message.

^{Reviewed by Cursor Bugbot for commit 985de76. Bugbot is set up for automated code reviews on this repo. Configure here.}

[vieveks]

Additional smoke test on a local RTX 5070 Ti:

env -u PYTHONPATH \
  PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \
  PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
  MODEL_ID=Qwen/Qwen3.5-2B \
  /home/vieveks/miniconda3/envs/trl-dev/bin/python \
  /home/vieveks/Vivek/opensource_contri/trl_real_vlm_tool_image_smoke.py

Result:

torch: 2.12.1+cu130
cuda available: True
gpu: NVIDIA GeForce RTX 5070 Ti
SMOKE PASSED

The smoke uses the local checkout by prepending /home/vieveks/Vivek/opensource_contri/trl to sys.path, so it exercises this PR's code path. It forces a tool call, returns a PIL image from the tool, and runs one GRPO training step with Qwen/Qwen3.5-2B using LoRA.