PaperlayerView exact revision on arXiv

Source-linked AI summary

Reading the Room: Implicit Confusion Encoding in Recurrent World Model States

Donald Aadithiyan

arXiv:2608.21582v1cs.LG

TL;DR

This paper asks whether a recurrent world-model state can encode its own recent confusion despite training only for prediction. It tests that signal against novelty and ensemble disagreement, characterizes and edits it, and finds a causally used signal whose geometry and closed form generalize across three control tasks, while practical routing helps two of three.

  • Problem

    The paper asks whether DreamerV3's recurrent hidden state tracks sustained prediction confusion, a property not explicitly trained and distinct from novelty or current reconstruction error.

  • Method

    The authors hold KL divergence fixed while varying reconstruction quality, probe h_t, model the probe with a discounted count of recent high-error steps, and test causality through hidden-state edits and real-value substitution.

  • Results

    The confusion probe reaches AUROC 0.72 while ensemble disagreement scores below chance; the signal lies 88 degrees from the top 50 principal components and carries only 9% of its variance there.

  • Takeaways & Limitations

    The hidden state can support decisions about when to check a real observation instead of trusting imagination, with routing useful on two of three tasks.

  • Takeaways & Limitations

    The experiments use a 256-dimensional GRU rather than full-scale DreamerV3, and the decisive dissociation test has one consistent exception across three tasks.

Abstract

from arXiv · show

World models built on the RSSM architecture, such as DreamerV3, keep a recurrent hidden state $h_t$ trained only to reduce prediction error. We show this state also tracks its own confusion, hiding in plain sight: nearly orthogonal to $h_t$'s directions of greatest variance, invisible to any variance-based method. It is functionally distinct from ensemble disagreement, which flags new inputs, and reconstruction error, which flags bad predictions right now. On a test holding prediction error fixed while confusion varies, a linear probe on $h_t$ finds the signal (AUROC 0.72, 5 runs), while an ensemble baseline scores below chance. A discounted count of recent high-error steps explains 80% of the probe's output ($R^2=0.80$). We confirm the signal is causally used, not merely present, by editing $h_t$ directly and watching behaviour change, including a check using real values from other trajectories instead of synthetic edits. Its geometry and closed form generalize across three control tasks; the decisive dissociation test itself holds cleanly on only one, and its practical use, deciding when to check reality instead of trusting imagination, generalizes to only two of the three tasks.

1 Introduction

The paper shows that a DreamerV3 recurrent state encodes accumulated confusion despite lacking an uncertainty objective, and distinguishes this signal from novelty and ensemble disagreement. The signal is geometrically hidden from dominant-variance analysis, causally relevant, and partly generalizes across tasks and uses.

  • Motivation: DreamerV3’s recurrent state h_t tracks recent confusion even though training only minimizes reconstruction error and KL divergence.The paper defines confusion as sustained surprise across consecutive steps, rather than a single current error.
  • Research questions: The paper tests whether confusion is separable from novelty, ensemble disagreement, causal correlation, and task-specific behavior.These tests are designed to distinguish a genuine signal from relabeling or a property limited to one architecture and task.
  • Dissociation: 0.72 AUROC: a linear probe on h_t detects confusion on a matched-KL test, while ensemble disagreement scores below chance.The two methods point in opposite directions on the same test, distinguishing familiar-but-confused states from novel inputs.
  • Geometry: 88 degrees from the top 50 components and only 9% of its own variance there: the confusion direction is nearly orthogonal to dominant PCA directions.Variance-based representation analysis would therefore discard the signal.
  • Generality and use: A discounted recent-error account remains stable across five seeds, causal editing supports load-bearing use, and direction, geometry, and closed form carry across three control tasks.The decisive dissociation test is clean on only one task, while the operational reality-checking use helps two of three tasks.

3 Setup and Method

The method trains a small recurrent world model, constructs evaluation sets that separate confusion from novelty, fits a hidden-state probe, and tests its causal role against empirical null directions. It also characterizes the probe as a discounted history of recent high-KL steps.

  • Pipeline: The pipeline trains once, builds three evaluation sets, fits a probe, tests causal effects against an empirical null, and repeats across tasks and seeds.The evaluation also checks two operational uses.
  • Model and data: The Mini-DreamerV3 uses a 256-dimensional GRU, 32 × 32 categorical latents, and about 12 million parameters, trained for 100,000 dm-control steps.Recorded quantities include h_t, z_t, KL_t, and reconstruction error.
  • Evaluation sets: Set C pools clean and noisy states, then splits them by reconstruction error within matched-KL bins, producing roughly a 9× error difference.A probe above chance in Set C must detect information beyond KL magnitude.
  • Probe evaluation: The probe is logistic regression on h_t predicting above-median KL, evaluated by AUROC with bootstrap intervals or across-seed mean and standard deviation.Sets A and B contain fresh in-distribution states and the same states with Gaussian observation noise.
  • Closed-form characterization: C_t is a running discounted count of past high-KL steps, with γ down-weighting older steps; γ is fitted by regressing probe scores across candidate values.A high-KL step is marked when KL exceeds the median.
  • Assumption: The closed-form comparison assumes h_t is sufficient for recent prediction-error history, but the paper tests this proposition rather than taking it for granted.Under that assumption, probes trained on binarised KL labels and directly on C_t should reach approximately the same Set C score.
  • Causal intervention: The causal test subtracts, adds, or substitutes the confusion-direction component of h_t, then compares later probe changes with 50 matched-scale random directions.The effect Δ_k measures immediate and persistent changes, and the real-state substitution check avoids synthetic edits.

