feat: GenEval reward support + unified trainer sampling pipeline

[Jayce-Ping]

Summary

• GenEval reward model: Mask2Former object detection + CLIP color classification for evaluating compositional T2I generation (object count, color, spatial position). Heavy deps (mmcv/mmdetection) are fully optional with lazy import guards.
• Unified sampling pipeline: Extract sample_batch() / generate_samples() / evaluate() into BaseTrainer, eliminating ~460 lines of duplicated code across all 6 trainers.
• Metadata passthrough: _inject_batch_metadata() automatically passes dataset-level metadata (e.g. geneval_metadata) from dataloader batches into samples' extra_kwargs, enabling metadata-dependent rewards without modifying adapters.
• Full GenEval dataset: 33,199 train / 553 test prompts (deduplicated from DiffusionNFT's 50K), stored as JSON strings to avoid Arrow heterogeneous struct serialization issues.

Key Design Decisions

Decision	Rationale
sample_batch() is public	Subclasses can override for custom per-batch pipeline
evaluate() removed from @abstractmethod	Concrete default in BaseTrainer; subclasses override only if needed
Metadata as JSON string in JSONL	Arrow cannot serialize heterogeneous nested structs (learned from DiffusionNFT experience)
GenEval deps fully optional	mmcv 1.x requires Python 3.10 + CUDA compilation; main framework stays Python 3.12 compatible
reward_type: "score" for training	Continuous 0-1 reward for policy gradient (aligned with DiffusionNFT); strict accuracy logged in extra_info

Files Changed

Area	Files
Core architecture	trainers/abc.py (+200), trainers/{grpo,nft,dpo,awm,crd,dgpo}.py (-461 total)
GenEval reward	rewards/geneval.py, rewards/registry.py
Dataset	dataset/geneval/{train,test}.jsonl
Config/Scripts	examples/grpo/lora/sd3_5/geneval.yaml, scripts/install_geneval_deps.sh, scripts/convert_geneval_dataset.py, pyproject.toml

Test plan

• Run existing GRPO training (e.g. examples/grpo/lora/sd3_5/default.yaml) — verify no regression from sampling pipeline refactor
• Run NFT training — verify generate_samples() delegation works
• Verify GenEval reward loads correctly with mmcv/mmdetection installed (Python 3.10 env)
• Verify graceful error message when GenEval deps are missing
• Run examples/grpo/lora/sd3_5/geneval.yaml end-to-end with GenEval dataset

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR adds a new GenEval reward model (Mask2Former detection + CLIP color classification) and refactors the trainer stack to use a unified sampling / evaluation pipeline in BaseTrainer, reducing duplicated code across trainers and enabling dataset metadata passthrough into reward-model kwargs.

Changes:

• Introduces GenEvalRewardModel and registers it in the reward registry, plus docs/config/scripts for installing optional deps and using the dataset.
• Extracts sample_batch(), generate_samples(), and a default evaluate() into BaseTrainer, and updates trainers to delegate sampling to the shared implementation.
• Adds a dataset metadata convention (metadata column → sample.extra_kwargs → reward __call__ kwargs) and ships GenEval JSONL + class list.

Reviewed changes

Copilot reviewed 21 out of 22 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
src/flow_factory/trainers/abc.py	Adds unified sampling (sample_batch/generate_samples) and default evaluation loop with metadata injection.
src/flow_factory/trainers/grpo.py	Delegates sampling to generate_samples() (keeps GRPO trajectory/logprob settings).
src/flow_factory/trainers/nft.py	Delegates sampling to generate_samples() (final-latent-only rollouts).
src/flow_factory/trainers/dpo.py	Delegates sampling to generate_samples() (final-latent-only rollouts).
src/flow_factory/trainers/dgpo.py	Delegates sampling to generate_samples() (final-latent-only rollouts).
src/flow_factory/trainers/crd.py	Delegates sampling to generate_samples() (final-latent-only rollouts).
src/flow_factory/trainers/awm.py	Delegates sampling to generate_samples() (final-latent-only rollouts).
src/flow_factory/rewards/geneval.py	Implements GenEval reward model with optional mmdet/mmcv + open_clip deps.
src/flow_factory/rewards/registry.py	Registers the geneval reward model name → implementation path.
src/flow_factory/rewards/abc.py	Clarifies that **kwargs comes from sample.extra_kwargs / dataset metadata.
src/flow_factory/rewards/my_reward.py	Documents how to receive metadata fields in custom reward signatures.
src/flow_factory/data_utils/dataset.py	Packs non-preprocess JSONL fields into a metadata column for passthrough.
guidance/rewards.md	Documents the dataset metadata convention and adds GenEval to the reward list.
README.md	Adds GenEval to the reward table and notes the optional dependency install path.
examples/grpo/lora/sd3_5/geneval.yaml	Provides an end-to-end example config for training/eval with GenEval.
scripts/install_geneval_deps.sh	Adds a helper script to install/compile mmcv + install mmdet/open_clip.
scripts/convert_geneval_dataset.py	Adds a conversion/dedup script for GenEval JSONL (store metadata as JSON strings).
pyproject.toml	Adds a geneval optional-dependency extra.
dataset/geneval/test.jsonl	Adds the GenEval test split in JSONL (include/exclude stored as JSON strings).
dataset/geneval/object_names.txt	Adds the COCO-ish class-name list used for label mapping.
AGENTS.md	Updates the documented reward list to include GenEval.

Comments suppressed due to low confidence (3)

src/flow_factory/rewards/geneval.py:219

• GenEvalRewardModel hard-codes CUDA device selection via cuda:{accelerator.local_process_index}. This ignores the reward config’s device (and BaseRewardModel’s self.device) and will break if the reward is configured for CPU or non-CUDA backends, and can also diverge from Accelerate’s chosen device in more complex setups. Prefer initializing both mmdet and open_clip using self.device (or self.device.type/index) rather than constructing a CUDA string manually.

        device_str = f"cuda:{self.accelerator.local_process_index}"
        self._detector = init_detector(
            detector_config,
            detector_checkpoint,
            device=device_str,
        )
        logger.info(f"Mask2Former loaded on {device_str}.")

src/flow_factory/rewards/geneval.py:185

• The ImportError guidance for missing GenEval deps is incomplete: Mask2Former inference requires mmcv ops in addition to mmdet/mmengine, and this repo provides a dedicated installer script / optional extra. Consider updating the error message to point users to pip install flow-factory[geneval] and/or bash scripts/install_geneval_deps.sh, so failures are actionable.

        try:
            from mmdet.apis import init_detector, inference_detector
        except ImportError:
            raise ImportError(
                "mmdet is required for GenEval reward. "
                "Install with: pip install mmdet mmengine"
            )

src/flow_factory/rewards/geneval.py:498

• torch.amp.autocast("cuda", enabled=False) is always entered, even if the reward model is configured to run on CPU. If you switch GenEval to use self.device, consider also keying this autocast guard off self.device.type (or using a no-op context on non-CUDA) to avoid backend mismatches.

        # Disable AMP: mmdet's ms_deform_attn CUDA kernel does not support bf16,
        # and CLIP weights are loaded as fp32. Without this guard, bf16 autocast
        # from the trainer's eval loop would cause runtime errors.
        with torch.amp.autocast("cuda", enabled=False):
            for i in range(batch_size):

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 21 out of 22 changed files in this pull request and generated 3 comments.