Reference-guided Policy Optimization for Molecular Optimization via LLM Reasoning

Xuan Li; Zhanke Zhou; Zongze Li; Jiangchao Yao; Yu Rong; Lu Zhang; Bo Han

cs.LGMar 6, 2026

Reference-guided Policy Optimization for Molecular Optimization via LLM Reasoning

Authors: Xuan Li, Zhanke Zhou, Zongze Li, Jiangchao Yao, Yu Rong, Lu Zhang, Bo Han

Organizations: Hong Kong Baptist University · CMIC, Shanghai Jiao Tong University · DAMO Academy, Alibaba Group · Hupan Lab

Abstract

Large language models (LLMs) benefit substantially from supervised fine-tuning (SFT) and reinforcement learning with verifiable rewards (RLVR) in reasoning tasks. However, these recipes perform poorly in instruction-based molecular optimization, where each data point typically provides only a single optimized reference molecule and no step-by-step optimization trajectory. We reveal that answer-only SFT on the reference molecules collapses reasoning, and RLVR provides sparse feedback under similarity constraints due to the model's lack of effective exploration, which slows learning and limits optimization. To encourage the exploration of new molecules while balancing the exploitation of the reference molecules, we introduce Reference-guided Policy Optimization (RePO), an optimization approach that learns from reference molecules without requiring trajectory data. At each update, RePO samples candidate molecules with their intermediate reasoning trajectories from the model and trains the model using verifiable rewards that measure property satisfaction under similarity constraints in an RL manner. Meanwhile, it applies reference guidance by keeping the policy's intermediate reasoning trajectory as context and training only the answer in a supervised manner. Together, the RL term promotes exploration, while the guidance term mitigates reward sparsity and stabilizes training by grounding outputs to references when many valid molecular edits exist. Across molecular optimization benchmarks, RePO consistently outperforms SFT and RLVR baselines (e.g., GRPO), achieving improvements on the optimization metric (Success Rate $\times$ Similarity), improving balance across competing objectives, and generalizing better to unseen instruction styles. Our code is publicly available at https://github.com/tmlr-group/RePO.

Figures & tables

Figure 1: Molecular optimization aims to optimize the given molecule by modifying its components while maintaining the structural similarity of the original molecule after modification. The molecule is presented as SMILES ( Weininger, 1988 ) , a sequence of symbols representing atoms and bonds.

Figure 6: Token partition and gradient application for SFT, GRPO, and RePO.

Task type	Objective	Metric	Base Model	Distill-SFT	SFT	GRPO	GRPO (SFT init)	RePO
Structure- based optimization	AddComponent	SR	0.086	0.100	0.238	0.005	0.246	0.307
		Sim	0.763	0.604	0.619	0.992	0.635	0.778
		SR $\times$ Sim	0.066	0.060	0.147	0.005	0.156	0.239
	DelComponent	SR	0.107	0.188	0.203	0.008	0.232	0.158
		Sim	0.864	0.682	0.755	0.994	0.759	0.887
		SR $\times$ Sim	0.092	0.128	0.153	0.008	0.176	0.140

Table 1: SuccessRate (SR), Similarity (Sim), and their product (SR

\times

Sim) on TOMG-Bench for structure- and property-based optimization tasks. The higher the better. The best results for each task are bolded , and the second-best results are underlined .

Task type	Objective	Metric	Base Model	Distill-SFT	SFT	GRPO	GRPO (SFT init)	RePO
Seen instruction	BDP	SR	0.052	0.078	0.398	0.156	0.088	0.206
		Sim	0.149	0.207	0.254	0.759	0.141	0.569
		SR $\times$ Sim	0.008	0.016	0.101	0.118	0.012	0.117
	BDQ	SR	0.034	0.022	0.319	0.082	0.022	0.160
		Sim	0.117	0.106	0.279	0.479	0.045	0.365
		SR $\times$ Sim	0.004	0.002	0.089	0.039	0.001	0.058

Table 2: Success Rate (SR), Similarity (Sim), and their product (SR

\times

Sim) on MuMOInstruct for seen and unseen instructions. The best results are bolded , and the second bests are underlined .

Panel of Figure 11 — Figure 11: Distribution of property gains on TOMG-Bench property-optimization tasks.

	Metric	Base Model	SFT	GRPO	GRPO (SFT init)	RePO
LogP	SR	0.237	0.287	0.197	0.192	0.360
	Sim	0.692	0.763	0.763	0.775	0.746
	SR $\times$ Sim	0.164	0.219	0.151	0.149	0.269
QED	SR	0.159	0.185	0.115	0.121	0.243
	Sim	0.722	0.809	0.805	0.811	0.783
	SR $\times$ Sim	0.115	0.150	0.093	0.098	0.190

