Original: C OMBINING I MITATION AND R EINFORCEMENT L EARNING WITH F REE E NERGY P RINCIPLE
Revised: C OMBINING I MITATION AND R EINFORCEMENT L EARNING WITH F REE E NERGY P RINCIPLE

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Imitation Learning (IL) and Reinforcement Learning (RL) from high dimensional sensory inputs are often introduced as separate problems, but a more realistic problem setting is how to merge the techniques so that the agent can reduce exploration costs by partially imitating experts at the same time it maximizes its return.
Revised: Imitation Learning (IL) and Reinforcement Learning (RL) from high dimensional sensory inputs are often introduced as separate problems, but a more realistic problem setting is how to merge the techniques so that the agent can reduce exploration costs by partially imitating experts at the same time it maximizes its return.

Original: Even when the experts are suboptimal (e.g. Experts learned halfway with other RL methods or human-crafted experts), it is expected that the agent outperforms the suboptimal experts’ performance.
Revised: Even when the experts are suboptimal (e.g. Experts learned halfway with other RL methods or human-crafted experts), it is expected that the agent outperforms the suboptimal experts’ performance.

Original: In this paper
Revised: In this paper

Original: , we propose to address the issue by using and theoretically extending Free Energy Principle a uniﬁed brain theory that explains perception, action and model learning in a Bayesian probabilistic way
Revised: , we propose to address the issue by using and theoretically extending Free Energy Principle a uniﬁed brain theory that explains perception, action and model learning in a Bayesian probabilistic way

Original: We ﬁnd that both IL and RL can be achieved based on the same free energy objective function
Revised: We ﬁnd that both IL and RL can be achieved based on the same free energy objective function

Original: Our results show that our approach is promising in visual control tasks especially with sparse-reward environments.
Revised: Our results show that our approach is promising in visual control tasks especially with sparse-reward environments.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Imitation Learning (IL) is a framework to learn a policy to mimic expert trajectories.
Revised: Imitation Learning (IL) is a framework to learn a policy to mimic expert trajectories.

Original: As the expert speciﬁes model behaviors, there is no need to do exploration or to design complex reward functions.
Revised: As the expert speciﬁes model behaviors, there is no need to do exploration or to design complex reward functions.

Original: 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.
Revised: 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.

Original: The key drawbacks of IL are that the policy never exceeds the suboptimal expert performance and that the policy is vulnerable to distributional shift.
Revised: The key drawbacks of IL are that the policy never exceeds the suboptimal expert performance and that the policy is vulnerable to distributional shift.

Original: Meanwhile, RL can achieve super-human performance and has potentials to transfer the policy to new tasks.
Revised: Meanwhile, RL can achieve super-human performance and has potentials to transfer the policy to new tasks.

Original: As real-world applications often needs high sample efﬁciency and little preparation (rough rewards and suboptimal experts), it is important to ﬁnd a way to effectively combine IL and RL.
Revised: As real-world applications often needs high sample efﬁciency and little preparation (rough rewards and suboptimal experts), it is important to ﬁnd a way to effectively combine IL and RL.

Original: When the sensory inputs are high-dimensional images as in the real world, behavior learning such as
Revised: When the sensory inputs are high-dimensional images as in the real world, behavior learning such as

Original: IL and RL would be difﬁcult without representation or model learning.
Revised: IL and RL would be difﬁcult without representation or model learning.

Original: Free Energy Principle (FEP), a uniﬁed 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.
Revised: Free Energy Principle (FEP), a uniﬁed 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.

Original: 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.
Revised: 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.

Original: By minimizing the Free Energy, the brain achieves model learning and behavior learning.
Revised: By minimizing the Free Energy, the brain achieves model learning and behavior learning.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Recent work in model-based reinforcement learning succeeds in latent planning from highdimensional image inputs by incorporating latent dynamics models.
Revised: Recent work in model-based reinforcement learning succeeds in latent planning from highdimensional image inputs by incorporating latent dynamics models.

Original: Behaviors can be derived either by imagined-reward maximization (Ha & Schmidhuber, 2018; Hafner et al., 2019a) or by online planning (Hafner et al., 2019b).
Revised: Behaviors can be derived either by imagined-reward maximization (Ha & Schmidhuber, 2018; Hafner et al., 2019a) or by online planning (Hafner et al., 2019b).

Original: Although solving high dimensional visual control tasks with modelbased methods is becoming feasible, prior methods have never tried to combine with imitation.
Revised: Although solving high dimensional visual control tasks with modelbased methods is becoming feasible, prior methods have never tried to combine with imitation.

Original: 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 ﬁrst place, then learns from sparsely speciﬁed reward functions to exceed the suboptimal expert performance.
Revised: 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 ﬁrst place, then learns from sparsely speciﬁed reward functions to exceed the suboptimal expert performance.

Original: The key contributions of this work are summarized as follows:
Revised: The key contributions of this work are summarized as follows:

Original: • Extension of Free Energy Principle: We theoretically extend Free Energy Principle, introducing policy prior and policy posterior to combine IL and RL.
Revised: • Extension of Free Energy Principle: We theoretically extend Free Energy Principle, introducing policy prior and policy posterior to combine IL and RL.

Original: 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.
Revised: 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.

Original: • 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).
Revised: • 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).

Original: We do not only use the default problem settings, we also set up problems with sparse rewards and with suboptimal experts.
Revised: We do not only use the default problem settings, we also set up problems with sparse rewards and with suboptimal experts.

Original: We demonstrate that our agent outperforms model-based RL using Recurrent State Space Model in sparse-reward settings.
Revised: We demonstrate that our agent outperforms model-based RL using Recurrent State Space Model in sparse-reward settings.

Original: We also show that our agent can achieve higher returns than Behavioral Cloning (IL) with suboptimal experts.
Revised: We also show that our agent can achieve higher returns than Behavioral Cloning (IL) with suboptimal experts.

Original:  B ACKGROUNDS ON F REE E NERGY P RINCIPLE 2. P ROBLEM SETUPS
Revised:  B ACKGROUNDS ON F REE E NERGY P RINCIPLE 2. P ROBLEM SETUPS

Original: We formulate visual control as a partially observable Markov decision process (POMDP) with discrete time steps t , observations o t , hidden states s t , continuous action vectors a t , and scalar rewards r t .
Revised: We formulate visual control as a partially observable Markov decision process (POMDP) with discrete time steps t , observations o t , hidden states s t , continuous action vectors a t , and scalar rewards r t .

