Disentangled Representations for Causal Cognition

TL;DR

This paper advances a unifying framework that connects causal cognition in natural agents with causal machine learning through disentangled representations. It formalizes three core dimensions—explicitness, sources, and integration of causal information—and defines weak and strong disentanglement to capture varying levels of causal structure. By articulating a mathematical base (disentanglement, structural causal models, and reinforcement learning) and a computational framework that uses trajectories and multi-source integration, the work lays groundwork for causal RL that better mirrors animal cognition and supports transfer, zero-shot learning, and interventions. The proposed synthesis aims to drive practical advances in AI and provide a principled lens for interpreting causal reasoning in natural agents, with concrete avenues for future benchmarks, algorithms, and cross-domain experiments.

Abstract

Complex adaptive agents consistently achieve their goals by solving problems that seem to require an understanding of causal information, information pertaining to the causal relationships that exist among elements of combined agent-environment systems. Causal cognition studies and describes the main characteristics of causal learning and reasoning in human and non-human animals, offering a conceptual framework to discuss cognitive performances based on the level of apparent causal understanding of a task. Despite the use of formal intervention-based models of causality, including causal Bayesian networks, psychological and behavioural research on causal cognition does not yet offer a computational account that operationalises how agents acquire a causal understanding of the world. Machine and reinforcement learning research on causality, especially involving disentanglement as a candidate process to build causal representations, represent on the one hand a concrete attempt at designing causal artificial agents that can shed light on the inner workings of natural causal cognition. In this work, we connect these two areas of research to build a unifying framework for causal cognition that will offer a computational perspective on studies of animal cognition, and provide insights in the development of new algorithms for causal reinforcement learning in AI.

Figures (9)

• Figure 1: An illustration of the phenomenon of backward blocking in rats.. Subjects are conditioned to elicit a response (salivation) to a stimulus (presence of food) by means of a compound cue (light + acoustic tone) as well as a single cue (light). When tested with the other cue (acoustic tone), the rats do not react as strongly as if they understood that in the compound-cue trials the only cause of the reward was the first cue (light).
• Figure 2: Causal cognition tasks.
• Figure 3: Disentanglement for a bouncing red ball. Following the example describe in the text, we sketch here (some possible) factors, observations and codes for a system disentangling factors from observations into codes. For such a system to understand what make a red ball bounce, we consider shape and colour as factors generating observations about different balls (round and punctured balls, and of different colours, red and blue), inside the dashed-line oval shape at the centre of the figure. Other observations can be part of the standard repertoire of observables for our system (say, dogs or trees), but their factors are not explicitly drawn as we only focus on (some aspects) of a system capable of disentangling factors that generated observations of different kinds of balls. A disentangled representation is one that "faithfully" maps factors to codes via the given observations, according to the assumptions provided in the main text, see \ref{['def:disentanglement']}.
• Figure 4: Variational autoencoder. An intuitive representation of a variational autoencoder, combining an encoder $q_{\boldsymbol{\phi}}(\mathbf{z}|\mathbf{x})$ taking observations $\mathbf{x} \in X$ as inputs and producing $\mathbf{z} \in Z$ as outputs, these outputs are then used by a decoder $p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z})$ providing reconstructions of observations $\mathbf{x} \in X'$ with the goal of making them "as close as possible" to the original observations.
• Figure 5: Weak disentanglement. Weak disentanglement as a variant of causal representation learning concerned only with learning causal codes $Z$ from high-dimensional inputs (observations $X$) generated by causal factors $S$, and the identification of causal mechanisms that map causal factors to observations (no causal mechanisms between causal codes). A generative model can be considered structurally approximate (with respect to the assumed generative process) if it fails to recover completely disentangled codes (some codes remain entangled, e.g. the code $Z_{1,2}$ suggestively standing in for two factors $S_1, S_2$) and/or if it leaves out causal mechanisms relating codes to observations. Note that in an actual implementation, like the VAE, entanglement might manifest as correlations among two or more codes in the latent representation, all capturing the same factors at the same time. In other words, the node $Z_{1,2}$ may in practice represent a groups of nodes with possibly bidirectional influences amongst each other.

Theorems & Definitions (8)

• Definition 4.1: Disentangled representations
• Definition 4.2: Structural causal model of the generative process
• Definition 4.3: Markov decision process (MDP)
• Definition 4.4: Partially observable Markov decision process (POMDP)
• Definition 4.5: Action policy
• Definition 4.6: Expected cumulative discounted reward
• Definition 5.1: Weak disentanglement
• Definition 5.2: Strong disentanglement