Table 3: Performance comparison using Llama-3.1-8B Instruct on TOMG-Bench property optimization tasks. Best per row is bolded; second-best is underlined.

Table 11

Figure 16: Comparison of RePO (left) and GRPO (right). Notably, RePO employs chemically sound reasoning with valid substitutions, whereas GRPO yields incorrect reasoning and modifications.

Figure 17: Inference-scaling effect on the success rate and similarity metrics.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Category	Reward Shaping	Knowledge-guided Exploration	RePO
Guidance Mechanism	Utilizes environmental signals or LLM knowledge to guide the policy model.	Integrates professional tools (e.g., chemistry toolkits ( M. Bran et al., 2024 ) , databases ( Wang et al., 2025a ) ) to inject expert priors into workflows.	Employs demonstration molecules, offering end-to-end supervision that promotes free-form intermediate reasoning.
Signal & Supervision	Introduces dense intermediate signals with manually designed heuristics ( Xie et al., 2024 ; Chan et al., 2024 ) .	Injects expert guidance into the agent’s workflow through domain-specific signals.	Avoids intermediate rewards, eliminating the need for manually designed heuristics.
Application Domain	Commonly used in settings with sparse rewards, such as robotics and games.	Focused on solving domain-specific problems by leveraging external expert sources with tool integration.	Designed for enhancing reasoning performance in chemical tasks by narrowing the search space.
Overall Role	Complements demonstration guidance by providing additional intermediate cues.	Serves as an additional strategy that can work in tandem with demonstration guidance, offering expert insights.	Acts as a complementary strategy by directly guiding LLM reasoning with demonstrative examples.

Appendix

Table 6: Comparison of Reward Shaping, Knowledge-guided Exploration, and RePO

Models	LogP			MR			QED
Models	SR	Sim.	SR $\times$ Sim.	SR	Sim.	SR $\times$ Sim.	SR	Sim.	SR $\times$ Sim.
Graph GA	0.509	0.125	0.064	0.509	0.122	0.062	0.493	0.140	0.064
REINVENT	0.465	0.125	0.058	0.595	0.116	0.069	0.558	0.115	0.058
MOLLEO	0.510	0.103	0.053	0.509	0.127	0.065	0.496	0.122	0.061
GPT-4o-mini	0.499	0.706	0.352	0.409	0.771	0.315	0.231	0.752	0.174
Mol-T5-large	0.424	0.102	0.043	0.450	0.107	0.048	0.465	0.119	0.055

Appendix

Table 7: Comparison of molecular optimization baselines and larger models on TOMG-Bench tasks (SR, Sim, SR

\times

Sim).

Figure 18: Training reward dynamics of GRPO and RePO on the property optimization task.

Objective	Baseline	40% demo guidance	RePO
QED	0.130	0.215	0.236
LogP	0.168	0.312	0.297
MR	0.173	0.297	0.294

Appendix

Table 8: Ablation on demonstration-quality robustness in TOMG-Bench property optimization (reported in SR

\times

Sim).

A conversation between User and Assistant. The user asks a question, and the Assistant solves it.

The assistant first thinks about the reasoning process in the mind and then provides the user

with the answer. The reasoning process and answer are enclosed within <think> and </think>

<answer> </answer> tags, respectively, i.e., <think> reasoning process here </think>

<answer> answer here </answer> . User: Task . Assistant:

Appendix

Table 9: System prompt adopted for training. Task will be replaced with the specific molecular optimization task.

Objective	Metric	Baseline	Distill-SFT	SFT	GRPO	GRPO (SFT init)	RePO
LogP	SR	$0.269\pm 0.022$	$0.246\pm 0.011$	$0.323\pm 0.026$	$0.179\pm 0.174$	$0.183\pm 0.026$	$0.381\pm 0.029$
	Sim	$0.628\pm 0.016$	$0.553\pm 0.023$	$0.700\pm 0.008$	$0.817\pm 0.025$	$0.897\pm 0.030$	$0.734\pm 0.016$
	SR $\times$ Sim	0.169	0.136	0.226	0.146	0.164	0.280
QED	SR	$0.196\pm 0.025$	$0.220\pm 0.010$	$0.317\pm 0.023$	$0.087\pm 0.044$	$0.169\pm 0.048$	$0.264\pm 0.042$
	Sim	$0.666\pm 0.039$	$0.585\pm 0.012$	$0.701\pm 0.004$	$0.905\pm 0.015$	$0.893\pm 0.026$	$0.797\pm 0.035$
	SR $\times$ Sim	0.131	0.129	0.222	0.079	0.151	0.210

