Train Qwen2-VL-2B with GRPO so it generalizes off the training distribution

You inherit a deliberately weak reinforcement-learning (GRPO) training project for the vision-language model Qwen2-VL-2B-Instruct. The starting point in /app/methods/main/ trains the model on the provided counting-VQA train split with a trivial reward and conservative hyper-parameters — it barely beats the untrained model. Going weak→strong is the task.

You submit a trained model (full merged weights written to /app/submission/), not code. A sealed verifier loads your model in a clean box and evaluates it on two hidden test sets it never lets you see, scoring counting accuracy. Maximize the mean accuracy across the two sets.

The two hidden sets are not in-distribution samples of your train split. One is a compositional-generalization split: the same visual domain, but attribute combinations that never co-occur in training. The other comes from a different distribution entirely (different renderer/scene statistics). Overfitting the in-domain data is therefore useless — only real generalization moves the score.

Your job (this is the research)

The inherited project in /app/methods/main/ is a GRPO trainer (open_r1/grpo.py + Qwen2VLGRPOTrainer), launched by train.sh. Design and run the training so the model generalizes. Two surfaces are yours:

1. The reward — the reward functions and their registry in /app/methods/main/open_r1/grpo.py. The baseline train.sh selects a single, deliberately weak reward term. What the model should be paid for is your design decision; new reward functions can be registered and referenced by name in --reward_funcs.
2. The training hyper-parameters — the flags in /app/methods/main/train.sh (optimization, GRPO group size, batching, step count, KL, vision-token budget, sequence lengths, and optionally LoRA via --use_peft --lora_r …). Keep per_device_train_batch_size 1 — batched training in this trainer is buggy.

Nothing here prescribes which changes work. Finding that out is the task.

Hardware & budget

Budget: 2 GPUs (48 GB each), 8h wall-clock. Use them (accelerate / torchrun / deepspeed are installed). This budget is meant for several rounds, not one long run: a LoRA GRPO round is ~40–60 min, so plan to train a version → run selfcheck.py → read the OOD val → change the reward/hyper-parameters → retrain → keep the best. Size the group count, batch, and steps so each round is short enough to iterate. Leave your best model in /app/submission/ before time runs out.

Hard constraints

1. The submitted model must be Qwen2-VL-2B-Instruct (same architecture/tokenizer/processor), fine-tuned. Write the full merged model (config + weights + tokenizer + processor) to /app/submission/ with save_pretrained(...) so the verifier can load it with Qwen2VLForConditionalGeneration.from_pretrained("/app/submission"). The trainer does full fine-tuning by default, so the saved checkpoint already is the full model; train.sh copies it to /app/submission. If you switch to a LoRA adapter, merge it before saving.
2. No internet at training time. The base model and the train data are already on disk / in the HF cache (see What you have). You may not download anything, and in particular you may not fetch any test data — the network is sealed off.
3. Train only on the provided train split. Do not hand-write or memorize answer tables, do not key behavior on the visible val ids.
4. The verifier uses a fixed evaluation protocol (prompt template + greedy decoding, below). Train your model to that protocol — there is no separate prompt you get to submit.
5. There is no submit step and no per-attempt feedback. Iterate against your local self-check, then leave your best model in /app/submission/; it is graded once at the end on the hidden sets.

What you have

• Base model (/app/models/Qwen2-VL-2B-Instruct/): the untrained instruct model to start from. This exact model, with no fine-tuning, is the reward-0 anchor — you must beat it.
• Train data (/app/data/train/): ~68k counting-VQA examples (columns image (PIL), problem, solution), baked to disk, with the validation set removed. The baseline train.sh points --dataset_name /app/data/train at it (loaded offline). Train only on this.
• In-distribution val (/app/data/val/): 2k examples held out from training, same distribution as train. A training-health check — is the model actually learning to count, is it under-/over-training. Do not train on it.
• OOD val (/app/data/ood_val/): ~1k examples that are out-of-distribution relative to train, and visible to you. This is the signal that tracks the sealed out-of-distribution test: it lets you see real generalization move as you tune. Tune for this (train a version → check the OOD val → adjust → retrain → keep the best), but do not train on it. It is disjoint from the sealed test data.
• Editable baseline (/app/methods/main/): train.sh (launcher + hyper-params), open_r1/ (trainer + reward functions you edit), zero3.json. See methods/main/README.md. Running bash /app/methods/main/train.sh trains and writes the model to /app/submission/.
• Self-check (/app/selfcheck.py): runs the exact same eval as the verifier (same prompt, extraction, matching) on both val sets against whatever model is in /app/submission/. This is your feedback loop — train, check, adjust, retrain, keep the best, within the budget. Free and unlimited.

What you submit

A directory /app/submission/ that Qwen2VLForConditionalGeneration.from_pretrained can load (full merged Qwen2-VL-2B weights + tokenizer + processor). That directory is the graded artifact.

Evaluation protocol (fixed — the verifier and your self-check both use it)

For every test example the model is prompted with the image and:

{question} First output the thinking process in <think> </think> and final answer (number) in <answer> </answer> tags.

Note: the trainer's training template says "Output …" while the eval template says "First output …". The grader uses the EVAL phrasing above. Train so your model is robust to it.

The model is generated greedily (do_sample=False, max_new_tokens=256). The grader extracts the first integer inside <answer>…</answer> (regex <answer>\s*(\d+)\s*</answer>) and counts the example correct only on exact integer equality with the gold count (no fuzzy matching). The metric is:

accuracy_pct = 100 * (correct / total)            # computed per test set
score = mean(accuracy_pct over the two hidden test sets)

How it is judged

The verifier loads /app/submission/ in a clean, network-isolated box and runs the protocol above on the two hidden test sets (5k examples each). Your reward increases monotonically with the mean accuracy: the untrained base model scores 0, a strong reference solution sits well above it, and the reward is not capped at that reference — beating it scores higher. A model that fails to load, is the wrong architecture, or produces unparseable output everywhere scores 0.

Common pitfalls

• Overfitting the visible in-domain split. A high in-distribution self-check score does not transfer to either hidden set. The OOD val is the honest signal.
• Reward hacking. GRPO will exploit whatever your reward actually pays for, including degenerate outputs that satisfy it without answering. Watch what the policy converges to.
• Output format drift. If the model stops emitting the <answer> tag, every example is parsed wrong regardless of reasoning quality — the grader has no fallback.
• Blowing the wall-clock. GRPO is slow (multiple generations per prompt). Size the group count, batch, and step count to finish inside the budget with the model saved.