🤖 machine learning

LatentGym: A Testbed For Cross-Task Experiential Learning With Controllable Latent Structure

The paper introduces LatentGym, a controllable testbed featuring ground-truth latent structures that enables the evaluation of cross-task experiential learning in agentic systems by separating exploration from exploitation, thereby facilitating the study of how and why LLMs adapt across sequential tasks.

Original authors: Daksh Mittal, Tommaso Castellani, Thomson Yen, Naimeng Ye, Fangyu Wu, Minghui Chen, Tiffany Cai, Emmanouil Koukoumidis, William Zeng, Hongseok Namkoong

Published 2026-06-16

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Daksh Mittal, Tommaso Castellani, Thomson Yen, Naimeng Ye, Fangyu Wu, Minghui Chen, Tiffany Cai, Emmanouil Koukoumidis, William Zeng, Hongseok Namkoong

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a very smart but slightly rigid robot to play a series of games. In the real world, we humans are great at "connecting the dots." If you play a game where the answer is always one of three numbers, and you play it ten times, you quickly realize, "Hey, the answer is always one of these three!" You stop guessing randomly and start guessing those three numbers. You learn from experience.

The paper argues that today's most advanced AI agents (like the ones powering chatbots) are surprisingly bad at this. They tend to treat every single game as if it's brand new, forgetting everything they learned in the previous nine games.

To prove this and figure out how to fix it, the authors built LatentGym.

The Playground: LatentGym

Think of LatentGym not as a single game, but as a factory for creating game sequences.

The "Latent" (The Hidden Rule): In a normal game, the rules are fixed. In LatentGym, there is a secret "hidden rule" that applies to all the games in a sequence. For example, in a number-guessing game, the hidden rule might be: "The number is always either 137 or 793." The AI doesn't know this at the start; it has to figure it out by playing.
The Control Knobs: The researchers can tweak the difficulty like a soundboard:
- The Prompt: How much of the secret rule do they tell the AI? (Nothing? A vague hint? The whole truth?)
- The Feedback: After a game, do they just say "You won/lost," or do they reveal the actual answer so the AI can learn?
- The Horizon: How many games are in the sequence? (5? 10? 20?)

This setup is crucial because, in most AI tests, if an AI fails, you don't know why. Did it not see the pattern? Did it see the pattern but forget to use it? In LatentGym, the researchers know the secret rule perfectly, so they can pinpoint exactly where the AI's brain short-circuited.

The Problem: How AI Fails

When the researchers tested top-tier AI models (like GPT-4o and Claude) on these simple sequences, they found three specific ways the AI failed to learn:

Adaptation Neglect (The "Reset Button" Syndrome): The AI ignores the fact that it has played before. Even if the answer is always "Red," it treats the second game as if it's the first time ever, starting with a random guess instead of checking "Red" first. It's like a student who forgets they took a math test yesterday and tries to re-derive the formula from scratch today.
Adaptation Breakdown (The "Squinting" Syndrome): The AI notices a pattern ("Wait, the numbers are getting smaller!"), but it doesn't know how to use that information. It sees the clue but keeps playing the old way. It's like seeing a "Wet Floor" sign but still walking normally instead of slowing down.
Adaptation Miscalibration (The "Over-Thinker" Syndrome): When the researchers explicitly told the AI the rule ("The answer is always 137 or 793"), the AI sometimes got confused and performed worse than when it wasn't told anything. It tried so hard to follow the rule that it ignored the actual game mechanics. It's like being told "Don't think about a pink elephant" and then spending the whole time thinking about pink elephants, forgetting to play the game.

The Solution: Cross-Task RL

The researchers tried to teach the AI to learn better using a method called Cross-Task Reinforcement Learning (RL).

Old Way (Single-Task RL): You train the AI on Game 1, then Game 2, then Game 3. It learns to win Game 1, then Game 2, but it doesn't learn to connect them.
New Way (Cross-Task RL): You train the AI on the entire sequence at once. You tell it, "Your goal isn't just to win Game 1; your goal is to win Game 10 by learning from Games 1 through 9."

The Results:
When they used this new training method, the AI actually learned to adapt.

It started looking for patterns early on.
It started using those patterns to win later games faster.
Crucially, this learning generalized. Even when they tested the AI on a new set of games with a new hidden rule it had never seen before, it was still better at figuring it out than the AI trained the old way. It learned a "meta-skill" of how to learn.

A Surprising Twist: Less Information is Better

One of the most interesting findings was about feedback.

Rich Feedback: Telling the AI the exact answer after every game.
Sparse Feedback: Just telling the AI "You won" or "You lost."

Counter-intuitively, the AI trained with Sparse Feedback (less information) actually generalized better to new situations than the one trained with Rich Feedback. It seems that when you give the AI the answer too easily, it becomes lazy or over-reliant on that specific hint. When it has to figure things out with just a "win/loss" signal, it builds a more robust strategy that works even when the hints change.

The Bottom Line

The paper introduces a new "gym" (LatentGym) to test if AI agents can learn from a sequence of related tasks. They found that current top AIs are bad at this, often failing to connect the dots between tasks. However, by training them on full sequences of tasks (Cross-Task RL) rather than isolated tasks, they can learn a general ability to adapt. This framework allows researchers to see exactly why an AI fails and how to fix it, moving us closer to agents that truly learn from experience.

