PaperlayerView exact revision on arXiv

Source-linked AI summary

WMLLM: Self-Evolving Optimization Agents via Predict-Then-Act World Modeling

Zhongzheng Li, Qingsong Ran, Shikun Feng, Nian Ran, Wenhao Li, Xiaoyuan Zhang, Yue Wang, Xiaoguang Zhao

arXiv:2609.01608v1cs.LGcs.AI

TL;DR

Black-box optimization is difficult in complex search spaces because direct generation and trial-and-error exploration waste limited evaluator queries. WMLLM uses a single LLM for predict-then-act optimization, combining agentic refinement, population search, and reinforcement learning. It improves sample efficiency and final performance, achieving state-of-the-art results on multi-objective molecular optimization under a limited evaluation budget.

  • Problem

    Existing black-box optimization methods often rely on direct generation or trial-and-error exploration, leading to poor sample efficiency when evaluations are expensive and feedback is sparse.

  • Method

    WMLLM uses a single LLM as an implicit world model and policy, predicting candidate outcomes before acting and refining through feedback, population search, and reinforcement learning.

  • Results

    WMLLM improves sample efficiency and final optimization performance, achieving state-of-the-art performance on multi-objective molecular optimization under a limited evaluation budget.

  • Takeaways & Limitations

    Prediction-observation discrepancy becomes a self-supervised signal that improves future decisions and supports closed-loop co-evolution of the policy and implicit action-outcome model.

  • Takeaways & Limitations

    Evidence mainly covers structured black-box optimization, while broader validation across diverse long-horizon agent environments remains future work.

Abstract

from arXiv · show

Black-box optimization problems remain challenging because of large, weakly structured, and high-dimensional search spaces. Existing methods often suffer from poor sample efficiency because they rely on direct candidate generation or trial-and-error refinement. A natural way to improve search efficiency is to use world modeling, which can help identify promising optimization directions before costly evaluation. Large language models can predict the outcomes of these candidates with nontrivial accuracy because of their implicit knowledge. Motivated by this observation, we propose WMLLM, a self-evolving optimization-agent framework based on predict-then-act world modeling. The agent first predicts promising directions and then acts to generate candidates. Combined with agentic multi-turn refinement, population-based search, and reinforcement learning, WMLLM refines both its implicit world model and its optimization strategy during search. Experiments on black-box optimization tasks, especially multi-objective molecular optimization, show that WMLLM improves sample efficiency and final optimization performance. On the multi-objective molecular optimization benchmark, WMLLM achieves state-of-the-art results under a limited evaluation budget.

1 Introduction

Black-box optimization requires efficient search under complex spaces, limited budgets, and sparse feedback. WMLLM addresses this with predict-then-act modeling, agentic refinement, population search, and reinforcement learning.

  • Existing methods often evaluate candidates only after execution, causing poor sample efficiency, unstable optimization, and premature convergence under expensive, sparse feedback.
  • World models can improve search efficiency by representing how actions influence future outcomes, but separate dynamics modules are difficult to scale in high-dimensional, discrete optimization.
  • WMLLM uses one LLM for outcome prediction and candidate generation, refining decisions from prediction errors and evaluator feedback.
  • Agentic multi-turn refinement adapts interaction depth, while population-based search preserves exploration across candidate solutions.
  • WMLLM improves sample efficiency and final optimization performance across molecular and additional structured black-box tasks, achieving state-of-the-art molecular optimization under a limited evaluation budget.
  • WMLLM combines predict-then-act reasoning with population-based search and reinforcement learning to improve its policy and implicit action-outcome model through interaction.

2 Related Work

Prior work uses agentic workflows and world models for iterative reasoning, optimization, and decision-making, but commonly separates prediction from action. WMLLM instead makes prediction a directly refined responsibility of the policy.

  • Agentic LLM workflows combine structured reasoning, tool use, and self-reflection to iteratively generate, execute, and refine solutions using environmental feedback.
  • LLM-evolutionary methods generate heuristics or candidates, while reflection, memory, and evaluator feedback support iterative search.
  • Most agentic workflows treat prediction quality as indirectly related to decisions, whereas WMLLM drives refinement through prediction residuals and closed-loop optimization.
  • Classical world-model approaches model environment dynamics separately from policy learning and action selection, often using planning or value estimation.
  • LLMs can approximate dynamics through next-state prediction, simulation, and counterfactual reasoning, but reliability and long-horizon consistency remain limited without additional structure or supervision.
  • WMLLM enforces predictive responsibility by requiring consequence prediction before action, with prediction error contributing directly to learning.

3 Method

