Training a Generally Curious Agent

Training a Generally Curious Agent

Speakers: Fahim Tajwar, Yiding Jiang, Abitha Thankaraj, Sumaita Rahman, Zico Kolter, Jeff Schneider, Russ Salakhutdinov

Conference: ICML 2025

YouTube: https://slideslive.com/39043905

Overview

In this compelling talk from ICML 2025, Yiding Jiang, a PhD student at Carnegie Mellon University, along with collaborators Fahim Tajwar, Abitha Thankaraj, Sumaita Rahman, Zico Kolter, Jeff Schneider, and Russ Salakhutdinov, presented groundbreaking research on training a generally curious agent capable of solving novel problems it has not encountered during its training phase. The core challenge addressed is the pervasive issue of generalization in machine learning models, particularly for autonomous agents deployed in dynamic, real-world environments. Such agents frequently face situations where crucial information is missing or problems are vaguely defined, necessitating active exploration and information gathering at test time—a distinct and more challenging paradigm than traditional train-time exploration.

The talk highlights the limitations of current state-of-the-art models, such as large language models (LLMs), in exhibiting truly general information-seeking behaviors. While LLMs can be fine-tuned or augmented with tools and chain-of-thought prompting for specific problems, these "patching" approaches often fall short of enabling genuine discovery or robust performance in entirely new contexts. The proposed solution, dubbed Paprika, introduces a novel pipeline designed to instill general, in-context, and task-agnostic information-seeking capabilities, effectively learning an amortized exploration strategy that can transfer across diverse tasks.

This work is highly significant for the future of AI agents, moving beyond mere task execution to genuine problem-solving in uncertain conditions. By demonstrating that general decision-making abilities can emerge from relatively small amounts of diverse data and strategic training, the researchers pave the way for more robust, adaptable, and intelligent agents capable of navigating the inherent complexities and unknowns of real-world deployment. The implications span various domains, from clarifying user intent in coding assistants to accelerating scientific discovery by guiding experimental design.

Background

▶ Watch: Introduction: Training a generally curious agent for new problems (0:00)

Generalization stands as one of the most critical, yet elusive, properties for any machine learning model, especially for intelligent agents designed to interact with the real world. The prevalent challenge is that models deployed in practical settings often encounter entirely new problems where the necessary information was simply not available during their training. These problems are frequently partially observable, meaning agents must actively explore their environment to gather new data and insights before a solution can be formulated. This need for exploration at test time fundamentally differs from the exploration typically performed during the training phase of reinforcement learning, where the environment is often known and stable.

The speakers illustrate this problem with several compelling examples. A coding assistant, for instance, might be asked to generate code for finding the shortest path in a weighted graph. A naive response might be Dijkstra's algorithm, yet this algorithm only works for non-negative weights—a crucial piece of information the user might not explicitly state or even be aware of. A truly intelligent agent should clarify this assumption by asking questions, thereby performing an act of information-seeking exploration. Another example is scientific discovery, an inherently exploratory process where new knowledge is sought through experiments. Given the high cost of such experiments, as exemplified by the Large Hadron Collider (LHC), efficient and strategic exploration is paramount.

The limitations of current powerful models are starkly demonstrated by a recent paper showing that GPT-4 struggles with even simple multi-armed bandit problems out-of-the-box, performing worse than a basic greedy strategy. While fine-tuning, tool use, or chain-of-thought (CoT) prompting can improve performance for known problems, these methods represent problem-specific engineering rather than a general capacity for information discovery. The authors argue that such "patching" approaches are unlikely to lead to truly general intelligence capable of discovering novel solutions or insights. Instead, the community needs to equip models with general, in-context, and task-agnostic information-seeking behaviors. This necessitates a paradigm shift towards training models not just to solve problems, but to effectively explore and learn how to learn in new, uncertain situations.

Key Findings

▶ Watch: Current LLM limitations: GPT-4 struggles with exploration (2:20)

The research introduces Paprika, a novel pipeline designed to train agents with general information-seeking behavior, enabling them to tackle new, unseen problems effectively. The key findings and contributions are:

• Amortized Exploration: Paprika successfully learns an amortized exploration strategy, allowing agents to leverage past experiences of good decision-making trajectories from diverse environments to explore more effectively in new tasks.
• Significant In-Distribution Generalization: When trained on a diverse set of 10 task groups and tested on held-out instances of these tasks, Paprika demonstrates substantial performance improvements. Specifically, it boosts the performance of Llama-31-8B by approximately 35% and Gemma-2-12B by about 10% (median improvement normalized by base model performance). These improvements were observed across 8 out of 10 tasks, with no degradation on the remaining two, which were already considered "solved."
• Positive Transfer and Out-of-Distribution Generalization: The study provides strong evidence for positive transfer and the ability to generalize to entirely unseen tasks.
• Single Task Group Training: Training on only one out of ten task groups resulted in a 15% drop in performance compared to training on all tasks, indicating that even unrelated tasks contribute significantly to a generalizable exploration strategy.
• Leave-One-Out Generalization: When trained on nine out of ten task groups and tested on the completely unseen tenth group, Paprika still achieved an average performance 11% better than the base model. This is a crucial indicator of the agent's learned general decision-making ability transferring to novel contexts.
• Effectiveness of Curriculum Learning: The integration of a curriculum learning strategy, guided by a "learning potential" metric (ratio of reward standard deviation to mean reward), was shown to be highly effective. In experiments involving a 20-question task, curriculum learning improved both average and Pass@4 performance across all difficulty levels without any observed downsides.
• Efficiency in Learning: The research highlights that surprisingly general decision-making abilities can emerge from a relatively small amount of data (22,000 trajectories in total) and training, suggesting potential for scalability.

Technical Deep Dive

▶ Watch: Call for general, task-agnostic information-seeking behavior (3:20)

The core of this research lies in the Paprika pipeline, a systematic approach to train agents for general, in-context, and task-agnostic information-seeking behavior. The pipeline is built upon four fundamental ideas:

1. Diverse Data and Task Construction: The initial step emphasizes creating as diverse a set of data and tasks as possible. The rationale is that deep learning models thrive on diverse data, which encourages the development of general decision-making capabilities rather than task-specific heuristics. For Paprika, the researchers constructed 10 distinct environments, some from prior work and some novel, with no consistent theme to maximize diversity. Crucially, all these environments were designed such that their tasks require active interaction and strategic exploration to reach a solution. Each environment is referred to as a "task group," capable of programmatically generating a large number of unique instances or subtasks.

1. Curriculum Learning for Task Difficulty: To ensure meaningful learning, Paprika incorporates a curriculum learning strategy that selects tasks at an appropriate level of difficulty. The intuition is that tasks that are either too easy (no room for improvement) or too hard (agent cannot succeed, leading to no useful experience) yield little learning. The challenge here is that determining task suitability typically requires expensive rollouts. To circumvent this, the paper defines "meaningful experience" by observing the variability of policy performance on a given task. Specifically, if a policy's performance on a task $\tau$ exhibits high variability, it implies there's more information to be learned. This is quantified by computing the ratio of the standard deviation of rewards to the mean reward across multiple trajectories, termed "learning potential." The standard deviation measures diversity in outcomes, while the mean normalizes it for cross-task comparability. To make this process efficient, the researchers assume that similar tasks can be grouped, and tasks within a group will have reasonably similar learning potentials, while different groups will have distinct potential statistics. This allows the problem of picking task groups to be modeled as a multi-armed bandit problem, where each "arm" is a task group, and the "reward" is the learning potential sampled from that group. This allows for efficient exploration of which task groups are most amenable to learning at any given stage.

1. Diverse Behavior Generation: For each selected task, the pipeline generates a diverse set of behaviors or trajectories. This is critical for two reasons: First, diversity in samples is essential for generalization, as deep learning models can only learn what they are exposed to ("you cannot learn what you cannot sample"). Second, a broader range of sampled behaviors increases the likelihood of discovering superior trajectories that are significantly better than average. To achieve this, the researchers employ Min-P sampling, a generalization of Top-P sampling. Unlike Top-P, which uses a fixed probability mass threshold, Min-P adjusts its sampling cutoff based on the maximum probability in the distribution. This technique enables sampling at a very high "temperature" (encouraging exploration) while still maintaining coherent and meaningful generations, preventing the agent from generating completely nonsensical actions. For data collection, 20 trajectories are generated per task instance.

