Source-linked AI summary

Training a Generally Curious Agent

Fahim Tajwar, Yiding Jiang, Abitha Thankaraj, Sumaita Sadia Rahman, J Zico Kolter, Jeff Schneider, Ruslan Salakhutdinov

arXiv:2502.17543v4cs.LGcs.AIcs.CL

TL;DR

LLMs often lack strategic information-gathering abilities for interactive decision-making, while real-world data collection is costly and risky. PAPRIKA fine-tunes models on synthetic trajectories from diverse tasks, uses curriculum sampling and preference tuning, and reports transfer to unseen tasks without task-specific fine-tuning. The approach improves decision-making across task groups, but its performance depends on the starting model and has been evaluated in relatively simple, controlled environments.

  • Problem

    LLMs need better strategic information gathering for interactive decision-making, while naturally occurring data lacks interaction structure and real-world collection can be costly and risky.

  • Method

    PAPRIKA fine-tunes LLMs on synthetic trajectories from diverse decision-making tasks, prioritizing high-learning-potential tasks and favoring successful trajectories with sequential preference optimization.

  • Results

    PAPRIKA transfers learned decision-making behavior to unseen tasks, improving Llama-3.1-8B-Instruct performance by 47% of its original success rate across 10 task groups after about 22,500 trajectories.

  • Takeaways & Limitations

    Synthetic data can teach LLMs generalizable in-context reinforcement learning for novel sequential decision-making tasks without task-specific fine-tuning.

  • Takeaways & Limitations

    PAPRIKA depends on a reasonably capable starting model and has been evaluated in relatively simple, controlled environments.

Abstract

from arXiv · show

Efficient exploration is essential for intelligent systems interacting with their environment, but existing language models often fall short in scenarios that require strategic information gathering. In this paper, we present Paprika, a fine-tuning approach that enables language models to develop general decision-making capabilities that are not confined to particular environments. By training on synthetic interaction data from different tasks that require diverse strategies, Paprika teaches models to explore and adapt their behavior on a new task based on environment feedback in-context without more gradient updates. Experimental results show that models fine-tuned with Paprika can effectively transfer their learned decision-making capabilities to entirely unseen tasks without additional training. Unlike traditional training, our approach's primary bottleneck lies in sampling useful interaction data instead of model updates. To improve sample efficiency, we propose a curriculum learning strategy that prioritizes sampling trajectories from tasks with high learning potential. These results suggest a promising path towards AI systems that can autonomously solve novel sequential decision-making problems that require interactions with the external world.

1. Introduction

PAPRIKA teaches LLMs to gather information strategically and perform in-context reinforcement learning across tasks, using synthetic interaction data and preference tuning to transfer decision-making behavior to unseen environments.

  • PAPRIKA targets LLM agents that must interact with external environments and gather information to achieve goals, addressing gaps in naturally occurring interaction data and the risks of real-world collection.Synthetic interaction data provides an alternative to direct deployment, which can be expensive and risky.
  • The method treats curiosity as task-relevant information gathering rather than intrinsic motivation, learning exploration strategies across environments for more efficient exploration on new problems.The paper describes this as amortized exploration and focuses on information needed for the task at hand.
  • PAPRIKA designs diverse decision-making tasks, generates interaction trajectories, scores them by task success, and applies sequential Direct Preference Optimization to favor successful trajectories.Its primary computational bottleneck is sampling useful interaction data rather than performing model updates.
  • A curriculum strategy improves sample efficiency by prioritizing trajectories from tasks with high learning potential.This addresses the difficulty of obtaining meaningful learning signals from tasks that are too difficult for the current model.
  • Training on task subsets improves performance on unseen tasks, indicating that synthetic data can teach in-context reinforcement learning without task-specific fine-tuning.The reported result supports transfer of decision-making capabilities across tasks.

2. Preliminary

The paper models each task as a partially observable sequential decision problem and evaluates agents across task groups whose environments may differ in observations, rewards, and termination conditions.

  • Each task τ is modeled as a partially observable Markov decision process, although the paper does not use the formal POMDP details.
  • A task group G contains related tasks that may share strategies without necessarily sharing the same optimal policy.From the agent’s perspective, each task is a black-box function mapping actions or interaction histories to observations.
  • An episode records an agent’s interaction trajectory and ends when the task objective is achieved or the interaction limit is reached.Tasks can use different observation spaces and termination conditions, with rewards either zero or inferred from observations.
  • The agent trains on groups in G_train and is evaluated on unseen groups in G_test.

3. PAPRIKA