Original: The goal is to develop an agent that maximizes expected return E [ (cid:80)
Revised: The goal is to develop an agent that maximizes expected return E [ (cid:80)

Original: Tt =1 r t ] .
Revised: Tt =1 r t ] .

Original: 2. F REE E NERGY P RINCIPLE
Revised: 2. F REE E NERGY P RINCIPLE

Original: Perception, action and model learning are all achieved by minimizing the same objective function,
Revised: Perception, action and model learning are all achieved by minimizing the same objective function,

Original: Free Energy (Friston et al., 2006; Friston, 2010).
Revised: Free Energy (Friston et al., 2006; Friston, 2010).

Original: In FEP, the agent is equipped with a generative model of the world, using a prior p ( s t ) and a likelihood p ( o t
Revised: In FEP, the agent is equipped with a generative model of the world, using a prior p ( s t ) and a likelihood p ( o t

Original: | s t ) .
Revised: | s t ) .

Original: [Equation]
Revised: [Equation]

Original: Perceptual Inference
Revised: Perceptual Inference

Original: Under the generative model, the posterior probability of hidden states given observations is calculated with Bayes’ theorem as follows.
Revised: Under the generative model, the posterior probability of hidden states given observations is calculated with Bayes’ theorem as follows.

Original: [Equation]
Revised: [Equation]

Original: Since we cannot compute p ( o t ) due to the integral, we think of approximating p ( s t | o t ) with a variational posterior q ( s t ) by minimizing KL divergence KL ( q ( s t )
Revised: Since we cannot compute p ( o t ) due to the integral, we think of approximating p ( s t | o t ) with a variational posterior q ( s t ) by minimizing KL divergence KL ( q ( s t )

Original: || p ( s t | o t )) .
Revised: || p ( s t | o t )) .

Original: [Equation]
Revised: [Equation]

Original: We deﬁne the Free Energy as (eq.4).
Revised: We deﬁne the Free Energy as (eq.4).

Original: Since p ( o t ) does not depend on s t , we can minimize (eq.3) w.r.t.
Revised: Since p ( o t ) does not depend on s t , we can minimize (eq.3) w.r.t.

Original: the parameters of the variational posterior by minimizing the Free Energy.
Revised: the parameters of the variational posterior by minimizing the Free Energy.

Original: Thus, the agent can infer the hidden states of the observations by minimizing F t .
Revised: Thus, the agent can infer the hidden states of the observations by minimizing F t .

Original: This process is called ’perceptual inference’ in FEP.
Revised: This process is called ’perceptual inference’ in FEP.

Original: Perceptual Learning
Revised: Perceptual Learning

Original: Free Energy is the same amount as negative Evidence
Revised: Free Energy is the same amount as negative Evidence

Original: Lower Bound
Revised: Lower Bound

Original: (ELBO) in variational inference often seen in machine learning as follows.
Revised: (ELBO) in variational inference often seen in machine learning as follows.

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: By minimizing F t w.r.t.
Revised: By minimizing F t w.r.t.

Original: the parameters of the prior and the likelihood, the generative model learns to best explain the observations.
Revised: the parameters of the prior and the likelihood, the generative model learns to best explain the observations.

Original: This process is called ’perceptual learning’ in FEP.
Revised: This process is called ’perceptual learning’ in FEP.

Original: Active Inference
Revised: Active Inference

Original: We can assume that the prior is conditioned on the hidden states and actions at the previous time step as follows.
Revised: We can assume that the prior is conditioned on the hidden states and actions at the previous time step as follows.

Original: [Equation]
Revised: [Equation]

Original: The agent can change the future by choosing actions.
Revised: The agent can change the future by choosing actions.

Original: Suppose the agent chooses a t when it is at s t , the prior can predict the next hidden state s t +1 .
Revised: Suppose the agent chooses a t when it is at s t , the prior can predict the next hidden state s t +1 .

Original: Thus, we can think of the Expected Free Energy
Revised: Thus, we can think of the Expected Free Energy

Original: G t +1 at the next time step t + 1 as follows (Friston et al., 2015).
Revised: G t +1 at the next time step t + 1 as follows (Friston et al., 2015).

Original: [Equation]
Revised: [Equation]

Original: Since the agent has not experienced time step t + 1 yet and has not received observations o t +1 , we take expectation over o t +1 using the likelihood p ( o t +

Revised: Since the agent has not experienced time step t + 1 yet and has not received observations o t +1 , we take expectation over o t +1 using the likelihood p ( o t +

Original: | s t +1 ) as (eq.7).
Revised: | s t +1 ) as (eq.7).

Original: In (eq.8), we approximate p ( o t +1 | s t +1 ) as q ( o t +1 | s t +1 ) and p ( s t +1 | o t +1 ) as q ( s t +1 | o t +1 ) .
Revised: In (eq.8), we approximate p ( o t +1 | s t +1 ) as q ( o t +1 | s t +1 ) and p ( s t +1 | o t +1 ) as q ( s t +1 | o t +1 ) .

Original: According to the complete class theorem (Friston et al., 2012), any scalar rewards can be encoded as observation priors using p ( o )
Revised: According to the complete class theorem (Friston et al., 2012), any scalar rewards can be encoded as observation priors using p ( o )

Original: ∝ exp r
Revised: ∝ exp r

Original: ( o ) and the second term in (eq.9) becomes a goal-directed value.
Revised: ( o ) and the second term in (eq.9) becomes a goal-directed value.

Original: This observation prior p ( o t +1 ) can also be regarded as the probability of optimality variable p ( O t +1 = 

Revised: This observation prior p ( o t +1 ) can also be regarded as the probability of optimality variable p ( O t +1 =

Original: | o t +1 ) , where the binary optimality variable
Revised: | o t +1 ) , where the binary optimality variable

Original: O t +1 = 1 denotes that time step t + 1 is optimal and O t +1 = 

Revised: O t +1 = 1 denotes that time step t + 1 is optimal and O t +1 =

Original: denotes that it is not optimal as introduced in the context of control as probabilistic inference(Levine, 2018).
Revised: denotes that it is not optimal as introduced in the context of control as probabilistic inference(Levine, 2018).

Original: The ﬁrst term in (eq.9) is called epistemic value that works as intrinsic motivation to further explore the world.
Revised: The ﬁrst term in (eq.9) is called epistemic value that works as intrinsic motivation to further explore the world.

