Personalizing MLLMs via a Reinforced Multimodal Reference Game

This repository contains the training and evaluation code for personalizing multimodal large language models (MLLMs) through a Lewis reference game optimized with GRPO (Group Relative Policy Optimization).

A speaker learns to describe a personalized concept (a specific object, pet, person, or place) so that a listener can pick the right image out of a pool of confusable distractors. The reward signal comes from whether the listener succeeds, which pushes the speaker toward concise, discriminative, identity-focused descriptions. Those learned descriptions are then used for several downstream personalization tasks.

Base models:Qwen2-VL / Qwen2.5-VL (2B and 7B Instruct). Training uses LoRA adapters + DeepSpeed ZeRO-3.

---

Table of Contents

• Method overview
• Repository layout
• Installation
• Datasets
• End-to-end workflow
  • 1. Data preparation
  • 2. GRPO training
  • 3. Generate descriptions
  • 4. Evaluation tasks
  • 5. Aggregate metrics
• Model types
• Environment variables
• Description quality analysis
• Troubleshooting
• Known limitations
• Citation
• License

---

Method overview

 ┌──────────────────────────────────────────┐
│ GRPO reference game │
└──────────────────────────────────────────┘
reference image ──► SPEAKER (policy) ──► description ──► LISTENER (reward model)
Qwen2-VL + LoRA │ │
│ picks image from
│ [target, distractor_1, …]
▼ │
GRPO policy update ◄─────────┘
reward = accuracy (listener) + format + length

• The speaker is the trained policy. It emits a structured description with <think>…</think> reasoning and an <answer>…</answer> caption (further parsed into coarse / detailed / state / location fields).
• The listener is served as a separate HTTP micro-service and scores how well a description identifies the correct image among distractors.
• Rewards combine listener accuracy, output format compliance, and a length term that encourages conciseness (soft_gated, binary, and soft_always reward modes are supported).
• The listener itself can also be trained with GRPO (selection + consistency tasks), with the speaker served as the helper service.

---

Repository layout

Lewis_Game/
├── README.md # This file
├── requirements.txt # Python dependencies
├── LICENSE # MIT
├── configs/
│ └── zero3.json # DeepSpeed ZeRO-3 config
├── scripts/ # Launch helpers
│ ├── run_listener_service.sh # Start the listener reward service
│ ├── run_speaker_service.sh # Start the speaker helper service (listener training)
│ ├── run_speaker_training.sh # GRPO-train the speaker
│ └── run_listener_training.sh # GRPO-train the listener
│
├── src/ # Data prep, inference & evaluation
│ ├── data_prepare/ # 01–05 data pipeline (+ ablation/, utils)
│ ├── inference_utils/ # Shared model / dataset / prompt / retriever code
│ ├── eval_utils/ # Metric aggregation + description processing
│ │ └── description_processing/# State/location evaluation & refinement (Qwen3)
│ ├── analysis/ # Hallucination / attribute / entropy analysis
│ ├── generate_descriptions.py # Build the description database
│ ├── personalize.py # Task A: retrieval-based identification
│ ├── personalize_skip_retrieval.py # Task A ablation (no retrieval, yes/no)
│ ├── recognition.py # Task B: binary same-object recognition
│ └── vqa.py # Task C: personalized visual QA
│
├── train_src/open_r1/ # GRPO training pipeline
│ ├── train_speaker_dist.py # Speaker GRPO entry point
│ ├── train_listener_dist.py # Listener GRPO entry point
│ ├── listener_service.py # Listener micro-service (single-image scoring)
│ ├── listener_service_ablation.py # Two-image (reference-matching) variant
│ ├── speaker_service.py # Speaker micro-service
│ ├── dist_helpers.py # torch.distributed gather/broadcast helpers
│ ├── logger.py # Prediction / W&B logging
│ ├── rouge_helpers.py # ROUGE overlap reward helpers
│ └── trainer/ # Custom GRPO trainers
│ ├── grpo_trainer.py # HF-generate GRPO trainer (default)
│ ├── vllm_grpo_trainer.py # vLLM-backed GRPO trainer (--use_vllm)
│ └── speaker_helpers.py # Listener-training ↔ speaker-service client
│
├── tests/ # Dataset / retrieval / LoRA sanity checks
├── visualize/ # Qualitative error-analysis scripts
└── find_qualitative_examples.py # Build paper-figure example triplets