Technical Summary: LatentGym

Problem Statement

Current training and evaluation frameworks for Large Language Model (LLM) agents lack the ability to rigorously study cross-task experiential learning: the capacity of an agent to improve over a sequence of related tasks by inferring and utilizing shared hidden structures (latents). Existing benchmarks often treat tasks independently or rely on implicit structures that are not controllable by the experimenter. Consequently, it is difficult to diagnose why an agent fails to adapt (e.g., whether it failed to detect the structure, failed to infer it, or failed to act on it) or to compare methods that aim to instill this capability.

Methodology: LatentGym

The authors introduce LatentGym, a controllable testbed designed to isolate and measure cross-task adaptation. The framework is built on three core components:

1. Controllable Environments

LatentGym organizes environments around an explicit, ground-truth latent variable ( $z$ ) that governs the regularity shared across a sequence of $N$ tasks. The framework factorizes each environment along five orthogonal axes:

Core Dynamics: The within-task decision problem (e.g., number guessing, bandits, word ladders).
Latent: The ground-truth structure shared across the sequence (e.g., a fixed set of target numbers, a rotating best arm in bandits).
Prompt Type: Controls information availability regarding the latent (No information, Vague hint, or Full information).
Feedback Type: Determines post-task information (Standard: success/failure only; Information: ground-truth outcome revealed).
Horizon ( $N$ ): The number of tasks in the sequence.

The suite includes seven environments built on text-based games (Number Guessing, Bandits, Secretary, Mastermind, Word Ladder, Wordle, Hangman), each with multiple latent configurations.

2. Diagnostics: Exploration vs. Exploitation

To move beyond aggregate reward metrics, LatentGym introduces a counterfactual agent-switching mechanism. At a switch point $K$ , the history gathered by an "explorer" agent is passed to an "exploiter" agent. This allows the separation of two distinct capabilities:

Exploration Efficiency: Whether the agent's actions in early tasks gather informative evidence about the latent.
Exploitation Efficiency: Whether the agent effectively utilizes that evidence to improve performance in later tasks.
Metrics include cumulative reward ( $R$ ), cross-task gain ( $r_N - r_1$ ), and final-task reward ( $r_N$ ).

3. Integrated RL Pipeline

The framework integrates with SkyRL, supporting standard algorithms (SFT, PPO, GRPO). It enables Cross-Task RL, where the policy is fine-tuned end-to-end on full sequences of related tasks using a cumulative reward objective ( $\sum r_i$ ). This contrasts with Single-Task RL, which optimizes for isolated tasks.

Key Contributions and Results

1. Failure Modes of Frontier Models

Evaluating frontier models (GPT-4o, Claude Sonnet 4.6, Gemini 2.5 Flash) on LatentGym reveals that they often fail to adapt even under simple latent structures. The authors identify three recurring failure modes:

Adaptation Neglect: Models repeatedly restart strategies (e.g., binary search) for each task, ignoring cross-task evidence even when hints are provided.
Adaptation Breakdown: Models recognize patterns but fail to act on them reliably, or act on them in ways that degrade performance (e.g., over-fitting to a pattern).
Adaptation Miscalibration: Providing explicit information about the latent can paradoxically degrade performance, as models may over-condition on the provided rule and violate operational constraints.

2. Cross-Task RL Induces General Adaptation

The authors demonstrate that Cross-Task RL (fine-tuning on full sequences) successfully instills cross-task adaptation capabilities that Single-Task RL fails to achieve.

Performance Gains: Cross-Task RL significantly outperforms both base models and Single-Task RL in cumulative reward, cross-task gain, and final-task reward across all tested environments.
Generalization: The learned adaptation strategies generalize to held-out latents (OOD-1) and held-out environments (OOD-2), suggesting the acquisition of a general meta-strategy rather than environment-specific cues.
Mechanism: Analysis via agent-switching shows that Cross-Task RL improves both exploration and exploitation efficiency, though the dominant factor varies by environment (e.g., exploitation dominates in Number Guessing, while exploration dominates in Hangman).

3. Impact of Training-Time Design Choices

The study investigates how training conditions affect generalization:

Feedback Density: Counterintuitively, training with standard (sparse) feedback transfers more robustly to deployment conditions than training with information-rich feedback. Models trained on full information suffer significant performance drops when evaluated under standard feedback.
Prompt Richness: Training with vague hints ("some info") yields models that are more uniformly transferable across prompt conditions compared to models trained with full information, which show a strong in-distribution advantage but poor cross-prompt transfer.

Significance

The paper positions LatentGym as a foundational instrument for the principled study of cross-task experiential learning in LLM agents. By providing a shared, controllable laboratory with ground-truth latents, it enables:

Diagnosis: Isolating specific failure modes (neglect, breakdown, miscalibration) that aggregate benchmarks obscure.
Algorithm Development: Testing whether training on sequences induces general adaptation strategies.
Design Insights: Revealing that "less is more" in training data (sparse feedback) for robust generalization.

The authors conclude that LatentGym offers a controlled foundation for designing agents that adapt more reliably in sequential, personalized, and interactive settings, moving beyond the current limitation of treating tasks in isolation.