Appendix

Table 10: TOMG-Bench property optimization over multiple random seeds. We report mean

\pm

std for SR and Sim, and SR

\times

Sim for overall quality.

Task type	Objective ( $\uparrow$ )	BioT5-base	MolT5-large	Baseline	SFT	GRPO	GRPO (SFT init)	RePO
Structure- based optimization	AddComponent	0.054	0.031	0.065	0.147	0.005	0.156	0.239
	DelComponent	0.027	0.027	0.092	0.154	0.008	0.176	0.140
	SubComponent	0.011	0.016	0.047	0.264	0.052	0.300	0.344
Property optimization	QED	0.080	0.055	0.130	0.207	0.123	0.193	0.236
	LogP	0.079	0.043	0.168	0.206	0.305	0.183	0.297
	MR	0.081	0.048	0.173	0.238	0.188	0.225	0.294

Appendix

Table 11: Comparison of different methods on TOMG-Bench target on structural and property optimization. The best results for each task are bolded, and the second-best is underlined.

Molecular Optimization

Molecular optimization computationally designs or edits chemical structures to improve desired properties while respecting chemical and practical constraints.

See all 103 papers

Active-GRPO: Adaptive Imitation and Self-Improving Reasoning for Molecular Optimization
Jul 1, 2026Xuefeng Liu, Mingxuan Cao, Qinan Huang +3Molecular Optimization Large Language Model-Guided Optimization
School of Medicine, Stanford University · 2Data Science Institute, University of Chicago · 3Pritzker School of Molecular Engineering, University of Chicago +2
C-MORAL: Controllable Multi-Objective Molecular Optimization with Reinforcement Alignment for LLMs
Apr 24, 2026Rui Gao, Youngseung Jeon, Swastik Roy +2RL for Language Models Molecular Optimization
University of California, Los Angeles · Amazon
LLM-Guided Dynamic Action Spaces for Synthesizable Molecular Optimization
Apr 9, 2026Tao Li, Kaiyuan Hou, Tuan Vinh +4Molecular Optimization Large Language Model-Guided Optimization
Department of Computer Science, Emory University · School of Computer Science, Carnegie Mellon University · MRC Brain Network Dynamics Unit, University of Oxford +3

Jul 1, 2026J/K move · Enter open · S save

Active-GRPO: Adaptive Imitation and Self-Improving Reasoning for Molecular Optimization

Xuefeng Liu, Mingxuan Cao, Qinan Huang, Thomas Brettin+2

School of Medicine, Stanford University · 2Data Science Institute, University of Chicago · 3Pritzker School of Molecular Engineering, University of Chicago+2

Molecular Optimization Large Language Model-Guided Optimization Group Relative Policy Optimization Policy Optimization Imitation Learning RL Fine-Tuning

Open paper AskPDF arXiv

Scientific reasoning is an increasingly important capability of large language models, yet improving the robustness and efficiency of training such reasoning remains a key open challenge. We study this problem in instruction-based molecular optimization, where answer-only supervised fine-tuning (SFT) collapses multi-step reasoning and reinforcement learning with verifiable rewards (RLVR) suffers from sparse feedback. Reference-guided Policy Optimization mitigates both by anchoring policy updates to dataset-provided references, but its effectiveness is tightly coupled to reference quality: weak or misaligned references impose a performance ceiling. To overcome this ceiling, we propose active reasoning, a paradigm in which the policy actively decides, on a per-instance basis, when to imitate a reference and when to reinforce its own discoveries, while continuously upgrading what it imitates. We instantiate this paradigm as Active Group Relative Policy Optimization (Active-GRPO), realized through two coupled mechanisms: active imitate-reinforce and active referencing. The former performs imitation learning when the reference still outperforms the policy's own candidates, and shifts to self-improvement via reinforcement learning once the policy has generated molecules that surpass the reference. The latter continuously upgrades the reference itself by replacing it with the best policy-generated candidate discovered so far, progressively raising the imitation target and ensuring that reference guidance remains informative-rather than restrictive-throughout training. Across TOMG-Bench MOLOPT, Active-GRPO improves average SRxSim from 0.0959 for GRPO and 0.1665 for RePO to 0.1773 under matched three-seed evaluation, with statistically significant gains on LogP, MR, and QED.

Abstract

Molecular Optimization

Explore similar work

Active-GRPO: Adaptive Imitation and Self-Improving Reasoning for Molecular Optimization

C-MORAL: Controllable Multi-Objective Molecular Optimization with Reinforcement Alignment for LLMs

LLM-Guided Dynamic Action Spaces for Synthesizable Molecular Optimization