CellFluxRL: Biologically-Constrained Virtual Cell Modeling via Reinforcement Learning

Dongxia Wu; Shiye Su; Yuhui Zhang; Elaine Sui; Emma Lundberg; Emily B. Fox; Serena Yeung-Levy

cs.LGMar 23, 2026

CellFluxRL: Biologically-Constrained Virtual Cell Modeling via Reinforcement Learning

Authors: Dongxia Wu, Shiye Su, Yuhui Zhang, Elaine Sui, Emma Lundberg, Emily B. Fox, Serena Yeung-Levy

Organizations: Stanford University, Stanford, CA

Abstract

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 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 $\beta$ . Each subplot shows reward sensitivity to $\beta$ 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

\bigl\{\hat{x}_{1}^{(j)}\bigr\}_{j=1}^{m}\;\sim\;\pi^{\mathrm{old}}(\cdot\mid x_{0},\,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.

Constrained Generative Modeling

Constrained generative modeling develops methods that produce samples satisfying structural, physical, semantic, or user-defined requirements.

See all 185 papers

CellRFT: Reinforcement Fine-Tuning for Single-Cell Perturbation Modeling
Sep 17, 2026Jie Yan, Li Liu, Hanze Guo +7Multi-Objective Reinforcement Learning Single-Cell Perturbation Modeling
Chinese Academy of Sciences · Peking University · Renmin University of China +2
PerturbBridge: Direct Phenotypic Supervision for Cellular Perturbation Generation
Oct 12, 2026Yucheng Yang, Xiaoxiao Sun, Huangqingbo Sun +3Single-Cell Perturbation Modeling Image Generation
Stanford University
CellScientist: From Execution Feedback to Auditable Model-Revision Trajectories for Cellular Perturbation Prediction
May 8, 2026Mengran Li, Bo Li, Jiaying Wang +7Single-Cell Perturbation Modeling Scientific Workflow Automation
Center for Artificial Intelligence and Robotics (CAIR), Hong Kong Institute of Science and Innovation (HKISI)

Sep 17, 2026J/K move · Enter open · S save

CellRFT: Reinforcement Fine-Tuning for Single-Cell Perturbation Modeling

Jie Yan, Li Liu, Hanze Guo, Jiaxin Hu+6

Chinese Academy of Sciences · Peking University · Renmin University of China+2

Multi-Objective Reinforcement Learning Single-Cell Perturbation Modeling RL Fine-Tuning

Open paper AskPDF arXiv

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.

Abstract

Constrained Generative Modeling

Explore similar work

CellRFT: Reinforcement Fine-Tuning for Single-Cell Perturbation Modeling

PerturbBridge: Direct Phenotypic Supervision for Cellular Perturbation Generation

CellScientist: From Execution Feedback to Auditable Model-Revision Trajectories for Cellular Perturbation Prediction