Concept based learning improves action selection for interactive AI agents
CB-Agent: Concept-Based In-Context Learning for Sequential Action Selection
Artificial IntelligenceMachine Learning
Summary
Choosing the next step in a task can be tricky for AI agents because they have limited memory for examples of how to act. The paper introduces CB-Agent, a new way to help AI pick actions based on key concepts learned from expert behavior without needing to remember full examples each time. This method scores possible actions by combining learned concepts and language model features to decide the best next move. Tests in multiple simulated environments showed that this approach outperforms other methods in selecting the right actions efficiently.
What this means in practice
- •For interactive ai developers: Improve AI agents' action choice in interactive environments by using concept-based scoring without large prompt contexts.
- •For game ai programmers: Create more efficient agents for text-based games by applying learned concept distributions to guide action selection.
Authors
Qingyue Zhang, Huaze Tang, Minyu Zhang, Yuxiang Yan, Wenbo Ding, Shao-Lun Huang
Abstract
In-context learning (ICL) allows large language models to use demonstrations without updating their parameters. For interactive agents, demonstrations are trajectories of reasoning, actions, and observations that illustrate how to complete a task. Reusing demonstrations in new situations is challenging: directly inputting full trajectories consumes limited prompt space and constrains the overall effectiveness, while natural-language summaries require the agent to translate general guidance into specific action choices. Therefore, we study whether expert action choices can directly supervise a reusable scoring rule. We propose CB-Agent, which extends concept-based ICL to sequential action selection. We extract environment-action decisions from expert trajectories and fit a concept vector by maximum likelihood to model the conditional distribution over candidate actions given the context. This model combines the concept with frozen language-model features and priors. During ReAct interaction, the agent uses this distribution to select environment actions without including demonstrations in the action-scoring context. Across TextWorld, ALFWorld, and ScienceWorld, CB-Agent achieves the highest cross-benchmark average among the evaluated methods at all three tested model sizes and ranks first on TextWorld and ScienceWorld at each size. ALFWorld experiments support the learned concept's contribution beyond prior-only scoring and show further gains when combined with Reflexion. Code is available at https://anonymous.4open.science/r/CB-Agent-DFF9/.