Original: Minimization of − KL ( q ( s t +1 |
Revised: Minimization of − KL ( q ( s t +1 |

Original: o t +1 )
Revised: o t +1 )

Original: ||
Revised: ||

Original: q ( s t +1 )) means that the agent tries to experience as different states s t +1 as possible given some imagined observations o t +1 .
Revised: q ( s t +1 )) means that the agent tries to experience as different states s t +1 as possible given some imagined observations o t +1 .

Original: By minimizing the Expected Free
Revised: By minimizing the Expected Free

Original: Energy, the agent can infer the actions that explores the world and maximize rewards.
Revised: Energy, the agent can infer the actions that explores the world and maximize rewards.

Original: This process is called ’active inference’.
Revised: This process is called ’active inference’.

Original:  D EEP F REE E NERGY N ETWORK (FEN ET )
Revised:  D EEP F REE E NERGY N ETWORK (FEN ET )

Original: Perceptual learning deals with learning the generative model to best explain the agent’s sensory inputs.
Revised: Perceptual learning deals with learning the generative model to best explain the agent’s sensory inputs.

Original: 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.
Revised: 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.

Original: Active inference deals with exploration and reward maximization, so it is compatible with reinforcement learning.
Revised: Active inference deals with exploration and reward maximization, so it is compatible with reinforcement learning.

Original: By minimizing the same objective function, the Free Energy, we can deal with both imitation and RL.
Revised: By minimizing the same objective function, the Free Energy, we can deal with both imitation and RL.

Original: In this section, we ﬁrst introduce a policy prior for imitation and a policy posterior for RL.
Revised: In this section, we ﬁrst introduce a policy prior for imitation and a policy posterior for RL.

Original: Second, we extend the Free Energy Principle to be able to accommodate these two policies in the same objective function, the Free Energy.
Revised: Second, we extend the Free Energy Principle to be able to accommodate these two policies in the same objective function, the Free Energy.

Original: Finally, we explain a detailed network architecture to implement the proposed method for solving image control tasks.
Revised: Finally, we explain a detailed network architecture to implement the proposed method for solving image control tasks.

Original: 3. I NTRODUCING A POLICY PRIOR AND A POLICY POSTERIOR Free Energy
Revised: 3. I NTRODUCING A POLICY PRIOR AND A POLICY POSTERIOR Free Energy

Original: We extend the Free Energy from (eq.4) so that actions are a part of sensory inputs that the generative model tries to explain.
Revised: We extend the Free Energy from (eq.4) so that actions are a part of sensory inputs that the generative model tries to explain.

Original: [Equation]
Revised: [Equation]

Original: We deﬁne p ( a t | s t ) as a policy prior.
Revised: We deﬁne p ( a t | s t ) as a policy prior.

Original: When the agent observes expert trajectories, by minimizing
Revised: When the agent observes expert trajectories, by minimizing

Original: F t , the policy prior will be learned so that it can best explain the experts.
Revised: F t , the policy prior will be learned so that it can best explain the experts.

Original: Besides the policy prior, we introduce and deﬁne a policy posterior q ( a t
Revised: Besides the policy prior, we introduce and deﬁne a policy posterior q ( a t

Original: | s t ) , which is the very policy that the agent samples from when interacting with its environments.
Revised: | s t ) , which is the very policy that the agent samples from when interacting with its environments.

Original: We explain how to learn the policy posterior in the following.
Revised: We explain how to learn the policy posterior in the following.

Original: Expected Free Energy for imitation
Revised: Expected Free Energy for imitation

Original: In a similar manner to active inference in Section 2.2, we think of the Expected Free Energy G t +1 at the next time step t +1 , but this time we take expectation over the policy posterior q ( a t
Revised: In a similar manner to active inference in Section 2.2, we think of the Expected Free Energy G t +1 at the next time step t +1 , but this time we take expectation over the policy posterior q ( a t

Original: | s t ) because G t +1 is a value expected under the next actions.
Revised: | s t ) because G t +1 is a value expected under the next actions.

Original: Note that in Section 2.2 a t was given as a certain value, but here a t is sampled from the policy posterior.
Revised: Note that in Section 2.2 a t was given as a certain value, but here a t is sampled from the policy posterior.

Original: We calculate the expected variational posterior at time step t + 1 as follows.
Revised: We calculate the expected variational posterior at time step t + 1 as follows.

Original: [Equation]
Revised: [Equation]

Original: We extend the Expected Free Energy from (eq.12) so that the variational posterior makes inference on actions as follows.
Revised: We extend the Expected Free Energy from (eq.12) so that the variational posterior makes inference on actions as follows.

Original: [Equation]
Revised: [Equation]

Original: In (eq.20)
Revised: In (eq.20)

Original: [Equation]
Revised: [Equation]

Original: the ﬁrst term is the entropy of the observation likelihood, and the second term is the KL divergence between the policy prior and the policy posterior.
Revised: the ﬁrst term is the entropy of the observation likelihood, and the second term is the KL divergence between the policy prior and the policy posterior.

Original: By minimizing G ILt +1 , the agent learns the policy posterior so that it matches the policy prior which has been learned through minimizing
Revised: By minimizing G ILt +1 , the agent learns the policy posterior so that it matches the policy prior which has been learned through minimizing

Original: F t to encode the experts’ behavior
Revised: F t to encode the experts’ behavior

Original: [Equation]
Revised: [Equation]

Original: Expected Free Energy for RL
Revised: Expected Free Energy for RL

Original: We can get the Expected Free Energy in a different way that has a reward component r ( o t +1 ) leading to the policy posterior maximizing rewards.
Revised: We can get the Expected Free Energy in a different way that has a reward component r ( o t +1 ) leading to the policy posterior maximizing rewards.

Original: We extend the
Revised: We extend the

Original: Expected Free Energy from (eq.8) so that the variational posterior makes inference on actions as follows.
Revised: Expected Free Energy from (eq.8) so that the variational posterior makes inference on actions as follows.

Original: [Equation]
Revised: [Equation]

Original: ≈
Revised: ≈

Original: [Equation]
Revised: [Equation]

Original: In a similar manner to active inference in Section 2.2, we use p ( o )
Revised: In a similar manner to active inference in Section 2.2, we use p ( o )

Original: ∝ exp r
Revised: ∝ exp r

Original: ( o ) in (eq.24).
Revised: ( o ) in (eq.24).

Original: The ﬁrst 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.
Revised: The ﬁrst 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.

Original: 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.
Revised: 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.

Original: Thus, the expected free energy can be regarded as one of entropy maximizing RL methods.
Revised: Thus, the expected free energy can be regarded as one of entropy maximizing RL methods.

Original: [Equation]
Revised: [Equation]

