Building virtual cells with generative models to simulate cellular behavior in silico is emerging as a promising paradigm for accelerating drug discovery. However, prior image-based generative approaches can produce implausible cell images that violate basic physical and biological constraints. To address this, we propose to post-train virtual cell models with reinforcement learning (RL), leveraging biologically meaningful evaluators as reward functions. We design seven rewards spanning three categories-biological function, structural validity, and morphological correctness-and optimize the state-of-the-art CellFlux model to yield CellFluxRL. CellFluxRL consistently improves over CellFlux across all rewards, with further performance boosts from test-time scaling. Overall, our results present a virtual cell modeling framework that enforces physically-based constraints through RL, advancing beyond "visually realistic" generations towards "biologically meaningful" ones.
Figures & tables
Figure 1: Failure of cell generation.
Figure 2: Motivation. Current generative models for simulating cellular perturbations can fail to produce physically plausible cell images. For example, nuclei may appear outside the cell membrane. We design a suite of biologically meaningful verifiers in three roles: (1) as evaluators to assess the biological correctness of generated images, (2) as reward signals to improve generation via reinforcement learning, and (3) as verification modules to enhance sample quality through test-time scaling.
Figure 3: CellFluxRL algorithm. RL post-training seeks to increase the likelihood of high-reward samples and decrease the likelihood of low-reward samples. Therefore, the core training loop of CellFluxRL consists of interleaved phases of sampling and training. (a) Sampling: we generate multiple rollouts from a fixed control image and perturbation condition, scoring each with the reward models. (b) Training: because exact likelihoods in flow matching are intractable, we construct positive and negative velocities from the batch of rollouts and optimize them contrastively to achieve this goal, following DiffusionNFT [ 51 ] .
Metric
PhenDiff [ 3 ]
IMPA [ 34 ]
CellFlux [ 50 ]
CellFluxRL
+TTS
Biological Rewards
MoA
0.18
0.12
0.26
0.34
0.56
Nuc-in-Cyto
0.91
0.79
0.88
0.96
0.97
Roundness
-0.24
-0.32
-0.34
-0.26
-0.19
NucSize
-0.88
-1.02
-2.21
-1.04
-0.38
CytoSize
-0.97
-1.59
-1.09
-0.65
-0.41
Table 1: Quantitative comparison across biological rewards and generative quality metrics. Each row represents a different evaluation metric, and each column represents a different method. TTS is CellFluxRL with test time scaling by best-of- N with N=4 , where the best sample is selected by the weighted total reward. Bold values indicate the best performance; underlined values indicate the second best.
Figure 4: Qualitative comparisons. CellFluxRL generates more biologically-grounded images, better capturing drug-induced morphological changes. In these examples, Etoposide-induced cell rounding, Demecolcine-driven microtubule destabilization, and AZ138-associated cell shrinkage are all more faithfully reproduced, and cell density more closely matches the ground truth for Cisplatin. Test-time scaling (+TTS) further refines these predictions toward the target images.
Figure 5: Test-time scaling by best-of- N further improves generation quality. The sample achieving the highest overall (combined) reward is selected from N rollouts, and each individual reward is plotted. RL (orange) consistently exhibits better scaling than the base model (blue) across all rewards.
Figure 6: Sensitivity analysis on KL weight β . Each subplot shows reward sensitivity to β after RL post-training.
Metric
MoA
Nuc-in-Cyto
Roundness
NucSize
CytoSize
NucCount
CytoCount
CellFluxRL
MoA
0.412
0.204
0.263
0.301
0.212
0.279
0.277
0.337
Nuc-in-Cyto
0.852
0.977
0.907
0.904
0.976
0.867
0.937
0.956
Roundness
-0.342
-0.332
-0.198
-0.298
-0.340
-0.271
-0.304
-0.264
NucSize
-3.319
-1.724
-2.328
-0.702
-1.784
-2.909
-2.853
-1.038
CytoSize
-1.198
-0.545
-0.919
-0.917
-0.497
-1.113
-0.801
-0.652
NucCount
-0.760
-0.904
-0.685
-0.818
-0.881
-0.476
-0.511
-0.528
Table 2: Single-reward optimization with reinforcement learning. Each column shows a model optimized for one reward only or CellFluxRL (all rewards), evaluated across the metrics in each row. Bold values indicate the best performance; underlined values indicate the second best; dashed-underlined values indicate the third best.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
{x^1(j)}j=1m∼πold(⋅∣x0,c)
Appendix
Algorithm 1 CellFluxRL : Biologically-Constrained RL Post-Training for Flow Matching
Compound
MoA
Cytochalasin B
Actin disruptors
Cytochalasin D
Actin disruptors
Latrunculin B
Actin disruptors
AZ258
Aurora kinase inhibitors
AZ841
Aurora kinase inhibitors
Mevinolin/Lovastatin
Cholesterol-lowering
Appendix
Table 3: Modes of action (MoA) for compounds in BBBC021.
Figure 7: MoA reward failure cases. The pretrained CellFlux baseline (left) generates images that do not match the expected morphological profile for the given perturbation, as measured by the MoA reward. CellFluxRL + TTS (right) corrects these failures by explicitly optimizing for MoA consistency during RL post-training. Ground-truth target images for the same perturbation conditions are shown in Figure 4 of the main paper.
Figure 8: Roundness reward failure cases. The pretrained CellFlux baseline (left) produces nuclei with irregular, implausible shapes that deviate from the MoA-conditioned ground-truth distribution. CellFluxRL + TTS (right) generates nuclei with roundness statistics consistent with real cells under the same perturbation condition. Ground-truth target images for the same perturbation conditions are shown in Figure 4 of the main paper.
Metric
CellFlux [ 50 ]
CellFluxRL
Biological Rewards
MoA
0.26
0.28
Nuc-in-Cyto
0.88
0.96
Roundness
-0.34
-0.25
NucSize
-2.21
-1.34
CytoSize
-1.09
-0.64
Appendix
Table 4: Quantitative comparison across biological evaluation metrics (rows) for CellFlux baselines and CellFluxRL variants (columns), where CellFluxRL is post-trained with min-max normalized rewards. Bold values indicate the best performance per metric.
MoA
MoA-CLIP
MoA-DINO
NuInCy
NuInCy-StarDist
CellFlux [ 50 ]
0.26
0.45
0.34
0.88
0.88
CellFluxRL
0.34
0.52
0.40
0.96
0.95
Appendix
Table 5: Evaluation with held-out reward functions unseen during RL training. MoA and NuInCy use the original training-time evaluators (Inceptionv3 classifier and Cellpose segmentation, respectively). MoA-CLIP and MoA-DINO use independently trained held-out MoA classifiers, and NuInCy-StarDist uses an independent segmentation pipeline. Bold values indicate the best performance per metric.
Predicting cellular responses to perturbations supports the study of gene function, disease mechanisms, and therapeutic strategies. Despite advances in single-cell perturbation modeling, existing models typically optimize surrogate losses that do not directly reflect the biological criteria used for evaluation, so better data fitting need not yield better biological predictions. To address this mismatch, we introduce CellRFT, a reinforcement fine-tuning framework that uses biological evaluation as direct training feedback. CellRFT uses policy-gradient optimization to learn from non-differentiable evaluations of generated cell populations and integrates multiple biological rewards through hierarchical reward aggregation. Comprehensive experiments demonstrate CellRFT's applicability across different pretrained models and effectiveness in improving perturbation prediction, reveal that optimizing one biological criterion can help or hinder others, and show that complementary rewards can improve criteria beyond those directly optimized, offering a way to probe how biological metrics shape model behavior, with the potential to inform evaluation design. Code will be made available.
Jie Yan, Li Liu, Hanze Guo +7
Chinese Academy of Sciences · Peking University · Renmin University of China +2
Virtual cell models should capture phenotypes induced by specific perturbations. However, image-based models are typically trained and evaluated to match real-image distributions, which alone does not explicitly prioritize perturbation-specific biological responses. Recent approaches add biological signals through reinforcement-learning post-training or inference-time steering, raising a central question: where should biological response information enter the generative pipeline? We introduce PerturbBridge, a stochastic data-to-data bridge model that directly supervises perturbation-specific phenotypes in predicted perturbed images during training. PerturbBridge jointly optimizes bridge regression and differentiable biological response losses, with biological gradients passing through a frozen image encoder. We further introduce perturbation-response evaluations based on downstream classification and retrieval. Across three microscopy benchmarks for perturbation profiling and subcellular imaging, PerturbBridge surpasses CellFlux in biological fidelity, image quality, and sampling efficiency. On BBBC021, mechanism-of-action classification improves by 5.66 and 7.5 percentage points with morphological and learned visual features, respectively. On JUMP and AllenCell, retrieval Hit@10 and Hit@1 improve by 1.77 and 7.15 percentage points, respectively. On BBBC021 and JUMP, PerturbBridge achieves lower FID with 80% fewer model evaluations than CellFlux. Overall, these results show that directly supervising perturbation responses during training improves phenotypic fidelity while enabling efficient cellular image generation. Code is available at https://github.com/yyc1215lll-creator/PerturbBridge.
Cellular perturbation-response modeling requires coordinated choices of representations, fusion mechanisms, objectives, and training procedures. Large language models (LLMs) can propose executable candidates, but unconstrained revision can produce invalid implementations, change task semantics, or discard useful components. We present CellScientist, a protocol-constrained workflow that converts execution and validation feedback into auditable model-revision trajectories. It records design states and outcomes, routes discrepancies to specific components, and applies local revisions under a fixed task contract. A matched-budget study fixes the candidate language, predictor, fitting, and evaluator: structured revision finds better held-out predictors at small budgets under two LLM backbones. Operational audits link history to fewer repeated proposals, contract checks to contained violations, and discrepancy routing to targeted repairs. Refits of two frozen designs on an independently acquired cohort evaluate external predictive utility. Open-workflow trajectories retain improvements, regressions, and failures, while transcriptomic and single-cell searches extend application to additional response spaces. CellScientist produces both a selected predictor and an inspectable record of its development. Project page: https://limengran98.github.io/CellScientist/.
Mengran Li, Bo Li, Jiaying Wang +7
Center for Artificial Intelligence and Robotics (CAIR), Hong Kong Institute of Science and Innovation (HKISI)