Improved Off-policy Reinforcement Learning in Biological Sequence Design

TL;DR

This paper introduces a conservative off-policy reinforcement learning method for biological sequence design that improves robustness and outperforms existing methods across various biological tasks.

Contribution

The paper proposes $oldsymbol{ extdelta}$-Conservative Search, a novel off-policy RL approach that adaptively restricts exploration based on proxy model confidence to enhance robustness in biological sequence design.

Findings

Outperforms existing methods in discovering high-score sequences.

Effectively adapts conservativeness based on proxy uncertainty.

Demonstrates success across DNA, RNA, protein, and peptide design tasks.

Abstract

Designing biological sequences with desired properties is challenging due to vast search spaces and limited evaluation budgets. Although reinforcement learning methods use proxy models for rapid reward evaluation, insufficient training data can cause proxy misspecification on out-of-distribution inputs. To address this, we propose a novel off-policy search, $\delta$-Conservative Search, that enhances robustness by restricting policy exploration to reliable regions. Starting from high-score offline sequences, we inject noise by randomly masking tokens with probability $\delta$, then denoise them using our policy. We further adapt $\delta$ based on proxy uncertainty on each data point, aligning the level of conservativeness with model confidence. Experimental results show that our conservative search consistently enhances the off-policy training, outperforming existing machine learning methods in discovering high-score sequences across diverse tasks, including DNA, RNA, protein, and peptide design.