WMLLM frames black-box optimization as sequential agentic search over evaluated candidates, combining an implicit world model with population-based refinement and reinforcement learning. It predicts candidate outcomes before acting, uses feedback to update decisions, and supports adaptive multi-turn optimization under limited evaluations.

  • Problem Formulation: The framework operates over structured discrete spaces with non-differentiable evaluators, vector objectives, scalar aggregation, and a limited evaluation budget.Every evaluator call during agentic refinement counts toward the budget.
  • Population-Based Search: An evolutionary outer loop samples parent subsets, generates offspring through multi-turn trajectories, and selects the next population using scalar scores and Pareto non-dominated sorting.The implementation selects half the population by scalar score and half by Pareto ranking.
  • Agentic Optimization: WMLLM treats candidate generation as a sequential decision process in which an LLM produces reasoning, predictions, proposals, and tool invocations.Context is updated after each action with generated content and returned tool feedback.
  • Predict-Then-Act World Modeling: Before submitting candidate x_t, the same LLM predicts outcome ˆo_t, observes o_t after evaluation, and measures discrepancy δ_t = D(ˆo_t, o_t).Prediction errors provide feedback for refining later decisions without a separate transition or reward model.
  • Trajectory-Level Reinforcement Learning: The trajectory reward combines evaluator and prediction signals, while R(τ) = max_t≤T r_t rewards discovering at least one strong candidate with reliable predictions.The best evaluator score max_t≤T s(x_t) is tracked separately for reporting and population updates.
  • Trajectory-Level Reinforcement Learning: GRPO samples K trajectories per prompt, normalizes their rewards into group-relative advantages, and optimizes the policy without a separate value model.This supports sparse rewards and comparison among alternative refinements from the same parent set.

4 Experiments

The experiments test WMLLM under fixed evaluation budgets on molecular and additional black-box optimization tasks, including ablations of prediction, refinement, warm start, scaling, and interaction. WMLLM-Evolve achieves stronger molecular optimization results and transfers gains across diverse non-molecular tasks.

  • Main Results on Molecular Optimization: Under a fixed budget of 3,000 unique molecule evaluations, WMLLM-Evolve achieves the best Avg. Top-1, Avg. Top-10, and Top-10 AUC.Top-10 AUC captures anytime performance across the search process, not only the final result.
  • Main Results on Molecular Optimization: WMLLM-Evolve improves over OpenEvolve from 4.123 to 4.385 in Avg. Top-1 and from 3.661 to 4.124 in Top-10 AUC.It remains competitive in uniqueness while accepting lower validity than some graph-based or domain-specific methods.
  • Ablation Study: Prediction raises Qwen3-8B Avg. Top-1 from 4.088 to 4.150, while Agent-Pred reaches 4.329 versus 4.071 for Agent-Qwen3-8B.These ablations indicate that agentic refinement is more effective when the agent predicts before acting.
  • Ablation Study: WMLLM-Evolve achieves the best overall ablation results without SFT, while reinforcement learning and population-based evolution support self-improvement from interaction.SFT-Agent-Pred slightly improves over Agent-Pred, showing that warm start can help the predictive agent.
  • Ablation Study: Larger models alone do not explain the gains: Qwen3-32B outperforms Qwen3-8B but remains below Agent-Pred and WMLLM-Evolve.The reported comparison attributes the main gain to predict-then-act optimization rather than model size alone.
  • Ablation Study: Prediction error decreases during interaction, while supervised warm start reduces median absolute prediction error across molecular properties and the overall score.Together with the ablations, this supports a link between prediction fidelity and decision quality.
  • Generalization to Additional Tasks: WMLLM improves the best score on CIRCLE PACKING, SUMS DIFFS, and HADAMARD DET under a 100-call evaluation budget.The tasks span continuous geometric layouts, discrete additive-combinatorial constructions, and structured matrix search.

5 Conclusion

WMLLM combines predict-then-act world modeling with agentic interaction, population search, and trajectory-level optimization. It improves optimization performance and learns from prediction errors as well as rewards.

  • WMLLM combines predict-then-act world modeling with agentic multi-turn interaction, population-based search, and trajectory-level optimization.
  • Explicit prediction improves sample efficiency and final performance on multi-objective molecular optimization beyond model scale or agentic refinement alone.
  • WMLLM learns from errors in its own predictions in addition to rewards, providing a practical direction for optimization agents.

6 Limitations

The evidence mainly covers structured black-box optimization, while broader long-horizon agent validation remains future work and evaluator feedback remains necessary.

  • Evidence mainly covers structured black-box optimization, so validation in more diverse long-horizon agent environments remains future work.
  • Although WMLLM improves sample efficiency, it still requires repeated evaluator feedback, leaving interaction-cost reduction important for extremely expensive domains.

A Additional Experimental Analysis

Additional analysis examines evaluator-call efficiency and the molecular-design workflow. Warm-started models use fewer calls, whereas reinforcement learning produces longer, more variable trajectories through deeper exploration.

  • Evaluator-call efficiency: RL-trained agents use more variable evaluator-call trajectories, reflecting deeper refinement and increased exploration.
  • Evaluator-call efficiency: SFT-initialized models consistently require fewer evaluator calls per conversation than the base model.This is attributed to more accurate predictive priors reducing redundant trial-and-error.
  • Molecular-design workflow: The optimization targets minimize sa and drd2, maximize qed and jnk3, and minimize gsk3b, with an overall score goal of at least 4.2.
  • Molecular-design workflow: The molecular-design assistant predicts six objective values before evaluator calls and compares predictions with actual outcomes afterward.The required values are sa, drd2, qed, gsk3b, jnk3, and overall.
  • Molecular-design workflow: The workflow forbids repeated molecule evaluations and requires immediate proposal of the next candidate after comparing prediction with evaluator output.