4 Results and Discussion

The hidden-state confusion signal is distinct from novelty, encoded outside dominant-variance directions, causally load-bearing, and broadly transferable, though evaluation and routing results vary by task.

  • 4.1 Confusion Is Not Novelty: 0.72 AUROC: the probe detects confusion on matched-KL Set C, while the ensemble baseline scores 0.550 near chance.On an unmatched novelty test, reconstruction error reaches 0.996, ensemble disagreement 0.966, and the probe 0.49, separating the signals in both directions.
  • 4.2 What the Probe Computes: R^2 = 0.798: a discounted recent-error count explains the probe score, versus R^2 = 0.519 for current KL alone.The discounted count uses γ = 0.95, corresponding to roughly a 13-step memory.
  • 4.3 Geometry, and Why It Is Not a Gate Artefact: 88.2 degrees: the confusion direction is nearly orthogonal to the top 50 principal components and retains only 9% of its variance there.This geometry persists when GRU gates are fixed between 0.5 and 0.99, ruling out gate saturation as the explanation.
  • 4.4 The Direction Is Causally Load-Bearing: −0.586: subtracting the confusion direction produces an extreme immediate probe-score change, replicated by real-state substitution at −0.761.Both effects reach the 100th percentile of the 50-direction empirical null; the effect decays across look-ahead rather than disappearing immediately.
  • 4.5 Generality Across Three Tasks, and Why One Metric Inverts: 88.0–89.4°: confusion geometry generalizes across cartpole, reacher, and pendulum, while Ct fit ranges from R^2 = 0.26 to 0.86.Pendulum’s Set C score inverts to 0.32–0.40, predicted by the within-matched-KL correlation between reconstruction error and Ct.
  • 4.6 An Operational Use: Deciding When to Check Reality: At a 30% query budget, probe routing beats reconstruction-error routing on cartpole and pendulum but loses on reacher.Against raw KL, the probe wins only on reacher (+0.273 ± 0.020), is unstable on cartpole (−0.006 ± 0.025), and loses on pendulum (−0.030).

5 Limitations

The results are constrained by the small model scale, task-specific exceptions, failed correction attempts, and limited evidence for routing and KL-based comparisons.

  • 256-dimensional GRU experiments are three orders of magnitude smaller than full DreamerV3, although geometry and causal-editing results hold at double width.The authors identify scaling to full size as the natural next step.
  • Reacher is the one consistent exception across the three checks because its reconstruction error is already unusually strong and confusion-aligned.
  • Two direct corrections to imagined-rollout value estimates failed, so the signal detects confusion but does not yet correct for it.
  • Routing beats reconstruction error on two of three tasks, while the probe beats KL alone only on reacher and should be checked on new tasks.

6 Conclusion

The work characterizes a small world model whose hidden state tracks recent confusion without explicit training, with a simple, causally load-bearing signal that transfers across structurally different tasks. It also reports that the experiments ran on a single laptop CPU and that routing can reserve real observations for when they are needed.

  • A DreamerV3 world model’s hidden state tracks its own recent confusion, distinct from novelty and ensemble disagreement, without being explicitly trained to do so.
  • The signal has a simple closed-form description, is causally load-bearing rather than merely correlated, and carries over to structurally different control tasks.
  • Every result ran on a single laptop CPU, lowering the compute barrier for causally validated interpretability findings.
  • Routing can reserve costly real observations only when needed, offering a direct efficiency payoff in resource-scarce settings.

A Full Experimental Setup

The experiments use small DreamerV3-style models across three control tasks and evaluate held-out, noisy, and KL-matched state sets with deterministic linear probes and ensemble baselines.

  • Model and training: Mini-DreamerV3 uses a 256-dimensional deterministic GRU state, a 32 × 32 categorical stochastic state, and about 12 million parameters.The GRU receives [z_t−1, a_t−1] at each step.
  • Tasks and collection: The study trains on cartpole-swingup, reacher-easy, and pendulum-swingup for 100,000 steps each, using uniform-random-policy analysis on frozen models.
  • Evaluation sets: Set A contains 10,000 fresh in-distribution states, while Set B reuses them with σ = 0.1 Gaussian observation noise.
  • Evaluation sets: Set C matches KL by construction at 22.9 versus 23.9 nats while reconstruction error differs 9×, testing separation beyond KL magnitude.
  • Evaluation and reporting: The same protocol is applied independently to all three tasks, with per-task ROC results reported for Sets A, B, and C.Set A reaches AUROC 0.872, Set B 0.807, and Set C 0.714 in the cited figures.

B Full Causal Validation

