CIG: Exploration via Conditional Information Gain

Technical Summary: CIG: Exploration via Conditional Information Gain

Problem Statement
In reinforcement learning (RL) with sparse or delayed extrinsic rewards, agents must rely on intrinsic motivation to explore. Effective exploration requires identifying transitions that are informative about the environment's dynamics. However, existing intrinsic reward methods face a fundamental trade-off in how they utilize context:

• Lifelong rewards (e.g., ensemble disagreement, prediction error) condition on the accumulated replay buffer. They assign credit to every transition probing a model gap, regardless of whether the current trajectory has already explored that gap. This leads to redundancy within a single rollout (e.g., repeatedly visiting the same uncertain corridor in a maze).
• Episodic rewards (e.g., count-based, reachability bonuses) condition on the current rollout prefix. They penalize repetition within a trajectory but ignore lifetime progress. In model-based RL, where policies optimize over short "imagined" rollouts ($T \approx 15$), the prefix is often too short to provide a rich signal for redundancy correction.
• Hybrid methods attempt to combine these signals via heuristic weights or rely on Gaussian Process (GP) dynamics, which do not scale to high-dimensional state spaces.
• Trajectory-level Information Gain (IG) offers a principled objective conditioning on both the buffer and the prefix simultaneously. However, computing the exact IG requires marginalizing over a posterior distribution of model parameters, which is intractable for deep neural networks.

Methodology: Conditional Information Gain (CIG)
The authors derive Conditional Information Gain (CIG), a tractable surrogate for trajectory-level information gain that scales to high-dimensional latent spaces while conditioning on both the replay buffer and the rollout prefix.

1. Formulation: The objective is to maximize the mutual information between the model parameters $w$ and a trajectory $s_{1:T}$ given the replay buffer $D$: $I(w; s_{1:T} | D)$. Using the chain rule, this decomposes into per-step terms $I(w; s_t | s_{<t}, D)$.
2. Approximations:
   • Ensemble Posterior (A1): The posterior $p(w|D)$ is approximated by a uniform distribution over $M$ independently trained ensemble members.
   • Markov-Gaussian Dynamics (A2): Each ensemble member predicts the next latent state as a Gaussian with a shared isotropic aleatoric covariance $\sigma^2 I_d$.
   • Moment-Matched Gaussian (A3): The mixture entropy of the ensemble predictions is upper-bounded by the entropy of a Gaussian with matching mean and covariance.
3. The Rank-Saturation Bottleneck: Maximizing the log-determinant of the full epistemic covariance matrix $\Sigma$ (size $Td \times Td$) fails in high dimensions because the rank of the covariance is limited by the ensemble size ($M-1$). For typical values ($M=5, d=200$), a single observation can saturate the reward signal, leaving no credit for subsequent steps.
4. Trace Reduction (Design Choice D1): To resolve the capacity bottleneck, the authors collapse each $d \times d$ block of the epistemic covariance into its trace. This yields a kernel $K \in \mathbb{R}^{T \times T}$ where $K_{jt} = \frac{1}{M} \sum_k (\delta^{(j)}_k)^\top \delta^{(t)}_k$, with $\delta$ representing the deviation of ensemble members from the mean. This raises the effective capacity from $M-1$ to $\min(T, (M-1)d)$.
5. Per-Step Decomposition: The trajectory objective $\tilde{J} = \log \det(K + \sigma^2 d I_T)$ is decomposed into causal per-step rewards using Cholesky factorization ($L L^\top = \tilde{K}$). The reward for step $t$ is:
   $$r_t = \log \left( K_{tt} + \sigma^2 d - k_{<t}^\top \tilde{K}_{<t}^{-1} k_{<t} \right)$$
   • Lifelong Term ($K_{tt}$): The standard ensemble disagreement (buffer conditioning).
   • Prefix Redundancy Term ($k_{<t}^\top \tilde{K}_{<t}^{-1} k_{<t}$): A quadratic form that subtracts the portion of the disagreement at step $t$ already explained by the preceding steps in the rollout.
   • Aleatoric Ridge ($\sigma^2 d$): A floor estimated post-hoc from training residuals, preventing rewards from vanishing due to irreducible noise.

Key Contributions

• CIG Reward: A tractable, per-step intrinsic reward derived from a log-determinant surrogate of trajectory-level information gain. It conditions simultaneously on the replay buffer (via the trained ensemble) and the rollout prefix (via Cholesky decomposition).
• Rank-Saturation Analysis: Identification of the capacity bottleneck in finite-ensemble surrogates and the derivation of the trace reduction technique to restore effective capacity without requiring full covariance estimation.
• Robustness to Stochastic Distractors: The method naturally filters action-dependent stochastic noise (Noisy-TV) through the ensemble's aleatoric scale, unlike count-based or raw prediction-error methods.

Experimental Results
The authors evaluated CIG on 12 tasks across discrete (MiniGrid) and continuous control (OGBench) domains, including clean and stochastic-distractor (Noisy-TV) variants. The evaluation used a reward-free protocol to isolate exploration quality.

• Performance: CIG outperformed or matched all baselines (including Plan2Explore, RND, ICM, APT, E3B, and E3B $\times$ P2E) across the majority of tasks. In aggregate statistics (Inter-Quartile Mean), CIG achieved the highest score with non-overlapping confidence intervals against all other methods.
• Noisy-TV Robustness: In environments with action-dependent stochastic distractors, methods relying on raw prediction error (ICM, RND, APT) collapsed to near-zero success rates. CIG and Plan2Explore (P2E) remained robust due to ensemble filtering, but CIG consistently outpaced P2E, demonstrating that the prefix redundancy correction provides additional value beyond simple disagreement.
• Ablations: Removing the trace reduction caused the reward to saturate immediately, halting exploration. Removing the prefix redundancy term reduced sample efficiency, confirming that the correction prevents overcounting correlated steps within a rollout.
• Entropy vs. Success: While methods like APT and E3B achieved higher state-visit entropy, this did not correlate with task success. CIG achieved high success with lower entropy, indicating that targeted uncertainty reduction (exploring informative frontiers) is more effective than uniform coverage.

Significance and Claims
The paper claims that CIG demonstrates that lifelong and episodic exploration signals need not be combined heuristically. Instead, they emerge jointly from a single information-theoretic objective. The Cholesky factorization of the log-determinant surrogate naturally produces causal, per-step rewards that retain both conditioning sets.

The authors emphasize that CIG is robust across problem structures: it filters irreducible noise in stochastic environments and provides structured rewards in deterministic ones. While CIG does not dominate every single task (e.g., APT performed better on clean ObstructedMaze), it is the only method that avoids degradation across all tested conditions. The work is limited to model-based RL with short imagined rollouts; the authors note that the $O(T^3)$ Cholesky cost and $T \times T$ kernel storage may become prohibitive for long-horizon model-free episodes, suggesting sliding-window or low-rank approximations as future work.