Context-Aware Policy-Guided Gradient Search for Offline Model-Based Optimization

Offline model-based optimization (MBO) aims to optimize complex and high-dimensional black-box functions to discover high performing designs by leveraging a static dataset, which can be widely used in both the science and engineering fields. Training a forward surrogate for the target function is an important method to solve offline MBO, since the learned forward model provides more possibilities for counterfactual queries and plentiful downstream tasks. A central challenge in this paradigm lies in determining the optimal step size during gradient-based design updates, as this hyperparameter is notoriously difficult to determine and often highly dependent on the data context. Inspired by policy-guided gradient search (PGS), which frames MBO as a sequential decision-making problem, we propose context-aware policy-guided gradient search (CA-PGS), a novel approach that dynamically adapts the gradient-based optimization process to the data and model context. Our approach introduces context-aware representation learning, which encodes model-specific and data-specific information into the MDP state space. This state representation is then refined and used to train a context-aware policy, which allows a subsequent design dreamer process that leverages ensemble forward models to generate high-quality candidate designs. Experiments on the Design-Bench demonstrate that CA-PGS outperforms existing methods. Furthermore, ablation studies highlight that the proposed key components bring significant advantages.