1. Training to Increase Likelihood of Good Behavior: The final stage involves training the model to increase the likelihood of generating these "good behaviors." While the choice of optimization algorithm is flexible, the paper specifically uses Direct Preference Optimization (DPO). DPO is an offline reinforcement learning algorithm that directly optimizes a policy to align with a preference dataset (in this case, trajectories deemed "good" based on their reward). The entire process, from task generation to data collection and model training, is iterative and can be repeated, forming the complete Paprika pipeline. This iterative self-improvement mechanism allows the agent to continuously refine its exploration strategies by learning from its own generated experiences. The overall pipeline can be viewed as learning amortized exploration, where the model learns general strategies from diverse past experiences to efficiently explore any new, unseen task.

Experimental Setup & Results

▶ Watch: Paprika pipeline: Learning amortized exploration strategies (5:00)

The experimental evaluation of Paprika focused on assessing both in-distribution and out-of-distribution generalization capabilities, as well as the efficacy of the curriculum learning component.

Base Models and Training Data:

The base large language models used for the experiments were Llama-31-8B and Gemma-2-12B. Paprika was trained on a relatively small dataset by modern standards, comprising approximately 22,000 trajectories collected across the diverse set of 10 task groups. For each task, 20 trajectories were generated using Min-P sampling to encourage diversity.

Evaluation Metrics:

Performance was aggregated and normalized, with the model's performance on a task divided by the base model's performance on the same task. The median normalized performance is reported. For specific tasks like the 20-question environment, Pass@4 (the percentage of tasks solved within 4 attempts) was also used.

In-Distribution Generalization:

The first set of experiments evaluated the model's ability to generalize to held-out tasks from the same distribution it was trained on (i.e., instances within the 10 task groups).

• Llama-31-8B: Paprika improved performance by approximately 35% over the base Llama model.
• Gemma-2-12B: Paprika achieved a 10% improvement over the base Gemma model.

These improvements were observed across 8 out of 10 tasks. For the remaining two tasks, which were deemed "too easy" and already largely solved by the base models, Paprika did not hurt performance, indicating robust learning without negative transfer.

Transfer and Out-of-Distribution Generalization:

To rigorously test the agent's ability to solve truly new problems, two critical experiments were conducted:

1. Single Task Group Training: In this setup, Paprika was trained using tasks from only one of the 10 task groups and then tested on held-out tasks from that same group. The results showed a 15% drop in performance compared to training on all task groups. This finding strongly suggests that training on diverse, even seemingly unrelated, tasks provides significant positive transfer, contributing to a more robust and generalizable exploration strategy.

1. Leave-One-Out Generalization: This experiment was designed to assess true out-of-distribution generalization. Paprika was trained on nine out of the ten task groups and then tested on the completely unseen tenth task group. Despite never having encountered any instances from this specific task group during training, Paprika achieved a performance that was approximately 11% better than the base model. This is a highly encouraging result, demonstrating that the agent learned a form of general decision-making ability that successfully transferred to a novel environment.

Effect of Curriculum Learning:

The impact of the proposed curriculum learning strategy was evaluated using a 20-question task environment, which was divided into three difficulty groups. Paprika was run for three rounds of curriculum learning. The results indicated that curriculum learning successfully improved both the average performance and the Pass@4 metric. More importantly, these improvements were observed across all difficulty levels, suggesting that the curriculum effectively guides the agent to learn from meaningful experiences without any discernible downsides.

Overall, the experimental results provide compelling evidence that Paprika is an effective method for training generally curious agents. It demonstrates significant performance gains on known tasks, crucial positive transfer, and a remarkable ability to generalize to entirely unseen problem domains, all while leveraging a relatively modest amount of training data.

Practical Implications

▶ Watch: Data collection details: Diverse trajectories via Min-P sampling (6:00)

