Compositional Perturbation Planning with a Generative Structural Causal World Model
Abstract
Multi-gene intervention planning is critical to achieve a desired final cell state from either an embryonic stem or an intermediary cell state, while simultaneously modeling epistatic gene interactions. Current perturbation foundation models are not reliably trained to predict multi-gene perturbations and are conditional-mean estimators rather than intervention operators, so composing perturbations under a reinforcement learning (RL) policy queries the model off-distribution. Here, we introduce a generative structural causal world model in which transcription factor (TF) perturbations are modeled as latent interventions with support for multi-perturbations where any k-factor combination is solved in closed form rather than rolled out, making the multi-perturbation solve closed-form and the model compositional under unseen combinations. On held-out TF overexpressions, the model matches or outperforms baselines on both single and combination perturbations, with cocktails of up to eight TFs, and predicts held-out gene-set epistasis. Epistatic interaction is explicitly learned through a low-rank interaction head trained on multi-TF perturbations by injecting set-dependent latent shifts through the causal graph, allowing the model to predict non-additive effects in held-out, experimentally validated TF combinations. Over this pre-trained world model, we built a policy to design perturbation cocktails that steer human embryonic stem cells (hESCs) to downstream differentiated fates including hematopoietic stem and progenitor cells (HSPC). We implement a goal-conditioned GFlowNet policy that proposes TF cocktails scored by the frozen world model's do-operator under an ensemble-pessimistic reward, penalizing cocktails whose value depends on simulator error. The policy generalizes to held-out fate signatures and recovers the driver TFs of a human embryonic stem cell to HSPC screen alongside other human fetal atlas annotated fates, matching model-predictive-control (MPC) search on broad recovery and exceeding it on top-rank metrics, while searching at roughly 36-fold lower inference cost and proposing far more diverse cocktails. Furthermore, since both the gene-sets that differentiate to the HSPC fate and the fate itself were held out in the world model and the policy, respectively, its zero-shot proposals for HSPC surface more validated answer-key drivers than the search and heuristic baselines at matched budget, nominating canonical blood-lineage regulators including LYL1, KLF1, and KLF2. Together, our work provides a framework for an exact, compositional causal world model allowing the implementation of combinatorial cell-fate design as an amortized, highly efficient policy.