PAPRIKA trains language models for strategic exploration and sequential decision-making by using diverse synthetic interaction tasks, preference tuning, and curriculum learning. Its learned decision-making behaviors can generalize to unseen task groups without additional training.

  • PAPRIKA: PAPRIKA trains models on diverse decision-making tasks requiring strategic information gathering, then improves their preference for higher-performing interaction trajectories.The framework designs complex multiturn, partially observable, text-based tasks and uses sequential DPO-style tuning on successful versus lower-scoring trajectories.
  • PAPRIKA: PAPRIKA's learned decision-making behaviors often generalize to unseen task groups without additional training.Training on different subsets of tasks improves performance on tasks excluded from training.
  • 3.1. Task Design: The benchmark contains 10 task groups spanning guessing games, code-like feedback tasks, dynamic text environments, Minesweeper, and multiturn bandit selection.These tasks require agents to use interaction outcomes to explore and exploit under partial observability.
  • 3.2. Dataset construction: Training data is generated with high-temperature Min-p sampling, while preference pairs use the highest-scoring trajectory against a randomly selected lower-scoring trajectory.Randomly selecting the lower-scoring trajectory preserves dataset diversity rather than always choosing the worst outcome.
  • 3.4. Scalable Online Curriculum Learning: PAPRIKA's curriculum prioritizes tasks with meaningful learning signals because trajectories from tasks that are too difficult may not provide useful training experience.The method uses trajectory variability relative to reward scale to compare task learning potential and addresses the expense of generating rollouts.
  • 3.4. Scalable Online Curriculum Learning: The curriculum's task-selection quantity is intended to make learning potential comparable across task groups with different reward scales.The paper notes that selection algorithms could be replaced by more sophisticated online learning methods.

4. Empirical Results

PAPRIKA improves multi-turn decision making and often transfers learned strategies to unseen task groups without additional training. Curriculum sampling further improves data efficiency, while standard evaluations show no noticeable degradation in regular capabilities.

  • Decision-making performance: PAPRIKA raises Llama-3.1-8B-Instruct’s average bandit success rate from 42.25% to 62.25% using trajectories from other task groups only.This demonstrates transfer of strategic exploration without synthetic trajectories from an optimal bandit algorithm.
  • Decision-making performance: Across 10 task groups, PAPRIKA increases Llama-3.1-8B-Instruct’s performance by 47% of its original success rate after training on about 22,500 trajectories.The improvement also occurs for Gemma-3-12B-It, according to the supplied results passage.
  • Generalization to unseen task groups: PAPRIKA (LOO) improves success rates on 9 of 10 task groups, while PAPRIKA (full) outperforms single-task-group training on 7 of 10 groups.LOO models can match or exceed group-specific training, although transfer varies and can be absent or negative on some groups.
  • Curriculum learning: Curriculum sampling improves twenty-questions training over uniform sampling by 1.4% in average success rate and 3.3% in pass@4 success rate.The comparison uses three rounds with Llama-3.1-8B-Instruct and reports success across easy, medium, and hard task groups.
  • Task efficiency and regular capabilities: PAPRIKA reduces the average turns needed to solve tasks, indicating more optimal intermediate actions, and does not cause noticeable degradation on standard benchmarks.The standard-capability comparison is reported against Llama-3.1-8B-Instruct.

5. Related Works

Related work covers alignment methods for multi-turn interaction, in-context reinforcement learning across environments, and curriculum learning for improving training efficiency. PAPRIKA differs by targeting general information seeking and cross-environment decision-making with the LLM itself as the agent.

  • LLM alignment: Unlike single-turn alignment and post-training methods, PAPRIKA studies multi-turn interaction in which an agent iteratively acts in an environment.Prior multi-turn environments and datasets include LMRL-Gym and related textual interaction settings.
  • In-context reinforcement learning: Prior in-context reinforcement learning studies commonly use grid worlds, bandits, or mazes and generate data with random policies or existing RL algorithms.PAPRIKA instead studies diverse environments and generalization to completely new environments.
  • In-context reinforcement learning: PAPRIKA targets general information seeking across tasks, unlike concurrent work focused on medical reasoning or using an LLM to assist a separate decision-making agent.PAPRIKA uses the LLM directly as the decision-making agent.
  • Curriculum learning in RL: Curriculum learning presents data nonuniformly, and easier tasks may scaffold learning of more difficult tasks in reinforcement learning.Related environment-design methods alter environment distributions or generate environments to maximize learning progress.
  • Curriculum learning in RL: Unlike a concurrent curriculum method requiring rollouts for every example, PAPRIKA designs its curriculum using grouping metadata and a constant number of rollouts per task group.The comparison concerns how task learnability is estimated for curriculum construction.

6. Discussion

