How to Stay Curious while Avoiding Noisy TVs using Aleatoric Uncertainty Estimation

TL;DR

The paper tackles exploration under sparse rewards by addressing the noisy TV problem through Aleatoric Mapping Agents (AMAs), which explicitly model aleatoric uncertainty and subtract it from prediction errors to focus on epistemic uncertainty. Using a double-headed network that predicts a mean $\hat{\mu}_{t+1}$ and a diagonal covariance $\hat{\Sigma}_{t+1}$, AMAs compute intrinsic rewards as $r^i_t = \|s_{t+1}-\hat{\mu}_{t+1}\|^2 - \eta \mathrm{Tr}(\hat{\Sigma}_{t+1})$ and minimize the likelihood-based loss $\mathcal{L}_{t+1} = (s_{t+1}-\hat{\mu}_{t+1})^T \hat{\Sigma}_{t+1}^{-1} (s_{t+1}-\hat{\mu}_{t+1}) + \lambda \log\det(\hat{\Sigma}_{t+1})$, with a separate policy. The approach is evaluated on Noisy MNIST, Minigrid, Mario/Space Invaders, and Bank Heist, demonstrating robustness to action-dependent noise and improved exploration compared to baselines like MSE-based curiosity, RND, and ensemble disagreements. The study also links to neuroscience via acetylcholine signaling, suggesting potential experimental tests and broader implications for understanding uncertainty in biological and artificial agents. Overall, AMAs advance robust exploration in stochastic environments and offer a framework for integrating aleatoric uncertainty estimation into curiosity-driven RL, with open-source code for replication.

Abstract

Exploration in environments with sparse rewards is difficult for artificial agents. Curiosity driven learning -- using feed-forward prediction errors as intrinsic rewards -- has achieved some success in these scenarios, but fails when faced with action-dependent noise sources. We present aleatoric mapping agents (AMAs), a neuroscience inspired solution modeled on the cholinergic system of the mammalian brain. AMAs aim to explicitly ascertain which dynamics of the environment are unpredictable, regardless of whether those dynamics are induced by the actions of the agent. This is achieved by generating separate forward predictions for the mean and variance of future states and reducing intrinsic rewards for those transitions with high aleatoric variance. We show AMAs are able to effectively circumvent action-dependent stochastic traps that immobilise conventional curiosity driven agents. The code for all experiments presented in this paper is open sourced: http://github.com/self-supervisor/Escaping-Stochastic-Traps-With-Aleatoric-Mapping-Agents.

Figures (15)

• Figure 1: AMAs can learn to ignore stochastic transitions. (a) Example transitions from the Noisy MNIST environment along with associated predictions. The top two rows show stochastic transitions where AMA's predicted variance is high in the majority of the image allowing intrinsic reward to be small despite the stochastic transition. (b) Two reward curves for MSE and AMA are plotted where stochastic is the $1\rightarrow\{2,...,9\}$ transitions and deterministic is the $0\rightarrow0$ transitions.
• Figure 2: AMA agents effectively explore sparse reward minigrid environments that contain action dependent stochastic traps. (a) and (b) panel show performance on the easiest environment, containing four rooms, while the (c) and (d) show performance on a more challenging environment with six rooms. AMA and MSE have similar exploration performance when the noisy TV is absent, outperforming a no-intrinsic-reward baseline---but when a noisy TV is present only the AMA curiosity approach is able to significantly explore the environment. Ensemble methods are robust to noisy TVs in this case, but not random network distillation. Panel (e) shows an example six room environment. Standard error represents seed variation. See appendix \ref{['sec:minigrid_extrinsic_rewards']} for plots of extrinsic rewards.
• Figure 3: Pixel AMA performs significantly better than all baselines with a noisy TV (a) and (c) and without the distracting noisy TV AMA nearly matches its most directly comparable method Pixel MSE (b) and (d). No extrinsic rewards were used for policy optimisation. In Mario distance covered and extrinsic reward are equivalent. The y-axis plots extrinsic rewards per episode. The no intrinsic reward baseline for Space Invaders is the mean performance of the random agent data provided by burda2018large. For Mario we were required to compute the no intrinsic reward baseline from scratch, which is the PPO agent from burda2018exploration with intrinsic rewards turned off without a noisy TV (but repeated in all panels for easy comparison).
• Figure 4: AMA is robust to Noisy TVs of a very random unifrom noise distribution (random pixels from 0-255) while other baselines are also trapped by this additional noisy TV. We verify in \ref{['mario']} the eventual collapse in Mario is a property of the curiosity methods tested and not a feature of AMA.
• Figure 5: IDF AMA (a) and Pixel AMA (b) avoids the natural trap in Bank Heist and so explores more of the maze on average than the IDF MSE. (b) shows the different seeds from Pixel AMA plotted individually due to the very long episodes of pixel AMA. An implementation detail means the values of the x-axis above are a close approximation to their true values see \ref{['bank_heist_x_explanation']} for details.