1 Introduction↩︎

The human brain is an experience machine [1]: a survival system that enables far-sighted planning by efficiently organizing complex multimodal sensory input. As the most advanced general intelligence we know, it is a natural inspiration for autonomous agents. Yet it remains unclear how the brain structures sensory stimuli into decision-relevant variables to solve control and adaptation problems so efficiently [2]. Reinforcement learning (RL) formalizes far-sighted planning, and recent advances enable web-scale learning of multimodal distributions for general-purpose action generation [3]–[5]. Still, we lack a satisfying account of how such rich representations are organized for efficient planning under tight computational and data constraints [6].

The active inference framework (AIF) offers a process-level theory of whole-brain information organization [7], [8]. It posits an action–perception cycle in which neural dynamics for perception and control evolve to minimize a single objective, the expected free energy (EFE) [9]. Formally, this amounts to variational Bayes on a controlled Markov process, yielding a form of model-predictive control. Its links to modern machine learning have motivated AIF-inspired adaptive control methods [10]–[15], but in practice these efforts have largely rederived familiar information-theoretic exploration heuristics (e.g., information gain) [16], [17]. To date, AIF has not delivered major practical gains in state-of-the-art RL.

Observing that the brain is an efficient information-processing system, we investigate whether AIF is useful in situations where high-fidelity forward simulation is infeasible. As distributional RL [18], [19] provides a computationally efficient alternative to model-based RL, we aim to incorporate AIF into this framework. We achieve this goal in three steps. Firstly, we analyze AIF strictly through the lens of variational Bayesian and causal inference, reconstructing its formulation by applying a mechanistic deduction chain over a coherent set of formal entities. By expressing the prior construction via do-calculus [20], we show that the standard objective admits a simpler equivalent of its commonly adopted form.

Secondly, we formalize a new theoretical framework, called push-forward RL, that explains the return distribution as pushing the trajectory measure of a specific policy-induced state transition kernel forward with a return functional. This recipe enables us to relate Bellman updates to the transition kernels implied by the resulting push-forwards of trajectory measures, thereby connecting model-based and model-free views of policy iteration. We use this connection to embed active inference into the distributional RL framework.

Lastly, we exploit our theoretical findings to curate a general-purpose policy optimization algorithm called Distributional Active Inference (DAIF). DAIF proves to be an intuitive and easily implementable extension of distributional RL. Specifically, DAIF performs temporal-difference quantile matching on a probabilistic embedding space constructed by a state-action amortized parametric distribution. Across a broad suite of tabular and continuous-control tasks, we find that this simple modification delivers substantial performance gains. The results support our view that AIF is particularly powerful when the agent has limited computational capabilities, mirroring the conditions of the biological brains it is intended to explain.

Notation. We denote probability measures by \(P\), and by \(\mathcal{P}_{\mathcal{X}}\) the set of all probability measures on the measurable space \((\mathcal{X},\sigma(\mathcal{X}))\) for a sigma-algebra \(\sigma(\mathcal{X})\). Marginalization of an intermediate random variable and conditioning on a point observation \(z\) are denoted as \(P_{X|Y} P_{Y|z} \triangleq \int P_{X|Y}(\cdot|Y)P_{Y|Z}(dY|z) = P_{X|z}\). For brevity, we suppress the conditioning on variables that are integrated out, e.g., \(P_X P_Y\) instead of \(P_{X|Y} P_Y\). By \(F_P\) and \(f_P\) we denote the cumulative distribution function and probability density function of probability measure \(P\), respectively. See 2 in the appendix for the details of our notational conventions.

Problem Statement. We consider controlled Markov processes defined by a transition kernel \(P(X'|X,A)\) where \(X,X' \in \mathcal{X}\) denote the current and next environment states, respectively, and \(A \in \mathcal{A}\) indicates the action taken by the agent at the current time step. The agent operates in a controlled Markov process with an unknown transition kernel \(P_*\). The agent has a reward function \(R:\mathcal{X} \times \mathcal{A} \rightarrow [0,R_{\text{max}}]\), a world model \(P_W(X,S)\) accompanied by auxiliary variables \(S\) in a latent embedding space \(\mathcal{S}\), and a policy \(\pi: \mathcal{X} \rightarrow \mathcal{A}\). The agent seeks to maximize its adaptation to the environment by updating \(P_W\) and \(\pi\) through its interactions with the environment which evolves according to \(P_\pi^*\). The agent updates its policy \(\pi\) by estimating the discounted sum of rewards, called the return functional, after taking action \(a_0\) in state \(x_0\) and subsequently following a fixed policy \(\pi\): \({\boldsymbol{G}}_{x_0,a_0}^\pi(\omega) \triangleq R(x_0,a_0)+\sum_{t=1}^\infty \gamma^t R(x_t, \pi(x_t))\) for some discount factor \(\gamma \in (0,1)\) and a trajectory of future states \(\omega \triangleq (x_1,x_2,\dots) \in \mathcal{X}^{\mathbb{N}_+}\).

2 A Rigorous Formulation of Active Inference↩︎

Active inference has substantial empirical support as a meta-level process theory for cognitive neuroscience [21]. However, its algorithmic details and probability-theoretic implications allow diverse interpretations. We will next derive the AIF training objective from the first principles of Bayesian and causal inference. The outcome yields important simplifications, which paves the way for the integration of AIF into distributional RL.

Let \(P_W(X,Y,S) \triangleq P_D(X|Y,S)P_0(Y,S)\) be the agent’s world model with a prior belief over the latent embeddings \(P_0(Y,S)\) and a decoding likelihood \(P_D(X|Y,S)\). Variable \(Y\) represents actions and variable \(S\) represents perceptions. The Evidence Lower Bound (ELBO) functional for observation \(x\) and approximate posterior \(P_Q(Y,S)\) follows from the standard machinery of variational inference \[\begin{align} &L(x,P_W,P_Q) = \\ &\quad \mathbb{E}_{y,s \sim P_Q}\Big [\log \big ( P_D(x|y,s) P_0(y,s) / P_Q(y,s) \big ) \Big ]. \end{align}\] Maximizing this functional with respect to \(Q\) minimizes the Kullback-Leibler divergence \(\mathrm{KL}(P_Q(Y,S)||P(Y,S|x))\), as this quantity equals the Jensen gap between \(\log P(X)\) and \(L(x,P_W,P_Q)\). AIF is defined as an extension of variational inference to control problems in which the transition kernel and the policy are jointly optimized with the same ELBO. AIF aims to solve \({P_Q := \arg \max_{\bar{P}_Q} L(x,P_{W},\bar{P}_Q)}\) with \[\begin{align} &f_{P_{0}}(Y:=y) \propto \exp(\mathbb{E}_{x' \sim P_{W|\mathrm{do}(Y \sim \delta_y, S \sim P_Q )}}[\\ &L(x',P_{W|\mathrm{do}(Y \sim \delta_y, S \sim P_Q)},P_Q)]) \end{align}\] where the \(\mathrm{do}(X \sim P_{\mathrm{new}})\) operator denotes a stochastic intervention [20], [22] to the causal graph of \(P_W\) such that the distribution of variable \(X\) is replaced by \(P_{\mathrm{new}}\). The intervention \(\mathrm{do}(S \sim P_Q)\) is intended to incorporate the effect of past observations on future predictions. A model equipped with such a prior would favor a world model whose predicted observations best fit it. This principle is known in cognitive neuroscience as predictive coding [23], [24].

AIF implements a particular way of predictive coding by intervening into the world model used in the ELBO with a desired state distribution \(P_R(X)\), which plays the role of the reward function in RL. The canonical AIF prior construction applies the product rule over the approximate world model in the reverse direction \[\begin{align} P_{W|\mathrm{do}(Y,S \sim P_Q)}(X,Y,S) &= P_D(X|Y,S)P_Q(Y,S) \\ &= \widetilde{P}(Y,S|X) \widetilde{P}(X) \end{align}\] which yields an approximate posterior \(\widetilde{P}(Y,S|X)\) and the corresponding prior \(\widetilde{P}(X)\). The prior \(\widetilde{P}(X)\) is then intervened by \(P_R\). We now point to an overlooked matter in the AIF literature. Due to the equality imposed by the product rule, the intervention on \(X\) yields (see 8.4): \[\begin{align} &P_{\widetilde{W}}(X,Y,S) \triangleq \label{eq:factorized95qz}\\ &\quad P_{W|\mathrm{do}(Y,S \sim P_Q, X \sim P_R)}(X,Y,S) = P_Q(Y,S) P_R(X)\nonumber \end{align}\tag{1}\] and disconnects the observable \(X\) from the latents \((Y,S)\). The key consequence is that the intervened distribution of \((Y,S)\) becomes independent from \(X\), eliminating the effort discussed in prior work about the inference of \(\widetilde{P}(Y,S|X)\) [25]. Substituting this outcome into the prior construction developed above, we obtain \[\begin{align} f_{P_{0}}(Y:=y) \propto \exp (\mathbb{E}_{x' \sim P_D P_{Q|y}} [ \log P_R(x') ] ) \end{align}\] which yields the complete AIF objective below \[\begin{align} \begin{aligned} &L(x,P_W,P_Q) = \\ &\quad \mathbb{E}_{y,s \sim P_Q}[\log P_D(x|y,s) ] + \mathbb{E}_{y,s \sim P_Q}[\log P_0(s|y)] \\ &\quad + \mathbb{E}_{x' \sim P_D P_Q} [ \log P_R(x') ]+ \mathbb{H}[P_Q]. \end{aligned}\label{eq:aif95elbo95small} \end{align}\tag{2}\] See 8.5 for the intermediate derivation steps.

Next, we connect the minimalist derivation presented above to the control setting. We exploit the factorization in Eq.@eq:eq:factorized95qz to simplify the ELBO further. The world model of an AIF agent comprises a policy variable \(\pi\) to represent actions and the latent embedding \((S,S')\) that represents the current perceptions by \(S\) and the future perceptions by \(S'\). AIF theory assumes that these variables factorize as: \[\begin{align} P_W(X', S', &S,\pi|X,A) \triangleq\\ &\quad P_D(X' | S') P_M(S'|S,A) P_E(S|X) P_0(\pi) \end{align}\] where \(P_M\) is a latent state transition kernel, \(P_E\) is an encoder that maps the sensory stimulus space to the latent perception space, \(P_D\) is a decoder that maps in the reverse direction, and \(P_A\) is a policy distribution. Let us assume an approximate posterior \[\begin{align} P_Q(S,S',\pi) \triangleq P_M(S'|S,A) P_E(S|X) P_A(\pi |X) \end{align}\] that shares the same transition kernel and encoder with the world model. Such a choice of approximate posterior is standard in the probabilistic state-space modeling literature [26] and is used to eliminate the latent overshooting term in the Dreamer model family [27], [28]. With this approximate posterior, the term \(\mathbb{E}_{y,s \sim P_Q}[\log P_0(s|y)]\) developed in Eq. 2 drops and we get: \[\begin{align} \begin{aligned} &L(x,\!a,\!x'\!,\! P_W,\!P_Q)\! =\!\mathbb{E}_{S' \sim P_M P_E |x,a} [\log P_D(x'|S')] \\ &\quad+ \mathbb{E}_{x'',\pi \sim P_D P_M P_E P_A | x'} \Big [ \! \log P_R(x''|x',\pi(x'))\!\Big ] \\ &\quad + \mathbb{H}[P_{A|x'}]. \end{aligned} \label{eq:full95aif95elbo} \end{align}\tag{3}\] In RL terms, this ELBO prescribes a Dyna-style model-based approach. The first term learns the state transition dynamics from observed triples \((x,a,x')\). The remaining two terms can be viewed as maximum-entropy policy search [29], [30] on the estimated transition model. The formulation can be straightforwardly extended to the case where \(x''\) is an arbitrarily long trajectory.

3 Push-forward Reinforcement Learning↩︎

We project AIF onto the distributional RL framework via a model-based extension of policy iteration presented in 1, which augments standard policy iteration with a model update step. It performs policy updates by Bellman backups using the following operator, where the backup transition kernel \(P\) comes from the estimated model \[\begin{align} \begin{aligned} (T^\pi_P Q) & (x,a) \triangleq \\ &R(x,a) +\gamma \mathbb{E}_{x' \sim P(\cdot|x,a)}[ Q(x',\pi(x')) ] \label{eq:generalized95bellman95operator} \end{aligned} \end{align}\tag{4}\] where \(Q:\mathcal{X} \times \mathcal{A}\rightarrow \mathbb{R}\) is an arbitrary function. The use of \(\mathcal{U}(\mathcal{X})\) in 1 is an idealized full-coverage assumption; in practice, \(x\) is drawn from the state distribution induced by the data (e.g., under a behavior policy). Viewed this way, MBPI is best seen as a template that highlights design choices shared by modern model-based RL. For instance, PSRL samples an MDP from the posterior and solves it [31]; PILCO combines model learning with policy search [32]; and Dreamer couples model learning with actor–critic updates [27]. We temporarily suppress the policy entropy maximization in Eq.@eq:eq:full95aif95elbo for notational brevity.

With finite samples and finite backups, the algorithm inherits the properties of Approximate Policy Iteration (API), which can be viewed as a two-error scheme [33]: each iteration incurs (i) a policy-evaluation error \(\varepsilon\) and (ii) an approximate greedification error \(\delta\), yielding asymptotic suboptimality on the order of \(\big(\delta + 2\gamma \varepsilon\big)/(1-\gamma)^2\). In API, \(\varepsilon\) and \(\delta\) are controlled from data; worst-case guarantees such as \(|Q(x,a) - \mathbb{E}[{\boldsymbol{G}}_{x,a}^\pi]|_\infty \le \varepsilon\) typically demand accurate estimates over a sufficiently large set of states, so the total sampling burden is dominated by the size of that set. For high-dimensional (vector) state spaces, \(|\mathcal{S}|\) can grow rapidly with dimension under discretization, motivating learned state abstractions that group states with similar transition dynamics in a latent embedding [34]–[36].

As the second step, we generalize the distributional RL framework to the case in which the transition dynamics are explicitly identified as an intermediate step and then pushed forward on the trajectory measures. Approaching RL from a probabilistic perspective, this decomposition will pave the way to integrate AIF into the distributional RL framework. Since the return functional \({\boldsymbol{G}}_{x,a}^\pi(\omega)\) operates on a state-action trajectory \(\omega\), its distribution (therefore expectation) can be expressed through a probability measure over trajectories. We construct such a measure below.

Let \(P_\pi(B|x) \triangleq P(B|x,\pi(x))\) be a Markov kernel induced by applying a policy \(\pi: \mathcal{X} \rightarrow \mathcal{A}\) to a controlled Markov process with transition kernel \(P\) for \(B \in \sigma(\mathcal{X})\). Define a cylinder set of length \(n\) for measurable sets \(B_1, B_2, \ldots, B_n \in \sigma(\mathcal{X})\) as \[C\! =\! \{ (x_1, x_2, \dots)\! \in\! \mathcal{X}^{{\mathbb{N}_+}}\!\!:\! x_1\! \in\! B_1, x_2\! \in\! B_2, \dots, x_n\! \in\! B_n \}.\] According to the Ionescu-Tulcea theorem [37], there exists a unique probability measure \(\mathbb{P}^{P_\pi}_{x_0,a_0}\) defined on the product \(\sigma\)-algebra \(\sigma(\mathcal{X})^{\otimes {\mathbb{N}_+}}\) such that for every \(n\in{\mathbb{N}_+}\) and every cylinder set \(C\) of the form \(B_1\times B_2\times\cdots\times B_n\times \mathcal{X}\times\cdots\), with \(B_t\in\sigma(\mathcal{X})\) for \(t\ge 1\), the measure is consistent with: \[\begin{align} \mathbb{P}_{x_0,a_0}^{P_\pi}&(C) = \int_{B_1} P(dx_1|x_0,a_0) \times \int_{B_2} P_\pi(dx_2 | x_1) \\ & \times \int_{B_3} P_\pi(dx_3 | x_2) \times \cdots \times \int_{B_n} P_\pi(dx_n | x_{n-1}). \end{align}\] This measure defines the canonical process \((X_n)_{n\ge 1}\) on the path space \(\mathcal{X}^{{\mathbb{N}_+}}\), where \(X_n(\omega) \triangleq \omega_n\) for all \(n\ge 1\). Here \(\omega_n\) indicates the \(n\)’th coordinate of the trajectory. The resulting triple \((\mathcal{X}^{{\mathbb{N}_+}},\;\sigma(\mathcal{X})^{\otimes {\mathbb{N}_+}},\;\mathbb{P}_{x_0,a_0}^{P_\pi})\) constitutes a valid probability space for the infinite-horizon controllable Markov process that starts from \(P(\cdot|x_0,a_0)\) and proceeds with the policy-induced kernel \(P_\pi\). We will refer to objects like \(\mathbb{P}_{x_0,a_0}^{P_\pi}\) as Markov process measures.

Let \((E,\sigma(E))\) be a new measurable space and \(f: \mathcal{X} \to E\) be a measurable map. We extend this to the infinite-horizon path space by defining the sequence-valued map \(\mathbf{F}: \mathcal{X}^{{\mathbb{N}_+}} \to E^{{\mathbb{N}_+}}\) as \(\mathbf{F}(\omega) = (f(x_1), f(x_2), \dots)\) where \(\omega = (x_1, x_2, \dots)\) is an element of the path space \(\mathcal{X}^{{\mathbb{N}_+}}\). We refer to objects like \(\mathbf{F}\) as push-forward process functionals. If we push a Markov process measure forward with this functional, we get the following measure on the observation path space \((E^{{\mathbb{N}_+}}, \sigma(E)^{\otimes {\mathbb{N}_+}})\): \[\begin{align} {\boldsymbol{F}}_{\#} \mathbb{P}_{x_0,a_0}^{P_\pi} &\triangleq \mathbb{P}_{x_0,a_0}^{P_\pi}(\mathbf{F}^{-1}(B)) \\ &= \mathbb{P}_{x_0,a_0}^{P_\pi} \{ \omega \in \mathcal{X}^{{\mathbb{N}_+}} : \mathbf{F}(\omega) \in B \} \end{align}\] for any measurable set \(B \in \sigma(E)^{\otimes {\mathbb{N}_+}}\). Under this definition, the measure \((\mathbf{G}_{x_0,a_0}^\pi)_{\#} \mathbb{P}_{x_0,a_0}^{P_\pi}\) corresponds precisely to the return distribution on the measurable space \((\mathbb{R}, \mathcal{B}(\mathbb{R}))\). For any Borel set \(B \in \mathcal{B}(\mathbb{R})\), this measure can be expressed as: \[\begin{align} \mathbb{P}_{x_0,a_0}^{P_\pi} &\Bigg( \Big\{ \omega \in \mathcal{X}^{{\mathbb{N}_+}} : \\ &R(x_0,a_0) + \sum_{t=1}^{\infty} \gamma^t R(x_t,\pi(x_t)) \in B \Big \} \Bigg ). \end{align}\] The return distributions induced by two different Markov kernels \(P\) and \(P'\) can then be expressed as \((\mathbf{G}_{x_0,a_0}^{\pi})_{\#} \mathbb{P}_{x_0,a_0}^{P_\pi}\) and \((\mathbf{G}_{x_0,a_0}^\pi)_{\#} \mathbb{P}_{x_0,a_0}^{P'_\pi}\), respectively.

We can define the distributional counterpart of the generalized Bellman operator introduced in Eq.@eq:eq:generalized95bellman95operator with a specific Markov kernel \(P\) as follows: \[\begin{align} {\boldsymbol{T}}_P^{\pi} \eta_{x,a} \triangleq R(x,a) + \gamma\mathbb{E}_{x' \sim P(x'|x,a)} [ \eta_{x',\pi(x')}] \end{align}\] for a real-valued probability measure \({\eta_{x,a}: \mathcal{X} \times \mathcal{A} \rightarrow \mathcal{P}_{\mathbb{R}}}\) conditioned on \((x,a) \in \mathcal{X} \times \mathcal{A}\). Define the maximal form of the \(p\)-Wasserstein distance \({\bar{\mathcal{W}}_p(P,\bar{P}) \triangleq \sup_{x,a} \mathcal{W}_p(P_{x,a},\bar{P}_{x,a})}\) where \(P_{x,a}\) and \(\bar{P}_{x,a}\) are two probability measures on the same space as \(\eta_{x,a}\). We can re-express the contraction property of the distributional Bellman operator [18] in terms of push-forwards of Markov process measures.

The key nuance here is that the contraction modulus is due to \(\gamma\) within the distributional Bellman operator, not to the discount factor applied by the return functional. From Banach’s fixed point theorem [38], it follows straightforwardly that there exists \(\widehat{P}_\pi \in \mathcal{P_X}\) that satisfies \(\bar{\mathcal{W}}_p ( {\boldsymbol{T}}_{P_*}^\pi {\boldsymbol{F}}_\# \mathbb{P}^{\widehat{P}_\pi}, {\boldsymbol{F}}_\# \mathbb{P}^{\widehat{P}_\pi} ) = 0\) for a sufficiently large set of admissible Markov kernels \(\mathcal{P_X}\). One can find this fixed point at a geometric rate (i.e., the error is \(O(\gamma^k)\)) by starting from an arbitrary \(P_\pi^0 \in \mathcal{P_X}\) and repeatedly applying \({\boldsymbol{T}}_{P_*}^\pi\). It is important to remark that this fixed point is unique in the space of Markov process measures after a push-forward with functional \({\boldsymbol{F}}\). There may be multiple Markov process measures almost surely equal to the fixed point \({\boldsymbol{F}}_{\#} \mathbb{P}^{{\widehat P}_\pi}_{x_0,a_0}\) after being pushed forward with \({\boldsymbol{F}}\). Obviously, \({\boldsymbol{F}}_{\#} \mathbb{P}^{P_\pi^*}_{x_0,a_0}\) also satisfies this condition. The Markov kernel of the Bellman operator plays the role of an anchor in the search process performed via Bellman backups.

The advantage of distributional RL over model-based RL appears in situations where capturing the push-forward of the Markov process measure with the return functional is sufficient for decision making, such as in risk-sensitive RL [39]–[42]. The equivalence of this measure within a set of transition kernels brings sample efficiency when it can be exploited by the learning algorithm. Model-based RL theory formulates such equivalence classes via state abstractions [34], [43], [44]. We will next develop some essential concepts to establish the link between state abstractions and distributional RL. First, we need to construct an embedding space on which state abstractions can be formulated.

Now consider the push-forward of the transition kernel \(P\) of a Markov process \(\mathbb{P}_{x,a}^P\) with an \(L_E\)-Lipschitz continuous function \(S\) to an embedding space \(\mathcal{S}\), which we denote by \(S_\# P\), and then mapping back to the state space by an \(L_D\)-Lipschitz Markov kernel \(P_D\). This composite operation can be viewed as the action of \(P_D\) on the measure \(S_{\#} P\). We can also construct the same outcome by transforming the input of \(P_D\), which then defines an operator on \(P\).

The key observation is that the action of \((KF~P)\) is identical to that of \(K(F_{\#}P)\). The following result establishes how the Lipschitz continuity of a Markov kernel is affected by an input transformation.

Our main theoretical result follows straightforwardly from the fact that transforming the transition kernels of two Markov processes with a point-wise fixed \(L\)-Lipschitz operator contracts their probability measures by a factor of \(L\) with respect to the \(p\)-Wasserstein distance. The consequence in our context is that performing distributional RL after an auto-encoding operation with an \(L_E\)-Lipschitz encoder \(S_\# P\) and an \(L_D\)-Lipschitz decoder \(P_D\) affects the contraction modulus by \(L_E \cdot L_D\).

When a narrow information bottleneck yields an efficient compression in the embedding space, \(L_D\) should be small as the returns will create large equivalence classes in the latent trajectory space (see 6 in the appendix), which will reduce the contraction modulus of the Bellman backups and speed up convergence. We can let distributional RL inherit these benefits of AIF by expressing an auto-encoding operation within that framework. Let us redefine \({\boldsymbol{G}}_\# \mathbb{P}^{ (P_D S) P_\pi}_{x_0,a_0}\) in terms of an encoding process defined on the latent space and a decoding process that maps back to the observation space. The encoding process on \(\mathcal{S}^{\mathbb{N}_+}\) pushes \(P^{P_\pi}_{x_0,a_0}\) forward through the sequence map \(\mathbf{S}(\omega) = (S(x_1), S(x_2), \dots)\), which leads to a new measure \(\mathbf{S}_{\#} \mathbb{P}_{x_0,a_0}^{P_\pi}\) on the space \((\mathcal{S}^{\mathbb{N}_+},\sigma(\mathcal{S})^{\otimes {\mathbb{N}_+}})\). For any sequence \(\mathbf{s} = (s_1, s_2, \dots) \in \mathcal{S}^{\mathbb{N}_{+}}\), the kernel \(P_D\) induces a conditional path measure \(\mathbb{P}^{P_D| \mathbf{s}}\) on \(\mathcal{X}^{\mathbb{N}_{+}}\) via Ionescu-Tulcea: \[\begin{align} \mathbb{P}^{P_D|\mathbf{s}}&(C) \triangleq \int_{B_1} P_D(dx_1|s_1) \times \cdots \times \int_{B_n} P_D(dx_n|s_n). \end{align}\] This represents the stochastic decoding of the representation path back into the state space. The return distribution can now be expressed as an integration over the intermediate measure \(\mathbf{S}_{\#} \mathbb{P}_{x_0,a_0}^{P_\pi}\) in the space \(\mathcal{S}^{\mathbb{N}_+}\). Defining the expected return of a fixed representation path \(\mathbf{s}\) as \[\begin{align} &(\mathbf{G}_{x_0,a_0}^\pi)_{\#} \mathbb{P}^{P_D|\mathbf{s}} (B) \triangleq \mathbb{P}^{P_D |\mathbf{s}} \Bigg( \\ & \Big\{ \omega \in \mathcal{X}^{\mathbb{N}_+} : R(x_0,a_0)+ \sum_{t=1}^\infty \gamma^t R(x_t,\pi(x_t)) \in B \Big\} \Bigg) \end{align}\] we can express the final return distribution as the integral of these local returns over the representation measure: \[\begin{align} &(\mathbf{G}_{x_0,a_0}^\pi)_{\#} \mathbb{P}^{ (P_D S) P_\pi}_{x_0,a_0}(B) = \label{eq:composite95return95measure}\\ &\int_{\mathcal{S}^{\mathbb{N}_+}}(\mathbf{G}_{x_0,a_0}^\pi)_{\#}\mathbb{P}^{P_D|{\boldsymbol{s}}} (B) \, \mathbf{S}_{\#} \mathbb{P}_{x_0,a_0}^{P_\pi}(d\mathbf{s}). \nonumber \end{align}\tag{5}\] The fact that the outcome is a composite measure will be instrumental in its implementation.

Let us next demonstrate how this result translates to a distributional RL algorithm. Although our result has more general implications, we will follow the established framework that performs Bellman residual minimization with respect to \(\mathcal{W}_2\) for policy evaluation and a policy update on the expectation of the fitted return distribution. Formally, the policy evaluation step minimizes the objective below \[\begin{align} J(P_\pi) \triangleq \mathbb{E}_{x \sim \mathcal{U}(\mathcal{X})} \big[\mathcal{W}_2^2( &(\mathbf{G}_{x_0,a_0}^\pi)_{\#} \mathbb{P}^{ P_\pi}_{x_0,a_0}, \\ &{\boldsymbol{T}}_{P_*}^\pi (\mathbf{G}_{x',\pi(x')}^\pi)_{\#} \mathbb{P}^{ P_\pi}_{x_0,a_0} )\big]. \end{align}\] Since \({\boldsymbol{T}}_{P_*}^\pi\) has a unique fixed point, we have \(J( \widehat{P}_\pi)=0\). When the Markov process is transformed by an auto-encoder \(P_D S\) we attain the following upper bound to this objective \[\begin{align} &J(S) \triangleq \mathbb{E}_{x \sim \mathcal{U}(\mathcal{X})} \big[\mathcal{W}_2^2( (\mathbf{G}_{x_0,a_0}^\pi)_{\#} \mathbb{P}^{ (P_D S) P_\pi}_{x_0,a_0}, \\ &{\boldsymbol{T}}_{P_*}^\pi (\mathbf{G}_{x',\pi(x')}^\pi)_{\#} \mathbb{P}^{ (P_D S) P_\pi}_{x_0,a_0} )\big]\\ &\leq \mathbb{E}_{x \sim \mathcal{U}(\mathcal{X})} \Big [\mathbb{E}_{{\boldsymbol{s}} \sim \mathbf{S}_{\#} \mathbb{P}_{x_0,a_0}^{P_\pi}} \Big [ \mathbb{E}_{\tau \in \mathcal{U}(0,1)} \Big [ \Big ( \\ & \Big ( F_{(\mathbf{G}_{x_0,a_0}^\pi)_{\#} \mathbb{P}^{P_D| \mathbf{s}}}^{-1}(\tau) - F_{{\boldsymbol{T}}_{P_*}^\pi (\mathbf{G}_{x',\pi(x')}^\pi)_{\#} \mathbb{P}^{P_D| \mathbf{s}}}^{-1}(\tau) \Big )^2 \Big ] \Big ] \Big ]. \end{align}\]

As the final statement is indexed only by \(\mathbf{S}_{\#} \mathbb{P}_{x_0,a_0}^{P_\pi}\), we can instead search directly the space of encoding measures \[\begin{align} \mathcal{Z}_{x,a}^\pi \triangleq \{Z_{x,a} : Z_{x,a}=\mathbf{S}_{\#} \mathbb{P}_{x_0,a_0}^{P_\pi}, P_\pi \in \mathcal{P}_\pi \} \label{eq:encoding95process95search} \end{align}\tag{6}\] for a state-action pair \((x,a) \in \mathcal{X} \times \mathcal{A}\). Using the expectation of the identified push-forward distributions in the policy improvement step, which is the Bayes-optimal way to make predictions with it, we arrive at 2. We call this algorithm template push-forward policy iteration and its practical applications push-forward reinforcement learning. This algorithm will deliver a policy sequence identical to Algorithm 1 whenever \((\mathbf{G}_{x_0,a_0}^\pi)_{\#} \mathbb{P}^{ P_\pi}_{x_0,a_0}\) is realizable under \(\mathcal{Z}_{x,a}^\pi\) for all \(x, a, \pi\), as this will make their policy improvement steps measurably identical. In practice, this realizability can be enforced with powerful function approximators. The gain in sample complexity depends on the trade-off between the capacity of these function approximators and the level of compression that the underlying transition dynamics permits via state abstractions.

4 Active Inference with Push-Forward RL↩︎

We next integrate AIF into the push-forward RL framework and remove the need to estimate the transition kernel of the underlying dynamical system. We will exploit the fact that in Eq.@eq:eq:full95aif95elbo , \(P_M\) and \(P_E\) always appear together in a marginalization process. Let us denote their marginal by \[\begin{align} P_{S|x,a}(S') \triangleq \mathbb{E}_{s \sim P_{E|x}} [P_M(S'|s,a)] . \end{align}\] This distribution can be interpreted as a transfer transition kernel from a source Markov chain defined on \(\mathcal{X}\) to a target Markov chain defined on \(\mathcal{S}\) [45]. Note that \(P_S\) can be viewed as a push-forward of an underlying \(P_\pi\) with some function \(S\).

Let us next define the next state \(S'\) as the infinitely-long trajectory of future latent states \({\boldsymbol{s}}\) simulated on the world model after taking action \(a\) at state \(x\) and following policy \(\pi\) afterwards. Let \(\omega\) be the corresponding states in the observation space. We can then construct the desired state distribution as the return of this trajectory: \(P_R(\omega) \propto \exp(({\boldsymbol{G}}_{x,a}^\pi)_{\#}(\omega))\). Placing this construction into the related term of the ELBO developed in 3 we get \[\begin{align} &\mathbb{E}_{\omega \sim P_D P_M P_E|x,a} [ ({\boldsymbol{G}}_{x,a}^\pi)_{\#} (\omega) ]. \\ &\qquad =\mathbb{E}_{\omega \sim \mathbb{P}_{x,a}^{ (P_D P_M P_E)_\pi}} [ ({\boldsymbol{G}}_{x,a}^\pi)_{\#}(\omega) ]\\ &\qquad =\mathbb{E} [ ({\boldsymbol{G}}_{x,a}^\pi)_{\#} \mathbb{P}_{x,a}^{ (P_D P_M P_E)_\pi} ]\\ &\qquad = \mathbb{E}_{{\boldsymbol{s}} \sim {\boldsymbol{S}}_{\#} \mathbb{P}_{x,a}^{P_\pi}} \left [ ({\boldsymbol{G}}_{x,a}^\pi)_{\#}\mathbb{P}^{ P_D | {\boldsymbol{s}}} \right ]\\ &\qquad = \mathbb{E}_{\omega \sim P_D P_{S|x,a}} \left [ ({\boldsymbol{G}}_{x',\pi(x')}^\pi)_{\#}(\omega) \right ] \end{align}\] where we omit the integral over \(P_A\) for notational brevity. The first and the fourth equations follow from definitions, the second from the Law of the Unconscious Statistician (LOTUS), and the third from Eq.@eq:eq:composite95return95measure . By learning an encoding process as in 6 and performing distributional RL on the latent trajectory space, we can capture \(({\boldsymbol{G}}_{x,a}^\pi)_{\#}\mathbb{P}^{ P_D | {\boldsymbol{s}}}\). Thus, we can inherit the benefits of learning a transition kernel as in 1 to simulate trajectories. Furthermore, we can quantify the effects of downstream calculations on these simulations in an infinitely long horizon.

We can implement this recipe by learning a state-action amortized encoder that maps to the latent space and performing quantile regression of returns on this space. To operate on a latent space, we need to cast the quantile regression problem as a maximum likelihood estimation instance and assign input-dependent priors to its parameters. The solution of the problem below for an i.i.d.sample of \(Y\) gives an unbiased estimate of the \(\tau\)’th quantile of a conditional distribution \(P(Y|X)\) with input \(X\) and output \(Y\): \[\begin{align} \widehat{g}_\tau \triangleq \mathop{\mathrm{arg\,min}}_{g_\tau} \mathbb{E}_{X,Y \sim P}[\ell_\tau(Y - g_\tau(X))] \label{eq:qr} \end{align}\tag{7}\] for a predictor \(g_\tau\) where \(\ell_\tau(u) \triangleq (|u| + (2 \tau -1 ) u)/2\) is the check loss. By extending the connection established by [46], we can frame \(\widehat{g}_\tau\) as the maximum likelihood estimate (MLE) of the Asymmetric Laplace Distribution (ALD) [47] \[\begin{align} f_\tau(Y|X; g_\tau, \sigma_\tau) = \frac{\tau(1-\tau)}{\sigma_\tau} \exp\left( -\frac{\ell_\tau(Y - g_\tau(X))}{\sigma_\tau} \right). \end{align}\] By setting \(\mu_\tau \triangleq g_\tau\) and treating both \(\mu_\tau\) and \(\sigma_\tau\) as random variables that follow a prior distribution with state and action dependent hyperparameters, we can perform distributional reinforcement learning on a latent space. We can capture the randomness caused by the auto-encoding step with a probability measure \(E_\phi(x,a,\tau)\), the parameters \(\phi\) of which are trainable with a high-capacity function approximator. In our context, \(\mu_\tau\) captures the mean quantile for these returns and \(\sigma_\tau\) their standard deviation. The regression target \(Y\) can be evaluated via Bellman backups.

The remaining two terms of the AIF ELBO perform a model update and policy entropy maximization. Whether to implement the latter is a design choice. We obtained better results with action-noise exploration than maximum entropy policy search (see Appendix 9.1.4, Table 6). This is expected, as the other steps of our uncertainty quantification pipeline already induce the necessary randomness for exploration. As the model update is a lower bound to the same step of 1, the same consequences of Lemma 1 apply.

With this we end up with 3, which we refer to as Distributional Active Inference (DAIF). This algorithm preserves the benefits of performing AIF whenever they are available in its native model-based version while eliminating the need to learn a transition kernel. A natural choice for \(E_\phi\) would be the Normal-Inverse Gamma distribution with parameterized inputs. However, this would not permit an analytical calculation of the expectation of \(\log p_\tau\) in the policy evaluation step. Having observed in our preliminary implementations that modeling the uncertainty around \(\mu_\tau\) does not improve performance, we instead assumed a normal prior on the mean with fixed variance and modeled an Inverse Gamma prior for \(\sigma_\tau\) (see 9.2.4).

5 The Prior Art of RL with Active Inference↩︎

AIF casts control as inference, selecting policies that minimize expected free energy (EFE), with epistemic terms inducing principled exploration [9], [48]. Most practical AIF agents are model-based: they learn an explicit world model and use variational inference to maintain beliefs and/or evaluate EFE, often with planning or lookahead [14], [49]–[54]. Complementary work analyzes links and limitations between AIF and RL/control-as-inference, including when EFE objectives recover RL-like optimality or reduce to standard intrinsic-motivation criteria [10], [55]–[57]. Closest to our work, [15] derive Bellman-style EFE recursions for POMDPs and develop actor–critic updates in belief space, but still rely on learned belief representations and a world model. DAIF sidesteps world-model learning for EFE evaluation by learning targets directly from sampled transitions. See 10 for an extended review.

6 Experiments↩︎

AIF postulates that the action-perception cycle of a biological agent is mediated by simulations on a latent space [8]. Since the latent space applies state abstractions to maximize simulation efficiency, the advantages of AIF should be observed in scenarios where state abstractions facilitate the learning task. We represent such a use case with a two-dimensional grid world environment whose transition probabilities share the dynamics of the RiverSwim [31] setup on a one-dimensional latent manifold unknown to the agent. See 9.1 and 6 for further details. 5 (right panel) demonstrates that both model-based RL and vanilla distributional RL fail to solve the task as the problem horizon grows, whereas our DAIF maintains its performance. As shown in the left panel, the models behave similarly in long problem horizons when state abstractions are not advantageous.¹

We also adapt DAIF to continuous control with deep actor-critics (9.2.4). As summarized in 1, DAIF provides consistent performance gains across various environments from three benchmark suites. Controlling DMC locomotors is difficult due to compounding function-approximation errors that cause the deadly triad [60]. DMC vision presents the additional challenge of controlling these platforms from the high-dimensional state space comprising raw pixels. EvoGym introduces control problems for soft robots with deformable bodies.

4 reports learning curves for one representative environment per suite. Catcher-v0 is classified as a hard task in EvoGym [61] and has the steepest learning curve within the interaction budget. Dog Run has the highest state dimensionality among all DMC robots, and running demands sustained high-speed control. Quadruped Run has the highest action dimensionality among the vision-based environments. DAIF outperforms state-of-the-art baselines by a clear margin throughout the entire learning process. DAIF’s pronounced performance boost corroborates our view that AIF is especially useful for controlling challenging environments with limited computational resources. The performance boost of DAIF requires about \(12\%\) more wall-clock computation time than distributional actor-critic methods. See 9.2 for full details of these experiments.

7 Takeaways and Open Questions↩︎

We provided a new theoretical framework that enables casting AIF as a simple extension of distributional RL. This implies that the improvements in the adaptation capabilities of an AIF agent should be observed in the distributional setting. We reported a comprehensive set of experiment results that support this claim. The deep RL implementation of DAIF improved the state of the art in multiple soft robotics benchmarks and challenging visual control tasks. DAIF inherits the convergence guarantees of distributional RL with an improved contraction modulus. A rigorous characterization of its finite-sample characteristics is an open question, as it is for the distributional RL field. Our theoretical framework can also lay a foundation for a formal analysis of AIF’s computational properties. A methodology for this could be extending the sample complexity analysis of linear quadratic control [62] to well-behaved non-linearities.

Impact Statement↩︎

This work makes both theoretical and algorithmic contributions to the general field of machine learning. The theoretical contributions are two-fold. The first is a formulation of the key concepts of distributional reinforcement learning using only stochastic process constructions and their transformations. The second is a derivation of the active inference machinery using variational and causal inference concepts in a rigorous and consistent manner. This derivation highlights a significant simplification in the resulting objective function. The algorithmic contribution is a distributional counterpart of active inference, which has thus far been practiced in computationally demanding model-based settings. We do not foresee negative societal impacts beyond those generally associated with increased automation.

8 Proofs and derivations↩︎

8.1 Proof of 1↩︎

For a fixed \((x,a)\) the Wasserstein distance is regular, homogeneous, and \(p\)-convex [19], i.e., \[\begin{align} \mathcal{W}_p^p\left( {\boldsymbol{T}}_{P_*}^{\pi} {\boldsymbol{F}}_\# \mathbb{P}^{P_\pi}_{x,a}, {\boldsymbol{T}}_{P_*}^{\pi} {\boldsymbol{F}}_\# \mathbb{P}^{\bar P_\pi}_{x,a} \right) &\leq \gamma^p {\mathcal{W}_p\left(\mathbb{E}_{x'\sim P^*_\pi(x'|x,\pi(x))}\left[{\boldsymbol{F}}_\# \mathbb{P}^{P_\pi}_{x',\pi(x)}\right], \mathbb{E}_{x'\sim P^*_\pi(x'|x,\pi(x))}\left[{\boldsymbol{F}}_\# \mathbb{P}^{\bar P_\pi}_{x',\pi(x)}\right]\right)} \\ &\leq \gamma^p\mathbb{E}_{x'\sim P^*_\pi(x'|x,\pi(x))}\left[ {\mathcal{W}_p\left({\boldsymbol{F}}_\# \mathbb{P}^{P_\pi}_{x',\pi(x)}, {\boldsymbol{F}}_\# \mathbb{P}^{\bar P_\pi}_{x',\pi(x)}\right)}\right] \\ &\leq \gamma^p \sup_{x',a'} \mathcal{W}_p^p\left({\boldsymbol{F}}_\# \mathbb{P}^{P_\pi}_{x',\pi(x)}, {\boldsymbol{F}}_\# \mathbb{P}^{\bar P_\pi}_{x',\pi(x)}\right) = \gamma^p \bar{\mathcal{W}}_p^p\left({\boldsymbol{F}}_\# \mathbb{P}^{P_\pi}, {\boldsymbol{F}}_\# \mathbb{P}^{\bar P_\pi}\right), \end{align}\] where the first inequality follows, via regularity and homogeneity, the second via \(p\)-convexity, the third via the definition of a supremum and finally the equality via the definition of \(\bar{\mathcal{W}}_p\). As this holds for every pair \((x,a)\) taking the \(p\)-th root and subsequently the supremum on the left hand side gives the desired inequality \(\square\)

8.2 Proof of Lemma 2↩︎

By definition, \((KF)(x, \cdot) = K(F(x), \cdot)\). Using the \(L\)-Lipschitz property of the kernel \(K\) on \(\mathcal{S}\), we have: \[\mathcal{W}_p(K(F(x), \cdot), K(F(x'), \cdot)) \leq L \cdot |F(x) - F(x')|.\] Applying the \(M\)-Lipschitz property of \(F\), it follows that \(|F(x) - F(x')| \leq M \cdot |x - x'|\). Substituting this into the inequality above gives: \[\mathcal{W}_p((KF)(x, \cdot), (KF)(x', \cdot)) \leq L \cdot M \cdot |x -x'|\square\]

8.3 Proof of 1↩︎

We apply a cascade of the following two simple results. The first is that passing a Markov kernel through an \(L\)-Lipschitz operator is a contraction with respect to \(p\)-Wasserstein distance.

The second is that passing the Markov kernels of two Markov processes through an \(L\)-Lipschitz operator contracts their \(p\)-Wasserstein distance by a factor \(L\).

8.4 Factorization of the intervened world model↩︎

Define \[\begin{align} \widetilde{P}(X,Y,S) \triangleq P_D(X|Y,S)P_Q(Y,S). \label{eq:joint95tilde} \end{align}\tag{8}\] Then by product rule \[\begin{align} \widetilde{P}(X,Y,S)_{\mathrm{do}(X \sim P_R)} &= \widetilde{P}(X|Y,S)_{\mathrm{do}(X \sim P_R)}\widetilde{P}(Y,S)_{\mathrm{do}(X \sim P_R)}\\ &=\widetilde{P}(X|Y,S)_{\mathrm{do}(X \sim P_R)}\widetilde{P}(Y,S)\\ &=P_R(X)\widetilde{P}(Y,S) \end{align}\] By definition we have \[\begin{align} \widetilde{P}(X,Y,S)_{\mathrm{do}(X \sim P_R)} &= (P_D)_{\mathrm{do}(X \sim P_R)}(X|Y,S) (P_Q)_{\mathrm{do}(X \sim P_R)}(Y,S)\\ &=P_R(X) P_Q(Y,S). \end{align}\] which yields \(\widetilde{P}(Y,S)=P_Q(Y,S)\) hence \[\begin{align} \widetilde{P}(X,Y,S)_{\mathrm{do}(X \sim P_R)}=P_R(X) P_Q(Y,S). \end{align}\] Therefore \[\begin{align} \widetilde{P}(Y,S|X)_{\mathrm{do}(X \sim P_R)}=P_Q(Y,S). \end{align}\]

8.5 Derivation of the active inference objective↩︎

Consider that \[\begin{align} P_{\widetilde{W}|\mathrm{do}(Y \sim \delta_y, S \sim P_Q)}(X,Y,S)& = P_{\widetilde{W}|\mathrm{do}(Y \sim \delta_y)}(X,Y,S)\\ &\quad = (P_Q(Y,S) P_R(X))_{\mathrm{do}(Y \sim \delta_y)}\\ &\quad = (P_Q(Y,S))_{\mathrm{do}(Y \sim \delta_y)} P_R(X)\\ &\quad = P_Q(y, S) P_R(X). \end{align}\] The first equality follows from the fact that \(P_{\widetilde{W}}\) has already been constructed with the intervention \(\mathrm{do}(S \sim P_Q)\) and the second is due to the result shown in 8.4. Now the related AIF posterior can be constructed as follows \[\begin{align} \log f_{P_0}(Y:=y) &+ \mathrm{const} = \mathbb{E}_{x',s \sim P_{W|\mathrm{do}(Y \sim \delta_y, S \sim P_Q)}} \Big [ L(x',P_{\widetilde{W}|\mathrm{do}(Y \sim \delta_y)}, P_Q \Big )\Big ]\\ &=\mathbb{E}_{x',s \sim P_{W|\mathrm{do}(Y \sim \delta_y, S \sim P_Q)}} \Big [ \mathbb{E}_{y',s' \sim P_Q}\Big [\log \Big ( P_Q(y,s') P_R(x') / P_Q(y',s') \Big ) \Big ] \Big ]\\ &=\mathbb{E}_{x',s \sim P_{W|\mathrm{do}(Y \sim \delta_y, S \sim P_Q)}} \Big [ \log P_R(x') + \mathbb{E}_{y',s' \sim P_Q}\Big [\log \Big ( P_Q(y,s') / P_Q(y',s') \Big ) \Big ] \Big ]\\ &= \mathbb{E}_{x' \sim P_D P_{Q|y}} \Big [ \log P_R(x') \Big ] + \mathbb{E}_{x' \sim P_D P_{Q|y}} \Big [ \mathbb{E}_{y',s' \sim P_Q} [ \log P_Q(y,s') ] - \mathbb{E}_{y',s' \sim P_Q} [ \log P_Q(y',s') ] \Big ] \end{align}\] where \(\mathrm{const}\) represents the normalization constant that does not depend on \(Y\). Now let us take the expectation with respect to the posterior on \(Y\): \[\begin{align} &\mathbb{E}_{y \sim P_Q}[ \log f_{P_0}(Y:=y) ] + \mathrm{const}\\ &\quad=\mathbb{E}_{y \sim P_Q} \Bigg [\mathbb{E}_{x' \sim P_D P_{Q|y}} \Big [ \log P_R(x') \Big ] + \mathbb{E}_{x' \sim P_D P_{Q|y}} \Big [ \mathbb{E}_{y',s' \sim P_Q} [ \log P_Q(y,s') ] - \mathbb{E}_{y',s' \sim P_Q} [ \log P_Q(s',y') ] \Big ] \Bigg ]\\ &\quad=\mathbb{E}_{x' \sim P_D P_{Q}} [ \log P_R(x') ] + \mathbb{E}_{y \sim P_Q} \Bigg [\mathbb{E}_{x' \sim P_D P_{Q|y}} \Big [ \mathbb{E}_{y',s' \sim P_Q} [ \log P_Q(y,s') ] - \mathbb{E}_{y',s' \sim P_Q} [ \log P_Q(y',s') ] \Big ] \Bigg ]\\ &\quad=\mathbb{E}_{x' \sim P_D P_{Q}} [ \log P_R(x') ] + \mathbb{E}_{y \sim P_Q} \Bigg [\mathbb{E}_{x' \sim P_D P_{Q|y}} \Big [ \mathbb{E}_{y',s' \sim P_Q} [ \log P_Q(y,s') ] \Big ] \Bigg ]- \mathbb{E}_{y \sim P_Q} \Bigg [\mathbb{E}_{y',s' \sim P_Q} [ \log P_Q(y',s') ] \Big ] \Bigg ]\\ &\quad=\mathbb{E}_{x' \sim P_D P_{Q}} [ \log P_R(x') ] + \mathbb{E}_{y \sim P_Q} \Bigg [\mathbb{E}_{x' \sim P_D P_{Q|y}} \Big [ \mathbb{E}_{y',s' \sim P_Q} [ \log P_Q(y,s') ] \Big ] \Bigg ] - \mathbb{E}_{y',s' \sim P_Q} [ \log P_Q(y',s') ] \Big ]\\ &\quad=\mathbb{E}_{x' \sim P_D P_{Q}} [ \log P_R(x') ] + \mathbb{E}_{x',y \sim P_D P_Q} \Big [ \mathbb{E}_{y',s' \sim P_Q} [ \log P_Q(y,s') ] \Big ] - \mathbb{E}_{y',s' \sim P_Q} [ \log P_Q(y',s') ] \Big ]\\ &\quad=\mathbb{E}_{x' \sim P_D P_{Q}} [ \log P_R(x') ] + \mathbb{E}_{x',y' \sim P_D P_Q} \Big [ \mathbb{E}_{y,s' \sim P_Q} [ \log P_Q(y,s') ] \Big ] - \mathbb{E}_{y',s' \sim P_Q} [ \log P_Q(y',s') ] \Big ]\\ &\quad=\mathbb{E}_{x' \sim P_D P_{Q}} [ \log P_R(x') ] + \mathbb{E}_{y,s' \sim P_Q} [ \log P_Q(y,s') ] - \mathbb{E}_{y',s' \sim P_Q} [ \log P_Q(y',s') ] \Big ]\\ &\quad=\mathbb{E}_{x' \sim P_D P_{Q}} [ \log P_R(x') ]. \end{align}\] Plugging this result into the ELBO, we get \[\begin{align} L(x,P_W,P_Q) &= \mathbb{E}_{y,s \sim P_Q}\Big [\log \Big ( P_D(x|y,s) P_0(y,s) / P_Q(y,s) \Big ) \Big ]\\ &=\mathbb{E}_{y,s \sim P_Q}[\log P_D(x|y,s) ] + \mathbb{E}_{y,s \sim P_Q}[\log P_0(y,s)] + \mathbb{H}[P_Q]\\ &=\mathbb{E}_{y,s \sim P_Q}[\log P_D(x|y,s) ]+ \mathbb{E}_{y,s \sim P_Q}[\log P_0(s|y)] + \mathbb{E}_{y \sim P_Q}[\log P_{0}(y)] + \mathbb{H}[P_Q]\\ &=\mathbb{E}_{y,s \sim P_Q}[\log P_D(x|y,s) ] + \mathbb{E}_{y,s \sim P_Q}[\log P_0(s|y)] + \mathbb{E}_{x' \sim P_D P_Q} [ \log P_R(x') ] + \mathbb{H}[P_Q]. \end{align}\]

8.6 Derivation of the upper bound to the push-forward RL objective↩︎

\[\begin{align} J(S) = \mathbb{E}_{x \sim \mathcal{U}(\mathcal{X})} &[d_2^2( ({\boldsymbol{G}}_{x_0,a_0}^\pi)_{\#} \mathbb{P}^{ (P_D S) P_\pi}_{x_0,a_0}, {\boldsymbol{T}}_{P_*}^{\pi} ({\boldsymbol{G}}_{x',\pi(x')}^\pi)_\# \mathbb{P}^{ (P_D S) P_\pi}_{x_0,a_0} )] \\ &\leq \mathbb{E}_{x \sim \mathcal{U}(\mathcal{X})} [\mathbb{E}_{{\boldsymbol{s}} \sim \mathbf{S}_{\#} \mathbb{P}_{x_0,a_0}^{P_\pi}}[ d_2^2(({\boldsymbol{G}}_{x_0,a_0}^\pi)_{\#} \mathbb{P}_{x_0,a_0}^{P_D|\mathbf{s}} (B), {\boldsymbol{T}}_{P_*}^{\pi} {\boldsymbol{G}}_{x',\pi(x')}^\pi)_\# \mathbb{P}_{x_0,a_0}^{P_D| \mathbf{s}} (B)) ]]\\ &=\mathbb{E}_{x \sim \mathcal{U}(\mathcal{X})} \Big [\mathbb{E}_{{\boldsymbol{s}} \sim \mathbf{S}_{\#} \mathbb{P}_{x_0,a_0}^{P_\pi}} \Big [ \int_0^1 \Big ( F_{({\boldsymbol{G}}_{x_0,a_0}^\pi)_{\#} \mathbb{P}_{x_0,a_0}^{P_D| \mathbf{s}}}^{-1}(\tau) - F_{{\boldsymbol{T}}_{P_*}^{\pi} {\boldsymbol{G}}_{x',\pi(x')}^\pi)_\# \mathbb{P}_{x_0,a_0}^{P_D| \mathbf{s}}}^{-1}(\tau) \Big )^2 d\tau \Big ] \Big ]\\ &=\mathbb{E}_{x \sim \mathcal{U}(\mathcal{X})} \Big [\mathbb{E}_{{\boldsymbol{s}} \sim \mathbf{S}_{\#} \mathbb{P}_{x_0,a_0}^{P_\pi}} \Big [ \mathbb{E}_{\tau \in \mathcal{U}(0,1)} \Big [ \Big ( \Big ( F_{({\boldsymbol{G}}_{x_0,a_0}^\pi)_{\#} \mathbb{P}_{x_0,a_0}{P_D| \mathbf{s}}}^{-1}(\tau) - F_{{\boldsymbol{T}}_{P_*}^{\pi} {\boldsymbol{G}}_{x',\pi(x')}^\pi)_\# \mathbb{P}^{P_D| \mathbf{s}}}^{-1}(\tau) \Big )^2 \Big ] \Big ] \Big ] \end{align}\] where the inequality follows from the composite nature of \(({\boldsymbol{G}}_{x_0,a_0}^\pi)_{\#} \mathbb{P}^{ (P_D S) P_\pi}_{x_0,a_0}\) established in Eq. 5 . The result follows straightforwardly from the role of couplings in the \(p\)-Wasserstein distance formulation and it is commonly used by prior work in other contexts such as Wasserstein Auto Encoders [63].

9 Experiment Details↩︎

9.1 Tabular environments↩︎

9.1.1 RiverSwim↩︎

The RiverSwim [31] environment is defined over a one-dimensional state space \(\mathcal{X} = \{1, \ldots, n\}\) with binary actions \(\mathcal{A} = \{-1, +1\}\). The transition dynamics are parameterized by probabilities \(p_{\texttt{forward}}\) and \(p_{\texttt{backward}}\) and are given by \[\begin{align} P(X' = k+1 | X = k, A = +1) &= p_{\texttt{forward}} \, \mathbf{1}(2 \leq k \leq n - 1) + \bigl(1 - (p_{\texttt{forward}} + p_{\texttt{backward}})\bigr)\, \mathbf{1}(k = 1), \\ P(X' = k | X = k, A = +1) &= \bigl(1 - (p_{\texttt{forward}} + p_{\texttt{backward}})\bigr)\, \mathbf{1}(k \geq 2) + (p_{\texttt{forward}} + p_{\texttt{backward}})\, \mathbf{1}(k = 1), \\ P(X' = k-1 | X = k, A = +1) &= p_{\texttt{backward}} \, \mathbf{1}(2 \leq k \leq n - 1) + (p_{\texttt{forward}} + p_{\texttt{backward}})\, \mathbf{1}(k = n), \\ P(X' = k+1 | X = k, A = -1) &= 0, \\ P(X' = k | X = k, A = -1) &= \mathbf{1}(k = 1), \\ P(X' = k-1 | X = k, A = -1) &= \mathbf{1}(k-1 \geq 1). \end{align}\]

The desired state distribution (reward) is given by \[\begin{align} P_R(X = n) = 0.99, \qquad P_R(X = 1) = 0.005, \qquad P_R(X \notin \{1, n\}) = \frac{0.005}{n - 2}. \end{align}\]

9.1.2 Latent RiverSwim↩︎

We introduce the Latent RiverSwim environment that extends RiverSwim to a two-dimensional observation space \[\mathcal{X} = \{1, \ldots, n\}^2,\] with action space \[\mathcal{A} = \{(+1,0), (-1,0), (0,+1), (0,-1)\},\] and a one-dimensional latent state space \[\mathcal{S} = \{1, \ldots, n\},\] which defines the latent manifold. The latent encoding and transition dynamics are defined as \[\begin{align} e_\alpha(i,j) &= \alpha i + (1 - \alpha) j, \qquad (i,j) \in \mathcal{X} && \text{(state encoder)}, \\ k &= \lfloor e_\alpha(i,j) \rfloor && \text{(latent state)}, \\ \bar{a} &= \text{sign}(e_\alpha(a_1, a_2)), \qquad (a_1,a_2) \in \mathcal{A} && \text{(latent action)}, \end{align}\] for \(\alpha \in (0,1)\). The latent transition and reward dynamics follow those of RiverSwim with the identification \(S \triangleq X\). The decoder maps a latent state \(k\) back to the observation space by uniformly sampling one of the states consistent with the encoding: \[\begin{align} P_D(X | S = k) = \mathcal{U}\bigl(\{ X : \lfloor e_\alpha(i,j) \rfloor = k,\;(i,j) \in \mathcal{X} \}\bigr). \label{eq:tabular95decoder} \end{align}\tag{9}\]

9.1.3 Tabular algorithms↩︎

We present the state of the art in tabular model-based reinforcement learning using a policy-iteration variant of posterior sampling for reinforcement learning (PSRL-PI), summarized in 7. As a representative of state-of-the-art distributional reinforcement learning in the tabular setting, we consider 8, which captures the key properties of implicit quantile networks. Our distributional active inference method can be implemented in tabular or discrete-action settings by extending the IQQL algorithm in 8. We refer to the resulting method as Distributional Active Inference (DAIF), which is presented in 9. For details, see 9.2.4.

9.1.4 Experiments↩︎

We evaluate the tabular methods described above on the RiverSwim and Latent RiverSwim environments using \(50\) independent repetitions. For RiverSwim, the agents interact with the environment for \(\num{5000}\) steps, while for Latent RiverSwim they interact for \(\num{10000}\) steps. The first \(10\%\) of the interactions are performed using a random policy to initialize exploration. For the IQQL and DAIF models, we use neural networks to approximate the value functions. In RiverSwim, we employ a single linear layer, whereas in Latent RiverSwim we use a multilayer perceptron with ReLU activations and a single hidden layer of width \(128\). 10 reports the results in terms of the frequency of visiting the most desired state within a \(100\)-step window for varying horizon.

9.2 Continuous control environments↩︎

9.2.1 EvoGym↩︎

EvoGym [61] provides continuous control tasks for soft robots with deformable bodies, which introduce non-stationarity into the control dynamics and make the learning problem more challenging. As our focus is on the control optimization problem rather than robot morphology design, we do not design or evolve robot structures. Instead, we adopt the highest-reward robot morphologies provided in the EvoGym repository.² These morphologies were identified by [61] through co-optimization of robot design and control, and represent structures that are well-suited to each task. By using these previously optimized designs, we isolate the control learning problem and ensure a fair comparison across algorithms without confounding effects from morphology variation. From the available environments, we choose seven tasks spanning different difficulty levels and task types to provide a comprehensive evaluation.

Locomotion tasks.

We include four locomotion environments: (1) Walker-v0 (easy), where the objective is to travel as far as possible in a single direction; (2) UpStepper-v0 (medium), where the robot must climb stairs of varying heights; (3) BidirectionalWalker-v0 (medium), where the target position changes dynamically during the episode, requiring the agent to learn locomotion in both directions; and (4) Traverser-v0 (hard), where the robot must cross a pit without sinking.

Object manipulation tasks.

We include three manipulation environments: (1) Thrower-v0 (medium), where a box is placed on top of the robot and must be thrown over obstacles of varying sizes; (2) Catcher-v0 (hard), where the robot must catch a rapidly falling and rotating box dropped from a height; and (3) Lifter-v0 (hard), where the robot must lift a box out of a hole.

Excluded environments.

From the original benchmark set, we exclude Carrier-v0, BridgeWalker-v0, Climber-v0, and BeamSlider-v0. We exclude Carrier-v0 because it is classified as easy, and our evaluation prioritizes more challenging tasks while retaining one easy task as a baseline. We exclude BridgeWalker-v0, Climber-v0, and BeamSlider-v0 because preliminary experiments showed minimal performance variance across algorithms in our initial seeds, limiting their discriminative value. Our final choice ensures balanced coverage across difficulty levels (one easy, three medium, three hard) and task types (four locomotion, three manipulation).

Inclusion of non-benchmark environment.

We include BidirectionalWalker-v0, which is not part of the original benchmark set, because it introduces a qualitatively different locomotion challenge: the agent must learn to move in both directions in response to a dynamically changing goal, testing adaptability beyond standard unidirectional locomotion.

11 shows the learning curves across EvoGym environments, while 3 reports the numerical results in terms of AULC and final return.

9.2.2 DeepMind Control Suite↩︎

DeepMind Control Suite (DMC) [64] provides a diverse set of continuous control tasks built on the MuJoCo physics engine [65] and is widely used as a benchmark for reinforcement learning algorithms. From the available environments, we choose eight tasks that span different robot morphologies and difficulty levels.

Dog tasks.

The dog robot has the largest state and action space dimensionality in DMC. We include dog-walk, dog-trot, and dog-run to evaluate performance across increasing difficulty on this challenging morphology.

Run tasks across robots.

We include cheetah-run, walker-run, quadruped-run, and humanoid-run. These robots range from planar (cheetah, walker) to 3D (quadruped, humanoid) and vary in state dimensionality and control complexity. We choose the run task as it is the most demanding locomotion variant for each robot.

12 presents the learning curves across DMC environments, while 4 reports the corresponding numerical results in terms of AULC and final return.

9.2.3 DMC Vision.↩︎

The environments and tasks are identical to those in the state-based DMC setting; however, the agent learns directly from pixel observations rather than proprioceptive state inputs. Following the experimental setup of [58], we adopt their task difficulty categorization and hyperparameter configurations.

We evaluate on five tasks: cheetah-run, quadruped-run, walker-run, reacher-hard, and finger-turn-hard. The three run tasks test locomotion from pixels across different morphologies, while reacher-hard and finger-turn-hard provide non-locomotion tasks requiring precise control. We exclude dog tasks because [58] did not evaluate on this environment and no reference configurations are available. We also exclude humanoid-run, which is categorized as hard and requires \(\num{30000000}\) frames, \(10 \times\) increase over the medium-difficulty tasks we consider.

We use DrQ-v2 as the backbone for our implementation, adopting DDPG-style exploration. The DTD3 variant corresponds to the distributional extension of DrQ-v2. We use the learning curves provided in the official repository³. 13 presents the learning curves across DMC Vision environments, while 5 reports the corresponding numerical results in terms of AULC and final return.

9.2.4 Implementation details↩︎

Given an observation tuple \((x, a, \tau)\), we define the input-dependent hyperpriors \[\begin{align} \mu \triangleq \mu_\phi(x, a, \tau), \qquad \alpha \triangleq \alpha_\phi(x, a, \tau), \qquad \beta \triangleq \beta_\phi(x, a, \tau), \end{align}\] where \(\mu_\phi\), \(\alpha_\phi\), and \(\beta_\phi\) are outputs of a neural network parameterized by \(\phi\). We denote \(\boldsymbol{m}\triangleq (\mu, \alpha, \beta)\). The corresponding generative model is defined as \[\begin{align} \sigma \sim \mathcal{IG}(\alpha, \beta), \qquad G | \mu, \sigma, \tau \sim \mathcal{ALD}(\mu, \sigma, \tau), \end{align}\] where \(\mathcal{ALD}\) denotes the asymmetric Laplace distribution with density \[\begin{align} f(G | \mu, \sigma, \tau) = \frac{\tau(1-\tau)}{\sigma} \exp\!\left( -\ell_\tau\!\left(\dfrac{G-\mu}{\sigma}\right) \right), \end{align}\] and the check function is defined as \(\ell_\tau(u) = \dfrac{|u| + (2\tau - 1)u}{2}\). The corresponding log-likelihood is given by \[\begin{align} \log f(G | \mu, \sigma, \tau) &= \log \tau(1-\tau) - \log \sigma - \frac{|G-\mu| + (2\tau - 1)(G-\mu)}{2\sigma}. \end{align}\]

Rather than drawing reparameterized samples from the inverse-gamma distribution, we analytically marginalize over \(\sigma\). Taking the expectation of the log-likelihood with respect to the inverse-gamma prior yields \[\begin{align} \mathbb{E}_\sigma\!\left[\log f(G | \mu, \sigma, \tau)\right] &= \log \tau(1-\tau) - \mathbb{E}_\sigma[\log \sigma] \\ & ~- \frac{1}{2} \left( |G-\mu| + (2\tau - 1)(G-\mu) \right) \mathbb{E}_\sigma\!\left[\frac{1}{\sigma}\right]. \end{align}\]

We first compute \(\mathbb{E}_\sigma[\log \sigma]\). For \(\sigma \sim \mathcal{IG}(\alpha, \beta)\), \[\begin{align} \mathbb{E}_\sigma[\log \sigma] &= \int_0^\infty \log \sigma\, \frac{\beta^\alpha}{\Gamma(\alpha)} \sigma^{-(\alpha+1)} \exp\!\left(-\frac{\beta}{\sigma}\right) d\sigma. \end{align}\] Applying the change of variables \(a = \beta/\sigma\) yields \[\begin{align} \mathbb{E}_\sigma[\log \sigma] &= \frac{1}{\Gamma(\alpha)} \int_0^\infty (\log \beta - \log a) a^{\alpha-1} \exp(-a)\, da \\ &= \log \beta - \psi(\alpha), \end{align}\] where \(\psi(\cdot)\) denotes the digamma function. Next, we compute \(\mathbb{E}_\sigma[1/\sigma]\). Using the same change of variables, \[\begin{align} \mathbb{E}_\sigma\!\left[\frac{1}{\sigma}\right] &= \frac{\beta^\alpha}{\Gamma(\alpha)} \int_0^\infty \sigma^{-(\alpha+2)} \exp\!\left(-\frac{\beta}{\sigma}\right) d\sigma = \frac{\alpha}{\beta}, \end{align}\] where we used the identity \(\Gamma(\alpha+1) = \alpha \Gamma(\alpha)\). Combining the above results, we obtain \[\begin{align} \mathbb{E}_\sigma\!\left[\log f(G | \mu, \sigma, \tau)\right] &= \log \tau(1-\tau) - \log \beta + \psi(\alpha) - \frac{\alpha}{2\beta} \left( |G-\mu| + (2\tau - 1)(G-\mu) \right). \label{eq:daif95nll} \end{align}\tag{10}\]

We use the negative of this expected log-likelihood as the critic training objective. To mitigate numerical instabilities, we enforce \(\alpha\) and \(\beta\) to be larger than \(10\) by adding a constant offset of \(10\) to the corresponding network outputs. In addition, we adopt the regularization strategy proposed by [66] who placed weak hyperpriors \[\begin{align} p(\mu) = \mathcal{N}(0, 1000^2), \qquad p(\alpha) = \mathcal{G}(10, 0.1), \qquad p(\beta) = \mathcal{G}(10, 0.1), \end{align}\] where the priors over \(\alpha\) and \(\beta\) are defined over the shifted parameters. We use this regularization with a regularization coefficient \(\xi=0.001\).

We adopt the design choices of [67] combined with action-noise exploration and delayed policy updates [68]. The full update procedure is presented in 15. The algorithm proceeds as follows: first, quantiles are randomly sampled and their midpoints are computed. Next, the target action is obtained from the lagged target actor with clipped Gaussian noise added. The parameters of the asymmetric Laplace distribution, \(\mu\) and \(\sigma\) are estimated for the current state-action pair and quantiles, where we model \(\sigma \sim \mathcal{IG}(\alpha, \beta)\). Temporal difference errors are then computed for each critic using min-clipping. The critic loss is the negative log-likelihood weighted by pairwise quantile regression weights following [67], [69]; specifically, line 23 weights each TD error by \((\tau_{i+1} - \tau_i)\), which corresponds to the quantile bin width and ensures that each quantile contributes proportionally to its probability mass under the target distribution. Regularization terms are computed using the priors, followed by parameter and Polyak updates for the critics. Every \(d^{th}\) step (where \(d=2\)), the policy is updated: quantiles and midpoints are recomputed, an action is sampled from the online actor, and the actor maximizes the estimated mean \(\mu\) over the quantile distribution, followed by parameter and Polyak updates. For DMC vision control environments, we also apply the random augmentation method of [58] to our baselines.

9.2.5 Training↩︎

9.2.5.1 Baselines.

We include DRND [70] as a representative exploration-driven model-free reinforcement learning baseline. We also consider DSAC [67], the most recent state-of-the-art distributional actor–critic algorithm for continuous control. In addition, we introduce DTD3, a distributional extension of TD3 [68], to directly assess whether DAIF improves over a distributional variant of TD3. For the DMC vision experiments, we include DrQ-v2 [58], which achieves the highest reported scores among model-free methods on DMC vision environments. While model-based approaches such as DreamerV3 [28] achieve strong performance, they require learning a world model, which is orthogonal to our focus on model-free policy optimization. Notably, at \(\num{1000000}\) environment steps (\(\num{2000000}\) frames), DAIF achieves competitive or superior performance compared to DreamerV3 on several tasks: DAIF outperforms DreamerV3 on quadruped-run (\(813\) vs.\(617\)), walker-run (\(780\) vs.\(684\)), and reacher-hard (\(950\) vs.\(862\)), while achieving comparable performance on cheetah-run (\(823\) vs.\(836\)) and finger-turn-hard (\(867\) vs.\(904\)). This demonstrates that our model-free approach can match or exceed the performance of state-of-the-art model-based methods without the additional complexity of world model learning.

9.2.5.2 Critic architecture.

We adopt the critic architecture of [67] for continuous control tasks, except for DMC Vision environments. The critic takes the state, action, and quantile fraction as inputs and outputs the parameters \((\mu, \alpha, \beta)\). The architecture follows a modular design with separate embeddings for the state–action pair and the quantile fraction, which are fused to produce the final quantile value estimates. Specifically, the state–action input is processed by a base network consisting of a linear layer with 256 hidden units, followed by layer normalization and a ReLU activation. The quantile fraction \(\tau\) is embedded using a separate quantile network that maps a 128-dimensional quantile encoding to a 256-dimensional representation via a linear layer, layer normalization, and a sigmoid activation. The resulting embeddings are combined through element-wise interaction and passed to an output network comprising an additional hidden layer with 256 units, layer normalization, and ReLU activation, followed by a final linear layer that outputs scalar values for \(\mu\), \(\alpha\), and \(\beta\).

For DMC Vision control tasks, in addition to the critic architecture described above, we employ the convolutional encoder of [58] to process high-dimensional observations. The encoder consists of four convolutional layers with 32 channels each. The first layer uses a stride of 2, while the remaining layers use stride 1. Each convolution is followed by a ReLU activation, and the resulting feature maps are flattened to obtain a compact latent representation. After the encoder, we use a Linear layer that maps the output of the encoder to a feature dimension of \(50\), followed by layer normalization and \(\tanh\) activation for both actor and critic.

9.2.5.3 Actor architecture.

We parametrize the actor by a multilayer perceptron that maps the state representation to a continuous action. The network consists of two hidden layers with 256 units each, each followed by a ReLU activation. A final linear layer outputs the action vector, which is subsequently passed through a \(\tanh\) activation to enforce bounded actions within the valid action range. This architecture provides sufficient expressive capacity for modeling complex continuous policies while remaining computationally efficient and stable in practice. For DMC Vision control tasks, the same actor architecture is used on top of the encoder, and a DDPG-style exploration scheme is applied following [58].

9.2.5.4 Optimization.

We use a learning rate of \(3\times10^{-4}\) for both the actor and the critic, except for DMC Vision control tasks, where a learning rate of \(1\times10^{-4}\) is used for the actor, critic, and encoder. We employ a replay buffer of size \(\num{1000000}\) and a batch size of \(256\), with \(\num{10000}\) environment steps for warm-up. The discount factor is set to \(0.99\), and Polyak averaging with a coefficient of \(0.005\) is used for target network updates. We use two critics and a single actor, with target networks maintained for both. Policy smoothing is applied using Gaussian noise with standard deviation \(0.1\), target policy noise with standard deviation \(0.2\), and target noise clipping of \(0.5\), together with a policy update delay of \(2\). We randomly sample \(8\) quantiles for value estimation, target computation, and actor training. The regularization coefficient is set to \(\xi = 0.001\). For DMC Vision tasks, we use a frame stack of \(3\) and an action repeat of \(2\), and all other hyperparameters follow [58]⁴.

9.2.5.5 Environment interactions.

We interact with the environments for \(\num{1000000}\) steps for DMC control tasks, except for dog environments, where we use \(\num{1500000}\) interaction steps. For EvoGym environments, we perform \(\num{2000000}\) interaction steps, while for DMC Vision tasks we train for \(\num{3000000}\) environment frames. All experiments are repeated \(10\) times, except for DMC Vision tasks, where we use \(5\) repetitions.

9.2.5.6 Ablation study on DAIF with SAC.

We also investigate the performance of DAIF with SAC version (DAIF-SAC) on three representative environments chosen from each suite. As shown in 6, DAIF-SAC performs well across all three environments, consistently outperforming its base algorithm DSAC on Catcher-V0 and Dog-Run, while achieving comparable performance on Quadruped-Run. These results demonstrate that DAIF can be effectively applied to distributional versions of actor-critic algorithms beyond TD3. It is worth noting that in the DMC Vision tasks, DTD3 and DAIF are based on DrQ-v2, which employs DDPG as its backbone, further highlighting the flexibility of our approach across different actor-critic frameworks.

9.2.5.7 Wall-clock time comparison.

We measure wall-clock training time for \(\num{10000}\) environment interaction steps on the Dog-Run task from the DMC suite, using a system equipped with a GeForce RTX 4090 GPU, an Intel Core i7-14700K CPU (\(5.6\) GHz), and \(96\) GB memory. DTD3 requires \(59.69\) seconds, DAIF \(66.81\) seconds, DSAC \(75.05\) seconds, and DRND \(81.69\) seconds. Relative to its baseline DTD3, DAIF incurs an overhead of approximately \(12\%\) in wall-clock time. These results indicate that the performance improvements achieved by DAIF come at only a modest additional computational cost over its underlying actor–critic architecture.

10 Extended Related Works↩︎

10.1 Distributional Reinforcement Learning↩︎

Distributional RL explicitly models the full return distribution rather than only its expectation. While it was initially introduced for value-based algorithms [18], [41], [69], [71], subsequent work extended these ideas to actor–critic architectures, particularly in continuous control settings [72]. Various distributional actor–critic variants have been proposed to address specific limitations, including training instability arising from categorical critics [73], [74], overestimation bias [75]–[77], and risk-sensitive decision making [67], with several approaches focusing on risk-aware locomotion control [78]–[80]. Our work generalizes the distributional reinforcement learning framework by working with push-forward mappings of trajectory measures, where the new push-forward RL formulation extends the distributional RL. This generalization allows us to integrate the active inference framework into distributional RL, which in turn lets us develop model-free RL algorithms for active inference. The term push-forward also appears in [81], but there it refers to transport-map parameterizations of action/return distributions from base noise, whereas our push-forward RL pushes forward trajectory measures within the Bellman theory. Therefore the overlap is primarily terminological. In our experiments, we compare our method against distributional RL baselines such as the distributional Soft Actor–Critic (DSAC) [67] and Implicit Quantile Q-Learning) [41]. To our knowledge, our method is the first to integrate active inference into the distributional reinforcement learning framework.

10.2 Directed Exploration in Reinforcement Learning↩︎

Directed exploration explicitly guides reinforcement learning agents toward uncertain or informative states, proving essential in non-stationary environments with changing dynamics and sparse-reward settings lacking frequent feedback. While directed exploration has been extensively studied in the context of classical reinforcement learning algorithms [16], [60], [70], [82]–[85], distributional reinforcement learning also provides a natural foundation for exploration. In particular, distributional losses have been shown to induce intrinsic exploration through uncertainty-aware regularization [86]. Beyond this implicit effect, several works explicitly design exploration strategies within the distributional reinforcement learning framework. [87] exploit epistemic uncertainty in return distributions for posterior sampling-like exploration; [88] introduce Decaying Left Truncated Variance to derive an optimistic bonus from upper-tail quantiles with decaying optimism; [89] propose Distributional Prediction Error, using Wasserstein distance between quantile distributions to construct exploration bonuses; [90] schedule risk levels from risk-seeking to risk-averse to balance exploration; and [91] randomize risk criteria through distributional perturbations to avoid biased optimism. While these methods explicitly design exploration bonuses, the exploration term in our distributional actor–critic objective arises naturally from variational free-energy minimization. For directed exploration baselines, we also consider DSAC, which has been shown to exhibit intrinsic exploration effects [86]. In addition, we include DRND [70], a recent state-of-the-art exploration method for actor–critic algorithms.

10.3 Active Inference and Reinforcement Learning↩︎

Active inference has increasingly been examined in relation to reinforcement learning as an alternative that unifies perception, action, and learning by minimizing expected free energy (EFE) without explicit rewards or value functions [9]–[11], [48]–[50], [52], [55], [56], [92]. A growing body of empirical work shows that active inference can solve standard reinforcement learning benchmarks, often achieving performance comparable to classical reinforcement learning algorithms [51], [53], [93]–[95]. Despite these advantages, particularly the intrinsic balancing of goal-seeking and information gain through epistemic value, most practical AIF agents remain model-based, relying on learned generative/world models, belief (variational) inference, and often explicit planning or lookahead to evaluate and optimize EFE [13], [14], [51], [52], [96]. This reliance contributes to the computational burden of multi-step EFE estimation and has historically kept many scalable demonstrations either in discrete settings or in continuous-control implementations with relatively short effective planning horizons [13], [57]. Closest in spirit to our work, [15] derive Bellman-style recursions for EFE and instantiate actor–critic updates; however, their formulation operates in belief space and still depends on learned belief representations and a learned world model for objective construction/evaluation, whereas our approach targets the model-free setting.

References↩︎