PAPRIKA provides a scalable fine-tuning method whose learned strategies can generalize zero-shot to unseen tasks, but its effectiveness depends on the starting model and currently relies on offline preference tuning and human-built environments.

  • Limitations: PAPRIKA depends on the starting model generating sufficiently good behaviors within a reasonable sampling budget, so weaker base models can limit performance.The method uses rejection sampling on self-generated data to teach better behaviors.
  • Limitations: The current approach uses offline preference tuning because of computational constraints, while online reinforcement learning remains a proposed direction.The authors expect online RL on diverse tasks could produce larger improvements in in-context RL capabilities.
  • Limitations: The environments require substantial human implementation effort, motivating future work on scalable automatic generation of suitable training tasks.The environments were designed with GPT-4o-mini assistance but still required considerable human effort.

Impact Statement

The work aims to improve strategic exploration and decision making in language-model agents, with potential real-world impact if agentic systems become widespread. Its societal effects remain open because experiments use relatively simple, controlled environments.

  • The approach targets better strategic exploration and decision making in language-model agents, with potential real-world impact if agentic systems become widespread.
  • Experiments use relatively simple and controlled environments, leaving the impacts of truly agentic systems on society as an open question.

B.1. Summary of Task Groups

The paper evaluates ten sequential, interactive task groups spanning information gathering, deduction, logical reasoning, exploration–exploitation, and decision making. These environments vary in their initial states, action spaces, observations, interaction formats, and environment implementations.

  • Twenty Questions and Guess My City test strategic questioning to identify hidden topics or cities within at most 20 questions.Twenty Questions restricts answers to yes/no-style interaction, whereas Guess My City permits broader questions that cannot directly reveal the city name.
  • Customer Service and Murder Mystery test directed information gathering for troubleshooting and culprit identification through interactive clues and actions.
  • Wordle and Mastermind test deductive reasoning by requiring agents to use structured feedback from guesses to refine subsequent attempts.Wordle allows six guesses for a five-letter word, while Mastermind allows twelve turns to identify a four-digit code.
  • Minesweeper, Cellular Automata, Bandit Best Arm Selection, and Battleship probe sequential logic, coding refinement, exploration–exploitation, and informative questioning.Bandit evaluation uses multi-turn reward feedback with randomized arm rewards and a best-arm objective rather than long-horizon regret.

C. Details of Training Dataset Construction

The training dataset is built from self-generated interaction trajectories, retaining successful behavior for supervised fine-tuning and pairing preferred and dispreferred trajectories for DPO. Minesweeper requires special caution because its environment depends on the agent’s first action.

  • For each training task, the authors generate 20 trajectories, except for Mastermind, where 100 are generated because the base model has a low success rate.Generation uses temperature 1.5 and Min-p 0.3; trajectories longer than 12,000 tokens are discarded.
  • Supervised fine-tuning retains successful trajectories with distinct turn counts for each task.Trajectories exceeding 12,000 total tokens are filtered mainly because of memory constraints from long contexts.
  • DPO uses each task’s fastest successful trajectory as preferred and randomly selects a lower-performing or failed trajectory as dispreferred.The construction uses one trajectory pair per task rather than multiple pairs.
  • Formatting violations cause trajectories to terminate and be filtered, while no other filtering mechanism is applied.
  • DPO or RPO is technically unsuitable for Minesweeper because random mine placement depends on the agent’s first action.The authors report no significant practical effect, but suggest unpaired preference learning such as KTO may be more appropriate.

D. Note about Task Environment Hacking

The paper addresses errors caused by LLM-simulated task environments that can produce incorrect observations or success rewards. A separate LLM judge reduces but does not eliminate these failures, which can therefore affect reported success rates.

  • LLM-simulated environments can be hacked into producing wrong intermediate observations or task success rewards.Strict string matching partly mitigates this issue for Twenty Questions, but open-ended tasks cannot fully prevent it.
  • A separate GPT-4o-mini judge filters trajectories mistakenly labeled successful by the simulated environment, but environment hacking remains possible.
  • 1.5% of 200 manually sampled successful trajectories showed environment hacking, and the authors caution that this error should be factored into success-rate comparisons.Using GPT-4o instead of GPT-4o-mini reduced observed mistakes, but GPT-4o-mini was retained because of API-call budget constraints.
  • Table 5 illustrates a Guess My City failure where the environment reveals the answer and the judge fails to detect the violation.

F. Additional Experimental Details

The experiments use dedicated GPU hardware for model training and inference, with project-wide data-generation and evaluation costs estimated at approximately $20,000.

  • Training uses one 8-GPU node: NVIDIA L40S GPUs for Llama-3.1-8B-Instruct and NVIDIA H100 GPUs for Gemma-3-12B-IT.
  • Inference and data generation use single NVIDIA A40 GPUs, while the full project’s estimated API cost is approximately $20,000.

G. Public Release of Code, Model and Dataset

