Skip to content

LLM-Driven Composite Neural Architecture Search for Multi-Source RL State Encoding

🕒 Published (v1): 2025-12-07 20:25 UTC · Source: Arxiv · Venue: NEURIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

LACER is an LLM-driven composite neural architecture search (NAS) framework that uses an LLM as a design agent to automatically discover multi-source state encoder architectures for reinforcement learning. By feeding the LLM richer performance signals (task metric, training reward, and feature mutual information/redundancy) beyond the standard task metric, LACER achieves better sample efficiency than both traditional and LLM-based NAS baselines on a mixed-autonomy traffic control benchmark.

Problem

Designing state encoders for RL environments with heterogeneous multi-source observations (sensors, time-series, images, text) is typically done by manual expert design. Existing NAS methods target single-modality supervised learning tasks, ignore the composite multi-module structure of RL encoders, and discard useful intermediate signals (e.g., representation quality of each source-specific module), leading to poor sample efficiency when evaluation requires thousands of costly RL simulator interactions.

Method

LACER formalizes multi-source RL state encoding as composite NAS: jointly searching over per-source encoder modules \(f_{\theta_i}\) (CNN, Transformer, GRU, FFN depending on modality) and a fusion module \(g_\phi\), optimizing the downstream RL performance \(\mathcal{M}(\pi \circ g_\phi \circ (f_{\theta_1},\ldots,f_{\theta_M}); \mathcal{E})\).

The LLM (Claude Sonnet 4.0, temperature=1.0) acts as an iterative architecture design agent in a closed loop: given a pruned conversation history containing the initial architecture, search space definition, and multi-signal performance feedback, the LLM proposes one (LACER-1) or five (LACER-5) new composite architecture candidates per iteration. Candidates are parsed from LLM output via regex, trained end-to-end with PPO for 200k steps, then evaluated for 50k steps. Performance feedback fed back to the LLM comprises three signals: (i) task metric (average traffic speed), (ii) average RL reward, and (iii) feature information — mutual information \(I(X;Y)\) and redundancy \(R(X;Y)\) between inputs/outputs of each encoder module, quantifying per-module representation quality.

Key Contributions

  • Formal definition of composite NAS for multi-source RL state encoding as a search over \(\Theta_1 \times \cdots \times \Theta_M \times \Phi\).
  • LACER pipeline: LLM-as-agent with enriched multi-signal feedback (task metric + RL reward + feature mutual information/redundancy), contrasting with GENIUS which uses only task metric.
  • Structured prompt engineering with strategic history pruning to mitigate LLM context/forgetting issues across many iterations.
  • Ablation study confirming each feedback component (feature information, average reward, initial architecture evaluation) is individually necessary.

Results

  • Both LACER variants significantly outperform expert-designed encoder, DARTS, ENAS, PEPNAS, and GENIUS on average traffic speed as a function of candidate evaluations (50 candidates total, 8 random seeds, mean ± 2× SE reported).
  • LLM query overhead is ~1% of total time cost; evaluation time dominates (>97%) across all methods, so LACER adds negligible overhead vs. baselines.
  • Ablation: removing feature information, then average reward, then initial architecture evaluation each causes significant performance degradation in LACER-1.
  • Claude Sonnet 4.0 at temperature=1.0 outperforms GPT-4 at both temperatures and Claude at temperature=0.0 on this task.

Limitations

  • Evaluated on a single benchmark (mixed-autonomy traffic control via SUMO); MiniGrid and ManiSkill are described in appendix setup but results are not reported in the paper.
  • Each candidate requires a full RL training run (200k steps), making the search computationally expensive despite improved sample efficiency.
  • Architecture spaces are manually designed per modality, requiring domain knowledge to define per-module search spaces.
  • The composite search space (~26M architectures for the traffic task) is not explored at scale; the budget is fixed at 50 candidates.
  • No theoretical guarantees on convergence or optimality of the LLM-guided search.

Relevance to Agentic AI / LLM Agents

This work instantiates the LLM-as-agent paradigm in a concrete optimization loop — the LLM perceives structured feedback, maintains conversation history, reasons over prior candidates, and proposes new actions (architectures) — directly relevant to agentic tool use and iterative decision-making research. The approach of providing the LLM with rich intermediate signals (not just final outcomes) connects to broader work on improving agent reasoning via richer environmental feedback. For researchers tracking agentic AI, LACER exemplifies how LLM agents can serve as meta-optimizers within compute-intensive loops, and the prompt pruning/history management strategies are applicable to any long-horizon agentic task with context limitations.