The Yokai Learning Environment: Tracking Beliefs Over Space and Time

This paper introduces the Yokai Learning Environment (YLE), a new open-source benchmark for zero-shot coordination that overcomes the saturation of the Hanabi Learning Environment by requiring agents to track moving cards and reason under ambiguous hints, thereby revealing that current state-of-the-art methods fail to maintain consistent internal models when paired with unseen partners.

The Gist

Imagine you are playing a card game with a stranger you've never met before. You can't talk, you can't text, and you can't even see their cards. You only see a few cards on the table, and every time you move a card, it changes the layout of the whole board.

Your goal? To sort the cards into color groups as fast as possible. But here's the catch: the faster you finish, the more points you get. If you wait too long to be "sure," you lose points. If you finish too early and guess wrong, you lose everything.

This is the core challenge of the Yōkai Learning Environment (YLE), a new "test track" for Artificial Intelligence researchers.

The Problem: The "Hanabi" Problem

For years, the gold standard for testing how well AI agents can cooperate without talking was a game called Hanabi. Think of Hanabi as a game where you hold your cards facing away from you, and your partner tells you exactly what they are (e.g., "You have a blue 3").

Recently, AI got too good at Hanabi. They figured out the rules so perfectly that they can play with any random version of themselves and win almost every time. It's like a student who memorized the entire textbook so well that they can pass any test, but they haven't actually learned how to think critically. The "Hanabi test" is no longer hard enough to tell us if AI is getting smarter.

The Solution: The "Yōkai" Test

The authors created a new game called Yōkai (inspired by a real board game) to be a much harder, more realistic test. Here is why it's different, using some analogies:

1. The Moving Target (Space & Time)

• Hanabi: Your cards are in fixed slots. If you know "Slot 1 is Blue," it stays Blue in Slot 1.
• Yōkai: The cards are like roaming animals in a forest. You see a blue card, you move it, and now it's next to a green card. You have to constantly update your mental map: "Wait, that blue card I saw five minutes ago is now three steps to the left."
• The AI Challenge: The AI has to track moving objects in its head, not just static slots.

2. The Ambiguous Whisper (Communication)

• Hanabi: If your partner points at a card, they are telling the truth by the rules of the game.
• Yōkai: Your partner can drop a "hint card" that says "Blue or Green." It's a riddle, not a fact. Maybe they mean "This card is Blue," or maybe they mean "The card next to this one is Blue."
• The AI Challenge: The AI has to guess what the other person meant, not just what they said. It has to read between the lines.

3. The High-Stakes Gamble (Early Termination)

• Hanabi: You play until the game ends naturally.
• Yōkai: You can shout "I'm done!" at any time. If you're right, you get a massive bonus. If you're wrong, you get zero.
• The AI Challenge: The AI has to decide: "Do I have enough shared understanding with my partner to finish now, or should I keep playing and risk losing points?" This requires Theory of Mind—the ability to think, "What does my partner know? Do they know that I know?"

What Happened When They Tested the AI?

The researchers took the smartest AI agents that were "World Champions" at Hanabi and dropped them into Yōkai.

The Result: They crashed.

• The "Self-Play" vs. "Stranger" Gap: When these AIs played against copies of themselves (Self-Play), they did great. But when they played against a different copy of themselves (Cross-Play), they failed miserably.
• The Analogy: Imagine two people who learned to speak a secret language with each other. They can understand each other perfectly. But if you swap one of them with a twin who learned the same language but with slightly different slang, they can't understand each other at all.
• The "Calibration" Failure: In Yōkai, the AIs were terrible at knowing when to stop. They either stopped too early (guessing wildly) or waited too long (missing the bonus points). They couldn't agree on a "common ground."

Why This Matters

This paper proves that being good at one game doesn't mean you're good at cooperation in general.

The current AI methods are like students who memorized the answers to a specific math test. When you give them a new type of problem that requires actual reasoning, tracking moving variables, and interpreting ambiguous hints, they fail.

The Takeaway:
The Yōkai Learning Environment is a new, tougher gym for AI. It forces robots to stop memorizing rules and start learning how to:

1. Keep a mental map of moving things.
2. Interpret vague hints and riddles.
3. Trust their partners enough to make a risky decision together.

If AI can master Yōkai, it will be much closer to being able to work alongside humans in the real world, where things move, hints are vague, and we have to make split-second decisions together without a rulebook.

Technical Summary

Here is a detailed technical summary of the paper "The Y¯okai Learning Environment: Tracking Beliefs Over Space and Time."

1. Problem Statement

The paper addresses the limitations of current Zero-Shot Coordination (ZSC) benchmarks, specifically the Hanabi Learning Environment (HLE). While HLE has been the standard for evaluating cooperative AI and Theory of Mind (ToM) reasoning, recent algorithms (e.g., High-Entropy IPPO, Other-Play, Off-Belief Learning) have achieved near-perfect inter-seed cross-play performance in Hanabi. This saturation limits the ability to track further algorithmic progress and risks overfitting to Hanabi's specific structural constraints.

The core challenge identified is that Hanabi lacks:

• Dynamic spatial reasoning: Beliefs are tied to fixed hand slots, not moving entities.
• Ambiguous communication: Hints in Hanabi are truthful by rule; real-world collaboration often involves ambiguous or non-truthful signals.
• High-stakes termination decisions: Agents must decide when to stop based on inferred shared knowledge, balancing risk and reward.

The authors propose that existing ZSC methods fail to generalize to environments requiring spatio-temporal belief tracking and reasoning under ambiguous, non-truthful hints.

