Uncertainty-aware reinforcement learning for chemical language models

In de novo molecular design, reinforcement learning (RL) agents often treat predictive scoring functions as deterministic oracles. This approach ignores the inherent uncertainty in property predictions, particularly when the agent explores regions of chemical space far from the training data. This paper investigates whether explicitly incorporating this predictive uncertainty can prevent the agent from over-exploiting unreliable, high-scoring regions, thereby improving the reliability of generated molecular candidates.

The authors propose two strategies to integrate uncertainty into the REINVENT4 framework:

1. Score Modulation (SM): Uncertainty is treated as an additional objective within the Multi-parameter Optimization (MPO) scoring function, forcing the agent to balance property optimization with prediction confidence.
2. Loss Modulation (LM): Uncertainty is used to weight the policy updates, reducing the influence of molecules whose predicted properties lie outside the model's confidence domain.

The authors evaluated these strategies across three settings: a synthetic model system with distance-dependent Gaussian noise, ChemProp models capturing aleatoric uncertainty, and Random Forest classifiers wrapped with Conformal Prediction.

The study demonstrates that uncertainty-aware RL enables more robust exploration of chemical space. By favoring regions with lower uncertainty, the agent avoids generating molecules that appear high-scoring due to model error rather than true activity. This approach improved the true hit rate by 0.25 (from 0.5 to 0.75) and nearly doubled the total number of true hits compared to standard, uncertainty-agnostic RL.