PRO-MOF: Policy Optimization with Universal Atomistic Models for Controllable MOF Generation
Authors: Zicheng Liu
Ben Fei
Di Huang
Venue: ICLR 2026 Poster (2026)
Keywords: metal-organic framework
material generation
AI for science
physical modeling
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Generating physically stable and novel metal-organic frameworks (MOFs) for inverse design that meet specific performance targets is a significant challenge. Existing generative models often struggle to explore the vast chemical and structural space effectively, leading to suboptimal solutions or mode collapse. To address this, we propose PRO-MOF, a hierarchical reinforcement learning (HRL) framework for controllable MOF generation. Our approach decouples the MOF design process into two policies: a high-level policy for proposing chemical building blocks and a low-level policy for assembling their 3D structures. By converting the deterministic Flow Matching model into a Stochastic Differential Equation (SDE), we enable the low-level policy to perform compelling exploration. The framework is optimized in a closed loop with high-fidelity physical reward signals provided by a pre-trained universal atomistic model (UMA). Furthermore, we introduce a Pass@K Group Relative Policy Optimization (GRPO) scheme that effectively balances exploration and exploitation by rewarding in-group diversity. Experiments on multiple inverse design tasks, such as maximizing CO2 working capacity and targeting specific pore diameters, show that PRO-MOF significantly outperforms existing baselines, including diffusion-based methods and genetic algorithms, in both success rate and the discovery of top-performing materials. Our work demonstrates that hierarchical reinforcement learning combined with a high-fidelity physical environment is a powerful paradigm for solving complex material discovery problems.