Original: Note that q ( o t +1 ) in (eq.24) can be calculated as follows.
Revised: Note that q ( o t +1 ) in (eq.24) can be calculated as follows.

Original: [Equation]
Revised: [Equation]

Original: By minimizing G RLt +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 F t .
Revised: By minimizing G RLt +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 F t .

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Figure]
Revised: [Figure]

Original: Y3Cpu7zRlGAtCGyTkoWi7WFLOAtoABpy2I0Gx73LacgdXmd
Revised: Y3Cpu7zRlGAtCGyTkoWi7WFLOAtoABpy2I0Gx73LacgdXmd

Original: [Figure]
Revised: [Figure]

Original: 6p0KyMLiDUURtH
Revised: 6p0KyMLiDUURtH

Original: [Figure]
Revised: [Figure]

Original: /xJzRkpohrpTHHe6IYl9GgDhWErLNCKwEyyAEU7TQieWNMJkgHvUUjTAPpV2MkmZ6gdK6epeKNQJQJ
Revised: cC5jGCQUlO8XxYlk4Ch2Zauej4GPqum1ynTpLJUBmWsQMPGUut6Hswux8pvWI7xZJRNSbQ/xJzRkpohrpTHHe6IYl9GgDhWErLNCKwEyyAEU7TQieWNMJkgHvUUjTAPpV2MkmZ6gdK6epeKNQJQJ

Original: [Figure]
Revised: [Figure]

Original: 3. I MITATION AND RL OBJECTIVES
Revised: 3. I MITATION AND RL OBJECTIVES

Original: To account for the long term future, the agent has to calculate the Expected Free Energy at t + 1 to ∞ .
Revised: To account for the long term future, the agent has to calculate the Expected Free Energy at t + 1 to ∞ .

Original: [Equation]
Revised: [Equation]

Revised: (cid:88)

Revised: [Equation]

Original: We deﬁne this curly F to be the objective that the Deep Free Energy Network should minimize.
Revised: We deﬁne this curly F to be the objective that the Deep Free Energy Network should minimize.

Original: Note that γ is a discount factor as in the case of general RL algorithms.
Revised: Note that γ is a discount factor as in the case of general RL algorithms.

Original: As it is impossible to sum over inﬁnity time steps, we introduce an Expected Free Energy Value function V ( s t +1 ) to estimate the cumulative Expected Free Energy.
Revised: As it is impossible to sum over inﬁnity time steps, we introduce an Expected Free Energy Value function V ( s t +1 ) to estimate the cumulative Expected Free Energy.

Original: Similarly to the case of Temporal Difference learning of Deep
Revised: Similarly to the case of Temporal Difference learning of Deep

Original: Q Network (Mnih et al., 2013), we use a target network V targ ( s t +2 ) to stabilize the learning process and deﬁne the loss for the value function as follows.
Revised: Q Network (Mnih et al., 2013), we use a target network V targ ( s t +2 ) to stabilize the learning process and deﬁne the loss for the value function as follows.

Original: [Equation]
Revised: [Equation]

Original: We made a design choice that the agent uses the value function only for RL, and not for imitation.
Revised: We made a design choice that the agent uses the value function only for RL, and not for imitation.

Original: In imitation, we use only the real value of the Expected Free Energy G t +1 at the next time step t +1 .
Revised: In imitation, we use only the real value of the Expected Free Energy G t +1 at the next time step t +1 .

Original: This is because imitation learning can be achieved without long term prediction as the agent is given the experts’ all time series data available.
Revised: This is because imitation learning can be achieved without long term prediction as the agent is given the experts’ all time series data available.

Original: On the other hand, in RL, using the value function to predict rewards in the long-term future is essential to avoid a local minimum and achieve the desired goal.
Revised: On the other hand, in RL, using the value function to predict rewards in the long-term future is essential to avoid a local minimum and achieve the desired goal.

Original: In conclusion, the objective functions of Deep Free Energy Network (FENet) for a data sequence ( o t , a t , r t , o t +1 ) are as follows.
Revised: In conclusion, the objective functions of Deep Free Energy Network (FENet) for a data sequence ( o t , a t , r t , o t +1 ) are as follows.

Original: [Equation]
Revised: [Equation]

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: The overall Free Energy calculation process is shown in Figure 1.
Revised: The overall Free Energy calculation process is shown in Figure 1.

Original: 3. N ETWORK A RCHITECTURE AND C ALCULATION
Revised: 3. N ETWORK A RCHITECTURE AND C ALCULATION

Original: For implementation, we made a design choice to use Recurrent State Space Model (Hafner et al., 2019b), a latent dynamics model with both deterministic and stochastic components.
Revised: For implementation, we made a design choice to use Recurrent State Space Model (Hafner et al., 2019b), a latent dynamics model with both deterministic and stochastic components.

Original: In this model, the hidden states s t are split into two parts: stochastic hidden states s t and deterministic hidden states h t .
Revised: In this model, the hidden states s t are split into two parts: stochastic hidden states s t and deterministic hidden states h t .

Original: The deterministic transition of h t is modeled using Recurrent Neural Networks (RNN) f as follows.
Revised: The deterministic transition of h t is modeled using Recurrent Neural Networks (RNN) f as follows.

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: We model the probabilities in Deep Free Energy Networks as follows.
Revised: We model the probabilities in Deep Free Energy Networks as follows.

Original: State prior
Revised: State prior

Original: [Equation]
Revised: [Equation]

Original: Observation likelihood
Revised: Observation likelihood

Original: [Equation]
Revised: [Equation]

Original: Reward likelihood
Revised: Reward likelihood

Original: [Equation]
Revised: [Equation]

Original: posterior
Revised: posterior

Original: [Equation]
Revised: [Equation]

Original: Policy
Revised: Policy

Original: [Equation]
Revised: [Equation]

Original: Value network
Revised: Value network

Original: [Equation]
Revised: [Equation]

Original: Target Value Network
Revised: Target Value Network

Original: [Equation]
Revised: [Equation]

Original: We model these probabilities as feedforward Neural Networks that output the mean and standard deviation of the random variables according to the Gaussian distribution.
Revised: We model these probabilities as feedforward Neural Networks that output the mean and standard deviation of the random variables according to the Gaussian distribution.

Original: The parameters θ, φ, ψ, ω are network parameters to be learned.
Revised: The parameters θ, φ, ψ, ω are network parameters to be learned.

Original: Using the network parameters, the objective loss functions can be written as follows.
Revised: Using the network parameters, the objective loss functions can be written as follows.

