Reinforcement Learning for Antibody Sequence Infilling

Abstract

We introduce a flexible framework for antibody sequence design that combines an infilling language model with reinforcement learning to optimize functional properties. Our approach leverages a pretrained infilling language model to generate specific antibody regions within full sequences, guided by reinforcement learning to improve desired biophysical characteristics. We implement a range of online learning strategies, exploring both vanilla REINFORCE and Proximal Policy Optimization with Kullback-Leibler (KL) regularization, and demonstrate that KL regularization is essential for maintaining a balance between score optimization and sequence plausibility. We also adapt Direct Reward Optimization to the protein domain by adding a value head to the infilling model, allowing it to learn directly from static (prompt, response, feedback) datasets using a mean-squared error objective. This formulation is particularly useful when only single-trajectory data is available, which is commonly the case for historically collected experimental assays. We evaluate both the online and offline methods across multiple antibody design tasks--including binding affinity, immunogenicity, and expression--and show that our framework improves alignment with measured biophysical properties while outperforming likelihood-only baselines. This integrated online/offline approach enables functionally driven antibody design and provides a scalable toolkit for therapeutic sequence engineering. Code and data are available at https://github.com/LLNL/protein_tune_rl.