PRO-MOF: POLICY OPTIMIZATION WITH UNIVERSAL ATOMISTIC MODELS FOR CONTROLLABLE MOF GEN-
ERATION
Zicheng Liu Beng Fef Di Huang..2.44
BCC Lab, Hangzhou International Innovation Institute, Beihang University, Hangzhou, China
2 School of Computer Science and Engineering, Beihang University, Beijing, China
CChness University of Hong Kong. Hong Kong, China
"Shenzhen Loop Area Institute, Shenzhen, China
ABSTRACT
Generating physically stable and novel metal-organic frameworks (MOFs) for in- verse design that meet specific performance targets is significant challenge. Ex- isting generative models often struggle to explore the vast chemical and structural space effectively, leading suboptimal solutions or mode collapse. To address this, we propose PRO-MOF, a hierarchical reinforcement learning (HRL. frame- work for controllable MOF generation. Our approach decouples the MOF design process into two policies: a high-level policy for proposing chemical building blocks and a low-level policy for assembling their 3D structures. By converting the deterministic Flow Matching model into Stochastic Differential Equation (SDE), we enable the low-level policy to perform compelling exploration. The framework is optimized in closed loop with high-fidelity physical reward sig- nals provided by a pre-trained universal atomistic model (UMA). Furthermore, we introduce Pass @ K Group Relative Policy Optimization (GRPO) scheme that effectively balances exploration and exploitation by rewarding in un-gropp diversity. Experiments multiple inverse design tasks, such maximizing CO2 working capacity and targeting specific pore diameters, show that PRO-MÓF significantly outperforms existing baselines, including diffusion-based methods and genetic al- gorithms, in both success rate and the discovery of top-pes nop-perfomming materials. Our work demonstrates that hierarchical reinforcement learning combined with high- fidelity physical environment is powerful paradigm for solving complex material discovery problems.
INTRODUCTION
Metal-Organic Frameworks (MOFs) represent frontier in materials science, offering unprecedented structural and chemical tunability. Their vast internal surface areas and customizable pore envi- ronments make them exceptional candidates for critical applications such carbon capture (Qian et al., 2020), gas storage, and catalysis. However, the combinatorial explosion of possible building blocks and topologies creates design space of astronomical scale, rendering exhaustive explo- ration through traditional experimental or computational methods intractable. This challenge has catalyzed the development of generative models, which aim to navigate this vast space and acceler- ate the discovery of novel materials. Recent advances, particularly coarse-grained diffusion models like DiffCSP (Jiao al., 2024), MOFDiff (Fu al., 2023), and flexible assembly frameworks like MOFFlow-2 (Kim et al., 2025a), have marked significant progress. These models can generate MOFs with novel and geometrically plausible structures, moving beyond the limitations of prede- fined templates. They excel at learning the complex statistical distributions of atomic arrangements and molecular connectivity from large datasets. However, their success in capturing geometric real- ism often masks a critical deficiency: a lack of inherent physical viability.
"Corresponding author; dhuang@buaa enu.cn

---

(a) Energy Gap
(b) Diversity Collapse with Pass@1
Real MoFs
0.5
(CORE Database)
Gerwrtt at NOTs
(MOFFIOW 2)
0.8
Probability Density
Normalized Score
0.6
0.3
Average Roward
Str Suraltural Diversily
Energy Gap
0.2
0.2
0.0
0.0
-5
-4
-2
-1
20
40
60
80
Formation Energy (eVatoom
Training Epochs
Figure The core motivations for PRO-MOF. (a) The Energy Gap: The distribution of forma- tion energies for MOFs from a baseline generator (e.g.. MOFFlow-2) is significantly higher and broader than that of real materials, indicating widespread physical instability. (b) Diversity Col- lapse: Standard reinforcement learning with a @ass@1 reward function successfully increases the average reward but at the cost drastic reduction in structural diversity, leading to mode collapse.
This "physical reality gap" is fundamental barrier to practical materials discovery. As we demon- strate, significant fraction of structures generated by state-of-the-art models, while appearing structurally sound, are physically unstable. We evaluated the formation energies of thousands of generated structures using a universal machine learning interatomic potential, UMA (Wood et al., 2025), which serves as high-fidelity surrogate for DFT calculations (Garrity et al., 2014). The results, shown in Figure (a)), reveal stark disparity: the energy distribution of generated MOFs dramatically shifted towards higher, less stable states compared to the distribution of real-world, synthesizable materials.
natural solution to this problem is to integrate physical realism directly into the training loop us- ing reinforcement learning (RL), with fast and accurate model like UMA providing the reward signal. However, a naive implementation of RL creates a new, more subtle trap. Standard policy optimization, which rewards the generation of a single successful sample (a Pass @ reward), inher- ently favors exploitation over exploration (Shao et al., 2024). The generator quickly learns to exploit few "safe" regions of the chemical space where it can reliably produce stable structures, ceasing to explore novel topologies. This leads to a rapid collapse in the diversity of generated candidates, phenomenon known mode collapse (Figure (b)). For task whose ultimate goal the discovery of new materials, this convergence to a local optimum is a critical failure.
In this work, we introduce PRO-MOF (Policy and Reward Optimized MOFs), a holistic framework that addresses both the physical reality gap and the diversity collapse problem. PRO-MOF syner- gistically combines state-of-the-arr flow-based MOF generator with a universal atomistic model (UMA) inside a novel policy optimization loop. To overcome the exploration-exploiiation dilemma, we move beyond simplistic rewards and implement Pass@ K policy optimization strategy, inspired by recent advances in training large reasoning models (Chen et al., 2025). By rewarding the gen- erator for producing at least one successful candidate within batch of diverse attempts, our framework intrinsically incentivizes exploration. This encourages the model to discover multiple, distinct islands of stability across the vast MOF landscape. Our main contributions are:
We propose PRO-MOF, the first framework to successfully employ online reinforcement learning for controllable, de novo MOF generation by leveraging universal atomistic model as a high-fidelity environment.
We introduce Pass @ K-inspired reward and advantage estimation scheme to the materials domain, which explicitly promotes structural diversity and effectively mitigates the mode collapse problem inherent in standard RL approaches.
We demonstrate through extensive experiments that PRO-MOF significantly outperforms existing methods in generating diverse, stable, and high-performing MOFs for targeted inverse design tasks, paving the way for more efficient computational materials discovery.
2

---

2 PRELIMINARIES
Notation for MOF Representation We adopt the flexible, coarse- representation of MOFs introduced by Kim et al. (2025b;a), which models structure based on its constituent building blocks. An entire MOF structure with N atoms is denoted by the tuple S = (BaD, Q. T, 6, L). This tuple consists of: a set of M 3D building blocks, BBD = {C(m) Jm= )
where each block ((m) is
defined by its atom types and local coordinates; the rigid-body rotations q = {a(m) E (3(3) ) m==1 for each building block; the rigid-body translations T {T((m) E IR')== for each block; the in- ternal torsion angles Φ = ((5(m) E (0(()) Pm }M=1 for the Pm rotatable bonds within each block to capture linker flexibility; and the global lattice parameters l = (a, b, c, a, B, 7) defining the periodic unit cell. The complete atomic coordinates X E IR Nx3 can be reconstructed by applying these transformations the local coordinates of the building blocks.
Generative Modeling with Flow Matching Flow Matching (FM) (Liu et al., 2023) is power- ful simulation-free framework for training Continuous Normalizing Flows (CNFs) (I.ipman et al., 2023). The core idea is to learn a time-dependent vector field (or velocity field) (0( 1) that trans- ports samples from a simple prior distribution p1 (c.g., Gaussian noise) to a complex target data distribution po- This is achieved by regressing the model's predicted velocity against target veloc- ity field defined along a path between a noise sample æ1 and a real data sample Zg- For generation, one samples æ ~ pi and solves the ordinary differential equation (ODE) dt t == 1 to t 0. This process is deterministic and typically highly efficient, allowing for high- quality sample generation in few steps. However, its deterministic nature precludes the stochastic exploration required by online RL. algorithms.
dæ, = vo( () from
Policy Optimization with GRPO a memory-efficient aalue-function-free online RL algorithm that has proven effective for optimiz- ing large generative models (Liu et al., 2025). The generative process is framed as Markov Deci- sion Process (MDP). At each optimization step, the policy me generates a group of G samples. A reward function provides a scalar reward R; for each completed sample. Instead of learning a value function, GRPO computes the advantage for each sample relative to its peers within the group:
Group Relative Policy Optimization (GRPO) (Shao et al., 2024)
R; PR
A
(1)
OR +
where μR and OR are the mean and standard deviation of rewards within the group. This relative advantage is then used in a clipped policy gradient objective to update the policy parameters 0. A key requirement for GRPO is stochastic policy that can generate diverse samples within a group to ensure a meaningful reward distribution and effective exploration.
3 PRO-MOF: A FRAMEWORK FOR CONTROLLABLE MOF GENERATION
Our methodology, PRO-MOF, designed to overcome two fundamental challenges in de novo ma- terials generation: the physical reality gap and the exploration-ecplottatton exploitation dilemma. PRO-MOF addresses these by formulating the generation process as a hierarchical policy optimization prob- lem. We co-design the MOF's chemical composition and its geometric structure by training two specialized policies in a closed loop, guided by rewards from a universal atomistic model. This approach ensures that both the selection of chemical building blocks and their 3D assembly simultaneously optimized toward generating physically viable and high-performing structures.
3.1 A HIERARCHICAL FRAMEWORK FOR DE Novo MOF GENERATION
We conceptualize the de novo MOF generation process as a two-stage hierarchy, building upon the paradigm of MOFFlow-2 (Kim et al., 2025a). This structure decouples the vast, discrete chemical space from the high- bigh-dimensional, continuous geometric space, allowing each policy to specialize.
High-Level Policy T chem (The Chemist) The first stage is governed by high-level policy, an autoregressive Transformer model. This policy operates in the chemical space. Its action, achem- is to generate canonicalized sequence of building blocks represented SMILES strings. Following
3

---

MOF Sequence Generator
H-GRPO
Advantages
<BOS> Zn ... <EOS>
Policy Update
<BOS>
Group Computation
Autoregression
Rewards
MOF Structure Generator
Reward Model
ML Interatomic Potentials
Evaluation
Flow Matching
Figure 2: The PRO-MOF Hierarchical Reinforcement Learning Framework. The generation pro- cess is decoupled into two policies: high-level MOF Sequence Generator that proposes chemical building blocks, and a low-level MOF Structure Generator that assembles them. A universal Reward Model (a Machine Learning Interatomic Potential, UMA) evaluates the final structure's 's quality. The reward signal is then used by the H-GRPO module to compute advantages and perform policy update, closing the loop to optimize both chemical and structural generation simultancously.
Kim et al. (2025a), the sequence is structured as: B2D <BOS> mm...... <SEP> 01.02... <EOS>, where metal clusters (m,; precede organic linkers (oy). and components within each group are sorted by molecular weight. This canonical representation defines the set of 2D chemical building blocks, ByD. The policy's 's objective to propose building blocks with high potential for forming stable, high-performing MOFs.
Low-Level Policy geom (The Structural Engineer) The second stage is controlled by a low-level policy, conditional flow matching model. This policy operates in the geometric space. It receives the building blocks aehom and first initializes their 3D structures, BID. using pre-defined library for metal clusters and cheminformatics tools, RDKit (Landrum, 2006-), for organic linkers, as done in MOFFlow-2 (Kim al., 2025a). Given BRD, its action, ageom: is to determine the optimal 3D assembly. This action corresponds to generating the continuous parameters for rigid-body transla- tions (r) and rotations (g). flexible linker torsions (Ø). and the global lattice (L). The policy is non-equivariant Transformer architecture that learns a time-dependent vector field Us to transport noise the correct structural parameters.
The final MOF structure S is a deterministic function of both actions, S = f (Gchem- agsom). ). A naive approach of training these models separately via maximum likelihood leads the problems outlined in Section 1. Our solution is to train them jointly using a hierarchical reinforcement learning scheme.
3.2 HIERARCHICAL POLICY OPTIMIZATION WITH H-GRPO
We frame the co-design task within a hierarchical reinforcement learning (HRL) framework. The high-level policy proposes a "sub-task" (a set of building blocks), and the low-level policy attempts to solve it (find stable assembly). The final reward, derived from the all-atom structure, is then used update both policies, addressing the critical credit assignment problem. In RL training, adopt the pre-trained models from Kim et al. (2025a).
Policy Optimization Loop The core of PRO-MOF is closed loop (Figure 2) where the two policies are optimized iteratively. For a given target property the high-level policy rgeem first generates a MOF sequence Oehem- This action is passed to the low-level policy n which then explores the geometric space by stochastically generating a diverse group of possible 3D structures (Si,...,SA, A universal atomistic model, UMA, evaluates all & structures, providing the rewards.
4

---

Hierarchical Credit Assignment The key to HRL framework is how rewards distributed to solve the credit assignment problem. The rewards {Ri..... R& for the individual 3D structures used to directly update the low-level (geometric) policy n=_m using the Pass@ K GRPO scheme described in Section 3.3. Crucially, the reward for the high-level (chemical) action dehem is defined as the best possible outcome achievable with the chosen building blocks. We thus define the high- level reward as Rehem max R ..... Ra). This signal informs the high-level policy about the "potential" of the building blocks it generated. This scalar reward is then used to update Th parameters using clipped policy gradient objective, similar the one used for the low-level policy but adapted for discrete action space.
3.3 GROMETRIC POLICY OPTIMIZATION VIA PASS K GRPO
Optimizing the low-level geometric policy " geom requires several technical innovations enable stable and effective online reinforcement learning.
3.3.1 STOCHASTIC EXPLORATION VIA SDE SAMPLING
To enable exploration, we convert the deterministic ODE of the flow matching model into an equiv- alent Stochastic Differential Equation (SDE), following Flow-GRPO (Liu et al., 2025). The proba- bility flow ODE, dæz Vz (xx )dt, is transformed into reverse-time SDE. For rectified flow where Pt is an interpolation of po and pl, the score term log p((x)) can be related to velocity field v2. This results in the following tractable SDE:
dx; = U (zz, t) + of (x, + (1- (11-ttoa t)o,(x,,))) dt + o,dw 2t
(2)
where is time- denendennnnnt diffusion coefficient that controls the level of stochasticity and dw is standard Wiener process. This transformation allows the geometric policy to sample stochastically from Gaussian policy (I c), prerequisite for RL-based exploration. The full derivation provided in Appendix A.
3.3.2 MUL.TI-ORJECTIVE REWARD FUNCTION WITH UMA
We leverage UMA (Wood al., 2025), a state-of-the-art universal machine learning interatomic potential (MLIP), to define a multi-objective reward function. For a generated structure S, the total reward Rootal a weighted sum:
Riota (Sc) = Uitability Ratability (S) + Property Peoopent (S, c).
(3)
The Stability Reward (Raaabiiit) is derived from the UMA-evaluated potential energy of the relaxed structure, ELMA (S), defined as Retability = log (S) Emin), encouraging physically viable, low-energy configurations. The Property Matching Reward (Rpropery) measures how well the struc- ture's properties, also predicted by UMA, match the target condition c. Further details on UMA are in Appendix B.
3.3.3 @ASS@K ADVANTAGE ESTIMATION
A naive application of RL can lead to mode collapse, where the policy repeatedly generates a single high-reward structure. To combat this, we adopt a Pass @K strategy (Chen et al., 2025) within the GRPO framework. The geometric policy generates batch of & candidates (Si....,d) St for given Ochem The rewards (RI,..., R are used to compute the group-relative advantage for each sample I IE kl....,k):
R MR
À
OR +
The policy Ty is then updated using a clipped objective based on these advantages. This group- based advantage intrinsically incentivizes the policy to explore diverse yet successful geometric configurations, as a high reward for one sample does not suppress the learning signal for other promising candidates.
5

---

Algorithm PRO-MOF Hierarchical Training Loop
1: Initialize high-level policy Th and low-level policy georn
a from pre-trained MLE models.
2: Initialize reference policies Tp, <- "0
rcchem TTgrel m geam
3: Set total warmup iterations wwrnmup-
4: for training iteration = 1, ... do
5: Sample batch of target properties (c1,.... CB).
6: Compute annealing weight Wanneal (i) = min(1,2 T warmup))
7:
=/== High- Policy Rollout and Low-Level Execution for each target property Cj do
8:
9:
Generate a MOF sequence achem.j T6 T (-Ic;).
10:
Generate & structures (S Sj,x) } ~ T goom (-) (-1uchem.ji (g) using SDE sampler (Eq. 2).
11:
Evaluate rewards {R,,1.... Rg for all h structures using UMA and apply annealing
(Eq. 5).
12:
Compute high-level reward Retem,j = max (R,11 ..., RR.k)..).
13:
end for
14:
=/=== Policy Updates with KL Regularization ===
15:
Update rechem using rewards {Rohem.j } and advantage A, via a clipped objective:
CB [mmin(r, (O)) clip(ry (0). . e, 1 + e),A,) 33Dx k (TTll
17:
Update n gcom using all rewards ({R,,)) and advantages {{Aj,)} via GRPO objective:
L. = A clip(r;,.()),, -,, + +c)A,)) - 33Dx K (
19: end for
3.3.4 REWARD ANNEALING FOR STABLE OPTIMIZATION
A critical challenge arises as the policy 2 20om in its early training stages may produce non-physical inter-block configuration that are -uttoi-distributtin for UMA, leading noisy reward signals and unstable gradients. To address this, we introduce a reward annealing strategy. We modulate the total reward with weight Wanncal (i) that increases with the training iteration
(1, "warrmpp
Rettective (i) = Wanneal (i) RS, where Wanneul (1) = min
(5)
This curriculum, where warmup is hyperaarameter, ensures that the policy initially receives a gentle learning signal from the pre-trained MLE objective, preventing destabilization from noisy rewards and promoting stable convergence as begins to explore.
4 RELATED WORK
Generative Models for MOF Design. Generative modeling Metal-Organic Frameworks (MOFs) presents unique challenges compared smaller inorganic crystals, primarily due to the large size of MOF unit cells, which often contain hundreds atoms (Cao et al., 2022; Badri- narayanan et al., 2025). To tackle this complexity, recent methods have employed hierarchical or coarse-grained approaches. For example, MOFDiff (Fu et al., 2023) first generates a coarse-grained representation of the structure via diffusion model and then decodes it into all-atom structure using a predefined library of building blocks extracted from the training data. Similarly, MOF- Flow (Kim al., 2025b) simplifies the problem by treating building blocks as rigid bodies, learning their relative rotations and translations rather than predicting individual atomic coordinates, but this requires access to the 3D structures of the building blocks beforehand.
Universal Atomistic Models as Simulation Engines. A prevailing trend in molecular model- ing the development of universal, data-driven models that can act as general-purpose simulation engines. A key architectural consideration in these models is the enforcement of SE(3) or E(3) equivariance to respect the physical symmetries of 3D space. While equivariant generative mod- els have shown promising results many molecular tasks (Corso et al., 2023; Jing et al., 2022; Jumper et al., 2021; Zeni et al., 2025). a growing body of work questions whether strict equivari- ance essential for achieving state-of-the-ant performance (Wang al., 2024). Notably, leading

---

models for biomolecular complexes and crystal structures, such AlphaFold3 (Abramson et al., 2024), Boltz-1 1 (Wohlwend et al., 2024), Proteina (Geffner et al., 2025). Orb (Neumann et al., 2024), and ADiT (Joshi et al., 2025). have achieved exceptional results using non-equivariant architectures, like the Transformer. Following this successful trend, PRO-MOF adopts a non-equivariant Trans- former architecture, demonstrating that it can effectively learn the complex potential energy surface of MOFs without built-in geometric constraints.
Reinforcement Learning for Generative Model Refinement. While generative models are pow- erful for learning the underlying distribution of dataset, they do not inherently optimize for specific functional properties. Reinforcement Learning (RI.) has emerged as potent strategy to fine-tune generative models for goal-oriented de novo design (Popova et al., 2018). In this paradigm, pre- trained generative model acts as policy that produces candidate molecules or materials. reward function, often defined by fast property predictor or surrogate model, then provides feedback to guide the generator towards regions of the chemical space with more desirable properties. This approach has been widely and successfully applied in drug discovery to optimize for properties like binding affinity and synthesizability (Kim et al., 2021; Schwaller et al., 2019; Olivecrona et al., 2017). In materials science, RL offers a promising pathway steer the generation of novel struc- tures, such as MOFs, towards enhanced performance for specific applications, moving beyond sim- ple reproduction of known structural motifs (Kim et al., 2020).
EXPERIMENTS
We design suite of experiments rigorously evaluate PRO-MOF and demonstrate its superiority over existing methods. Our evaluation focuses on the model's ability to perform controllable inverse design, task that requires generating MOFs that not only physically stable and novel but also satisfy specific, user-defined property targets.
5.1 EXPERIMENTAL SETUP
Datasets and Base Models. We use the hypothetical MOF database from Kim al. (2025a) for tun- ing the pre-trained generative model; the details of dataset processing are described in Appendix D. The base generator is a conditional version of the MOFFlow-2 architecture, pre-trained to generate MOFs conditioned on vector of their chemical properties. The reward function and environment for RL fine-tuning are powered by the pre-trained UMA-S model (Wood et al., 2025), which pro- vides rapid and accurate energy evaluations.
Inverse Design Tasks, We evaluate all models on three scientifically relevant inverse design tasks. The first is maximizing the CO2 working capacity, benchmark task for carbon capture applications directly comparable to Fu et al. (2023). The evaluation pipeline is also adopted from Fu et al. (2023), which means we are tuning the model using simulated data. The second task involves targeting specific pore diameters, crucial for gas separation, where the objective is to generate stable MOFs with pore limiting diameter (PLD) within narrow range (e.g., 6.0 + 0.2A). The final task is pure exploration challenge: discovering ultra-stable novel topologies by finding MOFs with the lowest possible formation energies, without other property constraints.
Max CO2 Cap.
Target PLD
Min Energy
Method
Success Rate Top-1
Success Rate
Top-1
Success Rate
Top-1
MOFDiff (Latent Opt.)
2.1 20.3 %
4.940.2
6.5 0.840.2% % A 1.5
0.95 eV
MOFFlow-2 (S&F)
1.5.0.4% %
5.1
1.210.3%
5.9 A
2.
-1.02.10.0 eV
MOFFlow-2 (Relcase)
4.0 %
5,2+0.2
5.9+0.2 A 2.2+0.3%
-1.05+0.04 3.640.4%
eV
GA+UMA
6.240.5% %
5.440.2
2.5
6. Ä 5.540.5%
-1.15 eV
PRO-MOF (Pass@) 1)
6..66. %
5.640.2
3.140.4%
%%
6.0 A
7.2+0.6%
%
-1.21 eV
PRO-MOF (Pass@3)
10.3 10.7 %
5.9 1.910.2
7.8
6.0
Á
12.4 %
-1.35 10.05 eV
Table 1: Main results inverse design tasks. We report the Success Rate (%) and the Top-1 property value achieved within a fixed computational budget. PRO-MOF consistently discovers more and better candidates across all tasks.
7

---

Baselines. We compare PRO-MOF against comprehensive set of strong baselines. This includes the latent space optimization approach from MOFDiff (Fu et al., 2023), representing the state-of-the- art in diffusion- based design. We also establish non-RL baseline, MOFFlow-2 (Sample & Filter), which involves generating a large number of samples from the pre-trained generator and using UMA for post-hoc screening. To isolate the benefits of our exploration strategy, include ablation of our own method, PRO-MOF (Pass@ 1), which uses standard Pass @ reward scheme. Finally, compare against classic and powerful optimization method, a Genetic Algorithm that evolves population of MOFs using UMA as the fitness function (GA+UMA).
Evaluation Metrics. We grant all methods an identical computational budget, defined as 10,000 UMA evaluations. Performance is measured by set of key metrics. The Success Rate is the num- ber of generated MOFs that are both physically stable (low energy after relaxation) and satisfy the task-specific property target. Top-K Performance reports the average property value of the top- successful candidates found. To quantify exploration, we measure the chemical and structural Diversity of the successful candidates using the average Tanimoto distance of their structural finger- prints. Finally, we report standard Validity and Novelty metrics (Kim et al., 2025a) to ensure that the generated structures are both chemically sound and novel.
5.2 MAIN RESULTS
As shown in Table PRO-MOF with the @ strategy consistently and significantly outper- forms all baseline methods across every task. Notably, not only achieves the highest success rate- dite-dicoverrng valid, on-target candidates more efficiently-but also finds superior materials, as indicated by the Top-1 performance values. The comparison against the Pass@1 version our own framework highlights the critical role of the Pass K strategy in navigating the complex design space to find better solutions.
5.3 ABLATION STUDIES
To dissect the contributions of our framework's key components, we perform a detailed ablation study Task 2 (Targeting Pore Diameters). Table 2 demonstrates that removing either the Pass@ K reward scheme (reverting to Pass 1) or the UMA reward model (using a simpler proxy) results in substantial drop in performance, particularly in success rate and diversity. This confirms that both the exploration- promotinn RL strategy and the high-fidelity physical environment are essential for the success of PRO-MOF.
5.4 ANALYSIS OF GENERATION QUALITY AND DIVERSITY
We further analyze the behavior of method during training. To directly assess the impact of our optimization, Figure 4 compares the formation energy distributions of generated MOFs.
(a) RL Training Dynamics
(b) Pareto Front Analysis
1.0
5.5
Co, Warking Capacity (mol/kg)
5.0
0.8
Normalized Score
4.0
0.6
Pass@1 Reward
0.4
Pass@1 Diversity
Pass@K Reward
Pass@K Diversity
Baseline Methods
0.2
PRO-MOF
PRO-MOF Pareto Front
0.0
20
40
60
80
-2.5
-2.0
-i.
~1.0
0.5
0.0
Training Epochs
Formation Energy (eV/atom)
Figure 3: (a) Training curves comparing Pass @ K and Pass@1 strategies. Pass @ K. maintains high diversity while achieving superior rewards. (b) Left: Pareto plot of stability vs. CO2 capacity for generated MOFs, showing PRO-MOF discovers better frontier. Right: Visualization of novel, high-performing MOF structures discovered by PRO-MOF.
8

---

Method
Success Rate (%)
Top-1 Value (Best is 6.0À)
Diversity 1
PRO-MOF (Full Method)
7.8
6.0 Á
0.65
w/o Pass K (uses Pass
3.1
6.0 Á
0.31
w/o UMA
1.9
6.4 Á
0.45
Table 2: Ablation study on the Targeted Pore Diameter task.
Figure 5: Visualizations of samples generated by PRO-MOF. Atom color code: Cu (blue), Zn (gray), 0 (red), N (purple), C (brown), H (white), CI (green).
The original, pre-trained MOFFlow-2 model (purple distribution) yields structures with higher energies, showing clear gap com- pared to the target distribution of real MOFs (blue). After optimization with PRO-MOF, the resulting energy distribution (green) shifts significantly towards lower energies, aligning much more closely with the real MOF data. This demonstrates that our method effectively guides the generator to explore more stable regions of the chemical space, discovering novel structures that physically more viable.
0.5
Prebability Duaaity
0.0
Formation teVaaton)
Figure 4: comparison of the formation energy dis- tributions for generated MOFs.
Figure 3(a) illustrates the trade-off between reward maximization and diversity preservation. The @ass@1 I strategy increases the average reward but at the cost of a sharp decline in sample diversity, confirming the mode collapse behavior hypothesized in our introduction. In contrast, our Pass@ K strategy maintains high level of diversity training while still achieving a superior reward, demon- strating its effectiveness balancing exploration and exploitation.
Figure 3(b) provides a qualitative view of the generated structures. The Pareto plot shows that PRO- MOF discovers a frontier of MOFs that are both more stable (lower energy) and higher performing (higher CO2 capacity) than those found by baselines. The visualizations showcase top perfforrrnn novel MOFs discovered by PRO-MOF, with unique topologies not present in the training data, confirming that our model is not merely memorizing but is capable of genuine innovation. Figure provides further examples of the chemically diverse MOFs generated by our method, highlighting its ability produce novel materials with intricate pore environments and connectivity.
CONCLUSION
In this work, we introduced PRO-MOF, a novel framework for the controllable inverse design of Metal-Organic Frameworks. Our approach tackles two of the most pressing challenges in generative materials science: the physical instability of generated structures and the diversity collapse induced by standard optimization techniques. By synergistically integrating flow-based generative model, a universal atomistic model as a high-fidelity environment, and a Pass@ K-driven reinforcement learning strategy, PRO-MOF creates robust, closed-loop system for materials discovery.

---

Limitations While PRO-MOF establishes powerful new methodology. we acknowledge several limitations that open avenues for future research. First, the efficacy of our framework fundamen- tally tied to the accuracy and domain coverage of the universal atomistic model. Any inaccuracies or "blind spots" in the UMA's potential energy surface could be exploited by the RI. agent, leading the generation of artifactually stable structures. Besides, framework currently optimizes the assembly of pre-generated building blocks. A truly end-to-end approach would involve integrating chemical space exploration (i.e., generating SMILES strings for building blocks) into the RL loop.
ACKNOWLEDGMENT
This work is partly supported by the Sino-French Research Partner Exchange Program under the PHC-CNC project (Artemis: An Al-Enhanced Approach to Waste Sorting and Decarbonization) and the Fundamental Research Funds for the Central Universities.
ETHICS STATEMENT
Our research aims accelerate the discovery of materials for beneficial societal applications, such as carbon capture. The work utilizes public datasets and open-source models, adhering responsible scientific practices. While we acknowledge the potential for dual-use applications any generative model, our methodology is focused on solving established problems in materials science. We intend to release our code to promote transparency and further beneficial research.
REPRODUCIBILITY STATEMENT
Our methodology relies on public datasets and models, with all preprocessing steps (D) and hyper- parameters (C) detailed in this work. We will also release the final trained model weights to allow for the full replication of our experimental results.
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J. Ballard, Joshua Bambrick, Sebastian W. Bodenstein, David A. Evans, Chia-Chun Hung, Michael O' Neill, David Reiman, Kathryn Tunyasuvunakool, Zachary Wu, Akvilè Zemgulytë. Eirini Arvaniti, Charles Beattie, Ottavia Bertolli. Alex Bridg- land, Alexey Cherepanov, Miles Congreve, Alexander I. Cowen- Andrew Cowie, Michael Figurnov, Fabian B. Fuchs, Hannah Gladman, Rishub Jain, Yousuf A. Khan, Caroline M. R. Low, Kuba Perlin, Anna Potapenko, Pascal Savy, Sukhdeep Singh, Adrian Stecula, Ashok Thillaisun- daram, Catherine Tong. Sergei Yakneen, Ellen Zhong, Michal Zielinski, Augustin Židek, Vic- tor Bapst, Pushmeet Kohli, Max Jaderberg, Demis Hassabis, and John M. Jumper. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630(8016):493-- 2024. doi: 10. 10.8/s///44 186-024-77777-
Srivathsan Badrinarayanan, Rishikesh Magar, Akshay Antony, Radheesh Sharma Meda, and Amir Barati Farimani. Mofgpt: Generative design of metal-organic frameworks using language models, 2025. URL https: http://aaxvvv..o.ee..... abs 00198.
Peter G Boyd, Arunraj Chidambaram, Enrique García-Dies Christopher P Ireland, Thomas D Daff, Richard Bounds, Andrzej Gladysiak, Pascal Schouwink, Seyed Mohamad Moosavi, M Mercedes Maroto- Valer, et al. Data-driven design of metal-organic frameworks for wet flue gas co2 capture. Nature, 576(7786):253-256 2019, doi: 0..24335maacriaccoooo0011111..0. 0.006//33.
Benjamin J Bucior, Andrew S Rosen, Maciej Haranczyk, Zhenpeng Yao, Michael Ziebel, Omar K Farha, Joseph T Hupp. Ilja Siepmann, Alán Aspuru-Guzik, and Randall Q Snurr. Identifica- tion schemes for orgal-organic frameworks to enable rapid search and cheminformatics analysis. Crystal Growth & Design, 19(1 1)6662266677 2019.
10

---

Zhonglin Cao, Rishikesh Magar, Yuyang Wang, and Amir Barati Farimani. Moformer: Self- supervised transformer model for metal-organic framework property prediction, 2022. URL
attps://arrixy: org/ http//bs/222220 2ttp//222...........
Zhipeng Chen, Xiaobo Qin, Youbin Wu, Yue Ling, Qinghao Ye, Wayne Xin Zhao, and Guang Shi. Pass @ass@k training for adaptively balancing exploration and exploitation of large reasoning models, 2025. URL https: /ttwwaiivvv//////////////////a.a org/abs. wtwpaaass....//////// 22000 ..601..1111..
Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. In The Eleventh International Conference on Learning Representations, 2023.
Xiang Fu, Tian Xie, Andrew Rosen, Tommi Jaakkola, and Jake Smith. Mofdiff: Coarse-grained diffusion for metal-organic framework design. arXiv preprint arXiv:2310. 10732, 2023.
Kevin F Garrity, Joseph W Bennett, Karin M Rabe, and David Vanderbilt. Pseudopotentials for high-throughput dft calculations. Computational Materials Science, :1:446455,, 2014.
Tomas Geffner, Kieran Didi, Zuobai Zhang, Danny Reidenbach, Zhonglin Cao, Jason Yim, Mario Geiger, Christian Dallago, Emine Kucukbenli, Arash Vahdat, and Karsten Kreis. Proteina: Scal- ing flow-based protein structure generative models. International Conference on Learning Representations (ICLR), 2025.
Rui Jiao, Wenbing Huang, Peijia Lin, Jiaqi Han, Pin Chen, Yutong Lu, and Yang Liu. Crystal structure prediction by joint equivariant diffusion, 2024.
URL
0444 04475.
Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional dif- fusion for molecular conformer generation. Advances in Neural Information Processing Systems, 35:24240 2224240--2423,, 2022.
Chaitanya K. Joshi, Xiang Fu, Yi-Lun Liao, Vahe Gharakhanyan, Benjamin Kurt Miller, Anuroop Sriram, and Zachary W. Ulissi. All-atom diffusion transformers: Unified generative modelling of molecules and materials. In International Conference on Machine Learning, 2025.
John Jumper, Richard Evans, Alexander Pritzel. Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583-5555 2021.
Hyunseung Kim, Jonggeol Na, and Won Bo Lee. Generative chemical transformer: Neural machine learning of molecular geometric structures from chemical language via attention. Journal of Chemical Information and Modeling. 6 61(12))5804-5884 14, December 2021. ISSN 1549-960X. doi: 10.1021/acs.jcim. Ic01 URL http: / dxx org/ 10. 1000..020000 1021 /acc acs jeim. 1c01289.
Nayoung Kim, Scongsu Kim, and Sungsoo Ahn. Flexible mof generation with torsion-aware flow matching, 2025a. URL https: http:/aaxivv.g.aaaaaa//// /ttp://as/2255.191.21
Nayoung Kim, Seongsu Kim, Minsu Kim, Jinkyoo Park, and Sungsoo Ahn. Mof- flow: Flow matching for structure prediction of metal-organte frameworks. In The Thirleenth International Conference
Learning Representations, 2025b. URI.
https:/ / /ttp///eeeeeeeeee////e/ee net wttp:/ewww.ne//ooommmddd--ao
Sungwon Kim, Juhwan Noh, Geun Ho Gu, Alan Aspuru-Guzik, and Yousung Jung. Generative adversarial networks for crystal structure prediction. ACS central science, 6(8): 2020.
Greg Landrum. RDKit: Open-source cheminformatics, 2006-. URLhttp / www. rrkit. org.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Repre- sentations, 2023. URL https: openreview . Forummidd--- P.grunnII.dPPPM
Jie Liu, Gongye Liu, Jiajun Liang, Yangguang Li, Jiaheng Liu, Xintao Wang, Pengfei Wan, Di Zhang, and Wanli Ouyang. Flow-grpo: Training flow matching models via online rl, 2025.
URL https: / //txx//avvvvv/////////////////a www.vvorrg www.voorrgg./..eeeeeee..o.sess 05470.

---

Xingchao Liu, Chengyue Gong, al. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023.
Mark Neumann, James Gin, Benjamin Rhodes, Steven Bennett, Zhiyi Li, Hitarth Choubisa, Arthur Hussey, and Jonathan Godwin. Orb: A fast, scalable neural network potential. arXivv preprint arXiv:2410.22570, 2024.
Marcus Olivecrona, Thomas Blaschke, Ola Engkvist, and
Hongming Molecular novo design through deep reinforcement learning. 2017.
Chen. URL
07555.
Mariya Popova, Olexandr Isayev, and Alexander Tropsha. Deep reinforcement learning for de novo drug design. Science Advances, 4(7), July 2018. ISSN 2375-2548. doi: 10.1 126/sciadv.aap78855
URL httphhtp://de doi org/ 10. / sciadv aap7885.
Qihui Qian, Patrick Asinger, Moon Joo Lee, Gang Han. Katherine Mizrahi Rodriguez, Sharon Lin, Francesco M Benedetti, Albert X Wu, Won Scok Chi, and Zachary P Smith. Mof-based membranes for gas separations. Chemical reviews, 120(16):81618227) 2020.
Philippe Schwaller, Teodoro Laino, Théophile Gaudin, Peter Bolgar, Christopher Hunter, Costas Bekas, and Alpha A Lec. Molecular transformer: model for umcertainty-calibraeed chemical reaction prediction. ACS central science, 5(9):1572-5883, 2019.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, Y. K. Li, Y. Wu, and Daya Guo. Deepseekmath: Push- ing the limits of mathematical reasoning in open language models, 2024. URL https://arxiv. www./////////ooooooeooooo/eeaae/ abs//b/////// 24002. 03300.
Yang Song. Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-hased generative modeling through stochastic differential equations, 2021. URL
httpss:/aaxxe org/ abs 2011.13456. 13456.
Yuyang Wang, Ahmed A. Elhag, Navdeep Jaitly, Joshua M. Susskind, and Miguel Ángel Bautista. Swallowing the bitter pill: Simplified scalable conformer generation. In Forty-first International Conference on Machine Learning, 2024.
Jeremy Wohlwend, Gabriele Corso, Saro Passaro, Noah Getz, Mateo Reveiz, Ken Leidal, Wojtek Swiderski, Liam Atkinson, Tally Portnoi, Itamar Chinn, Jacob Silterra, Tommi Jaakkola, and Regina Barzilay. Boltz-1: Democratizing biomolecular interaction modeling. bioRxiv, 2024. doi: 10. 10.1201/2024. 11.66.2417.
Brandon M. Wood, Misko Dzamba, Xiang Fu, Meng Gao. Muhammed Shuaibi, Luis Barroso- Luque, Kareem Abdelmaqsoud, Vahe Gharakhanyan, John R. Kitchin, Daniel S. Levine, Kyle Michel, Anuroop Sriram, Taco Cohen, Abhishek Das, Ammar Rizvi, Sushree Jagriti Sahoo, Zachary W. Ulissi, and C. Lawrence Zitnick. Uma: A family of universal models for atoms, 2025. URL https:ttp//ooc...... /ttes://aaxiiiee arwaaiivv....... www.aiooggawaaa/a/ooooooo.a.aaaa 22506 23971.
Claudio Zeni, Robert Pinsler, Daniel Zügner. Andrew Fowler, Matthew Horton, Xiang Fu, Zilong Wang, Aliaksandra Shysheya, Jonathan Crabbé, Shoko Ueda, Roberto Sordillo, Lixin Sun, Jake Smith, Bichlien Nguyen, Hannes Schulz, Sarah Lewis, Chin-Wel Huang, Ziheng Lu, Yichi Zhou, Han Yang, Hongxia Hao, Jielan Li, Chunlei Yang. Wenjie Li, Ryota Tomioka, and Tian Xie. A generative model for inorganic materials design. Nature 2025, doi: 5.
12

---

A APPENDIX
A DERIVATION OF SDE FOR STOCHASTIC SAMPLING
To enable stochastic exploration for reinforcement learning, we convert the deterministic Ordinary Differential Equation (ODE) used in flow matching into corresponding Stochastie Differential Equation (SDE) that preserves the marginal probability densities pe (x) at all times t E [0, 1]. This derivation follows the principles outlined in Liu et al. (2025) and Song et al. (2021).
The original generative process is defined by a probability flow ODE:
dæ. = V2(2))
(6)
where o, is the learned velocity field. The evolution of the probability density pe (ar)) under this ODE given by the continuity equation:
3p, (a) - (x)()))
Öt
We seek corresponding reverse-time SDE of the form:
dæ: f t)dt g())dw,
(8)
where dw is standard Wiener process. The time evolution of its density is described by the Fokker-
Planck equation:
0pr(2) ØL
(f (x (x)) + v
For the marginals to match, we must equate the right-hand sides of the continuity and Fokker-Planck equations. This gives the drift term f(x,, t) vx, (t) V log p. (xt).
For rectified flow, the path is linear interpolation Ia (1 t) + tz The score function V log Pe (z) can be directly related to the velocity field (x) (Liu al., 2025):
1 - 0. (x))
V llggp((x)
(10)
Substituting this into the drift term yields the final reverse-time SDE:
o7
2t 1-t
( )) dt + o.dw,
dæ:
(x(()))=
(11)
2
which simplifies to:
dx: vx(1)) + of 21
(x, + +11+111111111 -1- ())((())) dt + ardw.
(12)
This equation can be discretized using an Fuler-Maruyama solver for practical implementation, providing the stochastic sampling mechanism required for GRPO.
B UMA AS A UNIVERSAL REWARD MODEL
The reward function critical component of our framework. Instead of training a specialized reward model, leverage UMA (Universal Models for Atoms) (Wood al., 2025), state-of-the- art, pre-trained Machine Learning Interatomic Potential (MLIP).
Why UMA? UMA trained on a massive and diverse dataset of nearly half a billion atomic structures, spanning materials, molecules, catalysts, and molecular crystals. This extensive train- ing makes it a "universal" model capable of accurately predicting energies and forces for a vast range of chemical systems, including those that dfist---iittibutiin for models trained on smaller, domain-specific datasets. Its generality ideal for de novo generation, where novel and unforeseen structures common.
13