---

Installation

Requirements: Python ≥ 3.10, CUDA ≥ 12.1, and one or more NVIDIA GPUs (A100 40GB or better recommended for the 7B models).

git clone https://github.com/Deepayan137/Lewis_Game.git
cd Lewis_Game
python -m venv .venv &&source .venv/bin/activate
# Install PyTorch matching your CUDA version first, e.g.:
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
pip install -r requirements.txt
# Optional accelerators:
pip install flash-attn --no-build-isolation # Flash-Attention 2
pip install vllm # only for the vLLM GRPO trainer

Run all commands from the repository root — several scripts add src/ / train_src/ to sys.path relative to the current working directory.

---

Datasets

The pipeline is built around PerVA (multi-category personalization) and also supports the single-category personalization benchmarks YoLLaVA, MyVLM, and DreamBooth. Datasets are not distributed here — download them from their original sources and arrange them as follows.

PerVA (explicit train/test split layout):

data/PerVA/
├── train_/<category>/<concept>/*.jpg # reference images
│ └── <concept>/laion/*.jpg # (optional) LAION hard negatives
└── test_/<category>/<concept>/*.jpg # query images

YoLLaVA / MyVLM / DreamBooth (single "all" category, auto-split per concept):

data/YoLLaVA/<concept>/*.jpg

Concept → category mappings for each dataset live in src/inference_utils/common.py (DATASET_CATEGORY_MAPS). For VQA, the YoLLaVA/MyVLM *-visual-qa.json files are expected under data/<dataset>/.

---

End-to-end workflow

1. Data preparation

Scripts run in numeric order (all under src/data_prepare/, run from repo root).

# 1) Scan images into a catalog (one JSON with train/test/negative lists)
python src/data_prepare/01_build_image_catalog.py \
--data_root data/PerVA \
--out manifests/PerVA/catalog.json \
--num_train 5 --seed 23
# 2) Split concepts into train/test "combined" sets (seeded, reproducible)
python src/data_prepare/02_create_concept_splits.py \
--input_json manifests/PerVA/catalog.json \
--out_dir manifests/PerVA \
--concept_frac 0.65 --min_concepts_threshold 8 --seed 23
# → train_combined_concepts_seed_23.json# → test_combined_concepts_seed_23.json# → train_test_combined_metadata_seed_23.json# 3) CLIP hard-negative mining per category (builds FAISS indices + retrieval pools)
python src/data_prepare/03_build_retrieval_per_category.py \
--category clothe \
--catalog_file manifests/PerVA/train_combined_concepts_seed_23.json \
--distractors 2 --seed 23 --out_dir outputs/PerVA
# → outputs/PerVA/<category>/seed_23/retrieval_top3.json (+ *.faiss, class_mappings.json)# 3b) (optional) inject cross-category distractors as noise
python src/data_prepare/add_cross_category_distractors.py \
--input outputs/PerVA/clothe/seed_23/retrieval_top3.json \
--cross_cat_prob 0.3 --seed 23
# 4) Combine per-category retrieval files (optionally subsample)
python src/data_prepare/04_combine_retrieval_data.py \
--input_dir outputs/PerVA \
--catalog manifests/PerVA/train_combined_concepts_seed_23.json \
--input_filename retrieval_top3.json --seed 23 --num_samples 30
# 5) Convert to a HuggingFace DatasetDict for training
python src/data_prepare/05_convert_to_hf_dataset.py \
--input_filename outputs/PerVA/all/seed_23/retrieval_top3_subset_30.json \
--K 3 --seed 23 --task speaker # or --task listener# → share_data/PerVA_speaker_train_seed23_K3/

Step 3 is per-category and embarrassingly parallel — launch one job per category. The --with_negative (LAION) sampling path is not implemented in this release; use --random_negative or step 3b for distractor variety.

2. GRPO training

Training needs two processes: a reward micro-service and the trainer.

Terminal 1 — start the listener reward service:

CUDA_VISIBLE_DEVICES=0 ./scripts/run_listener_service.sh 9000 Qwen/Qwen2-VL-7B-Instruct
# → serves POST /batch_score at http://<hostname>:9000

Terminal 2 — GRPO-train the speaker (uses the remaining GPUs):

./scripts/run_speaker_training.sh $(hostname -s) 9000 4 PerVA 23
# <listener_host> <port> <epochs> <dataset> <seed>

The launch scripts wire up LISTENER_URL, LoRA flags, DeepSpeed ZeRO-3 (configs/zero3.json), and reward functions. Hyper-parameters in the scripts are reasonable defaults — override them via the environment variables noted inline (LR, LORA_RANK, NUM_GENERATIONS, MAX_PIXELS, …) or edit the scripts directly. Adapters are saved to share_models/….

To train the listener instead, start the speaker service (run_speaker_service.sh) and run run_listener_training.sh.

The trainer entry points accept the full TRL GRPOConfig / ModelConfig flag set plus:

Flag	Default	Meaning
--reward_funcs	accuracy format	Any of accuracy, format, length, overlap, accuracy_ablation
--listener_reward_mode	soft_gated	binary / soft_gated / soft_always (speaker)
--weighted_consistency	false	Weighted consistency reward (listener)
--lo_rank / --lo_alpha / --lo_dropout	64 / 128 / 0.0	LoRA config
--max_pixels / --min_pixels	12845056 / 3136	Image token budget
--use_vllm	false	Use the vLLM-backed trainer

3. Generate descriptions

Build the description database for a set of reference images. This is the bridge between training and all evaluation tasks.

python src/generate_descriptions.py \
--data_name PerVA \
--model_type sp_concise_soft_gated \
--seed 23
# → outputs/PerVA/all/seed_23/descriptions_<model_type>.json# → outputs/PerVA/all/seed_23/database_<model_type>.json

Extra flags: --num_return_sequences, --analyze, --copy_to_rap.

4. Evaluation tasks

All tasks read the descriptions/database from outputs/… and write per-concept results under results/<dataset>/<category>/<concept>/seed_<seed>/.

Task A — Personalized identification (CLIP retrieval + multiple choice):

python src/personalize.py \
--data_name PerVA \
--model_type original_7b \
--db_type sp_concise_soft_gated \
--k_retrieval 3 --seed 23
# --mode analysis → also compute answer-probability / entropy / margin# --gt_present → force the ground-truth concept into the candidate pool

Task A (ablation) — skip retrieval, pooled yes/no matching:

python src/personalize_skip_retrieval.py \
--data_name PerVA --model_type original_7b \
--db_type sp_concise_soft_gated --seed 23

Task B — Binary recognition (same object, yes/no):

python src/recognition.py \
--data_name PerVA --model_type original_7b \
--db_type sp_concise_soft_gated --seed 23 \
--use_description # include the reference description in the prompt

Task C — Personalized VQA:

python src/vqa.py \
--data_name YoLLaVA --model_type original_7b \
--db_type original_7b --seed 23 \
--qa_file data/YoLLaVA/yollava-visual-qa.json

• --model_type = the MLLM used to answer the query.
• --db_type = which description database to condition on (i.e. which speaker produced the reference descriptions).

5. Aggregate metrics

Per-concept result files are aggregated into dataset-level metrics:

# Identification (macro precision / recall / F1)
python src/eval_utils/aggregate_identification.py \
--dataset PerVA --model_type original_7b \
--db_type sp_concise_soft_gated --k 3 --seed 23
# Recognition (yes/no accuracy, weighted accuracy)
python src/eval_utils/aggregate_recognition.py \
--dataset PerVA --model_type original_7b \
--db_type sp_concise_soft_gated --seed 23
# Skip-retrieval ablation (micro + macro P/R/F1)
python src/eval_utils/aggregate_skip_ret.py \
--dataset PerVA --model_type original_7b \
--db_type sp_concise_soft_gated --seed 23

Identification aggregation reads the test-concept list from OSC_subset_seed_<seed>.txt (PerVA, produced by data_prepare/ablation/create_perva_concepts.py) or <dataset>_concept_list.txt (YoLLaVA/MyVLM) in the working directory.

---

Model types

Model identifiers are resolved centrally in src/inference_utils/common.py (MODEL_CONFIGS). LoRA paths are built from $SHARE_MODELS_DIR (default ./share_models) and the --seed.

model_type / db_type	Description	PEFT
original_2b	Qwen/Qwen2-VL-2B-Instruct	–
original_7b	Qwen/Qwen2-VL-7B-Instruct	–
sp_concise_soft_gated	Speaker, soft-gated accuracy + length reward	✓
sp_concise_binary	Speaker, binary listener reward	✓
sp_concise_only	Speaker, length reward only	✓
sp_accuracy_only	Speaker, accuracy reward only	✓
ls_original_7b	Listener trained vs. original_7b descriptions	✓
ls_soft_gated	Listener trained vs. soft-gated speaker	✓
ls_*_no_wt	Listener variants without weighted consistency	✓

The LoRA model paths are currently PerVA-specific. To evaluate a checkpoint trained on another dataset, add an entry (or adjust the path_template) in MODEL_CONFIGS.

---

Environment variables

Variable	Default	Used by
SHARE_MODELS_DIR	./share_models	Model path resolution
LISTENER_URL	http://127.0.0.1:9000/batch_score	Speaker training
LISTENER_TIMEOUT	30	Speaker training
LISTENER_REWARD_MODE	soft_gated	Speaker training
SPEAKER_URL	http://127.0.0.1:9000/batch_describe	Listener training
LISTENER_BATCH_SIZE / SPEAKER_BATCH_SIZE	5	Services
INFERENCE_CONCURRENCY	1	Services
DEBUG_MODE	false	Verbose per-step debug logs
WANDB_PROJECT / WANDB_MODE	–	Logging
NO_PROXY	–	Bypass proxy for the reward service host
DASHSCOPE_API_KEY	–	src/analysis attribute extraction (OpenAI-compatible)

Security: API keys are read from the environment only — never hard-code credentials in source. src/analysis/attribute_extraction.py and analyze_hallucination.py read DASHSCOPE_API_KEY.

---

Description quality analysis

src/eval_utils/description_processing/ (driven by src/eval_utils/eval_with_qwen.py) uses Qwen3-8B to measure and optionally remove undesirable state (pose/action) and location attributes from descriptions — the properties that hurt identity-based personalization.

# Evaluate: fraction of descriptions mentioning state / location, mean length
python src/eval_utils/eval_with_qwen.py \
--input outputs/PerVA/all/seed_23/descriptions_original_7b.json \
--out results/ --refine none --batch-size 2
# Refine: strip attributes (state | location | location_and_state)
python src/eval_utils/eval_with_qwen.py \
--input outputs/PerVA/all/seed_23/descriptions_original_7b.json \
--out results/ --refine state --batch-size 2

Additional analyses live in src/analysis/ (hallucination, attribute extraction, entropy) and visualize/ (qualitative error comparisons).

---

Troubleshooting

• ImportError / module not found — run from the repository root; the scripts insert src/ and train_src/ into sys.path relative to the CWD.
• Listener connection refused — confirm the service is up (curl http://<host>:<port>/) and that the host is in NO_PROXY.
• CUDA OOM — lower --per_device_train_batch_size, raise --gradient_accumulation_steps, reduce --max_pixels, or keep DeepSpeed ZeRO-3 enabled.
• Wrong processor for a checkpoint — the processor is now loaded from the model / LoRA base automatically; make sure the adapter's adapter_config.json points at the correct base model.
• Debug logs — set DEBUG_MODE=true and inspect debug_files/debug_speaker_*_job_*.txt.

---

Known limitations

• LoRA MODEL_CONFIGS paths are PerVA-specific (see Model types).
• --with_negative (LAION hard-negative sampling) in 03_build_retrieval_per_category.py is not implemented and raises a clear error if requested.
• VQA ships default question files only for YoLLaVA and MyVLM.
• The vLLM GRPO trainer supports Qwen2-VL / Qwen2.5-VL; the default (HF-generate) trainer is the more thoroughly exercised path.

---

Citation

If you use this code, please cite (update with the final publication details):

@inproceedings{lewis_game_personalization,
title = {Personalizing MLLMs via a Reinforced Multimodal Reference Game},
author = {Das, Deepayan and collaborators},
booktitle = {European Conference on Computer Vision (ECCV)},
year = {2026}
}

---

License

Released under the MIT License.

Acknowledgments

Built on Qwen2-VL / Qwen2.5-VL, TRL, DeepSpeed, and HuggingFace Transformers. The GRPO training code derives from the open-r1 project structure.