GRPO (VLM)

.
└── dataset_dir/
    ├── images/
    │   ├── /path/to/img_1
    │   ├── ...
    │   └── /path/to/img_i
    ├── train.jsonl
    └── reward_models.py

train_dataset = [
  {
    "prompt": "...",
    "image": "...",
    "answer": "..."
  },
]

{"prompt": "...","image": "...","answer": "..."}
{"prompt": "...","image": "...","answer": "..."}
{"prompt": "...","image": "...","answer": "..."}
{"prompt": "...","image": "...","answer": "..."}

[
  {
    "role": "user",
    "content": [
      { "type": "image" },
      {
        "type": "text",
        "text": "When a spring does work on an object, we cannot find the work by simply multiplying the spring force by the object's displacement. … Provide your reasoning between <think> and </think> and then your final answer between <answer> and (put a float here) </answer>."
      }
    ]
  }
]

images/cumin_canister_spring.jpg

1.2

train_dataset = [
  {
    "prompt": [
      {
        "role": "user",
        "content": [
          { "type": "image" },
          {
            "type": "text",
            "text": (
              "When a spring does work on an object, … provide your reasoning between <think> and </think> "
              "and then your final answer between <answer> and (put a float here) </answer>"
            )
          }
        ]
      }
    ],
    "image": "images/cumin_canister_spring.jpg", 
    "answer": "1.2",
  },
]

from your_tags import reasoning_start, reasoning_end, solution_start, solution_end
import re
from typing import Callable, List

def formatting_reward_func(completions, **kwargs) -> List[float]:
    """
    Rewards the presence of both reasoning and answer tags in the model's output.
    """
    thinking_pattern = f'{reasoning_start}(.*?){reasoning_end}'
    answer_pattern   = f'{solution_start}(.*?){solution_end}'
    scores = []

    for completion in completions:
        text    = completion[0]['content']
        score   = 0.0
        if len(re.findall(thinking_pattern, text, re.DOTALL)) == 1:
            score += 1.0
        if len(re.findall(answer_pattern, text, re.DOTALL)) == 1:
            score += 1.0
        scores.append(score)

    return scores


def correctness_reward_func(prompts, completions, answer, **kwargs) -> List[float]:
    """
    Rewards exact match of the numeric answer within the <answer> tags.
    """
    answer_pattern = f'{solution_start}(.*?){solution_end}'
    responses = [
        re.findall(answer_pattern, comp[0]['content'], re.DOTALL)
        for comp in completions
    ]

    # Example debug print
    q     = prompts[0][-1]['content']
    a     = answer[0]
    resp0 = completions[0][0]['content']
    print("-"*20, f"Q: {q}\nA: {a}\nR: {resp0}")

    return [
        2.0 if len(r)==1 and r[0].strip()==str(a) else 0.0
        for r in responses
    ]

reward_functions: List[Callable] = [formatting_reward_func, correctness_reward_func]