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.