C Cost Analysis

WMLLM is evaluated under a fixed 3,000-molecule budget using eight A100 GPUs. Training takes 3.53 ± 0.24 hours on average, indicating moderate computational overhead in this setting.

  • Experimental setting: 3,000 unique molecules define the fixed evaluation budget used throughout the experiments.
  • Experimental setting: 8 NVIDIA A100 GPUs with 80GB memory each support training and optimization under the reported experimental configuration.
  • Computational cost: 3.53 ± 0.24 hours is the average wall-clock training time across multiple runs.The reported variance is relatively low despite stochastic trajectory sampling and adaptive multi-turn interaction.
  • Computational cost: The authors characterize the resulting computational overhead as moderate for black-box optimization with expensive evaluations.
  • Broader impacts: Potential deployment risks include harmful design spaces, sensitive decision-making domains, and misleading search from inaccurate predictions.The passage recommends safety filters, evaluator checks, and human oversight.

E Details of Additional Black-Box Optimization Tasks

The section defines three additional black-box optimization tasks spanning continuous geometry, additive combinatorics, and structured matrix search. Each task uses external evaluation to verify candidates and compute objective scores, while WMLLM predicts task-specific outcomes before acting.

  • All three tasks use external evaluators that verify candidate validity and recompute the true objective score.The evaluator rejects invalid constructions or checks required structure before scoring valid candidates.
  • CIRCLE PACKING 26: Circle Packing 26 places 26 non-overlapping circles inside a unit square while maximizing their total radius.Candidates specify each circle’s center and nonnegative radius, subject to boundary and collision constraints.
  • CIRCLE PACKING 26: WMLLM predicts radius sum, validity, and possible failure causes before submitting circle layouts for evaluation.The task tests geometric priors about boundary constraints, pairwise collisions, and local layout improvements.
  • SUMS DIFFS: SUMS DIFFS constructs a finite integer set and optimizes a score based on the sizes of its sumset and difference set.The evaluator recomputes |A + A|, |A − A|, and C(A), so reported objective values cannot substitute for valid constructions.
  • SUMS DIFFS: WMLLM predicts sumset size, difference-set size, score, and validity before receiving deterministic evaluator feedback.This probes prediction-guided decisions when candidate quality depends on global set statistics rather than local syntax alone.
  • HADAMARD DET 29: Hadamard Det 29 constructs a 29 × 29 {−1,+1} matrix maximizing absolute determinant through a normalized determinant ratio.The evaluator checks shape and entries, then computes the determinant and returns the normalized ratio as the score.
  • HADAMARD DET 29: WMLLM predicts determinant ratio and validity before matrix evaluation in a search where single-entry changes have global effects.The task tests transfer of predict-then-act modeling to high-dimensional discrete matrix search involving row correlations and conditioning.

E.4 Evaluation Protocol

Across the three additional tasks, optimization treats evaluation as a black-box feedback loop. WMLLM adds an explicit prediction step and prediction-error learning signal before proposing each candidate, unlike the otherwise matched baseline.

  • The evaluator verifies each proposed candidate and recomputes its true score, after which search continues using returned feedback.The protocol applies this black-box interaction across all three tasks.
  • WMLLM predicts the candidate’s task-specific evaluation outcome before proposing it, whereas w/o Pred. removes prediction and prediction-error learning.Both methods otherwise follow the same optimization protocol.

F Implementation Details and Hyperparameters

The implementation combines population-based molecular optimization, hybrid score and Pareto selection, and online GRPO training under fixed evaluation budgets. Main molecular results use 3000 unique molecule evaluations and five random seeds.

  • Population construction: WMLLM maintains a population of N = 50 evaluated candidates and generates offspring through single-round or multi-turn refinement.The model receives three parent molecules as context, and offspring count varies with the number of valid candidates produced.
  • Hybrid population update: Population updates select half the next population by scalar score and half by Pareto ranking.The 0.5/0.5 hybrid preserves high-performing candidates while retaining diverse objective trade-offs.
  • Warm start: The full WMLLM-EVOLVE method uses no supervised warm start, relying on interaction, population evolution, and GRPO training.The 700-example warm-start dataset is used only by SFT-AGENT-PRED.
  • GRPO training: GRPO samples 8 trajectories per prompt and weights evaluator reward at 0.8 and prediction reward at 0.2.Training uses prediction weight α = 0.2 and runs for 300 RL iterations.
  • Trajectory and budget limits: Main molecular optimization evaluates all methods under a budget of 3000 unique molecule evaluations.Every evaluator call during variable-depth refinement counts toward this same budget.
  • Random seeds: Main results are averaged over five random seeds, 42 to 46, with mean and standard deviation reported.
Loading 2609.01608v1…