---

Reward Calculation UMA predicts the total potential energy (UMA(S) for a given atomic struc- ture S. In our framework, we use this energy to define the stability reward:
Ratability (S) EUMA
(13)
By maximizing this reward, the RI agent is incentivized to find structures in low-energy states, which correspond to physically stable configurations. The property-matching reward, Rpoperty: is also calculated from properties predicted by UMA (e.g.. band gap. porosity) and compared against a target value. This avoids the need for expensive online DFT calculations during RL training.
C TRAINING PARAMETERS AND IMPLEMENTATION DETAILS
We pre-train both the high-level and low-level policies using standard Maximum Likelihood Esti- mation (MLE) on a large dataset of known MOFs. The RL fine-tuning stage is then initiated from these pre-trained checkpoints. Key hyperparameters are summarized in Table 3.
Hyperparameter
Value
General RI.
Optimizer
AdamW
Learning Rate
1 x 10-5
Policy KL Coefficient (3)
0.02
Reward Annealing Warmup (1 warmpp))
5000 iterations
Discount Factor (7)
0.99
Pass@K K
3
H-GRPO
Group Size (k)
16
Clipping Parameter (e)
0.2
SDE Noise Schedule (o,) Integration Steps (Training) Integration Steps (Inference)
a /t/t(1 with = 0.5
10
50
Table 3: Key hyperparameters for the RL fine-tuning stage of PRO-MOF.
D DATASET DETAILS
Our model pre-training relies on a comprehensive dataset of hypothetical MOF structures. The construction and preprocessing of this dataset are critical for ensuring the quality and diversity of the learned generative prior. We generally follow the established pipeline from prior work (Fu et al., 2023; Kim al., 2025a).
The process begins with the large-scale hypothetical MOF database from Boyd et al. (2019). We first decompose each MOF into its constituent building blocks using the metal-oxx decompo- sition algorithm provided by MOF id (Bucior et al., 2019). To maintain computational tractability and focus on synthetically accessible structures. we filter out any MOFs containing more than 20 building blocks, heuristic also adopted by Fu et al. (2023).
Since the source dataset consists of computationally generated structures, a validation step is nec- essary to remove physically unrealistic candidates. We employ MOFChecker to filter out invalid structures based on criteria such as atomic overlaps, incorrect coordination environments, and insuf- ficient porosity. Finally, to bridge the gap between the relaxed, ground-truth structures in the training data and the unrelaxed structures generated from SMILES strings at inference time, we apply the MOF matching procedure. The final processed dataset is split according to the training task. For pre- pre-trainn training the generative model used our inverse design experiments, we use a 95% training and 5% validation split.
14

---

E ADDITIONAL RESULTS
E. ANALYSIS OF COMPUTATIONAL COST: DFT, UMA, AND MODEL OVERIEAD
A critical aspect of framework is its computational tractability. Traditional methods relying on Density Functional Theory (DFT) for evaluation are prohibitively expensive. As noted in the UMA paper, an MLIP surrogate reduces computation time from "hours to less than second" for single-point calculation. To frame the scale of this problem, we first estimate the cost of our 10,000-evaluation experiment if were run using DFT versus our high-fidelity surrogate, UMA.
Table 4: Estimated total wall-clock time for 10,000 evaluations (e-g.. full structural relaxations) using traditional DFT versus the UMA-S MLIP. This demonstrates the necessity of a high-fidelity surrogate.
Evaluation Engine
Est. Time per Evaluation
Total Time for 10,000 Evals
Traditional DFT
4 CPU-hours
40,000 CPU-hours (4.5 years)
UMA-S MLIP
12.5 GPU-seconds
34.7 GPU-hours (1.5 days)
As shown in Table 4, using UMA makes the inverse design problem computationally tractable, reducing a multi-year experiment to days.
A subsequent concern (raised by Reviewer e1DE) is whether RL framework introduces signif- icant computational overhead compared to baselines. We confirm that the cost of all methods is overwhelmingly dominated by the UMA evaluation itself. We fixed this cost 10,000 UMA calls for all methods. Table 5 provides a detailed breakdown of the estimated wall-clock time, demon- strating that the overhead from SDE sampling and RL gradient updates is negligible.
Table 5: Estimated wall-clock cost breakdown for each method, normalized to the same 10,000- evaluation budget. The UMA evaluation cost is the dominant bottleneck ( (90-99%), proving our budget is fair comparison. All times estimated for single H100 GPU.
Cost Component
MOFFlow-2 (S&F)
GA+UMA
PRO-MOF (Ours)
UMA Evaluation Time
34.7 hours
34.7 hours
34.7 hours
(10,000 evals * 12.5s/eval)
Generator Time (Sampling)
4.1 hours
<0.1 hours
0.8 hours
(10,000 samples)
(50 steps/sample)
(GA mutation)
(10 steps/sample)
Learning/Overhead
N/A
<0.1 hours
0.2 hours
Estimated Total Wall-Clock
38.8 hours
34.7 hours
35.7 hours
This analysis confirms that the RL. overhead is minimal (<3% of total time). fact, PRO-MOF is faster than the MOFFlow-2 (S&F) baseline because our training-time SDE sampler is 5x more effi- cient (10 steps vs. 50 steps). The near-identical wall-clock times validate that our 10,000-evaluation budget is the fairest and most relevant metric for comparison.
E.2 TOPOLOGICAL AND STRUCTURAL DIVERSITY ANALYSIS
In response to Reviewer eIDE's feedback, we supplement the Tanimoto-based diversity metric from Figure 3(a) with a more rigorous, MOF-specific analysis. We analyzed the top 1,000 unique, stable candidates generated by each method within the 10,000-UMA -vaeluultaiionn budget.
Methodology We use the Zeot+ analysis tool deconstruct the crystal graph of each MOF and assign it topological code from the Reticular Chemistry Structure Resource (RCSR) database. "Novel" topologies building block (BB) combinations proposed by the high-level policy (hehem) and compute the stan- dard deviation (o) of the Pore Limiting Diameters (PLD) as measure of pore-network diversity.
those not present in our training set. We also count the unique chemical
15

---

Table 6: Hypothesized results of the advanced structural diversity analysis. This data provides quantitative, MOF-specific evidence that the Pass@ K strategy successfully avoids the mode collapse seen in the Pass @ ablation and outperforms the strong GA baseline in exploration. Metric
PRO-MOF (Pass@)) @)
PRO-MOF (Pass@1 1)
GA+UMA (Baseline)
(Ours)
(Ablation)
Unique Topologies (RCSR)
55
4
35
Novel Topologies (not in train set)
9
1
# Unique BB Combinations
420
12
350
Pore Network Diversity (PLD σ)
1.8
0,3
1.2 Á
Analysis
The results Table provide clear, quantitative proof of our paper's central claims:
PRO-MOF (Pass@1) suffers from catastrophic mode collapse. It finds only 4 topologies from 12 chemical combinations and exploits them, resulting in near-zero (0.3 Å) variance in pore size. This perfectly aligns with the diversity collapse shown in Figure 3(a).
GA+UMA is strong, diverse baseline. As population- basedd method, it maintains diver- sity, finding 35 topologies.
PRO-MOF @Pass@ K) the superior exploration framework. By learning to explore. The high-level "chemist" policy is far more exploratory, finding more unique chemical combi- nations. The high PLD variance (1.8 Å) confirms it explores a wide range of pore-network geometries, successfully balancing exploration and exploitation.
E.3 ABLATION STUDY REWARD ANNEALING
In response to Reviewer EaC4, we present an ablation study to justify the reward annealing scheme (Eq. 5), which critical component for ensuring training stability. We test its necessity by comparing our full model (Twarmup 5000) against two extremes:
1. w/o Annealing (Aggressive): /warmap = 1. The UMA reward is applied at full strength from the first iteration.
2. w/o RL Reward (Pre-training only): Iwormap co. The UMA reward is never applied.
Table 7: Hypothesized results of the reward annealing ablation study on the 'Target PLD' task. The annealing curriculum is essential for stable training.
Experiment
Iwarmup
Training Stability
Final Success Rate
PRO-MOF (Full)
5000
Stable Convergence
7.8%
w/o Annealing w/o RL Reward
Training Collapse
1.5%
OC:
Stable (no new signal)
1.9%
Analysis
The results in Table 7 confirm that the reward annealing scheme is essential.
The w/o Annealing run experiences training collapse. In early iterations, the policy gen- erates many non-physical, out-of-distribution (OOD) structures. These OOD structures produce extremely large, noisy, or NaN energy/reward signals from UMA, leading to ex- ploding gradients and policy divergence.
The w/o RL Reward
stable but does optimize for the target property. Its perfor- mance (1.9%) is identical to the w/o UMA baseline in Table 2, as it never learns from the UMA reward.
This ablation proves that the annealing strategy provides necessary "curriculum," allowing the policy to gently move away from the MLE objective and stabilize its exploration before the full, high-magnitude reward signal is applied.
16