The authors release code, training datasets, model checkpoints, and supporting examples for reproducing PAPRIKA and generating task data.

  • The codebase for reproducing PAPRIKA’s results is publicly available.
  • The released resources include supervised fine-tuning and preference-training datasets, plus a PAPRIKA-trained Llama-3.1-8B-Instruct checkpoint.
  • The appendix provides an example conversation and procedure for generating twenty-questions difficulty labels from estimated guessing difficulty and turns.The procedure generates 20 trajectories per task, treats failed trajectories as taking the environment’s maximum of 20 turns, and uses turns as a proxy for reward.
  • The main comparisons report pass@4 and average success rate across task groups, alongside average turns required to solve tasks.The turn metric excludes bandit best-arm selection because its interaction budget is fixed at 20.

I.3. PAPRIKA Imporoves Task Success Rate on Gemma-3

On Gemma-3-12B-IT, PAPRIKA improves or matches success rates across all task groups and often reaches GPT-4o-mini performance.

  • PAPRIKA matches or outperforms GPT-4o-mini on 7 of 10 task groups by pass@4 success rate.
  • PAPRIKA reduces the average number of turns required by Gemma-3-12B-IT across all task groups, indicating improved information-seeking behavior.Bandit best-arm selection is excluded from turn comparisons because its number of turns is fixed at 20.

I.5. More Performance Metrics

Additional evaluations show that PAPRIKA transfers across sampling budgets, datasets, task groups, and model-training variants, while its gains depend on structured exploration data and the finetuning stage.

  • I.5. More Performance Metrics: PAPRIKA-finetuned Llama-3.1-8B-Instruct and Gemma-3-12B-IT outperform their regular instruct counterparts for pass@k values from 1 through 4.
  • I.7. Evaluation on LMRL-Gym split: PAPRIKA’s gains on guess my city and twenty questions transfer to the disjoint LMRL-Gym split.
  • I.8. Experiments on Modified Wordle to Further Test Generalization: On modified Wordle with non-five-letter words, PAPRIKA outperforms the starting Llama-3.1-8B-Instruct without training on that task group.The evaluation averages success over 1,000 tasks with three random seeds.
  • I.9. Ablation Study over Different Finetuning Stages of PAPRIKA: RPO improves performance beyond SFT alone on all six ablated task groups.
  • I.10. Finetuning on regular multiturn data does not help: Fine-tuning on 100,000 randomly sampled WildChat trajectories causes performance degradation across task groups, unlike PAPRIKA’s structured exploration data.The authors speculate that WildChat prioritizes coherence rather than information gathering.
  • I.11. Performance comparison between different starting models: The reported model comparison is limited because the main experiments use Llama-3.1-8B-Instruct, and reasoning-model evaluations did not produce final answers within 10,000 tokens.The authors also hypothesize that PAPRIKA’s gains depend on the base model’s quality and diversity.

J. Limitations of PAPRIKA: Evaluation on Standard Bandit

PAPRIKA improves empirical regret over Llama-3.1-8B-Instruct on unseen bandit tasks with few arms, but this advantage disappears as the number of arms increases.

  • Evaluation on Standard Bandit: The bandit evaluation defines empirical regret from the optimal arm reward and the reward obtained by the policy’s chosen arm over timesteps.The metric is averaged across 100 trials.
  • Evaluation on Standard Bandit: PAPRIKA improves empirical regret over Llama-3.1-8B-Instruct only when the bandit task has a small number of arms.The evaluation reports empirical regret averaged across 100 trials.
  • Evaluation on Standard Bandit: As the number of arms grows, the performance gap between PAPRIKA and Llama-3.1-8B-Instruct vanishes.This result limits the transfer of PAPRIKA’s learned decision-making behavior to larger bandit problems.

K. Example Trajectories

Qualitative trajectories show PAPRIKA asking more concise, higher-quality questions and making better guesses than Llama-3.1-8B-Instruct, although the evaluation is limited.

  • Twenty Questions: In the orca task, PAPRIKA asks more concise questions and reaches the final topic quicker than Llama-3.1-8B-Instruct.The comparison is based on qualitative behavior trajectories.
  • Twenty Questions: In the primary-numbers task, PAPRIKA guesses the concept in 8 turns, while Llama-3.1-8B-Instruct fails within 20 turns across all 4 attempts.PAPRIKA also avoids redundant questions that appear in the baseline trajectories.
  • Wordle: In Wordle, PAPRIKA produces more concise, higher-quality guesses, whereas Llama-3.1-8B-Instruct reasons incorrectly against the correct answer.The baseline reaches the correct answer but then rejects it and makes a wrong guess.
  • Overall Assessment: Overall, the limited trajectory evaluation indicates quantitative and qualitative improvement over Llama-3.1-8B-Instruct.The paper leaves a detailed study of model behavior and generation quality for future work.
Loading 2502.17543v4…