Original: [Equation]
Revised: [Equation]

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: 43) when
Revised: 43) when

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Algorithm 1 in Appendix shows overall calculations using these losses.
Revised: Algorithm 1 in Appendix shows overall calculations using these losses.

Original: The agent minimizes F IL for expert data D E and the agent minimizes F RL for agent data D A that the agent collects on its own.
Revised: The agent minimizes F IL for expert data D E and the agent minimizes F RL for agent data D A that the agent collects on its own.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We evaluate FENet on three continuous control tasks from images.
Revised: We evaluate FENet on three continuous control tasks from images.

Original: We compare our model with model-based RL and model-based imitation RL in dense and sparse reward setting when optimal expert is available.
Revised: We compare our model with model-based RL and model-based imitation RL in dense and sparse reward setting when optimal expert is available.

Original: Then we compare our model with imitation learning methods when only suboptimal experts are available.
Revised: Then we compare our model with imitation learning methods when only suboptimal experts are available.

Original: Finally, we investigate the merits of combining imitation and RL as an ablation study.
Revised: Finally, we investigate the merits of combining imitation and RL as an ablation study.

Original: Control tasks
Revised: Control tasks

Original: We used Cheetah-run, Walker-walk, and Quadruped-walk tasks, image-based continuous control tasks of DeepMind Control Suite (Tassa et al., 2018) shown in Figure 6.
Revised: We used Cheetah-run, Walker-walk, and Quadruped-walk tasks, image-based continuous control tasks of DeepMind Control Suite (Tassa et al., 2018) shown in Figure 6.

Original: The agent gets rewards ranging from 0 to 1 .
Revised: The agent gets rewards ranging from 0 to 1 .

Original: Quadruped-walk is the most difﬁcult as it has more action dimensions than the others.
Revised: Quadruped-walk is the most difﬁcult as it has more action dimensions than the others.

Original: Walker-walk is more challenging than Cheehtah-run because an agent ﬁrst has to stand up and then walk, meaning that the agent easily falls down on the ground, which is difﬁcult to predict.
Revised: Walker-walk is more challenging than Cheehtah-run because an agent ﬁrst has to stand up and then walk, meaning that the agent easily falls down on the ground, which is difﬁcult to predict.

Original: The episode length is 1000 steps starting from randomized initial states.
Revised: The episode length is 1000 steps starting from randomized initial states.

Original: We use action repeat R =4 for the Cheetah-run task, and R =
Revised: We use action repeat R =

Original: 4. P ERFORMANCE IN STANDARD VISUAL CONTROL TASKS
Revised: 4. P ERFORMANCE IN STANDARD VISUAL CONTROL TASKS

Original: We compare the performance of FENet to PlaNet (RL) and ”PlaNet with demonstrations” (imitation
Revised: We compare the performance of FENet to PlaNet (RL) and ”PlaNet with demonstrations” (imitation

Original: RL) in standard visual control tasks mentioned above.
Revised: RL) in standard visual control tasks mentioned above.

Original: We use PlaNet as a baseline method because PlaNet is one of the most basic methods using Recurrent State Space Model, on top of which
Revised: We use PlaNet as a baseline method because PlaNet is one of the most basic methods using Recurrent State Space Model, on top of which

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: we build our model.
Revised: we build our model.

Original: As FENet uses expert data, we create ”PlaNet with demonstrations” for fair comparison.
Revised: As FENet uses expert data, we create ”PlaNet with demonstrations” for fair comparison.

Original: This variant of PlaNet has an additional experience replay pre-populated with expert trajectories and minimize a loss calculated from the expert data in addition to PlaNet’s original loss.
Revised: This variant of PlaNet has an additional experience replay pre-populated with expert trajectories and minimize a loss calculated from the expert data in addition to PlaNet’s original loss.

Original: Figure 2 shows that ”PlaNet with demonstrations” is always better than PlaNet and that FENet is ranked higher as the difﬁculty of tasks gets higher.
Revised: Figure 2 shows that ”PlaNet with demonstrations” is always better than PlaNet and that FENet is ranked higher as the difﬁculty of tasks gets higher.

Original: In Cheetah-run, FENet gives competitive performance with PlaNet.
Revised: In Cheetah-run, FENet gives competitive performance with PlaNet.

Original: In Walker-walk, FENet and ”PlaNet with demonstrations” are almost competitive, both of which are substantially better than PlaNet thanks to expert knowledge being leveraged to increase sample efﬁciency.
Revised: In Walker-walk, FENet and ”PlaNet with demonstrations” are almost competitive, both of which are substantially better than PlaNet thanks to expert knowledge being leveraged to increase sample efﬁciency.

Original: In Quadruped-walk, FENet is slightly better than the other two baselines.
Revised: In Quadruped-walk, FENet is slightly better than the other two baselines.

Original: 4. P ERFORMANCE IN SPARSE - REWARD VISUAL CONTROL TASKS
Revised: 4. P ERFORMANCE IN SPARSE - REWARD VISUAL CONTROL TASKS

Original: In real-world robot learning, it is demanding to craft a dense reward function to lead robots to desired behaviors.
Revised: In real-world robot learning, it is demanding to craft a dense reward function to lead robots to desired behaviors.

Original: It would be helpful if an agent could acquire desired behaviors simply by giving sparse signals.
Revised: It would be helpful if an agent could acquire desired behaviors simply by giving sparse signals.

Original: We compare the performance of FENet to PlaNet and ”PlaNet with demonstrations” in sparse-reward settings, where agents do not get rewards less than 0.5 per time step (Note that in the original implementation of Cheetah-run, Walker-walk and Quadruped-walk, agents get rewards ranging from 0 to 1 per time step).
Revised: We compare the performance of FENet to PlaNet and ”PlaNet with demonstrations” in sparse-reward settings, where agents do not get rewards less than 0.5 per time step (Note that in the original implementation of Cheetah-run, Walker-walk and Quadruped-walk, agents get rewards ranging from 0 to 1 per time step).

Original: Figure 3 shows that FENet outperforms PlaNet and ”PlaNet with demonstrations” in all three tasks.
Revised: Figure 3 shows that FENet outperforms PlaNet and ”PlaNet with demonstrations” in all three tasks.

Original: In Cheetah-run, PlaNet and ”PlaNet with demonstrations” are not able to get even a single reward.
Revised: In Cheetah-run, PlaNet and ”PlaNet with demonstrations” are not able to get even a single reward.

Original: [Figure]
Revised: [Figure]

Original: 4. P ERFORMANCE WITH SUBOPTIMAL EXPERTS
Revised: 4. P ERFORMANCE WITH SUBOPTIMAL EXPERTS

Original: In real-world robot learning, expert trajectories are often given by human experts.
Revised: In real-world robot learning, expert trajectories are often given by human experts.

Original: It is natural to assume that expert trajectories are suboptimal and that there remains much room for improvement.
Revised: It is natural to assume that expert trajectories are suboptimal and that there remains much room for improvement.

Original: We compare the performance of FENet to Behavioral Cloning imitation methods.
Revised: We compare the performance of FENet to Behavioral Cloning imitation methods.

Original: We use two types of networks for behavioral cloning methods: recurrent policy and recurrent decoder policy.
Revised: We use two types of networks for behavioral cloning methods: recurrent policy and recurrent decoder policy.

Original: The recurrent policy π R ( a t | o t ) is neural networks with one gated recurrent unit cell and three dense layers.
Revised: The recurrent policy π R ( a t | o t ) is neural networks with one gated recurrent unit cell and three dense layers.

Original: The recurrent decoder policy π R ( a t , o t +1 | o t ) is neural networks with one gated recurrent unit cell and four dense layers and deconvolution layers as in the decoder of PlaNet.
Revised: The recurrent decoder policy π R ( a t , o t +1 | o t ) is neural networks with one gated recurrent unit cell and four dense layers and deconvolution layers as in the decoder of PlaNet.

Original: Both networks does not get raw pixel observations but take observations encoded by the same convolutional encoder as PlaNet’s.
Revised: Both networks does not get raw pixel observations but take observations encoded by the same convolutional encoder as PlaNet’s.

Original: Figure 4 shows that while imitation methods overﬁt to the expert and cannot surpass the suboptimal expert performance, FENet is able to substantially surpass the suboptimal expert’s performance.
Revised: Figure 4 shows that while imitation methods overﬁt to the expert and cannot surpass the suboptimal expert performance, FENet is able to substantially surpass the suboptimal expert’s performance.

Original: 4. L EARNING STRATEGIES
Revised: 4. L EARNING STRATEGIES

Original: Figure 5 compares learning strategies of FENet in Cheetah-run and Walker-walk (ablation study).
Revised: Figure 5 compares learning strategies of FENet in Cheetah-run and Walker-walk (ablation study).

Original: ’Imitation RL’ is the default FENet agent that does imitation learning and RL at the same time, minimizing F IL + F RL . ’
Revised: ’Imitation RL’ is the default FENet agent that does imitation learning and RL at the same time, minimizing F IL + F RL . ’

Original: Imitation-pretrained RL’ is an agent that ﬁrst learns the model only with imitation (minimizing
Revised: Imitation-pretrained RL’ is an agent that ﬁrst learns the model only with imitation (minimizing

Original: F IL )
Revised: F IL )

Original: and then does RL using the pre-trained model (minimizing F RL ). ’
Revised: and then does RL using the pre-trained model (minimizing F RL ). ’

Original: RL only’ is an agent that does RL only, minimizing
Revised: RL only’ is an agent that does RL only, minimizing

Original: F RL . ’
Revised: F RL . ’

Original: Imitation only’ is an agent that does imitation only, minimizing
Revised: Imitation only’ is an agent that does imitation only, minimizing

Original: F IL .
Revised: F IL .

Original: While ’imitation only’ gives the best performance and ’imitation
Revised: While ’imitation only’ gives the best performance and ’imitation

Original: RL’ gives the second best in Cheetah-run, ’imitation RL’ gives the best performance and ’imitation only’ gives the worst performance in Walker-walk.
Revised: RL’ gives the second best in Cheetah-run, ’imitation RL’ gives the best performance and ’imitation only’ gives the worst performance in Walker-walk.

Original: We could say ’imitation RL’ is the most robust to the properties of tasks.
Revised: We could say ’imitation RL’ is the most robust to the properties of tasks.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Active Inference
Revised: Active Inference

Original: Friston, who ﬁrst proposed Active Inference, has evaluated the performance in simple control tasks and a low-dimensional maze (Friston et al., 2012; 2015).
Revised: Friston, who ﬁrst proposed Active Inference, has evaluated the performance in simple control tasks and a low-dimensional maze (Friston et al., 2012; 2015).

Original: Ueltzhoffer implemented Active Inference with Deep Neural Networks and evaluated the performance in a simple control task (Ueltzh¨offer, 2018).
Revised: Ueltzhoffer implemented Active Inference with Deep Neural Networks and evaluated the performance in a simple control task (Ueltzh¨offer, 2018).

Original: Millidge proposed a Deep Active Inference framework with value functions to estimate the correct Free Energy and succeeded in solving Gym environments (Millidge, 2019).
Revised: Millidge proposed a Deep Active Inference framework with value functions to estimate the correct Free Energy and succeeded in solving Gym environments (Millidge, 2019).

Original: Our approach extends Deep Active Inference to combine imitation and RL, solving more challenging tasks.
Revised: Our approach extends Deep Active Inference to combine imitation and RL, solving more challenging tasks.

Original: RL from demonstration
Revised: RL from demonstration

Original: Reinforced Imitation Learning succeeds in reducing sample complexity by using imitation as pre-training before RL (Pfeiffer et al., 2018).
Revised: Reinforced Imitation Learning succeeds in reducing sample complexity by using imitation as pre-training before RL (Pfeiffer et al., 2018).

Original: Adding demonstrations into a replay buffer of off policy RL methods also demonstrates high sample efﬁciency (Vecerik et al., 2017; Nair et al., 2018;
Revised: Adding demonstrations into a replay buffer of off policy RL methods also demonstrates high sample efﬁciency (Vecerik et al., 2017; Nair et al., 2018;

Original: Paine et al., 2019).
Revised: Paine et al., 2019).

Original: Demo Augmented Policy Gradient mixes the policy gradient with a behavioral cloning gradient (Rajeswaran* et al., 2018).
Revised: Demo Augmented Policy Gradient mixes the policy gradient with a behavioral cloning gradient (Rajeswaran* et al., 2018).

Original: Deep Q-learning from
Revised: Deep Q-learning from

Original: Demonstrations (DQfD) not only use demonstrations for pre-training but also calculates gradients
Revised: Demonstrations (DQfD) not only use demonstrations for pre-training but also calculates gradients

Original: [Figure]
Revised: [Figure]

Original: from demonstrations and environment interaction data (Hester et al., 2018).
Revised: from demonstrations and environment interaction data (Hester et al., 2018).

Original: Truncated HORizon
Revised: Truncated HORizon