2. Methodology: The Y¯okai Learning Environment (YLE)

The authors introduce YLE, a new open-source, multi-agent Reinforcement Learning (RL) benchmark based on the cooperative board game Yokai.

Environment Design

• Game Mechanics: Players collaborate to group face-down cards by color on a grid without direct communication.
• Partial Observability: Agents privately observe only two cards per turn. They must track the location and color of cards that move across the grid.
• Ambiguous Communication: Players can place "hint cards" that may have multiple colors. Crucially, hints do not have to be truthful. A hint might indicate "Blue or Red," and the receiver must infer the actual color based on context and partner behavior.
• Termination Mechanic: Players can choose to end the game early.
  • Reward Structure: The score heavily penalizes unused hint cards. To maximize reward, agents must end the game as early as possible once they have established sufficient "common ground" (shared belief).
  • Risk: Ending too early (premature termination) based on incorrect beliefs results in a score of zero.
• Implementation: Built in JAX (via JaxMARL) for end-to-end GPU training, supporting hundreds of thousands of steps per second. It models the game as a Decentralized Partially Observable Markov Decision Process (Dec-POMDP) using graph representations for efficient legal move computation.

Metrics

Beyond standard Return ($R$), the paper introduces Successful Early Ending (SEE):

• EE (Early Ending Rate): Frequency of early termination.
• WEE (Calibration): Probability of winning given an early termination.
• Significance: High WEE indicates accurate belief formation and successful coordination; low WEE indicates guessing or poor common ground.

3. Experimental Setup & Adaptations

The authors evaluated leading ZSC algorithms adapted for YLE:

1. High-Entropy IPPO (HE): A state-of-the-art method for Hanabi.
2. Other-Play (OP): Optimizes for symmetry invariance. The authors extended OP to handle YLE's richer symmetries (color permutations + spatial grid rotations).
3. Off-Belief Learning (OBL): Constructs a hierarchy of policies and belief models. This is the first implementation of OBL outside of Hanabi.
   • Challenge: OBL must model beliefs over unobserved card colors and treat hints as unreliable evidence, unlike Hanabi where hints are truthful.
   • Cost: Requires longer fictitious rollouts ($K=4|N|$) compared to Hanabi ($K=2$), making it computationally expensive.

Baselines:

• Self-Play (SP) vs. Cross-Play (XP) performance.
• Human performance baseline (5 participants, 25 games).
• Ablation studies on entropy coefficients, symmetry groups, and neural architectures (CNN-GRU vs. GNNs).

4. Key Results

The experiments reveal that algorithms performing well in Hanabi struggle significantly in YLE.

• Persistent SP–XP Gaps: Unlike in Hanabi where SP and XP returns are nearly identical, YLE exhibits large gaps. For example, in 2-player 9C settings, the best SP return was ~7.5, while XP dropped to ~4.8 (a 43.9% degradation).
• Method Ranking Reversal:
  • Hanabi: HE > OBL > OP.
  • YLE: OP > HE > OBL.
  • High-Entropy IPPO (HE) collapses in multi-player settings, often learning to terminate immediately (game length = 1) to avoid penalties, failing to coordinate.
• Calibration Failure: In cross-play, the WEE (calibration) metric drops significantly, often to chance levels. This indicates agents cannot maintain consistent internal models of their partners' beliefs, leading to premature and failed game endings.
• Symmetry Sensitivity: OP performance degrades if only a subset of symmetries (e.g., color only) is enforced. Full spatial rotation invariance is required but difficult to enumerate perfectly.
• Memory Limitations: In imperfect memory settings (standard YLE), standard GRU-based policies fail completely. Transformer-XL (TrXL) improves performance but still falls short of human levels and perfect-memory agents.
• Belief Probing: Linear probes on hidden states show that while agents encode card colors well in Self-Play, this representation degrades significantly in Cross-Play, confirming that belief representations are not robust to unseen partners.

5. Key Contributions

1. YLE Benchmark: Introduction of the Y¯okai Learning Environment, an open-source JAX-based benchmark that forces agents to track beliefs over space and time, handle ambiguous hints, and make high-stakes termination decisions.
2. First OBL Implementation Outside Hanabi: A novel adaptation of Off-Belief Learning to a non-truthful hint environment, highlighting the computational and modeling challenges of this approach in complex settings.
3. Empirical Evidence of Benchmark Saturation: Demonstration that state-of-the-art ZSC methods (HE, OP, OBL) fail to generalize from Hanabi to YLE, proving that progress in ZSC is currently benchmark-specific.
4. Diagnostic Metrics: Introduction of Successful Early Ending (SEE) and Calibration (WEE) as critical metrics for evaluating ToM reasoning and common ground formation, moving beyond simple return scores.

6. Significance

This paper fundamentally challenges the assumption that ZSC is "solved" based on Hanabi results. It demonstrates that:

• Structural Diversity Matters: Algorithms optimized for truthful hints and static slots (Hanabi) do not transfer to environments with dynamic entities and ambiguous communication (YLE).
• Theory of Mind is Harder: The requirement to infer beliefs about moving objects and unreliable signals creates a much higher barrier for coordination than previously tested.
• Future Directions: The results suggest that future ZSC research must focus on robust belief tracking, handling ambiguity, and generalizing across diverse environmental structures rather than just tuning entropy or symmetry invariance for a single game.

The YLE serves as a necessary "stress test" to prevent overfitting and to drive the development of AI agents capable of true, flexible human-AI collaboration.