Combining Imitation and Reinforcement Learning with Free Energy Principle

1 INTRODUCTION . Imitation Learning ( IL ) is a framework to learn a policy to mimic expert trajectories . As the expert specifies model behaviors , there is no need to do exploration or to design complex reward functions . Reinforcement Learning ( RL ) does not have these features , so RL agents have no clue to realize desired behaviors in sparse-reward settings and even when RL succeeds in reward maximization , the policy does not necessarily achieve behaviors that the reward designer has expected . The key drawbacks of IL are that the policy never exceeds the suboptimal expert performance and that the policy is vulnerable to distributional shift . Meanwhile , RL can achieve super-human performance and has potentials to transfer the policy to new tasks . As real-world applications often needs high sample efficiency and little preparation ( rough rewards and suboptimal experts ) , it is important to find a way to effectively combine IL and RL . When the sensory inputs are high-dimensional images as in the real world , behavior learning such as IL and RL would be difficult without representation or model learning . Free Energy Principle ( FEP ) , a unified brain theory in computational neuroscience that explains perception , action and model learning in a Bayesian probabilistic way ( Friston et al. , 2006 ; Friston , 2010 ) , can handle behavior learning and model learning at the same time . In FEP , the brain has a generative model of the world and computes a mathematical amount called Free Energy using the model prediction and sensory inputs to the brain . By minimizing the Free Energy , the brain achieves model learning and behavior learning . Prior work about FEP only dealt with limited situations where a part of the generative model is given and the task is very low dimensional . As there are a lot in common between FEP and variational inference in machine learning , recent advancements in deep learning and latent variable models could be applied to scale up FEP agents to be compatible with high dimensional tasks . Recent work in model-based reinforcement learning succeeds in latent planning from highdimensional image inputs by incorporating latent dynamics models . Behaviors can be derived either by imagined-reward maximization ( Ha & Schmidhuber , 2018 ; Hafner et al. , 2019a ) or by online planning ( Hafner et al. , 2019b ) . Although solving high dimensional visual control tasks with modelbased methods is becoming feasible , prior methods have never tried to combine with imitation . In this paper , we propose Deep Free Energy Network ( FENet ) , an agent that combines the advantages of IL and RL so that the policy roughly learns from suboptimal expert data without the need of exploration or detailed reward crafting in the first place , then learns from sparsely specified reward functions to exceed the suboptimal expert performance . The key contributions of this work are summarized as follows : • Extension of Free Energy Principle : We theoretically extend Free Energy Principle , introducing policy prior and policy posterior to combine IL and RL . We implement the proposed method on top of Recurrent State Space Model ( Hafner et al. , 2019b ) , a latent dynamics model with both deterministic and stochastic components . • Visual control tasks in realistic problem settings : We solve Cheetah-run , Walker-walk , and Quadruped-walk tasks from DeepMind Control Suite ( Tassa et al. , 2018 ) . We do not only use the default problem settings , we also set up problems with sparse rewards and with suboptimal experts . We demonstrate that our agent outperforms model-based RL using Recurrent State Space Model in sparse-reward settings . We also show that our agent can achieve higher returns than Behavioral Cloning ( IL ) with suboptimal experts . 2 BACKGROUNDS ON FREE ENERGY PRINCIPLE . 2.1 PROBLEM SETUPS . We formulate visual control as a partially observable Markov decision process ( POMDP ) with discrete time steps t , observations ot , hidden states st , continuous action vectors at , and scalar rewards rt . The goal is to develop an agent that maximizes expected return E [ ∑T t=1 rt ] . 2.2 FREE ENERGY PRINCIPLE . Perception , action and model learning are all achieved by minimizing the same objective function , Free Energy ( Friston et al. , 2006 ; Friston , 2010 ) . In FEP , the agent is equipped with a generative model of the world , using a prior p ( st ) and a likelihood p ( ot|st ) . p ( ot , st ) = p ( ot|st ) p ( st ) ( 1 ) Perceptual Inference Under the generative model , the posterior probability of hidden states given observations is calculated with Bayes ’ theorem as follows . p ( st|ot ) = p ( ot|st ) p ( st ) p ( ot ) , p ( ot ) = ∫ p ( ot|st ) p ( st ) ds ( 2 ) Since we can not compute p ( ot ) due to the integral , we think of approximating p ( st|ot ) with a variational posterior q ( st ) by minimizing KL divergence KL ( q ( st ) ||p ( st|ot ) ) . KL ( q ( st ) ||p ( st|ot ) ) = ln p ( ot ) +KL ( q ( st ) ||p ( ot , st ) ) ( 3 ) Ft = KL ( q ( st ) ||p ( ot , st ) ) ( 4 ) We define the Free Energy as ( eq.4 ) . Since p ( ot ) does not depend on st , we can minimize ( eq.3 ) w.r.t . the parameters of the variational posterior by minimizing the Free Energy . Thus , the agent can infer the hidden states of the observations by minimizing Ft . This process is called ’ perceptual inference ’ in FEP . Perceptual Learning Free Energy is the same amount as negative Evidence Lower Bound ( ELBO ) in variational inference often seen in machine learning as follows . p ( ot ) ≥ −Ft ( 5 ) By minimizing Ft w.r.t . the parameters of the prior and the likelihood , the generative model learns to best explain the observations . This process is called ’ perceptual learning ’ in FEP . Active Inference We can assume that the prior is conditioned on the hidden states and actions at the previous time step as follows . p ( st ) = p ( st|st−1 , at−1 ) ( 6 ) The agent can change the future by choosing actions . Suppose the agent chooses at when it is at st , the prior can predict the next hidden state st+1 . Thus , we can think of the Expected Free Energy Gt+1 at the next time step t+ 1 as follows ( Friston et al. , 2015 ) . Gt+1 = KL ( q ( st+1 ) ||p ( ot+1 , st+1 ) ) = Eq ( st+1 ) [ ln q ( st+1 ) − ln p ( ot+1 , st+1 ) ] = Eq ( st+1 ) p ( ot+1|st+1 ) [ ln q ( st+1 ) − ln p ( ot+1 , st+1 ) ] ( 7 ) = Eq ( st+1 ) p ( ot+1|st+1 ) [ ln q ( st+1 ) − ln p ( st+1|ot+1 ) − ln p ( ot+1 ) ] ≈ Eq ( ot+1 , st+1 ) [ ln q ( st+1 ) − ln q ( st+1|ot+1 ) − ln p ( ot+1 ) ] ( 8 ) = Eq ( ot+1 ) [ −KL ( q ( st+1|ot+1 ) ||q ( st+1 ) ) − ln p ( ot+1 ) ] ( 9 ) Since the agent has not experienced time step t+ 1 yet and has not received observations ot+1 , we take expectation over ot+1 using the likelihood p ( ot+1|st+1 ) as ( eq.7 ) . In ( eq.8 ) , we approximate p ( ot+1|st+1 ) as q ( ot+1|st+1 ) and p ( st+1|ot+1 ) as q ( st+1|ot+1 ) . According to the complete class theorem ( Friston et al. , 2012 ) , any scalar rewards can be encoded as observation priors using p ( o ) ∝ exp r ( o ) and the second term in ( eq.9 ) becomes a goal-directed value . This observation prior p ( ot+1 ) can also be regarded as the probability of optimality variable p ( Ot+1 = 1|ot+1 ) , where the binary optimality variableOt+1 = 1 denotes that time step t+1 is optimal andOt+1 = 0 denotes that it is not optimal as introduced in the context of control as probabilistic inference ( Levine , 2018 ) . The first term in ( eq.9 ) is called epistemic value that works as intrinsic motivation to further explore the world . Minimization of −KL ( q ( st+1|ot+1 ) ||q ( st+1 ) ) means that the agent tries to experience as different states st+1 as possible given some imagined observations ot+1 . By minimizing the Expected Free Energy , the agent can infer the actions that explores the world and maximize rewards . This process is called ’ active inference ’ . 3 DEEP FREE ENERGY NETWORK ( FENET ) . Perceptual learning deals with learning the generative model to best explain the agent ’ s sensory inputs . If we think of not only observations but also actions given by the expert as a part of the sensory inputs , we can explain imitation leaning by using the concept of perceptual learning . Active inference deals with exploration and reward maximization , so it is compatible with reinforcement learning . By minimizing the same objective function , the Free Energy , we can deal with both imitation and RL . In this section , we first introduce a policy prior for imitation and a policy posterior for RL . Second , we extend the Free Energy Principle to be able to accommodate these two policies in the same objective function , the Free Energy . Finally , we explain a detailed network architecture to implement the proposed method for solving image control tasks . 3.1 INTRODUCING A POLICY PRIOR AND A POLICY POSTERIOR . Free Energy We extend the Free Energy from ( eq.4 ) so that actions are a part of sensory inputs that the generative model tries to explain . Ft = KL ( q ( st ) ||p ( ot , st , at ) ) = KL ( q ( st ) ||p ( ot|st ) p ( at|st ) p ( st|st−1 , at−1 ) ) ( 10 ) = Eq ( st ) [ ln q ( st ) p ( ot|st ) p ( at|st ) p ( st|st−1 , at−1 ) ] ( 11 ) = Eq ( st ) [ − ln p ( ot|st ) − ln p ( at|st ) + ln q ( st ) − ln p ( st|st−1 , at−1 ) ] ( 12 ) = Eq ( st ) [ − ln p ( ot|st ) − ln p ( at|st ) ] +KL ( q ( st ) ||p ( st|st−1 , at−1 ) ) ( 13 ) We define p ( at|st ) as a policy prior . When the agent observes expert trajectories , by minimizing Ft , the policy prior will be learned so that it can best explain the experts . Besides the policy prior , we introduce and define a policy posterior q ( at|st ) , which is the very policy that the agent samples from when interacting with its environments . We explain how to learn the policy posterior in the following . Expected Free Energy for imitation In a similar manner to active inference in Section 2.2 , we think of the Expected Free EnergyGt+1 at the next time step t+1 , but this time we take expectation over the policy posterior q ( at|st ) becauseGt+1 is a value expected under the next actions . Note that in Section 2.2 at was given as a certain value , but here at is sampled from the policy posterior . We calculate the expected variational posterior at time step t+ 1 as follows . q ( st+1 ) = Eq ( st ) q ( at|st ) [ p ( st+1|st , at ) ] ( 14 ) q ( ot+1 , st+1 , at+1 ) = Eq ( st+1 ) [ p ( ot+1|st+1 ) q ( at+1|st+1 ) ] ( 15 ) We extend the Expected Free Energy from ( eq.12 ) so that the variational posterior makes inference on actions as follows . GILt+1 = Eq ( ot+1 , st+1 , at+1 ) [ − ln p ( ot+1|st+1 ) − ln p ( at+1|st+1 ) + ln q ( st+1 , at+1 ) − ln p ( st+1|st , at ) ] ( 16 ) = Eq ( ot+1 , st+1 , at+1 ) [ − ln p ( ot+1|st+1 ) − ln p ( at+1|st+1 ) + ln q ( at+1|st+1 ) ] +KL ( q ( st+1 ) ||p ( st+1|st , at ) ) ( 17 ) = Eq ( ot+1 , st+1 ) [ − ln p ( ot+1|st+1 ) +KL ( q ( at+1|st+1 ) ||p ( at+1|st+1 ) ) ] +KL ( q ( st+1 ) ||p ( st+1|st , at ) ) ( 18 ) = Eq ( ot+1 , st+1 ) [ − ln p ( ot+1|st+1 ) +KL ( q ( at+1|st+1 ) ||p ( at+1|st+1 ) ) ] + 0 ( 19 ) = Eq ( st+1 ) [ H [ p ( ot+1|st+1 ) ] +KL ( q ( at+1|st+1 ) ||p ( at+1|st+1 ) ) ] ( 20 ) In ( eq.20 ) , the first term is the entropy of the observation likelihood , and the second term is the KL divergence between the policy prior and the policy posterior . By minimizing GILt+1 , the agent learns the policy posterior so that it matches the policy prior which has been learned through minimizing Ft to encode the experts ’ behavior . Expected Free Energy for RL We can get the Expected Free Energy in a different way that has a reward component r ( ot+1 ) leading to the policy posterior maximizing rewards . We extend the Expected Free Energy from ( eq.8 ) so that the variational posterior makes inference on actions as follows . GRLt+1 = Eq ( ot+1 , st+1 , at+1 ) [ ln q ( st+1 , at+1 ) − ln p ( at+1|st+1 ) − ln q ( st+1|ot+1 ) − ln p ( ot+1 ) ] ( 21 ) = Eq ( ot+1 , st+1 ) [ ln q ( st+1 ) − ln q ( st+1|ot+1 ) +KL ( q ( at+1|st+1 ) ||p ( at+1|st+1 ) ) − ln p ( ot+1 ) ] ( 22 ) = Eq ( ot+1 ) [ −KL ( q ( st+1|ot+1 ) ||q ( st+1 ) ) − ln p ( ot+1 ) ] + Eq ( st+1 ) [ KL ( q ( at+1|st+1 ) ||p ( at+1|st+1 ) ) ] ( 23 ) ≈ Eq ( ot+1 ) [ −KL ( q ( st+1|ot+1 ) ||q ( st+1 ) ) − r ( ot+1 ) ] + Eq ( st+1 ) [ KL ( q ( at+1|st+1 ) ||p ( at+1|st+1 ) ) ] ( 24 ) In a similar manner to active inference in Section 2.2 , we use p ( o ) ∝ exp r ( o ) in ( eq.24 ) . The first KL term is the epistemic value that lets the agent explore the world , the second term is the expected reward under the action sampled from the policy posterior , and the last KL term is the KL divergence between the policy prior and the policy posterior . The last KL term can be written as follows ( eq.25 ) , meaning that minimizing this term leads to maximizing the entropy of the policy posterior at the same time the policy posterior tries to match the policy prior . Thus , the expected free energy can be regarded as one of entropy maximizing RL methods . KL ( q ( at+1|st+1 ) ||p ( at+1|st+1 ) ) = −H [ q ( at+1|st+1 ) ] − Eq ( at+1|st+1 ) [ ln p ( at+1|st+1 ) ] ( 25 ) Note that q ( ot+1 ) in ( eq.24 ) can be calculated as follows . q ( ot+1 ) = Eq ( st+1 ) [ p ( ot+1|st+1 ) ] ( 26 ) By minimizing GRLt+1 , the agent learns the policy posterior so that it explores the world and maximizes the reward as long as it does not deviate too much from the policy prior which has encoded experts ’ behavior through minimizing Ft .

Summary: This paper extends and explains how to apply the "free energy principle" and active inference to RL and imitation learning. They implement a neural network approximation of losses derived this way and test on some control tasks. Importantly the tasks focus on here are imitation + control tasks. That is, there is both a reward signal but also demonstration trajectories. The demonstrations may be suboptimal. The compare against PLaNet, a latent planning based approach.