Original: Policy Search uses demonstrations to shape rewards so that subsequent planning can achieve superior performance to RL even when experts are suboptimal (Sun et al., 2018).
Revised: Policy Search uses demonstrations to shape rewards so that subsequent planning can achieve superior performance to RL even when experts are suboptimal (Sun et al., 2018).

Original: Soft Q Imitation
Revised: Soft Q Imitation

Original: Learning gives rewards that encourage the agent to return to demonstrated states in order to avoid policy collapse (Reddy et al., 2019).
Revised: Learning gives rewards that encourage the agent to return to demonstrated states in order to avoid policy collapse (Reddy et al., 2019).

Original: Our approach is similar to DQfD in terms of mixing gradients calculated from demonstrations and from environment interaction data.
Revised: Our approach is similar to DQfD in terms of mixing gradients calculated from demonstrations and from environment interaction data.

Original: One key difference is that FENet concurrently learns the generative model of the world so that it can be robust to wider environment properties.
Revised: One key difference is that FENet concurrently learns the generative model of the world so that it can be robust to wider environment properties.

Original: Control with latent dynamics model
Revised: Control with latent dynamics model

Original: World Models acquire latent spaces and dynamics over the spaces separately, and evolve simple linear controllers to solve visual control tasks (Ha & Schmidhuber, 2018).
Revised: World Models acquire latent spaces and dynamics over the spaces separately, and evolve simple linear controllers to solve visual control tasks (Ha & Schmidhuber, 2018).

Original: PlaNet learns Recurrent State Space Model and does planning with Model Predictive
Revised: PlaNet learns Recurrent State Space Model and does planning with Model Predictive

Original: Control at test phase (Hafner et al., 2019b).
Revised: Control at test phase (Hafner et al., 2019b).

Original: Dreamer, which is recently built upon PlaNet, has a policy for latent imagination and achieved higher performance than PlaNet (Hafner et al., 2019a).
Revised: Dreamer, which is recently built upon PlaNet, has a policy for latent imagination and achieved higher performance than PlaNet (Hafner et al., 2019a).

Original: Our approach also uses Recurrent State Space Model to describe variational inference, and we are the ﬁrst to combine imitation and RL over latent dynamics models to the best of our knowledge.
Revised: Our approach also uses Recurrent State Space Model to describe variational inference, and we are the ﬁrst to combine imitation and RL over latent dynamics models to the best of our knowledge.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We present FENet, an agent that combines Imitation Learning and Reinforcement Learning using
Revised: We present FENet, an agent that combines Imitation Learning and Reinforcement Learning using

Original: Free Energy objectives.
Revised: Free Energy objectives.

Original: For this, we theoretically extend the Free Energy Principle and introduce a policy prior that encodes experts’ behaviors and a policy posterior that learns to maximize expected rewards without deviating too much from the policy prior.
Revised: For this, we theoretically extend the Free Energy Principle and introduce a policy prior that encodes experts’ behaviors and a policy posterior that learns to maximize expected rewards without deviating too much from the policy prior.

Original: FENet outperforms model-based RL and imitation RL especially in visual control tasks with sparse rewards and FENet also outperforms suboptimal experts’ performance unlike Behavioral cloning.
Revised: FENet outperforms model-based RL and imitation RL especially in visual control tasks with sparse rewards and FENet also outperforms suboptimal experts’ performance unlike Behavioral cloning.

Original: Strong potentials in sparse environment with suboptimal experts are important factors for real-world robot learning.
Revised: Strong potentials in sparse environment with suboptimal experts are important factors for real-world robot learning.

Original: Directions for future work include learning the balance between imitation and RL, i.e. Free Energy and Expected Free Energy so that the agent can select the best approach to solve its confronting tasks by monitoring the value of Free Energy.
Revised: Directions for future work include learning the balance between imitation and RL, i.e. Free Energy and Expected Free Energy so that the agent can select the best approach to solve its confronting tasks by monitoring the value of Free Energy.

Original: It is also important to evaluate FENet in real-world robotics tasks to show that our method is effective in more realistic settings that truly appear in the real world.
Revised: It is also important to evaluate FENet in real-world robotics tasks to show that our method is effective in more realistic settings that truly appear in the real world.

Original: A PPENDIX
Revised: A PPENDIX

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: See Algorithm 1.
Revised: See Algorithm 1.

Original: A. I MPLEMENTATION
Revised: A. I MPLEMENTATION

Original: To stabilize the learning process we adopt burn-in, a technique to recover initial states of RNN’s hidden variables h t (Kapturowski et al., 2019).
Revised: To stabilize the learning process we adopt burn-in, a technique to recover initial states of RNN’s hidden variables h t (Kapturowski et al., 2019).

Original: As shown in Algorithm 1, the agent calculates the
Revised: As shown in Algorithm 1, the agent calculates the

Original: Free Energy with mini batches sampled from the expert or agent experience replay buffer D , which means that h t is initialized randomly in every mini batch calculation.
Revised: Free Energy with mini batches sampled from the expert or agent experience replay buffer D , which means that h t is initialized randomly in every mini batch calculation.

Original: Since the Free Energy heavily depends on h t , it is crucial to estimate the accurate hidden states.
Revised: Since the Free Energy heavily depends on h t , it is crucial to estimate the accurate hidden states.

Original: We set a burn-in period when a portion of the replay sequence is used only for unrolling the networks to produce initial states.
Revised: We set a burn-in period when a portion of the replay sequence is used only for unrolling the networks to produce initial states.

Original: After the period, we update the networks only on the remaining part of the sequence.
Revised: After the period, we update the networks only on the remaining part of the sequence.

Original: We use PyTorch (Paszke et al., 2017) to write neural networks and run experiments using NVIDIA
Revised: We use PyTorch (Paszke et al., 2017) to write neural networks and run experiments using NVIDIA

Original: GeForce GTX 1080 Ti / RTX 

Revised: GeForce GTX 1080 Ti / RTX

Original: Ti / Tesla V100 GPU (1 GPU per experiment).
Revised: Ti / Tesla V100 GPU (1 GPU per experiment).

Original: The training time for our FENet implementation is about 24 hours on the Control Suite environment.
Revised: The training time for our FENet implementation is about 24 hours on the Control Suite environment.

