Add Vision / VLM models to environments and GRPO trainer

pyproject.toml

@@ -36,6 +36,7 @@ dependencies = [
     "requests",
     "rich",
     "textual",
+    "pillow>=10.0.0",
     "pydantic>=2.11.9",
     "prime-sandboxes>=0.1.0",
 ]

tests/test_environment.py

@@ -255,6 +255,8 @@ def apply_template(conversation, tokenize=False, add_generation_prompt=True):
         (
             prompt_ids,
             prompt_mask,
+            prompt_image_grid, 
+            prompt_pixel_value,
             completion_ids,
             completion_mask,
             completion_logprobs,
@@ -300,6 +302,8 @@ def test_process_completion_format(self, mock_openai_client, sample_dataset):
         (
             prompt_ids,
             prompt_mask,
+            prompt_image_grid, 
+            prompt_pixel_value,
             completion_ids,
             completion_mask,
             completion_logprobs,

verifiers/envs/environment.py

@@ -4,8 +4,7 @@
 from abc import ABC, abstractmethod
 from concurrent.futures import ThreadPoolExecutor
 from copy import deepcopy
-from typing import TYPE_CHECKING, Literal
-
+from typing import TYPE_CHECKING, Literal, Union
 from datasets import Dataset
 from openai import AsyncOpenAI, BadRequestError, OpenAI
 
@@ -27,6 +26,7 @@
     SamplingArgs,
     State,
 )
+from verifiers.utils.processor_utils import encode_text_with_processor, encode_chat_with_processor
 from verifiers.utils.message_utils import (
     cleanup_messages,
     get_overlong_prompt_dummy_response,
@@ -35,10 +35,9 @@
 
 if TYPE_CHECKING:
     from transformers.tokenization_utils_base import (  # type: ignore
-        PreTrainedTokenizerBase,
+        PreTrainedTokenizerBase, ProcessorMixin
     )
-
-
+
 class Environment(ABC):
     """
     Base class for all environments.
@@ -69,7 +68,6 @@ def __init__(
             self.logger.warning(
                 "The parser and rubric parser are different. This may cause unexpected behavior."
             )
-
         if self.message_type == "chat":
             if dataset is not None:
                 self.dataset = self.format_dataset(
@@ -228,6 +226,7 @@ async def get_model_response(
         ):
             sampling_args.pop("max_completion_tokens")
         clean_sampling_args = {k: v for k, v in sampling_args.items() if v is not None}
+
         try:
             if message_type == "chat":
                 assert isinstance(prompt, list)
@@ -444,6 +443,7 @@ async def a_generate(
             reward=[],
             metrics={},
         )
+
         n = len(results.prompt)
 
         # Resolve concurrency knobs
@@ -593,6 +593,7 @@ def generate(
     ) -> GenerateOutputs:
         if isinstance(client, OpenAI):
             client = AsyncOpenAI(api_key=client.api_key, base_url=client.base_url)
+
         coro = self.a_generate(
             inputs,
             client,
@@ -819,9 +820,9 @@ def process_chat_format_vllm(
         prompt: list[ChatMessage],
         completion: list[ChatMessage],
         state: State,
-        processing_class: "PreTrainedTokenizerBase",
+        processing_class: Union["PreTrainedTokenizerBase", "ProcessorMixin"],
         mask_env_responses: bool = False,
-    ) -> tuple[list[int], list[int], list[int], list[int], list[float]]:
+    ) -> tuple[list[int], list[int], list[int], list[int], list[int], list[int], list[float]]:
         """
         Process chat format conversations using incremental prefixes.
         """
@@ -836,10 +837,13 @@ def process_chat_format_vllm(
                 zipped.append((turn, None))
         assert len(responses) == responses_idx, "Responses not fully consumed"
         assert len(zipped) == len(completion), "Length mismatch"
-        prompt_ids: list[int] = processing_class.apply_chat_template(
-            conversation=prompt,  # type: ignore
+
+        prompt_ids, prompt_image_grid, prompt_pixel_value = encode_chat_with_processor(
+            conversation=prompt,
+            processing_class=processing_class,
             add_generation_prompt=True,
         )
+
         messages_consumed = [m for m in prompt]
         prompt_mask: list[int] = [0] * len(prompt_ids)
         completion_ids: list[int] = []
@@ -900,13 +904,15 @@ def deserialize_tool_call(tool_call) -> dict:
                 while j < len(zipped) and zipped[j][0]["role"] != "assistant":
                     consecutive_messages.append(zipped[j][0])
                     j += 1
-                token_prefix: list[int] = processing_class.apply_chat_template(
-                    conversation=messages_consumed  # type: ignore
+                token_prefix, token_prefix_image_grid, token_prefix_pixel_values = encode_chat_with_processor(
+                    conversation=messages_consumed, # type: ignore
+                    processing_class=processing_class,
+                    add_generation_prompt=False,
                 )
-                token_prefix_with_turn: list[int] = (
-                    processing_class.apply_chat_template(
-                        conversation=messages_consumed + consecutive_messages,  # type: ignore
-                    )
+                token_prefix_with_turn, token_prefix_with_turn_image_grid,token_prefix_with_turn_pixel_values  = encode_chat_with_processor(
+                    conversation=messages_consumed + consecutive_messages,  # type: ignore
+                    processing_class=processing_class,
+                    add_generation_prompt=False,
                 )
                 assert token_prefix_with_turn[: len(token_prefix)] == token_prefix, (
                     f"Token prefix mismatch. Token prefix: {token_prefix}, token prefix with turn: {token_prefix_with_turn}"
@@ -916,6 +922,7 @@ def deserialize_tool_call(tool_call) -> dict:
                     completion_turn_mask = [0] * len(completion_turn_ids)
                 else:
                     completion_turn_mask = [1] * len(completion_turn_ids)
+
                 completion_turn_logprobs = [0.0] * len(completion_turn_ids)
                 completion_ids.extend(completion_turn_ids)
                 completion_mask.extend(completion_turn_mask)
@@ -925,6 +932,8 @@ def deserialize_tool_call(tool_call) -> dict:
         return (
             prompt_ids,
             prompt_mask,
+            prompt_image_grid, 
+            prompt_pixel_value,
             completion_ids,
             completion_mask,
             completion_logprobs,
@@ -935,9 +944,9 @@ def process_completion_format_vllm(
         prompt: str,
         completion: str,
         state: State,
-        processing_class: "PreTrainedTokenizerBase",
+        processing_class: Union["PreTrainedTokenizerBase", "ProcessorMixin"],
         mask_env_responses: bool = False,
-    ) -> tuple[list[int], list[int], list[int], list[int], list[float]]:
+    ) -> tuple[list[int], list[int], list[int], list[int], list[int], list[int], list[float]]:
         """
         Process completion format conversations using incremental prefixes.
         """
@@ -958,12 +967,16 @@ def process_completion_format_vllm(
             idx = response_start_idx + len(response_text)
         assert idx == len(completion), "Completion not fully consumed"
 
-        prompt_ids: list[int] = processing_class.encode(prompt)
+        prompt_ids, prompt_image_grid, prompt_pixel_value = encode_text_with_processor(
+            text=prompt,
+            processing_class=processing_class,
+        )
         rollout_consumed = prompt
         prompt_mask: list[int] = [0] * len(prompt_ids)
         completion_ids: list[int] = []
         completion_mask: list[int] = []
         completion_logprobs: list[float] = []
+
         i = 0
         while i < len(zipped):
             text, response = zipped[i]
@@ -1000,6 +1013,8 @@ def process_completion_format_vllm(
         return (
             prompt_ids,
             prompt_mask,
+            prompt_image_grid,
+            prompt_pixel_value,
             completion_ids,
             completion_mask,
             completion_logprobs,
@@ -1011,7 +1026,7 @@ def process_env_results_vllm(
         completions: list[Messages],
         states: list[State],
         rewards: list[float],
-        processing_class: "PreTrainedTokenizerBase",
+        processing_class: Union["PreTrainedTokenizerBase", "ProcessorMixin"],
         max_seq_len: int = -1,
         mask_env_responses: bool = False,
         mask_truncated_completions: bool = False,
@@ -1024,6 +1039,8 @@ def process_env_results_vllm(
 
         all_prompt_ids = []
         all_prompt_masks = []
+        all_prompt_image_grid = []
+        all_prompt_pixel_value = []
         all_completion_ids = []
         all_completion_masks = []
         all_completion_logprobs = []
@@ -1037,6 +1054,8 @@ def process_env_results_vllm(
                 (
                     prompt_ids,
                     prompt_mask,
+                    prompt_image_grid,
+                    prompt_pixel_value,
                     completion_ids,
                     completion_mask,
                     completion_logprobs,
@@ -1048,6 +1067,8 @@ def process_env_results_vllm(
                 (
                     prompt_ids,
                     prompt_mask,
+                    prompt_image_grid,
+                    prompt_pixel_value,
                     completion_ids,
                     completion_mask,
                     completion_logprobs,
@@ -1080,16 +1101,22 @@ def process_env_results_vllm(
             )
             all_prompt_ids.append(prompt_ids)
             all_prompt_masks.append(prompt_mask)
+            all_prompt_image_grid.append(prompt_image_grid)
+            all_prompt_pixel_value.append(prompt_pixel_value)
             all_completion_ids.append(completion_ids)
             all_completion_masks.append(completion_mask)
             all_completion_logprobs.append(completion_logprobs)
+
             if zero_truncated_completions and is_truncated:
                 all_rewards.append(0)
             else:
                 all_rewards.append(reward)
+
         return ProcessedOutputs(
             prompt_ids=all_prompt_ids,
             prompt_mask=all_prompt_masks,
+            image_grid_thw = all_prompt_image_grid,
+            pixel_values= all_prompt_pixel_value,
             completion_ids=all_completion_ids,
             completion_mask=all_completion_masks,
             completion_logprobs=all_completion_logprobs,

verifiers/trainers/async_batch_generator.py

@@ -4,14 +4,13 @@
 import threading
 import time
 from collections import deque
-from typing import Any
+from typing import Any, Optional
 
 from pydantic import BaseModel, Field
 
 from verifiers import GenerateOutputs
 from verifiers.types import ProcessedOutputs
 
-
 class BatchRequest(BaseModel):
     """Request for batch generation"""
 
@@ -38,6 +37,7 @@ class BatchResult(BaseModel):
         default_factory=list
     )  # Store completions for logging
     prompts: list[Any] = Field(default_factory=list)  # Store prompts for logging
+    answers : Optional[list[Any]]
 
 
 class AsyncBatchGenerator:
@@ -264,6 +264,7 @@ async def _generate_batch_async(self, request: BatchRequest) -> BatchResult:
         """
         # Call environment generation
         self.is_generating = True
+
         env_results = await self.env.a_generate(
             request.env_inputs,
             client=self.client,
@@ -272,6 +273,7 @@ async def _generate_batch_async(self, request: BatchRequest) -> BatchResult:
             score_rollouts=True,
             max_concurrent=request.max_concurrent,
         )
+
         self.is_generating = False
 
         # Extract all reward-related keys
@@ -281,6 +283,7 @@ async def _generate_batch_async(self, request: BatchRequest) -> BatchResult:
         for k in env_results.metrics:
             all_reward_dict[k] = env_results.metrics[k]
 
+
         # Process results
         processed_results = self.env.process_env_results_vllm(
             prompts=env_results.prompt,
@@ -300,6 +303,7 @@ async def _generate_batch_async(self, request: BatchRequest) -> BatchResult:
             all_reward_dict=all_reward_dict,
             completions=env_results.completion,
             prompts=env_results.prompt,
+            answers=request.env_inputs.get("answer")
         )
 
     async def _evaluate_async(self, num_samples: int = -1) -> GenerateOutputs: