Contrastive Active Inference

Active inference is a unifying theory for perception and action resting upon the idea that the brain maintains an internal model of the world by minimizing free energy . From a behavioral perspective , active inference agents can be seen as selfevidencing beings that act to fulfill their optimistic predictions , namely preferred outcomes or goals . In contrast , reinforcement learning requires human-designed rewards to accomplish any desired outcome . Although active inference could provide a more natural self-supervised objective for control , its applicability has been limited because of the shortcomings in scaling the approach to complex environments . In this work , we propose a contrastive objective for active inference that strongly reduces the computational burden in learning the agent ’ s generative model and planning future actions . Our method performs notably better than likelihood-based active inference in image-based tasks , while also being computationally cheaper and easier to train . We compare to reinforcement learning agents that have access to human-designed reward functions , showing that our approach closely matches their performance . Finally , we also show that contrastive methods perform significantly better in the case of distractors in the environment and that our method is able to generalize goals to variations in the background . 1 Introduction . Deep Reinforcement Learning ( RL ) has led to successful results in several domains , such as robotics , video games and board games [ 42 , 36 , 2 ] . From a neuroscience perspective , the reward prediction error signal that drives learning in deep RL closely relates to the neural activity of dopamine neurons for reward-based learning [ 44 , 3 ] . However , the reward functions used in deep RL typically require domain and task-specific design from humans , spoiling the generalization capabilities of RL agents . Furthermore , the possibility of faulty reward functions makes the application of deep RL risky in real-world contexts , given the possible unexpected behaviors that may derive from it [ 10 , 29 , 38 ] . Active Inference ( AIF ) has recently emerged as a unifying framework for learning perception and action . In AIF , agents operate according to one absolute imperative : minimize their free energy [ 15 ] . With respect to past experience , this encourages to update an internal model of the world to maximize evidence with respect to sensory data . With regard to future actions , the inference process becomes ‘ active ’ and agents select behaviors that fulfill optimistic predictions of their model , which are represented as preferred outcomes or goals [ 17 ] . Compared to RL , the AIF framework provides a more natural way of encoding objectives for control . However , its applicability has been limited because of the shortcomings in scaling the approach to complex environments , and current implementations have focused on tasks with either low-dimensional sensory inputs and/or small sets of discrete actions [ 12 ] . Moreover , several experiments in the literature have replaced the agent ’ s preferred outcomes with RL-like rewards from the environment , downplaying the AIF potential to provide self-supervised objectives [ 13 , 34 , 49 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . One of the major shortcomings in scaling AIF to environments with high-dimensional , e.g . imagebased , environments comes from the necessity of building accurate models of the world , which try to reconstruct every detail in the sensory data . This complexity is also reflected in the control stage , when AIF agents compare future imaginary outcomes of potential actions with their goals , to select the most convenient behaviors . In particular , we advocate that fulfilling goals in image space can be poorly informative to build an objective for control . In this work , we propose Contrastive Active Inference , a framework for AIF that aims to both reduce the complexity of the agent ’ s internal model and to propose a more suitable objective to fulfill preferred outcomes , by exploiting contrastive learning . Our method provides a self-supervised objective that constantly informs the agent about the distance from its goal , without needing to reconstruct the outputs of potential actions in high-dimensional image space . The contributions of our work can be summarised as follows : ( i ) we propose a framework for AIF that drastically reduces the computational power required both for learning the model and planning future actions , ( ii ) we combine our method with value iteration methods for planning , inspired by the RL literature , to amortize the cost of planning in AIF , ( iii ) we compare our framework to state-of-the-art RL techniques and to a non-contrastive AIF formulation , showing that our method compares well with reward-based systems and outperforms non-contrastive AIF , ( iv ) we show that contrastive methods work better than reconstruction-based methods in presence of distractors in the environment , ( v ) we found that our contrastive objective for control allows matching desired goals , despite differences in the backgrounds . The latter finding could have important consequences for deploying AIF in real-world settings , such as robotics , where perfectly reconstructing observations from the environment and matching them with high-dimensional preferences is practically unfeasible . 2 Background The control setting can be formalized as a Partially Observable Markov Decision Process ( POMDP ) , which is denoted with the tupleM = { S , A , T , Ω , O , γ } , where S is the set of unobserved states , A is the set of actions , T is the state transition function , also referred to as the dynamics of the environment , Ω is the set observations , O is a set of conditional observation probabilities , and γ is a discount factor ( Figure 1 ) . We use the terms observations and outcomes interchangeably throughout the work . In RL , the agent has also access to a reward function R , mapping state-action pairs to rewards . Active Inference . In AIF , the goal of the agent is to minimize ( a variational bound on ) the surprisal over observations , − log p ( o ) . With respect to past observations , the upper bound leads to the variational free energy F , which for timestep t is : F = Eq ( st ) [ log q ( st ) − log p ( ot , st ) ] ≥ − log p ( ot ) ( 1 ) where q ( st ) represents an approximate posterior . The agent hence builds a generative model over states , actions and observations , by defining a state transition function p ( st|st−1 , at−1 ) and a likelihood mapping p ( ot|st ) , while the posterior distribution over states is approximated by the variational distribution q ( st|ot ) . The free energy can then be decomposed as : FAIF = DKL [ q ( st|ot ) ||p ( st|st−1 , at−1 ) ] ︸ ︷︷ ︸ complexity −Eq ( st|ot ) [ log p ( ot|st ) ] ︸ ︷︷ ︸ accuracy . ( 2 ) This implies that minimizing variational free energy , on the one hand , maximizes the likelihood of observations under the likelihood mapping ( i.e . maximizing accuracy ) , whilst minimizing the KL divergence between the approximate posterior and prior ( i.e . complexity ) . Note that for the past we assume that outcomes and actions are observed , hence only inferences are made about the state st. Also note that the variational free energy is defined as the negative evidence lower bound as known from the variational autoencoder framework [ 39 , 27 ] . For future timesteps , the agent has to make inferences about both future states and actions q ( st , at ) = q ( at|st ) q ( st ) , while taking into account expectations over future observations . Crucially , in active inference the agent has a prior distribution p̃ ( ot ) on preferred outcomes it expects to obtain . Action selection is then cast as an inference problem , i.e . inferring actions that will yield preferred outcomes , or more formally that minimize the expected free energy G : G = Eq ( ot , st , at ) [ log q ( st , at ) − log p̃ ( ot , st , at ) ] , ( 3 ) where p̃ ( ot , st , at ) = p ( at ) p ( st|ot ) p̃ ( ot ) is the agent ’ s biased generative model , and the expectation is over predicted observations , states and actions q ( ot , st , at ) = p ( ot|st ) q ( st , at ) . If we assume the variational posterior over states is a good approximation of the true posterior , i.e . q ( st|ot ) ≈ p ( st|ot ) , and we also consider a uniform prior p ( at ) over actions [ 35 ] , the expected free energy can be formulated as : GAIF = −Eq ( ot ) [ DKL [ q ( st|ot ) ||q ( st ) ] ] ︸ ︷︷ ︸ intrinsic value −Eq ( ot ) [ log p̃ ( ot ) ︸ ︷︷ ︸ extrinsic value ] − Eq ( st ) [ H ( q ( at|st ) ) ] ︸ ︷︷ ︸ action entropy . ( 4 ) Intuitively , this means that the agent will infer actions for which observations have a high information gain about the states ( i.e . intrinsic value ) , which will yield preferred outcomes ( i.e . extrinsic value ) , while also keeping its possible actions as varied as possible ( i.e . action entropy ) . Full derivations of the equations in this section are provided in the Appendix . Reinforcement Learning . In RL , the objective of the agent is to maximize the discounted sum of rewards , or return , over time ∑∞ t γ trt . Deep RL can also be cast as probabilistic inference , by introducing an optimality variable Ot which denotes whether the time step t is optimal [ 30 ] . The distribution over the optimality variable is defined in terms of rewards as p ( Ot = 1|st , at ) = exp ( r ( st , at ) ) . Inference is then obtained by optimizing the following variational lower bound − log p ( Ot ) ≤ Eq ( st , at ) [ log q ( st , at ) − log p ( Ot , st , at ) ] = −Eq ( st , at ) [ r ( st , at ) ] − Eq ( st ) [ H ( q ( at|st ) ) ] , ( 5 ) where the reward-maximizing RL objective is augmented with an action entropy term , as in maximum entropy control [ 20 ] . As also highlighted in [ 35 ] , if we assume log p̃ ( ot|st ) = log p ( Ot|st ) , we can see that RL works alike AIF , but encoding optimality value in the likelihood rather than in the prior . In order to improve sample-efficiency of RL , model-based approaches ( MBRL ) , where the agent relies on an internal model of the environment to plan high-rewarding actions , have been studied . Contrastive Learning . Contrastive representations , which aim to organize the data distinguishing similar and dissimilar pairs , can be learned through Noise Contrastive Estimation ( NCE ) [ 19 ] . Following [ 37 ] , an NCE loss can be defined as a lower bound on the mutual information between two variables . Given two random variables X and Y , the NCE lower bound is : I ( X ; Y ) ≥ INCE ( X ; Y ) , E [ 1 K K∑ i=1 log ef ( xi , yi ) 1 K ∑K j=1 e f ( xi , yj ) ] , ( 6 ) where the expectation is over K independent samples from the joint distribution : ∏ j p ( xj , yj ) and f ( x , y ) is a function , called critic , that approximates the log density ratio log p ( x|y ) p ( x ) . Crucially , the critic can be unbounded , as in [ 50 ] , where the authors showed that an inner product of transformated samples from X and Y , namely f ( x , y ) = h ( x ) T g ( y ) , with h and g functions , works well as a critic .

Summary: This work proposes an active inference implementation that scales in high dimensional observations and continuous actions spaces. To achieve this the authors propose to replace the reconstruction required to optimize the free energy objective with the contrastive loss which aims to maximize the mutual information between the state and the observation at each timestep and is more robust under distractors and low capacity models. They show comparable performance to state of the art model-based RL baselines (dreamer, CVRL) and they demonstrate robustness under distractors setting. Also, the provided code is great and that would make a very nice contribution as well (the existing implementations of AIF are too difficult to parse I think)