Original: As for the hyper parameters, we use the convolutional encoder and decoder networks from (Ha & Schmidhuber, 2018) and Recurrent State Space Model from (Hafner et al., 2019b) and implement all other functions as three dense layers of size 200 with ReLU activations (Nair & Hinton, 2010).
Revised: As for the hyper parameters, we use the convolutional encoder and decoder networks from (Ha & Schmidhuber, 2018) and Recurrent State Space Model from (Hafner et al., 2019b) and implement all other functions as three dense layers of size 200 with ReLU activations (Nair & Hinton, 2010).

Original: We made a design choice to make the policy prior, the policy posterior, and the observation likelihood, the reward likelihood deterministic functions while the state prior and the state posterior are stochastic.
Revised: We made a design choice to make the policy prior, the policy posterior, and the observation likelihood, the reward likelihood deterministic functions while the state prior and the state posterior are stochastic.

Original: We use the batch size B = 25 for ’imitation RL’ with FENet, and B = 50 for other types and baseline methods.
Revised: We use the batch size B = 25 for ’imitation RL’ with FENet, and B = 50 for other types and baseline methods.

Original: We use the chunk length L = 50 , the burn-in period 20.
Revised: We use the chunk length L = 50 , the burn-in period 20.

Original: We use seed episodes
Revised: We use seed episodes

Original: S = 40 , expert episodes N = 10000 trained with PlaNet (Hafner et al., 2019b), collect interval
Revised: S = 40 , expert episodes N = 10000 trained with PlaNet (Hafner et al., 2019b), collect interval

Original: C = 100 and action exploration noise Normal(0, 0.3).
Revised: C = 100 and action exploration noise Normal(0, 0.3).

Original: We use the discount factor γ = 0 .
Revised: We use the discount factor γ = 0 .

Original: 99 and the
Revised: 99 and the

Original: Algorithm
Revised: Algorithm

Original: Input:
Revised: Input:

Original: Seed episodes S
Revised: Seed episodes S

Original: Collect interval C
Revised: Collect interval C

Original: Batch size B
Revised: Batch size B

Original: Chunk length L
Revised: Chunk length L

Original: Expert episodes N
Revised: Expert episodes N

Original: Target smoothing rate ρ
Revised: Target smoothing rate ρ

Original: Learning rate α
Revised: Learning rate α

Original: State prior p θ ( s t | h t ) State posterior q φ ( s t | h t , o t ) Policy prior p θ ( a t | s t , h t ) Policy posterior q ψ ( a t | s t , h t )
Revised: State prior p θ ( s t | h t ) State posterior q φ ( s t | h t , o t ) Policy prior p θ ( a t | s t , h t ) Policy posterior q ψ ( a t | s t , h t )

Original: Likelihood p θ ( o t | s t , h t ) , p θ ( r t − 1 | s t , h t )
Revised: Likelihood p θ ( o t | s t , h t ) , p θ ( r t − 1 | s t , h t )

Original: Value function V ω ( s t ) Target value function V ω targ ( s t )
Revised: Value function V ω ( s t ) Target value function V ω targ ( s t )

Original: Initialize expert dataset D E with N expert trajectoriesInitialize agent dataset D A with S random episodesInitialize neural network parameters θ, φ, ψ, ω randomly while not converged do for update step c = 1 ..
Revised: Initialize expert dataset D E with N expert trajectoriesInitialize agent dataset D A with S random episodesInitialize neural network parameters θ, φ, ψ, ω randomly while not converged do for update step c = 1 ..

Original: C do
Revised: C do

Original: // Imitation Learning
Revised: // Imitation Learning

Original: Draw expert data { ( o t , a t , r t , o t +1 )
Revised: Draw expert data { ( o t , a t , r t , o t +1 )

Original: k
Revised: k

Original: + L t = k } Bi =1 ∼
Revised: + L t = k } Bi =1 ∼

Original: D E Compute Free Energy F IL from equation
Revised: D E Compute Free Energy F IL from equation

Original:  // Reinforcement Learning
Revised:  // Reinforcement Learning

Original: Draw agent data { ( o t , a t , r t , o t +1 )
Revised: Draw agent data { ( o t , a t , r t , o t +1 )

Original: k
Revised: k

Original: + L t = k } Bi =1 ∼
Revised: + L t = k } Bi =1 ∼

Original: D A Compute Free Energy F RL from equation
Revised: D A Compute Free Energy F RL from equation

Original: Compute V function’s Loss L from equa- tion
Revised: Compute V function’s Loss L from equa- tion

Original: // Update parameters
Revised: // Update parameters

Original: θ ← θ − α ∇ θ ( F IL + F RL )
Revised: θ ← θ − α ∇ θ ( F IL + F RL )

Original: φ ← φ − α ∇ φ ( F IL + F RL ) ψ ← ψ − α ∇ ψ ( F IL + F RL ) ω ← ω − α ∇ ω L ω targ ← ρω targ + (1 − ρ ) ω end for // Environment interaction o 1 ← env.reset() for time step t = 1 ..
Revised: φ ← φ − α ∇ φ ( F IL + F RL ) ψ ← ψ − α ∇ ψ ( F IL + F RL ) ω ← ω − α ∇ ω L ω targ ← ρω targ + (1 − ρ ) ω end for // Environment interaction o 1 ← env.reset() for time step t = 1 ..

Original: T do
Revised: T do

Original: Infer hidden states s t ← q φ ( s t | h t , o t )
Revised: Infer hidden states s t ← q φ ( s t | h t , o t )

Original: Calculate actions a t ← q ψ ( a t | s t , h t )
Revised: Calculate actions a t ← q ψ ( a t | s t , h t )

Original: Add exploration noise to actions r t , o t +1 ← env.step ( a t ) end for
Revised: Add exploration noise to actions r t , o t +1 ← env.step ( a t ) end for

Original: D A ← D A ∪ { ( o t , a t , r t , o t +1 )
Revised: D A ← D A ∪ { ( o t , a t , r t , o t +1 )

Original: Tt =1 } end while target smoothing rate ρ = 0 .
Revised: Tt =1 } end while target smoothing rate ρ = 0 .

Original: We use Adam (Kingma & Ba, 2014) with learning rates α = 10 −and scale down gradient norms that exceed 1000.
Revised: We use Adam (Kingma & Ba, 2014) with learning rates α = 10 −and scale down gradient norms that exceed 1000.

Original: We scale the reward-related loss by 100, the policy-prior-related loss by 10.
Revised: We scale the reward-related loss by 100, the policy-prior-related loss by 10.

Original: We clip KL loss between the hidden states below 3 free nats and clip
Revised: We clip KL loss between the hidden states below 3 free nats and clip

Original: KL loss between the policies below 0.6.
Revised: KL loss between the policies below 0.6.