Causal validation edits the confusion direction in held-out states and compares the effects with empirical nulls, later look-ahead steps, real-value substitutions, routing, and forward-dynamics controls. The effects are extreme, decay over time as predicted by the closed form, and replicate unevenly across seeds.

  • Setup: 600 held-out states are used for causal edits, disjoint from probe training, while the setup is summarized in Figure 7.
  • Synthetic ablation: −0.586 immediate probe-score change lies at the 100th percentile of a 50-direction null, with z ≈−22; the effect remains extreme through k = 10.
  • Decay and closed form: The ablation decay follows γ^k, linking the causal effect’s time constant to the discounted-history account of the probe.
  • Behavioural consequences: 0.817 routing separation versus a 0.247 ± 0.019 null average reaches the 100th percentile, while confusion correlates with imagined-real drift and reward over-estimation.
  • Specificity controls: Forward-dynamics drift and next-observation reconstruction do not separate from random-direction nulls, unlike the probe readout and routing decision.
  • Real-value validation: Real-value substitution changes only the confusion-direction projection and produces Δprobe = −0.761 and routing-flip rate 0.868, matching the synthetic causal pattern.
  • Replication: Across five models, probe-decay effects are extreme in all five, routing flips in three of five, and next-step prediction error never separates from its null.
  • Gate and geometry controls: The full probe remains at AUROC 1.0 across forced gate values from 0.5 to 0.99, whereas magnitude-only effects track overwrite fraction and collapse as overwriting vanishes.

C Why Set C Inverts on Pendulum

Pendulum’s Set C inversion arises from how reconstruction error tracks the confusion count within a fixed KL bin: its negative within-bin correlation reverses the label direction. This is not explained by kinematic hardness, while cartpole’s positive trend supports a working Set C test.

  • Pendulum’s Set C scores were below chance across all four runs, with mean 0.396±0.063.The per-seed scores were 0.322, 0.395, 0.496, and 0.371.
  • Pendulum’s inversion is not caused by dynamically hard instants such as states near the top of the swing.Reconstruction error correlates more with Ct than with any single kinematic feature.
  • Reconstruction error correlates negatively with Ct on pendulum, producing the inverted Set C direction.The within-bin correlation is −0.12, matching pendulum’s Set C result of 0.32.
  • Cartpole has a positive reconstruction-error–Ct trend and a working recon-based Set C label despite kinematics dominating reconstruction error.Its within-bin correlation is +0.39, matching a Set C result of 0.72.
  • Relabelling pendulum directly by Ct only recovers chance-level performance because KL already explains most of Ct there.The relabelled Set C score is 0.53 at matched KL.

D Scale Check

A doubled-width GRU preserves the study’s two load-bearing geometric and causal findings, but the check remains partial rather than a full XL-scale validation.

  • At 512 dimensions, the confusion direction remains near-orthogonal to the top principal components.The double-width table reports preservation of the near-orthogonal geometry.
  • At doubled width, ablating the confusion direction collapses the probe beyond a random-direction control.The table describes both findings as load-bearing and replicated in the width check.
  • The doubled-width result is partial evidence against an XS-specific finding, not a resolution of the full XL-scale question.

E Two Negative Results

Two interventions fail to improve imagined-rollout value estimates: down-weighting by confusion worsens rewards, while early stopping is only weakly positive and unstable.

  • Down-weighting imagined-rollout value estimates by the probe output worsens performance, with ∆r = −0.526 versus no weighting.Using continuous Ct performs worse still, with ∆r = −0.576.
  • Early stopping after confusion crosses the 90th percentile gives only a small unstable benefit across five samples.It truncates 24% of rollouts; mean ∆r = +0.005 ± 0.009, with sign instability.
  • Both approaches support confusion for detecting problematic imagined rollouts, not for correcting their computed value estimates.Early stopping at least does not make things worse.

F Routing: Probe-A Versus a Ct-Direct Router and a KL-Only Baseline

Routing experiments compare Probe-A with a Ct-direct router and a KL-only threshold under the same 30% query budget. Ct-direct does not rescue the dissociation, while only reacher shows a reliable probe advantage over KL.

  • Routing setup: All routers are evaluated at a 30% query budget against the same reconstruction-error baseline on all three tasks.The KL-only router uses KLt−1, matching Probe-A’s pre-decision information regime; same-step KLt is excluded as tautological.
  • Ct-direct router: The Ct-direct router tracks Probe-A closely and is consistently slightly worse, with recall-vs-budget AUC lower by 0.01–0.03 on every task.
  • KL-only baseline: On pendulum, a zero-parameter KL threshold beats the trained probe, 0.808 versus 0.778 across 3 of 3 seeds.The routing gain over reconstruction error remains +0.30, but KL redundancy accounts for it.
  • KL-only baseline: On cartpole, the probe has no reliable advantage over KL, with −0.006 ± 0.025 and an unstable sign across five seeds.
  • KL-only baseline: Reacher is the only task where the probe adds a replicated advantage over KL, +0.273 ± 0.020 in 3 of 3 seeds.Its confusion fit is weakest there, with R2 = 0.26 and the shortest memory.
Loading 2608.21582v1…