The development of a generally curious agent with the Paprika pipeline carries significant practical implications for various stakeholders in the AI/ML ecosystem, from practitioners and infrastructure teams to model builders and deployers.

For practitioners and model builders, Paprika offers a pathway to construct more robust and adaptable AI agents. The ability to learn amortized exploration means that agents can be deployed into novel, partially observable environments and effectively gather information to solve problems without requiring extensive, task-specific fine-tuning or explicit programming for every new scenario. This could dramatically reduce the engineering overhead for real-world applications where uncertainty and novelty are common, such as complex robotic tasks, interactive coding assistants, or conversational AI that needs to clarify user intent. The concept of learning how to explore rather than just what to do is a paradigm shift that enables more intelligent and autonomous systems.

Infrastructure teams and deployers will find value in the efficiency and generalizability demonstrated. The fact that general decision-making abilities can emerge from a relatively small dataset (22,000 trajectories) suggests that the data collection and training costs for such agents might be more manageable than for traditional, heavily data-dependent models. This is particularly relevant for in-context reinforcement learning, where collecting high-quality multi-turn human preference data can be prohibitively expensive. Paprika's reliance on programmatically generated, synthetic data offers a cheaper and more scalable alternative for developing robust agents. The positive transfer observed means that investments in training on a diverse set of tasks can yield benefits across a wide range of future, unseen applications, improving the return on investment for model development.

However, several tradeoffs and limitations must be considered. The talk acknowledges that this is a preliminary study, and many aspects of the pipeline, such as the environment design and the specific choices of algorithms, can be further improved. While the current results are encouraging, scaling up the number of environments and the complexity of tasks will inevitably introduce new challenges. The open-ended question of "how to automate environment design" is crucial, as manually crafting diverse and challenging environments for general exploration can itself be a bottleneck. The current approach, while effective, still relies on pre-defined task groups and programmatic generation, which might not fully capture the unbounded diversity of real-world problems. Furthermore, while Min-P sampling helps maintain coherence, ensuring the quality and meaningfulness of synthetic trajectories remains an ongoing area of research. Despite these limitations, Paprika represents a significant step towards building agents that are genuinely curious and capable of autonomous discovery, pushing the boundaries of generalization in AI.

Key Takeaways

• Generalization through Test-Time Exploration: For AI agents to succeed in the real world, they must learn to actively explore and gather new information at test time to solve novel, partially observable problems.
• Paprika Pipeline for Amortized Exploration: The proposed Paprika pipeline effectively trains agents to learn an amortized exploration strategy by leveraging diverse data, curriculum learning, and diversified behavior generation.
• Significant Performance Gains and Transfer: Paprika substantially improves in-distribution generalization (e.g., 35% for Llama-31-8B) and, critically, demonstrates positive transfer and out-of-distribution generalization, achieving 11% better performance on unseen tasks compared to base models.
• Curriculum Learning is Crucial: A curriculum learning strategy, guided by a "learning potential" metric (reward variability), is highly effective in improving performance across all difficulty levels without downsides, and is expected to become more important with scaling.
• Efficiency of Learning General Abilities: Surprisingly, general decision-making and information-seeking abilities can emerge from a relatively small amount of diverse training data (22,000 trajectories), suggesting a cost-effective path for developing adaptable agents.
• In-Context RL with Synthetic Data: The approach focuses on in-context reinforcement learning using synthetically generated data, offering a cheaper and more scalable alternative to human preference-based fine-tuning for multi-turn interaction.

About the Speaker(s)

The research presented was a collaborative effort. Yiding Jiang, a PhD student at Carnegie Mellon University, delivered the talk and is a key contributor to this work. His research focuses on developing intelligent agents capable of generalization and effective exploration in novel environments. The team also includes Fahim Tajwar, Abitha Thankaraj, Sumaita Rahman, and senior researchers and faculty members Zico Kolter, Jeff Schneider, and Russ Salakhutdinov, all affiliated with Carnegie Mellon University. Their collective expertise spans machine learning, reinforcement learning, and the development of advanced AI systems, contributing to this significant advancement in training generally curious agents.