Memento No More: Coaching AI Agents to Master Multiple Tasks via Hints Internalization

TL;DR

This work addresses the challenge of training AI agents to master multiple tasks without relying on ever-growing prompts. It introduces Memento No More (MNM), an iterative coaching framework where humans provide task hints that are internalized into model weights via context distillation, using LoRA adapters to accumulate knowledge across rounds. Across ToolQA and OfficeBench, a Llama-3.1-70B agent achieves near- and sometimes surpasses GPT-4o and DeepSeek-V3 after three rounds, with significant inference speedups and reduced prompt length. The results demonstrate that structured human feedback, efficiently managed with automated reviewers and data balancing, can produce scalable, high-performing multi-task agents while mitigating prompting overhead and preserving core reasoning and tool-use capabilities.

Abstract

As the general capabilities of artificial intelligence (AI) agents continue to evolve, their ability to learn to master multiple complex tasks through experience remains a key challenge. Current LLM agents, particularly those based on proprietary language models, typically rely on prompts to incorporate knowledge about the target tasks. This approach does not allow the agent to internalize this information and instead relies on ever-expanding prompts to sustain its functionality in diverse scenarios. This resembles a system of notes used by a person affected by anterograde amnesia, the inability to form new memories. In this paper, we propose a novel method to train AI agents to incorporate knowledge and skills for multiple tasks without the need for either cumbersome note systems or prior high-quality demonstration data. Our approach employs an iterative process where the agent collects new experiences, receives corrective feedback from humans in the form of hints, and integrates this feedback into its weights via a context distillation training procedure. We demonstrate the efficacy of our approach by implementing it in a Llama-3-based agent that, after only a few rounds of feedback, outperforms advanced models GPT-4o and DeepSeek-V3 in tasksets requiring correct sequencing of information retrieval, tool use, and question answering.

Figures (4)

• Figure 1: The proposed iterative process in which a human expert coaches an AI agent to master multiple tasks.
• Figure 2:
• Figure 3: Overview of the proposed training procedure. Left: In Round 1, a student agent is trained to internalize initial hints (tool documentation and general best practices) by learning to distill the outputs (actions $a_t$) of a teacher agent that has access to these hints, while the student itself only sees a minimal task description and the execution history (actions $a_{1:t-1}$, environment observations $o_{1:t-1}$). Right: In Round 2 and subsequent rounds, the behavior of the trained student is further refined based on human feedback in the form of mistake descriptions and corrective hints, each addressing a specific kind of mistake the agent exhibits. These hints can be inserted only at steps where this mistake occurs, by implementing an automated reviewer (AI or a script) that locates the mistake in past trajectories. The student is then further trained to distill the outputs of a teacher that is conditioned on this corrective feedback, while the student is not.
• Figure 4: Example of how a hint is used to correct a mistake in the agent's behavior.