Add Vision / VLM models to environments and GRPO trainer

[UlrickBL]

Description

This Pull Request introduces support for Vision-Language Models (VLMs) into the environments and the GRPO trainer. This functionality is implemented by tracking pixel values and image grids as the base inputs, and then transforming images into Base64 format to comply with VLLM/OpenAI chat formats. The implementation is adapted to work with both standard text tokenizers and multimodal/mixin processors. It also adds Image and Answer login in wanDB table to simplify data analysis during training.

The motivation for adding VLM support is strategic: I believe Vision-Language environments are critical for advancing AGI and Reinforcement Learning (RL) research. This feature was necessary to begin testing several promising, high-value environments.

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Test improvement

Testing

• All existing tests pass
• New tests have been added to cover the changes
• Tests have been run locally with uv run pytest

It was end to end tested with 3 Prime-RL envs :

OCR VL with Qwen VL 2.5 3B and 7B : https://app.primeintellect.ai/dashboard/environments/ulrick-bl/ocr-vl (single turn image)

Rebus VL Thinking with Qwen VL 2.5 7B : https://app.primeintellect.ai/dashboard/environments/ulrick-bl/rebus-vl-thinking (single turn image)

Semantix with Qwen 2.5 0.5B : https://app.primeintellect.ai/dashboard/environments/ulrick-bl/semantic (multi turn text)

Test Coverage

• Current coverage: 33%
• Coverage after changes: 29%

Checklist

• My code follows the style guidelines of this project
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• Any dependent changes have been merged and published

[UlrickBL]

Sounds good, I'll work on that !

[UlrickBL]

Hello @willccbb ,

For point 2, I managed to clean up the problematic dependencies, handle lazy imports, and reorganize the code into utils/image_utils.py and utils/processing_utils.py. Is it ok with the current pattern ?

For point 1, I tried to find a task that was challenging enough to demonstrate the relevance of the training, but not too expensive to run.

I used an OCR environment I set up on Prime Hub (ocr-vl): https://app.primeintellect.ai/dashboard/environments/ulrick-bl/ocr-vl
.

I trained Qwen 2.5 VL 3B on the "hi" (Hindi) scope, since the model doesn’t perform very well on this task. The reward is mainly based on format and CER.

Warning that there are some issues in the dataset I use as the base for the env such as this type of data where the screen fails because of popup and on which my small training setup was very sensitive :

I trained with the following setup :
Qwen 2.5 VL 3B with Lora rank 16 on 2xA100 40GB

args = vf.grpo_defaults(run_name="ocr-vl")
args.per_device_train_batch_size = 8
args.num_generations = 16
args.gradient_accumulation_steps = 2
args.max_steps = 1000
args.eval_strategy = "steps"
args.eval_steps = 2
args.max_tokens = 1024
args.vllm_server_port= 8000
args.fp16 = True
args.temperature = 0.4
args.learning_rate = 1e-5
args.lr_scheduler_type = "cosine"
args.warmup_steps = 10

I would say the training is stable and started off very well. The first slowdown in reward progression was due to a series of poor-quality images in the data, like the ones I showed earlier. Nevertheless, we can observe the model improving and maintaining stable training performance on the task, which highlights the relevance of the implementation.

If needed, I can spend some time cleaning the dataset and retraining it.

Let’s keep in touch if there’s anything else to adjust, test, or adapt.

[anaszil]

+1 @willccbb any idea when this MR might be merged or available