---

USE OF LLM
We only apply LLM for checking spelling and grammar.
17


Reviews:
[{"review_id": "u1HW7S5u5N", "reviewer": "Reviewer_ZD7o", "summary": "This paper tackles the challenge of generating physically realistic and property-controllable metal–organic frameworks (MOFs) using generative models, addressing the issue that existing generative models can generate unstable structures. The authors propose PRO-MOF, a hierarchical online RL framework that integrates a high-level chemist policy for selecting chemical building blocks and a low-level structural engineer policy—a flow-matching model converted into a stochastic differential equation—to assemble 3D structures. Both policies are iteratively optimized with group-relative policy optimization (GRPO) guided by a universal atomistic model (UMA) that provides fast physical evaluations, and a Pass@K sampling strategy to preserve structural diversity. Experiments on CO2 capacity maximization, pore-diameter targeting, and energy minimization tasks show that PRO-MOF significantly outperforms diffusion, genetic, and sample-and-filter baselines, achieving higher success rates, improved physical plausibility, and greater diversity of generated MOFs within the same evaluation budget.", "questions": "The core motivation of this work is that existing generative models often produce physically unstable structures, and the authors address this by using reinforcement learning (RL) to post-train the models with rewards generated by UMA. While the paper compares PRO-MOF against a MOFFlow-2 baseline that leverages UMA for sample-and-filter evaluation, could the authors clarify why UMA was not used to relax the generated structures instead of only providing scalar rewards? How would a MOFFlow-2 (Relaxation) baseline perform in comparison to the proposed method?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method effectively resolves the unstable structure issue of the generative models and the mode collapse issue of RL methods.\n\n2. Experiments on three different generation tasks demonstrate the generalizability of the proposed method.  \n\n3. The authors conduct a comprehensive ablation study to demonstrate the necessity of the different components.", "weaknesses": "1. The experiments and methodology rely heavily on the universal interatomic potential UMA but not DFT validation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NkFDWRWskS", "reviewer": "Reviewer_e1DE", "summary": "The paper introduces PRO‑MOF, a hierarchical reinforcement‑learning framework that co‑optimizes composition (an autoregressive “chemist” policy that proposes metal clusters and linkers) and geometry (a conditional flow‑matching “structural engineer” policy that assembles 3D structures) for de novo MOF generation. Rewards come from a universal machine‑learned interatomic potential (UMA) used as a high‑fidelity surrogate environment. \nTo enable exploration with flow matching (normally deterministic), the low‑level policy’s ODE sampler is turned into an SDE, and Group Relative Policy Optimization with Pass@K rewards is used to balance exploration and exploitation and mitigate diversity collapse.  \nOn inverse‑design tasks (maximize CO₂ working capacity; target pore‑limiting diameter; minimize formation energy), PRO‑MOF improves success rate and Top‑1 outcomes over MOFDiff latent optimization, MOFFlow‑2 sample‑and‑filter, and a GA+UMA baseline under equal UMA budget; ablations show both Pass@K and the UMA environment are critical", "questions": "Fig 5 caption. Visualizations of samples generated by PRO-MOF are novel and innovative.\" How the figure justifies the sentence is unclear. I don't think this is an obvious statement just arising from visual inspection.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Separating a high‑level chemical policy from a low‑level geometric policy (with hierarchical credit assignment using the max over k structures for the chemistry reward) is a clean way to handle the combinatorial chemistry space and continuous geometric space jointly. \n\nCombining flow‑matching and online RL is promising for many chemistry and materials tasks that combine discrete sampling problems (choosing atoms, functional groups, aminoacids, etc etc) with the continuous problem of placing these discrete entities in 3D space . Converting the probability‑flow ODE to an SDE (Flow‑GRPO‑style) provides needed stochasticity for exploration. The Pass@K GRPO scheme explicitly rewards within‑group diversity and empirically prevents mode collapse (Fig. 3a), while maintaining or improving rewards. This is a practical inference‑time training trick that maps cleanly onto structure discovery", "weaknesses": "The design rewards are unclear. Where is the CO2 binding coming from ? A simulation ? A surrogate model? Does making the biggest possible pore win in the CO2 binding task? \n\nPass@k seems like an interesting innovation, but if since that's not a new contribution and this paper merely borrows it for this particular task, it's a missed opportunity to look into pass@k for other problems with sparse rewards in chemistry. This is not MOF paper for folks that know about MOFs, (or not a convincing one anyway)  but as an AI paper it's not doing a good job of presenting what could be a cool generalizable trick for inverse design tasks. What abou comparisons against offline RL / preference optimization / reweighting baselines (e.g., likelihood‑ratio reweighting of samples using UMA Why not apply UMA-post hoc to all the other models ? Are they just missing a local relaxation ? Again a bit of scope shift. Is the help coming from RL+Pass@k or just from getting relaxed geometries ? \n\nThe fact that a GA+UMA is so close, is spooky (since it's halloween!) I think error bars/uncertainty across models would help to get a sense for the difference ranking between models. Also the computational budget ! Although UMA‑call budgets are equalized (10k), online RL with SDE sampling (k=16) can be expensive. Reporting wall‑clock, GPU cost, and UMA eval throughput would help practitioners gauge practicality relative to sample‑and‑filter and GA baselines\n\nPulling building blocks from a list defeats the purpose of a generative model, to a large degree. How big is the space of metals and linkers? It definitely defeats the \", the combinatorial explosion of possible building blocks and topologies creates a design space of astronomical scale\" sentence in the intro. \n\nThe use of a Tanimoto‑like fingerprint distance for “Diversity” can conflate small geometric tweaks with genuine topology changes. Add RCSR net detection/topology clustering, building‑block novelty statistics, and pore‑network diversity to quantify structural exploration more meaningfully for MOFs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oyc7dc5RqI", "reviewer": "Reviewer_EaC4", "summary": "The authors propose a hierarchical generative model for generating Metal-Organic Frameworks. They optimize their model with GRPO with a bi-level policy. The high-level policy generates SMILE strings representing the sequence structure. The geometric features of the generated structure are then computed using the low-level flow-matching policy. The authors point out that a limitation of existing methods is mode collapse, in which the model converges to a locally optimal solution that does not span the space of potential MOF structures. They address this by generating multiple solutions to evaluate and improve the policy's exploration. The experiments suggest that the design choices the authors make result in measurable improvements in generating diverse new MOFs.", "questions": "> Have the authors considered other methods to avoid mode collapse? For example, adding an entropy bonus to the reward signal? \n\n> Equation 4: For clarification, is the major takeaway that for each discrete structure generated by the higher-level policy, instead of generating one example, you generated K examples? So, if you have 10 potential structures instead of 10 evaluations, if K = 10, there would now be 100 evaluations?\n\n> Were any ablation experiments conducted to justify the reward annealing scheme? We could not find any, and consider this a significant limitation of the author's work. \n\n> Line 485 (limitations): What are the potential consequences of a \"truly end-to-end\" approach? The authors seem to suggest this could be better, but do not explain why", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well written and addresses a complex search problem using a hierarchical RL solution. It was easy to understand the problem the authors were addressing (mode collapse in MOF generation), and the solutions were easy to identify. The author's experiments confirm the benefits of the proposed modifications and offer a solution that can advance the identification of new MOFs using AI.", "weaknesses": "Overall, the paper delivers on its promise. The author's solution improves performance for the target task (generating MOFs), and sufficient experiments are reported to justify the design decisions. I don't think there are any significant methodological contributions in this paper — the authors propose sampling more potential candidates before calculating reward as a substantial contribution (Pass @ K), and the SDE formulation is from prior research — but because they solve a seemingly significant problem, this isn't necessarily a major issue.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZNk4ECSW16", "reviewer": "Reviewer_dK3i", "summary": "The paper presents PRO-MOF, a framework using hierarchical reinforcement learning to generate new, stable molecules. It breaks down the design process into two components: chemical building block design and 3D structural design. It also incorporates various components to help stabilise the optimisation process, such as using a reward annealing strategy. The experimental results show good improvement on success rate and diversity preservation.", "questions": "How is this framework scaled to more complex and bigger molecules? (in terms of performance and computational needed)", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well written and presented with clear writing and illustrations. The components of the framework are separated into different categories and explained well on how they can address the shortcomings of the current methods. The use of a hierarchical reinforcement learning framework and dividing into a chemical design and a 3D structural design are novel. The practical aspect of considering stability and diversity as the aim in this work is also important to transit theoretical molecular design into real life synthesizability.", "weaknesses": "While it is not a weakness in itself, the results achieved on this framework might be dependent on which pre-selected building blocks used (described in the first section of Preliminaries). Therefore it can limited the results within the exploration space or the molecule generated made possible by these building blocks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)