Original: Solipsistic Reinforcement Learning
Revised: Solipsistic Reinforcement Learning

Original: Abstract
Revised: Abstract

Original: We present a model-based reinforcement learning framework that aims to tackle environments with high dimensional state spaces.
Revised: We present a model-based reinforcement learning framework that aims to tackle environments with high dimensional state spaces.

Original: In contrast to traditional approaches, agents under our framework learn a low dimensional internal representation of the environment while avoiding the need to learn a generative model of the environment itself.
Revised: In contrast to traditional approaches, agents under our framework learn a low dimensional internal representation of the environment while avoiding the need to learn a generative model of the environment itself.

Original: This solipsistic representation is trained to encode a belief that is consistent with the dynamics of the environment and is then exploited for eﬀective planning.
Revised: This solipsistic representation is trained to encode a belief that is consistent with the dynamics of the environment and is then exploited for eﬀective planning.

Original: We present speciﬁc cases of our framework with choices of model and corresponding planning algorithms that can deal with both discrete and continuous state environments.
Revised: We present speciﬁc cases of our framework with choices of model and corresponding planning algorithms that can deal with both discrete and continuous state environments.

Original: We demonstrate empirically gains in eﬃciency over existing model-free methods when learning directly from pixels and analyze the properties of our learned representations
Revised: We demonstrate empirically gains in eﬃciency over existing model-free methods when learning directly from pixels and analyze the properties of our learned representations

Original:  Introduction
Revised:  Introduction

Original: The real world is complex and a learning agent must be able to recognize relevant signals to decide what actions to take towards reaching a goal.
Revised: The real world is complex and a learning agent must be able to recognize relevant signals to decide what actions to take towards reaching a goal.

Original: The focus of our work is to form environment representations for model-based planning and reward prediction without having to learn a generative model of the potentially complex environment.
Revised: The focus of our work is to form environment representations for model-based planning and reward prediction without having to learn a generative model of the potentially complex environment.

Original: To motivate our approach, Pavlov’s dog (Dennis and Mitterer, 2004) learns to associate the sound of a bell with the eventual reward of food, despite sensory distractions.
Revised: To motivate our approach, Pavlov’s dog (Dennis and Mitterer, 2004) learns to associate the sound of a bell with the eventual reward of food, despite sensory distractions.

Original: The dog learns to associate the of with the eventual arrival of food salivating due to this expectation
Revised: The dog learns to associate the of with the eventual arrival of food salivating due to this expectation

Original: There are classically two interpretations: (a) a modelfree interpretation is that the dog learns a value v ( x t ) (expectation of eventual food reward) as a function of the environment state x t at time t ; (b) a standard model-based interpretation is that the dog models the environment and can use that to predict the future p ( x t + k | x t ) and any eventual reward.
Revised: There are classically two interpretations: (a) a modelfree interpretation is that the dog learns a value v ( x t ) (expectation of eventual food reward) as a function of the environment state x t at time t ; (b) a standard model-based interpretation is that the dog models the environment and can use that to predict the future p ( x t + k | x t ) and any eventual reward.

Original: In contrast to these standard approaches, we posit an alternative model-based interpretation in which the dog forms an internal
Revised: In contrast to these standard approaches, we posit an alternative model-based interpretation in which the dog forms an internal

Original: 1 low-dimensional representation s t as a function of the external environment state x t and forms a predictive model p ( s t + k | s t ) of the representation, without learning a model of the environment itself.
Revised: 1 low-dimensional representation s t as a function of the external environment state x t and forms a predictive model p ( s t + k | s t ) of the representation, without learning a model of the environment itself.

Original: This representation is useful if the dog is able to accurately predict eventual reward r t + k using the solipsistic transition p ( s t + k | s t ) and reward model p (
Revised: This representation is useful if the dog is able to accurately predict eventual reward r t + k using the solipsistic transition p ( s t + k | s t ) and reward model p (

Original: [Equation]
Revised: [Equation]

Original: In the Reinforcement Learning (RL) setting an agent observes state x t from the environment at time-step t , takes action a t , and subsequently observes x t +1 and reward r t +1 .
Revised: In the Reinforcement Learning (RL) setting an agent observes state x t from the environment at time-step t , takes action a t , and subsequently observes x t +1 and reward r t +1 .

Original: The goal of the agent is to learn (through interactions with the environment) how to take actions that result in favorable long-term rewards (Sutton and Barto, 2018).
Revised: The goal of the agent is to learn (through interactions with the environment) how to take actions that result in favorable long-term rewards (Sutton and Barto, 2018).

Original: A standard
Revised: A standard

Original: RL assumption is that there is an underlying Markov Decision Process with transition p ( x t +1 | x t , a t ) (Deisenroth and Rasmussen, 2011; Gal et al., 2016; Amos et al., 2018; Chua et al., 2018); rewards are functions of the observed state and the goal is usually to take actions that maximize cumulative reward, see Figure 2a 2 .
Revised: RL assumption is that there is an underlying Markov Decision Process with transition p ( x t +1 | x t , a t ) (Deisenroth and Rasmussen, 2011; Gal et al., 2016; Amos et al., 2018; Chua et al., 2018); rewards are functions of the observed state and the goal is usually to take actions that maximize cumulative reward, see Figure 2a 2 .

Original: The action a t depends on the state x t − 1 , meaning that the state x t − 1 is revealed before the action a t is decided.
Revised: The action a t depends on the state x t − 1 , meaning that the state x t − 1 is revealed before the action a t is decided.

Original: In model-based RL, we attempt to learn the model of the transition dynamics.
Revised: In model-based RL, we attempt to learn the model of the transition dynamics.

Original: Compared to model-free approaches, model-based RL can be signiﬁcantly more sample eﬃcient (Deisenroth and Rasmussen, 2011; Gal et al., 2016;
Revised: Compared to model-free approaches, model-based RL can be signiﬁcantly more sample eﬃcient (Deisenroth and Rasmussen, 2011; Gal et al., 2016;

Original: Amos et al., 2018; Chua et al., 2018).
Revised: Amos et al., 2018; Chua et al., 2018).

Original: However, for environments with high-dimensional states (such as an image pixels) the complexity and potential redundancy in the observations can make learning the environment dynamics using a model diﬃcult and potentially unnecessary (Ha and Schmidhuber, 2018; Hafner et al., 2019b).
Revised: However, for environments with high-dimensional states (such as an image pixels) the complexity and potential redundancy in the observations can make learning the environment dynamics using a model diﬃcult and potentially unnecessary (Ha and Schmidhuber, 2018; Hafner et al., 2019b).

Original: A recent trend is to learn a lower dimensional representation s t that is used to model relevant dynamics and reward prediction, such as PlaNet Hafner et al.
Revised: A recent trend is to learn a lower dimensional representation s t that is used to model relevant dynamics and reward prediction, such as PlaNet Hafner et al.

Original: (2019b)
Revised: (2019b)

Original: , World Model (Ha and Schmidhuber, 2018) and other variants (Chiappa et al., 2017; Ha and Schmidhuber, 2018), see Figure 2b.
Revised: , World Model (Ha and Schmidhuber, 2018) and other variants (Chiappa et al., 2017; Ha and Schmidhuber, 2018), see Figure 2b.

Original: The usual strategy is to train a latent variable generative model p ( x t ) =
Revised: The usual strategy is to train a latent variable generative model p ( x t ) =

Original: (cid:82) p ( x t | s t ) p ( s t ) ds t with an encoding-decoding structure (Kingma and Welling, 2014), while jointly training a dynamics model
Revised: (cid:82) p ( x t | s t ) p ( s t ) ds t with an encoding-decoding structure (Kingma and Welling, 2014), while jointly training a dynamics model

Original: [Equation]
Revised: [Equation]

Original: , a t ) in the lower dimensional latent space (Hafner et al., 2019b; Chiappa et al., 2017; Ha and Schmidhuber, 2018; Hafner et al., 2020).
Revised: , a t ) in the lower dimensional latent space (Hafner et al., 2019b; Chiappa et al., 2017; Ha and Schmidhuber, 2018; Hafner et al., 2020).

Original: The learned model infers a latent representation s t given a state observation x t (for example a sample from the posterior p ( s t | x t ) ∝ p ( x t | s t ) p ( s t ) ) that can be used by the dynamics model for eﬃcient planning.
Revised: The learned model infers a latent representation s t given a state observation x t (for example a sample from the posterior p ( s t | x t ) ∝ p ( x t | s t ) p ( s t ) ) that can be used by the dynamics model for eﬃcient planning.

Original: Arguably, a limitation of these recent approaches is that they spend signiﬁcant computational eﬀort on learning a generative model of the high-dimensional state x t – however, this generative model is not used directly during the planning phase.
Revised: Arguably, a limitation of these recent approaches is that they spend signiﬁcant computational eﬀort on learning a generative model of the high-dimensional state x t – however, this generative model is not used directly during the planning phase.

Original: The learned representation s t in these approaches is therefore likely encoding redundant information about the environment x t .
Revised: The learned representation s t in these approaches is therefore likely encoding redundant information about the environment x t .

Original: An alternative is to use model-free algorithms such as Deep Q-Networks (DQN), Proximal
Revised: An alternative is to use model-free algorithms such as Deep Q-Networks (DQN), Proximal

Original: Policy Optimization (PPO), amongst others (Mnih et al., 2015; Schulman et al., 2017; 2015;
Revised: Policy Optimization (PPO), amongst others (Mnih et al., 2015; Schulman et al., 2017; 2015;

Original: Babaeizadeh et al., 2017).
Revised: Babaeizadeh et al., 2017).

Original: The upside is that these approaches avoid creating a generative model of the environment by learning a policy from pixels to state values.
Revised: The upside is that these approaches avoid creating a generative model of the environment by learning a policy from pixels to state values.

Original: A potential downside is that they suﬀer from poor sample eﬃciency compared to model based approaches.
Revised: A potential downside is that they suﬀer from poor sample eﬃciency compared to model based approaches.

Original: The question we therefore study here, similar to other recent research work in this area
Revised: The question we therefore study here, similar to other recent research work in this area

Original: Hafner et al.
Revised: Hafner et al.

Original: (2019a), is whether it is possible to perform model-based RL without making a generative model of the environment.
Revised: (2019a), is whether it is possible to perform model-based RL without making a generative model of the environment.

Original: If this were possible, we could potentially reap the beneﬁts of the sample-eﬃciency of model-based RL, without the need to model complex high-dimensional observations.
Revised: If this were possible, we could potentially reap the beneﬁts of the sample-eﬃciency of model-based RL, without the need to model complex high-dimensional observations.

Original:  Solipsistic Representations
Revised:  Solipsistic Representations

Original: A solipsistic representation s t of an observation x t is one that is consistent (the predicted next solipsistic state s t +1 given action a t matches the observed next solipsistic state s t +1 ) and informative (one can predict the reward well using s t ).
Revised: A solipsistic representation s t of an observation x t is one that is consistent (the predicted next solipsistic state s t +1 given action a t matches the observed next solipsistic state s t +1 ) and informative (one can predict the reward well using s t ).

Original: A solipsistic Markov model is depicted in Figure 3a, where we remove the arrow from s to x in the latent variable model and instead introduce a recognition distribution p ( s t | x t ) , the purpose of which is to encode only the information in x t that is needed to eﬀectively learn the dynamics and reward.
Revised: A solipsistic Markov model is depicted in Figure 3a, where we remove the arrow from s to x in the latent variable model and instead introduce a recognition distribution p ( s t | x t ) , the purpose of which is to encode only the information in x t that is needed to eﬀectively learn the dynamics and reward.

Original: We will use throughout a toy ‘MNIST game’ to help build intuition, see Figure 4.
Revised: We will use throughout a toy ‘MNIST game’ to help build intuition, see Figure 4.

Original: Each observation x t is a 28 × 28 MNIST image representing a digit from 0 to 9.
Revised: Each observation x t is a 28 × 28 MNIST image representing a digit from 0 to 9.

Original: The agent has two possible actions: ‘minus 1’ or ‘plus 1’; the environment shows the resulting digit’s image.
Revised: The agent has two possible actions: ‘minus 1’ or ‘plus 1’; the environment shows the resulting digit’s image.

Original: The digit will stay the same when taking action ‘minus 1’ from digit 0 and ‘plus 1’ from digit rrrxxxaa(
Revised: The digit will stay the same when taking action ‘minus 1’ from digit 0 and ‘plus 1’ from digit rrrxxxaa(

Original: [Figure]
Revised: [Figure]

Original: Observed Markov Model xrxrxrsssaa(b) Latent Variable Model.
Revised: Observed Markov Model xrxrxrsssaa(b) Latent Variable Model.

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: xrrrsssaassxx
Revised: xrrrsssaassxx

Original: [Figure]
Revised: [Figure]

Original: The game is initialized at digit 4.
Revised: The game is initialized at digit 4.

Original: The reward r t is 1 if the state is an image of digit 9; otherwise the reward r t = 0 .
Revised: The reward r t is 1 if the state is an image of digit 9; otherwise the reward r t = 0 .

Original: Whilst the observation x t is a 784 dimensional image, clearly the underlying dynamics is representable by an integer s t ∈ { 0 , . .
Revised: Whilst the observation x t is a 784 dimensional image, clearly the underlying dynamics is representable by an integer s t ∈ { 0 , . .

Original: [Figure]
Revised: [Figure]

Original: Solipsistic Consistency We wish to ensure that, for a given recognition distribution p rec ( s t | x t ) , the dynamics of the solipsistic model are consistent with the dynamics of the true environment when training the model under sampled trajectories.
Revised: Solipsistic Consistency We wish to ensure that, for a given recognition distribution p rec ( s t | x t ) , the dynamics of the solipsistic model are consistent with the dynamics of the true environment when training the model under sampled trajectories.

Original: Thus, in the setting of Pavlov’s dog, if the model maps the current environment x t to the internal state x t → s t that represents hearing the bell s t = bell , and the dog predicts s t → s t +1 from this that s t +1 = food , then we must have that the next external state x t +1 → s t +1 indeed maps to food .
Revised: Thus, in the setting of Pavlov’s dog, if the model maps the current environment x t to the internal state x t → s t that represents hearing the bell s t = bell , and the dog predicts s t → s t +1 from this that s t +1 = food , then we must have that the next external state x t +1 → s t +1 indeed maps to food .

Original: This ensures that solipsistic transitions are eﬀective for planning.
Revised: This ensures that solipsistic transitions are eﬀective for planning.

Original: Similarly, in the MNIST game, we need to force the solipsistic model to predict a digit that is consistent with the image that would appear in the next time step under the true environment transition.
Revised: Similarly, in the MNIST game, we need to force the solipsistic model to predict a digit that is consistent with the image that would appear in the next time step under the true environment transition.

Original: More speciﬁcally, we assume a Markov transition distribution p tran ( s t +1 | s t , a t ) which takes the current solipsistic state and action as input and gives the distribution for the next solipsistic state s t +1 .
Revised: More speciﬁcally, we assume a Markov transition distribution p tran ( s t +1 | s t , a t ) which takes the current solipsistic state and action as input and gives the distribution for the next solipsistic state s t +1 .

Original: Given x t and a t , the solipsistic state distribution at time t + 1 can be predicted using p
Revised: Given x t and a t , the solipsistic state distribution at time t + 1 can be predicted using p

Original: [Equation]
Revised: [Equation]

Original: We want this predicted distribution to be consistent with the recognition distribution p rec ( s t +1 | x t +1 ) given the observation x t +1 from the next time step.
Revised: We want this predicted distribution to be consistent with the recognition distribution p rec ( s t +1 | x t +1 ) given the observation x t +1 from the next time step.

Original: To achieve this we introduce an agreement objective, for example the KL divergence,
Revised: To achieve this we introduce an agreement objective, for example the KL divergence,

Original: [Equation]
Revised: [Equation]

Original: Reward Consistency For eﬀective planning, we want the solipsistic representation to be useful for reward prediction p rew ( r t | s t ) .
Revised: Reward Consistency For eﬀective planning, we want the solipsistic representation to be useful for reward prediction p rew ( r t | s t ) .

Original: Given a solipsistic state s t and observed reward r t we use maximum likelihood to learn the reward model p rew ( r t | s t ) .
Revised: Given a solipsistic state s t and observed reward r t we use maximum likelihood to learn the reward model p rew ( r t | s t ) .

Original: This is equivalent to min- imizing the KL divergence KL (˜ p ( r t ) || p rew ( r t | s t )) between the empirical reward distribution ˜ p ( r t ) that places all mass in the observed reward r t and the model p rew ( r t | s t ) .
Revised: This is equivalent to min- imizing the KL divergence KL (˜ p ( r t ) || p rew ( r t | s t )) between the empirical reward distribution ˜ p ( r t ) that places all mass in the observed reward r t and the model p rew ( r t | s t ) .

Original: Solipsistic Contrast Finally, the intention of the recognition distribution p rec ( s t | x t ) is to ﬁlter out redundant information when producing the solipsistic representation.
Revised: Solipsistic Contrast Finally, the intention of the recognition distribution p rec ( s t | x t ) is to ﬁlter out redundant information when producing the solipsistic representation.

Original: To achieve this, we need it to learn to distinguish if the state information is useful or not for both reward prediction and transition dynamics.
Revised: To achieve this, we need it to learn to distinguish if the state information is useful or not for both reward prediction and transition dynamics.

Original: For example in the MNIST game, the recognition distribution may just focus on the image backgrounds which are stationary over time and the solipsistic transition function could just learn the identity mapping.
Revised: For example in the MNIST game, the recognition distribution may just focus on the image backgrounds which are stationary over time and the solipsistic transition function could just learn the identity mapping.

Original: Although this forms a consistent solipsistic representation, it would be useless for planning.
Revised: Although this forms a consistent solipsistic representation, it would be useless for planning.

Original: Furthermore, having a reward objective is not suﬃcient to avoid this behavior.
Revised: Furthermore, having a reward objective is not suﬃcient to avoid this behavior.

Original: Features of the state which inform reward prediction may not be the same as those which inform a useful dynamics model.
Revised: Features of the state which inform reward prediction may not be the same as those which inform a useful dynamics model.

Original: For example in the MNIST game, the recognition distribution may just keep the features which are relevant to distinguish if an image is 9 or not for predicting the reward, but ignore other useful information about the system dynamics.
Revised: For example in the MNIST game, the recognition distribution may just keep the features which are relevant to distinguish if an image is 9 or not for predicting the reward, but ignore other useful information about the system dynamics.

Original: One solution is to additionally encourage the agent’s solipsistic dynamics to be inconsistent with trajectories which are not observed in reality.
Revised: One solution is to additionally encourage the agent’s solipsistic dynamics to be inconsistent with trajectories which are not observed in reality.

Original: In other words, we want an objective that forces the predictive distribution p pred ( s
Revised: In other words, we want an objective that forces the predictive distribution p pred ( s

Original: [Equation]
Revised: [Equation]

Original: 1 | x , a t ) to be diﬀerent from p rec ( s i | x i ) for i (cid:54) =
Revised: 1 | x , a t ) to be diﬀerent from p rec ( s i | x i ) for i (cid:54) =

Original: t + 1 3 .
Revised: t + 1 3 .

Original: We thus maximize the expected KL divergence (for convenience dropping henceforth the “pred”, “rec”, “rew” and “tran” subscripts)
Revised: We thus maximize the expected KL divergence (for convenience dropping henceforth the “pred”, “rec”, “rew” and “tran” subscripts)

Original: [Equation]
Revised: [Equation]

Original: This solipsistic contrast term is similar to a contrastive loss used for representation learning or self-supervised learning (Chopra et al., 2005;
Revised: This solipsistic contrast term is similar to a contrastive loss used for representation learning or self-supervised learning (Chopra et al., 2005;

Original: Hadsell et al., 2006; Chen et al., 2020).
Revised: Hadsell et al., 2006; Chen et al., 2020).

Original: Since the KL divergence is unbounded, we use a positive constant cap m .
Revised: Since the KL divergence is unbounded, we use a positive constant cap m .

Original: Overall Objective: For a trajectory of length T the overall objective is to minimize,
Revised: Overall Objective: For a trajectory of length T the overall objective is to minimize,

Original: [Equation]
Revised: [Equation]

Original: with respect to the parameters of the recognition distribution p ( s t | x t ) , transition distribution p ( s t +1 | s t , a t ) and reward distribution p ( r t | s t ) ; λ s , λ r and m are user chosen hyper-parameters.
Revised: with respect to the parameters of the recognition distribution p ( s t | x t ) , transition distribution p ( s t +1 | s t , a t ) and reward distribution p ( r t | s t ) ; λ s , λ r and m are user chosen hyper-parameters.

Original: In our experience, the results are not particularly sensitive to the choice of these parameters, see section 3.
Revised: In our experience, the results are not particularly sensitive to the choice of these parameters, see section 3.

Original: For tasks that require accurate long-term planning, prediction errors using the simple solipsistic Markov model may accumulate during roll-out 4 .
Revised: For tasks that require accurate long-term planning, prediction errors using the simple solipsistic Markov model may accumulate during roll-out 4 .

Original: In such cases we consider a Solipsistic Memory model(SMM) whose transition depends on all past solipsistic states and actions – see Figure 3b.
Revised: In such cases we consider a Solipsistic Memory model(SMM) whose transition depends on all past solipsistic states and actions – see Figure 3b.

Original: Within our overall objective, we replace p ( s t +1 | x t , a t ) with p ( s t +1 | x t , h t ) where h t = { s 1: t , a 1: t } and use a recurrent neural network (Hochreiter and Schmidhuber, 1997; Chung et al., 2014) to learn the transition dynamics of the solipsistic model.
Revised: Within our overall objective, we replace p ( s t +1 | x t , a t ) with p ( s t +1 | x t , h t ) where h t = { s 1: t , a 1: t } and use a recurrent neural network (Hochreiter and Schmidhuber, 1997; Chung et al., 2014) to learn the transition dynamics of the solipsistic model.

Original: 2. Acting in the Environment and Planning
Revised: 2. Acting in the Environment and Planning

Original: For a given solipsistic model, we take a sequence of actions using a re-planning procedure: we observe x 1 and determine the ﬁrst solipsistic state distribution using the recognition process p ( s 1 | x 1 ) .
Revised: For a given solipsistic model, we take a sequence of actions using a re-planning procedure: we observe x 1 and determine the ﬁrst solipsistic state distribution using the recognition process p ( s 1 | x 1 ) .

Original: We then determine a ∗ 1 using a planning procedure (e.g. dynamic programming, or sampling from a trained parameterized policy, as we discuss below) and take this action in the environment, observing the resulting x 2 .
Revised: We then determine a ∗ 1 using a planning procedure (e.g. dynamic programming, or sampling from a trained parameterized policy, as we discuss below) and take this action in the environment, observing the resulting x 2 .

Original: We repeat this recognition and planning procedure until time T , observing x t at each step and then planning the best next action.
Revised: We repeat this recognition and planning procedure until time T , observing x t at each step and then planning the best next action.

Original: This process of re-planning at every time step helps prevent the accumulation of prediction errors from our model and is eﬃcient since our solipsistic representation is low dimensional.
Revised: This process of re-planning at every time step helps prevent the accumulation of prediction errors from our model and is eﬃcient since our solipsistic representation is low dimensional.

Original: How planning and acting in the environment is folded into the overall RL process of model and policy training is described in detail in algorithm 1,
Revised: How planning and acting in the environment is folded into the overall RL process of model and policy training is described in detail in algorithm 1,

Original: appendix A.
Revised: appendix A.

Original: 2.1. Planning in the Solipsistic Markov Model
Revised: 2.1. Planning in the Solipsistic Markov Model

Original: Given the observation x 1 , we would like to predict the expected rewards we would obtain by taking a sequence of subsequent actions a 1: T − 1 .
Revised: Given the observation x 1 , we would like to predict the expected rewards we would obtain by taking a sequence of subsequent actions a 1: T − 1 .

Original: The recognition distribution enables us to determine the distribution for the ﬁrst solipsistic state p ( s 1 | x 1 ) .
Revised: The recognition distribution enables us to determine the distribution for the ﬁrst solipsistic state p ( s 1 | x 1 ) .

Original: Then, given x 1 and a 1: T − 1 , the state-action trajectory can be described by the distribution
Revised: Then, given x 1 and a 1: T − 1 , the state-action trajectory can be described by the distribution

Original: [Equation]
Revised: [Equation]

Original: For planning, the goal is to maximize the cumulative reward by choosing a sequence of actions.
Revised: For planning, the goal is to maximize the cumulative reward by choosing a sequence of actions.

Original: For a discrete solipsistic state space, the objective is
Revised: For a discrete solipsistic state space, the objective is

Original: [Equation]
Revised: [Equation]

Original: where we use the reward function deﬁned as R ( s t )
Revised: where we use the reward function deﬁned as R ( s t )

Original: ≡ E
Revised: ≡ E

Original: [Equation]
Revised: [Equation]

Original: t | s t ) [ r t ] .
Revised: t | s t ) [ r t ] .

Original: If we interpret the solipsistic model as a Markov Decision Process (MDP) 5 then the optimal action sequence has value
Revised: If we interpret the solipsistic model as a Markov Decision Process (MDP) 5 then the optimal action sequence has value

Original: [Equation]
Revised: [Equation]

Original: This is readily solved by dynamic programming, which we describe below.
Revised: This is readily solved by dynamic programming, which we describe below.

Original: Value Estimation We let V ( s T ) =
Revised: Value Estimation We let V ( s T ) =

Original: R ( s T ) .
Revised: R ( s T ) .

Original: For t = T − 1 , . . .
Revised: For t = T − 1 , . . .

Original: , 2 and each state of s t we calculate
Revised: , 2 and each state of s t we calculate

Original: [Equation]
Revised: [Equation]

Original: The ﬁrst optimal action can be computed using the value function
Revised: The ﬁrst optimal action can be computed using the value function

Original: [Equation]
Revised: [Equation]

Original: We then take the ﬁrst optimal action a ∗ 1 in the real environment to get observation x 2 , and do re-planning based on the new solipsistic distribution p ( s 2 | x 2 ) .
Revised: We then take the ﬁrst optimal action a ∗ 1 in the real environment to get observation x 2 , and do re-planning based on the new solipsistic distribution p ( s 2 | x 2 ) .

Original: The general procedure is to repeat: (1) take the action a ∗ t in the environment and get the new observation x t +1 , (2) use the recognition function to compute p ( s t +1 | x t +1 ) , (3) compute the next action a ∗ t +1 using
Revised: The general procedure is to repeat: (1) take the action a ∗ t in the environment and get the new observation x t +1 , (2) use the recognition function to compute p ( s t +1 | x t +1 ) , (3) compute the next action a ∗ t +1 using

Original: [Equation]
Revised: [Equation]

Original: In situations where a full action sequence is desired before interacting with the environment, see appendix C. By using this re-planning scheme, we can sequentially decide the optimal action sequence under our model, which we demonstrate in our MNIST game in section 3.1.
Revised: In situations where a full action sequence is desired before interacting with the environment, see appendix C. By using this re-planning scheme, we can sequentially decide the optimal action sequence under our model, which we demonstrate in our MNIST game in section 3.1.

Original: For continuous solipsistic states and non-linear transition dynamics, exact dynamic programming is usually not available.
Revised: For continuous solipsistic states and non-linear transition dynamics, exact dynamic programming is usually not available.

Original: See (Bertsekas, 1995) for alternative approximate dynamic programming techniques.
Revised: See (Bertsekas, 1995) for alternative approximate dynamic programming techniques.

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

Original: 2.1. Planning in the Solipsistic Memory Model
Revised: 2.1. Planning in the Solipsistic Memory Model

Original: In the SMM dynamic programming is problematic and we instead learn a policy p w ( a t | s t ) , parameterized by a neural network.
Revised: In the SMM dynamic programming is problematic and we instead learn a policy p w ( a t | s t ) , parameterized by a neural network.

Original: Here we assume that the state s t contains suﬃcient information to determine the best action; however, in order to accurately track long term behavior, we need to use the SMM dynamics to track the state.
Revised: Here we assume that the state s t contains suﬃcient information to determine the best action; however, in order to accurately track long term behavior, we need to use the SMM dynamics to track the state.

Original: This is typically the case in pixel-based planning in which the best action is readily determinable from the current state; however, keeping track of the long-term consequences of a sequence of actions requires using a history of states (due to discretization errors in the image rendering).
Revised: This is typically the case in pixel-based planning in which the best action is readily determinable from the current state; however, keeping track of the long-term consequences of a sequence of actions requires using a history of states (due to discretization errors in the image rendering).

Original: The resulting objective to maximize is
Revised: The resulting objective to maximize is

Original: [Equation]
Revised: [Equation]

Original: T da
Revised: T da

Original: [Equation]
Revised: [Equation]

Original: T − 1 .
Revised: T − 1 .

Original: (11) where, for discrete actions, the integral over a 1: T − 1 is replaced by summation.
Revised: (11) where, for discrete actions, the integral over a 1: T − 1 is replaced by summation.

Original: Previous work has demonstrated that Variational Optimization (VO) style algorithms (Staines and Barber, 2012; Salimans et al., 2017) have an advantage over policy gradients for environments with long time horizons.
Revised: Previous work has demonstrated that Variational Optimization (VO) style algorithms (Staines and Barber, 2012; Salimans et al., 2017) have an advantage over policy gradients for environments with long time horizons.

Original: We therefore use a standard VO algorithm to learn the policy parameters w , see appendix D.
Revised: We therefore use a standard VO algorithm to learn the policy parameters w , see appendix D.

Original:  Experiments
Revised:  Experiments

Original: Rather than showing state-of-the-art across a range of RL challenges, the goal of the experiments is to conﬁrm our hypothesis that model-based RL can be achieved without requiring a generative model of the observations x t .
Revised: Rather than showing state-of-the-art across a range of RL challenges, the goal of the experiments is to conﬁrm our hypothesis that model-based RL can be achieved without requiring a generative model of the observations x t .

Original: We discuss the simple MNIST game and
Revised: We discuss the simple MNIST game and

Original: RL from pixels using continuous s t for a version of the Cartpole benchmark from OpenAI gym
Revised: RL from pixels using continuous s t for a version of the Cartpole benchmark from OpenAI gym

Original: (Brockman et al., 2016).
Revised: (Brockman et al., 2016).

Original: Full details of the models and training are given in appendix E.
Revised: Full details of the models and training are given in appendix E.

Original: 3. MNIST Game
Revised: 3. MNIST Game

Original: For the solipsistic state we assume we know the true number of states s t ∈ { 1 , . . .
Revised: For the solipsistic state we assume we know the true number of states s t ∈ { 1 , . . .

Original: The recognition distribution p ( s t | x t ) is parameterized by a convolutional network.
Revised: The recognition distribution p ( s t | x t ) is parameterized by a convolutional network.

Original: We param- eterize the transition p ( s t +1 | s t , a t ) with two normalized 10 × 10 matrices p ( s t +1 | s t , a t = 0) and p ( s t +1 | s t , a t = 1) .
Revised: We param- eterize the transition p ( s t +1 | s t , a t ) with two normalized 10 × 10 matrices p ( s t +1 | s t , a t = 0) and p ( s t +1 | s t , a t = 1) .

Original: The reward distribution p ( r t | s t ) is parameterized by a normalized 10 × 2 matrix.
Revised: The reward distribution p ( r t | s t ) is parameterized by a normalized 10 × 2 matrix.

Original: At each episode we select random actions a 1:14 and collect a single trajectory x 1:15 to add to our memory 6 .
Revised: At each episode we select random actions a 1:14 and collect a single trajectory x 1:15 to add to our memory 6 .

Original: Then during model training, we iteratively sample batches of consecutive states from memory with batch size 64 and update the model using equationand ADAM (Kingma and Ba, 2015).
Revised: Then during model training, we iteratively sample batches of consecutive states from memory with batch size 64 and update the model using equationand ADAM (Kingma and Ba, 2015).

Original: This process then repeats at each episode.
Revised: This process then repeats at each episode.

Original: We can see from Figure 5a that the 10 diﬀerent images are assigned to 10 diﬀerent solipsistic states with high probability.
Revised: We can see from Figure 5a that the 10 diﬀerent images are assigned to 10 diﬀerent solipsistic states with high probability.

Original: In Figure 5b, we see successful reward predictions, with the image state 9 having high probability of reward 1.
Revised: In Figure 5b, we see successful reward predictions, with the image state 9 having high probability of reward 1.

Original: At test time, we randomly initialize the image state and set the horizon for planing to T = 10 , since for any given state at t = 1 , the agent can reach the goal state within 10 steps.
Revised: At test time, we randomly initialize the image state and set the horizon for planing to T = 10 , since for any given state at t = 1 , the agent can reach the goal state within 10 steps.

Original: In Figure 5c we plot the value of each state from t = 1 , . . .
Revised: In Figure 5c we plot the value of each state from t = 1 , . . .

Original: The optimal action at time 1 is ‘plus 1’ for each initial state; similarly the subsequent optimal action is always ‘plus 1’, thus correctly solving the problem.
Revised: The optimal action at time 1 is ‘plus 1’ for each initial state; similarly the subsequent optimal action is always ‘plus 1’, thus correctly solving the problem.

Original: 3. Gym Control : Binary Action Cartpole
Revised: 3. Gym Control : Binary Action Cartpole

Original: In Cartpole, a pole is attached by an un-actuated joint to a cart that moves along a frictionless track, which can be controlled by applying a force of +1 or − 1 to the cart at each time step.
Revised: In Cartpole, a pole is attached by an un-actuated joint to a cart that moves along a frictionless track, which can be controlled by applying a force of +1 or − 1 to the cart at each time step.

Original: A reward of 1 is received at every time step that the pole remains upright and the episode terminates when it falls over or the cart moves too far from the center.
Revised: A reward of 1 is received at every time step that the pole remains upright and the episode terminates when it falls over or the cart moves too far from the center.

Original: A maximum horizon of 200 time steps is generally used.
Revised: A maximum horizon of 200 time steps is generally used.

Original: Instead of using the low-dimensional states provided by the OpenAI Gym, we use the rendered image frames as the state observations.
Revised: Instead of using the low-dimensional states provided by the OpenAI Gym, we use the rendered image frames as the state observations.

Original: We ﬁrst gray scale pre-process each video frame and down-sample to produce a 64 × 64 frame f t at time t ; three consecutive frames are then stacked to represent each state observation, x t = { f t , f t − 1 , f t − 2 } for t ≥ 3 .
Revised: We ﬁrst gray scale pre-process each video frame and down-sample to produce a 64 × 64 frame f t at time t ; three consecutive frames are then stacked to represent each state observation, x t = { f t , f t − 1 , f t − 2 } for t ≥ 3 .

Original: This provides higher-order dynamics information like speed, which are hidden in a single frame.
Revised: This provides higher-order dynamics information like speed, which are hidden in a single frame.

Original: For simplicity we assume a stationary initial position (initial pole position is upright) and deﬁne x
Revised: For simplicity we assume a stationary initial position (initial pole position is upright) and deﬁne x

Original: 1 = { f 1 , f 1 , f 1 } and x 2 = { f 2 , f 1 , f 1 } .
Revised: 1 = { f 1 , f 1 , f 1 } and x 2 = { f 2 , f 1 , f 1 } .

Original: We use the SMM since this gives accurate long- term prediction, despite discretization error from the pixel rendering of the true continuous underlying dynamics.
Revised: We use the SMM since this gives accurate long- term prediction, despite discretization error from the pixel rendering of the true continuous underlying dynamics.

Original: We follow algorithm 1 (see appendix A) to interact with the environment and learn the policy and model.
Revised: We follow algorithm 1 (see appendix A) to interact with the environment and learn the policy and model.

Original: We set s t ∈ R 16 and choose p ( s t | x t ) =
Revised: We set s t ∈ R 16 and choose p ( s t | x t ) =

Original: δ ( s t − g θ ( x t )) , p ( s t +1 | h t ) =
Revised: δ ( s t − g θ ( x t )) , p ( s t +1 | h t ) =

Original: δ ( s t +1 − f θ
Revised: δ ( s t +1 − f θ

Original: ( h t )) .
Revised: ( h t )) .

Original: The recognition function g θ and transition function f
Revised: The recognition function g θ and transition function f

Original: θ are a convolutional neural network and recurrent network respectively.
Revised: θ are a convolutional neural network and recurrent network respectively.

Original: The reward distribution p θ ( r t | s t ) is a Bernoulli distribution with the probability parameterized using a small neural network with a sigmoid output.
Revised: The reward distribution p θ ( r t | s t ) is a Bernoulli distribution with the probability parameterized using a small neural network with a sigmoid output.

Original: All model parameters are trained jointly using the ADAM optimizer.
Revised: All model parameters are trained jointly using the ADAM optimizer.

Original: Since the KL divergence between two delta distributions is not formally deﬁned for equation 4, we use the spread
Revised: Since the KL divergence between two delta distributions is not formally deﬁned for equation 4, we use the spread

Original: KL divergence (Zhang et al., 2020b) with ﬁxed Gaussian spread noise that has variance 0.5, resulting in a square loss objective.
Revised: KL divergence (Zhang et al., 2020b) with ﬁxed Gaussian spread noise that has variance 0.5, resulting in a square loss objective.

Original: [Figure]
Revised: [Figure]

Original: The policy, section 2.1.2, p θ ( a t | s t ) is a Bernoulli distribution with the probability parameterized using a small feed-forward network with a sigmoid output and trained using VO - see
Revised: The policy, section 2.1.2, p θ ( a t | s t ) is a Bernoulli distribution with the probability parameterized using a small feed-forward network with a sigmoid output and trained using VO - see

Original: appendix E.
Revised: appendix E.

Original: When evaluating the trained policy in the real environment we take the most likely action at each replanning step.
Revised: When evaluating the trained policy in the real environment we take the most likely action at each replanning step.

Original: We compare to DQN and PPO (Mnih et al., 2015; Schulman et al., 2017) following their standard open source implementations adapted for acting directly in pixel space – see appendix E.3 for details.
Revised: We compare to DQN and PPO (Mnih et al., 2015; Schulman et al., 2017) following their standard open source implementations adapted for acting directly in pixel space – see appendix E.3 for details.

Original: In Figure 6 we report the average reward over 5 diﬀerent runs for all methods, with each run carried out using diﬀerent random parameter seeds, but the same ﬁxed initial position of the environment.
Revised: In Figure 6 we report the average reward over 5 diﬀerent runs for all methods, with each run carried out using diﬀerent random parameter seeds, but the same ﬁxed initial position of the environment.

Original: Compared to PPO and DQN our solipsistic model-based approach is signiﬁcantly more sample eﬃcient, learning to balance the pole for over 150 time steps, after seeing onlytrajectories.
Revised: Compared to PPO and DQN our solipsistic model-based approach is signiﬁcantly more sample eﬃcient, learning to balance the pole for over 150 time steps, after seeing onlytrajectories.

Original: This demonstrates that our solipsistic approach is accurately modeling both the relevant environment dynamics and the reward.
Revised: This demonstrates that our solipsistic approach is accurately modeling both the relevant environment dynamics and the reward.

Original: For this setting, we found PPO and
Revised: For this setting, we found PPO and

Original: DQN struggled to balance the pole for over 50 time steps on average after 2500 sampled trajectories from the environment - see Figure 10 in appendix E.3.
Revised: DQN struggled to balance the pole for over 50 time steps on average after 2500 sampled trajectories from the environment - see Figure 10 in appendix E.3.

Original: 3.2. Learned Representations
Revised: 3.2. Learned Representations

Original: Of interest for a qualitative analysis on the recognition function and learned solipsistic representations are: (1) has the recognition function learned to ﬁlter out redundant information?
Revised: Of interest for a qualitative analysis on the recognition function and learned solipsistic representations are: (1) has the recognition function learned to ﬁlter out redundant information?

Original: (2) are the solipsistic trajectories consistent with their corresponding observation trajectories?
Revised: (2) are the solipsistic trajectories consistent with their corresponding observation trajectories?

Original: To answer (1) we extract ﬁlters from the ﬁrst layer of the CNN recognition function, which act as attention maps over the pixels.
Revised: To answer (1) we extract ﬁlters from the ﬁrst layer of the CNN recognition function, which act as attention maps over the pixels.

Original: In Figure 7 we see that diﬀerent ﬁlters are attending to diﬀerent physical attributes, for example pole position, pole speed and cart position.
Revised: In Figure 7 we see that diﬀerent ﬁlters are attending to diﬀerent physical attributes, for example pole position, pole speed and cart position.

Original: In contrast, no clear physical interpretation was apparent for samples from PPO’s convolutional policy and value networks, which we discuss further in appendix E.4.
Revised: In contrast, no clear physical interpretation was apparent for samples from PPO’s convolutional policy and value networks, which we discuss further in appendix E.4.

Original: In support of (2), in Figure 8d we illustrate that the solipsistic trajectories have smooth transitions and are disentangled in accordance with their corresponding trajectories in pixel space.
Revised: In support of (2), in Figure 8d we illustrate that the solipsistic trajectories have smooth transitions and are disentangled in accordance with their corresponding trajectories in pixel space.

Original: We select 4 ﬁlters’ activations (there are 8 ﬁlters in the ﬁrst layer of recognition network) and use a sigmoid function to create these grey-scale images.
Revised: We select 4 ﬁlters’ activations (there are 8 ﬁlters in the ﬁrst layer of recognition network) and use a sigmoid function to create these grey-scale images.

Original: The ﬁrst ﬁlter appears to encode spatial information of the cart and pole (the activated pixels are consistent with the position of both throughout).
Revised: The ﬁrst ﬁlter appears to encode spatial information of the cart and pole (the activated pixels are consistent with the position of both throughout).

Original: The second and third ﬁlters we believe represent velocity information for the pole (given the activated pixels can be interpreted as providing a ﬁnite diﬀerence type estimate of velocity by encoding the position of the pole in the ﬁrst and last frames), while the fourth can be interpreted as encoding information about the velocity of the cart (given more pixels become activated in the region of the cart the faster it goes).
Revised: The second and third ﬁlters we believe represent velocity information for the pole (given the activated pixels can be interpreted as providing a ﬁnite diﬀerence type estimate of velocity by encoding the position of the pole in the ﬁrst and last frames), while the fourth can be interpreted as encoding information about the velocity of the cart (given more pixels become activated in the region of the cart the faster it goes).

Original:  Related Work
Revised:  Related Work

Original: A traditional approach to reduce the complexity of the state is to apply state-aggregation methods, such as non-parametric dimensionality reduction techniques (Powell, 2007), or hand-coded features, to obtain lower-dimensional state representations – see (Mahadevan, 2009) and (Bertsekas, 2018) for a review, These methods require strong prior knowledge about the environment and are not generally useful in situations with complex state spaces.
Revised: A traditional approach to reduce the complexity of the state is to apply state-aggregation methods, such as non-parametric dimensionality reduction techniques (Powell, 2007), or hand-coded features, to obtain lower-dimensional state representations – see (Mahadevan, 2009) and (Bertsekas, 2018) for a review. These methods require strong prior knowledge about the environment and are not generally useful in situations with complex state spaces.

Original: Further the representations are divorced from the modeling process, which can hinder overall performance when utilized for planning (Kuvayev and Sutton, 1996).
Revised: Further the representations are learned separately from the modeling process, which can hinder overall performance when utilized for planning (Kuvayev and Sutton, 1996).

Revised: In contrast, our solipsistic approach jointly learns the dynamics model and representations.

Original: Recent work (Gelada et al., 2019; Zhang et al., 2020a) also propose to use deep neural networks to learnthe state abstraction without reconstructing the original state.
Revised: Recent work (Gelada et al., 2019; Zhang et al., 2020a) also propose to use deep neural networks to learn state representations without reconstructing the original state.

Original: However, they only show that the representation can beneﬁt the  model-free learning whereas we show how to dolong-term planning using the learned representation and the model dynamics.
Revised: However, they only demonstrate the beneﬁts of this approach in model-free learning whereas we show how to do planning using the learned representation and the model dynamics.

Original: [Figure]
Revised: [Figure]

Original: Contrastive learning is widely used in the ﬁeld of representation learning (Bengio et al., 2013).
Revised: Contrastive learning is widely used in the ﬁeld of representation learning (Bengio et al., 2013).

Original: The aim is to encourage similar datapoints to have similar representations (Chopra et al., 2005;
Revised: The aim is to encourage similar datapoints to have similar representations (Chopra et al., 2005;

Original: Hadsell et al., 2006; Chen et al., 2020).
Revised: Hadsell et al., 2006; Chen et al., 2020).

Original: In RL, this idea has been used to improve the data eﬃciency of model-free algorithms by treating contrastive learning as an auxiliary task (Oord et al., 2018; Srinivas et al., 2020; Jaderberg et al., 2017).
Revised: In RL, this idea has been used to improve the data eﬃciency of model-free algorithms by treating contrastive learning as an auxiliary task (Oord et al., 2018; Srinivas et al., 2020; Jaderberg et al., 2017).

Original: The contrastive objective we use (section 2) instead aims to prevent the recognition function from learning a trivial solution to ensure the solipsistic representation is useful for planning.
Revised: The contrastive objective we use (section 2) instead aims to prevent the recognition function from learning a trivial solution to ensure the solipsistic representation is useful for planning.

Original: Most closely related to our work is Hafner et al.
Revised: Most closely related to our work is Hafner et al.

Original: (2019a) which uses a Recurrent State Space
Revised: (2019a) which uses a Recurrent State Space

Original: Model without a decoder component, akin to our Solipistic Markov Model setup.
Revised: Model without a decoder component, akin to our Solipistic Markov Model setup.

Original: They use the variational information bottleneck (VIB) principle (Tishby et al., 2000; Alemi et al., 2016) to derive the following regularizing term for their learning objective:
Revised: They use the variational information bottleneck (VIB) principle (Tishby et al., 2000; Alemi et al., 2016) to derive the following regularizing term for their learning objective:

Original: [Equation]
Revised: [Equation]

Original: where the summation is over the observations in the current sequence batch.
Revised: where the summation is over the observations in the current sequence batch.

Original: The paper shows that this term keeps s t predictable from the current image, whilst also keeping the latent representations diverse.
Revised: The paper shows that this term keeps s t predictable from the current image, whilst also keeping the latent representations diverse.

Original: Although this work also achieves model based RL in the representation space without reconstructing the original image state, their regularized VIB objective is diﬀerent to our Solipsistic consistent-contrast objective (equation 4).
Revised: Although this work also achieves model based RL in the representation space without reconstructing the original image state, their regularized VIB objective is diﬀerent to our Solipsistic consistent-contrast objective (equation 4).

Original: We leave detailed comparisons in both theory and practice to future work.
Revised: We leave detailed comparisons in both theory and practice to future work.

Original:  Summary
Revised:  Summary

Original: We introduced Solipsistic Reinforcement Learning, a new model-based reinforcement learning framework that learns useful latent representations of the environment for planning and reward prediction, without constructing a generative model of the environment.
Revised: We introduced Solipsistic Reinforcement Learning, a new model-based reinforcement learning framework that learns useful latent representations of the environment for planning and reward prediction, without constructing a generative model of the environment.

Original: Our work is consistent with the recent general trend away from modeling the dynamics of high dimensional spaces and towards learning models that more directly solve the task at hand.
Revised: Our work is consistent with the recent general trend away from modeling the dynamics of high dimensional spaces and towards learning models that more directly solve the task at hand.

Original: Whilst modelbased reinforcement learning is arguably preferable to model-free alternatives, particularly because of data eﬃciency, most previous approaches do not learn in an end-to-end fashion and also require an explicit model of the environment.
Revised: Whilst modelbased reinforcement learning is arguably preferable to model-free alternatives, particularly because of data eﬃciency, most previous approaches do not learn in an end-to-end fashion and also require an explicit model of the environment.

Original: We hope therefore that we have shown that there is scope to solve reinforcement learning problems in a model-based way, but without the downsides of requiring complex models of the environment.
Revised: We hope therefore that we have shown that there is scope to solve reinforcement learning problems in a model-based way, but without the downsides of requiring complex models of the environment.

Original: A Overall Learning Procedure
Revised: A Overall Learning Procedure

Original: Here we provide the general procedure we follow for tackling RL problems with solipsistic reinforcement learning, where we assume the use of the memory model and a parameterized policy.
Revised: Here we provide the general procedure we follow for tackling RL problems with solipsistic reinforcement learning, where we assume the use of the memory model and a parameterized policy.

Original: ) for time step policy do
Revised: ) for time step policy do

Original: Sample a t ∼ p w
Revised: Sample a t ∼ p w

Original: h t ←
Revised: h t ←

Original: Predict R
Revised: Predict R

Original: Sample s t +1 ∼ p
Revised: Sample s t +1 ∼ p

Original: Compute E
Revised: Compute E

Original: ; // See section D
Revised: ; // See section D

Original: Compute µ n +1 with { E ( w j
Revised: Compute µ n +1 with { E ( w j

Original: } Jj using VO update equation 18 and ADAM
Revised: } Jj using VO update equation 18 and ADAM

Original: Update policy parameters w
Revised: Update policy parameters w

Original: + µ n
Revised: + µ n

Original: B Best Trajectory Planning
Revised: B Best Trajectory Planning

Original: In section 2.1, our goal is to maximize the cumulative reward objective
Revised: In section 2.1, our goal is to maximize the cumulative reward objective

Original: [Equation]
Revised: [Equation]

Original: We then assume the state s t is only revealed after we take action a t − 1 and form the problem as a Markov Decision Process in equation 7.
Revised: We then assume the state s t is only revealed after we take action a t − 1 and form the problem as a Markov Decision Process in equation 7.

Original: An alternative objective is to directly optimize equation 13 with respect to the action sequence a 1 , . .
Revised: An alternative objective is to directly optimize equation 13 with respect to the action sequence a 1 , . .

Original: . , a t − 1 – we refer to this as the ‘best trajectory planning’ objective.
Revised: . , a t − 1 – we refer to this as the ‘best trajectory planning’ objective.

Original: Although the optimization problem is no longer tractable, an approximate but eﬃcient Expectation Maximization (EM) style algorithm can be used following the lines of (Furmston and Barber, 2011).
Revised: Although the optimization problem is no longer tractable, an approximate but eﬃcient Expectation Maximization (EM) style algorithm can be used following the lines of (Furmston and Barber, 2011).

Original: When applied to the setting of only a ﬁnal reward, we didn’t ﬁnd any signiﬁcant improvement over MDP planning; however, we leave a more detailed discussion of this approach for future work.
Revised: When applied to the setting of only a ﬁnal reward, we didn’t ﬁnd any signiﬁcant improvement over MDP planning; however, we leave a more detailed discussion of this approach for future work.

Original: C Full Trajectory Planning Using the Model
Revised: C Full Trajectory Planning Using the Model

Original: In section 2.1 we discussed a planning approach using our model where we do re-planning each time we sample an action.
Revised: In section 2.1 we discussed a planning approach using our model where we do re-planning each time we sample an action.

Original: For the situation that the environment needs an ‘immediate’ response from the agent (with no time for re-planning), we may decide a sequence of actions purely based on the model predictions given observation x 1 .
Revised: For the situation that the environment needs an ‘immediate’ response from the agent (with no time for re-planning), we may decide a sequence of actions purely based on the model predictions given observation x 1 .

Original: Subsequently, we deploy that action sequence in the environment.
Revised: Subsequently, we deploy that action sequence in the environment.

Original: We describe the method below.
Revised: We describe the method below.

Original: We ﬁrst compute the value function using equation 8 and decide the ﬁrst optimal action based on
Revised: We ﬁrst compute the value function using equation 8 and decide the ﬁrst optimal action based on

Original: [Equation]
Revised: [Equation]

Original: For t = 2 , . . .
Revised: For t = 2 , . . .

Original: , T − 1 , we sequentially compute
Revised: , T − 1 , we sequentially compute

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: for t ≥ 3 .
Revised: for t ≥ 3 .

Original: The potential drawback of this approach is that the error will accumulate during the long term predictions using the model.
Revised: The potential drawback of this approach is that the error will accumulate during the long term predictions using the model.

Original: We implemented this method for the MNIST game, but found little diﬀerence in the results over re-planning.
Revised: We implemented this method for the MNIST game, but found little diﬀerence in the results over re-planning.

Original: In that setting the model provides very accurate long term planning.
Revised: In that setting the model provides very accurate long term planning.

Original: We leave it to future work to assess the performance of this approach in more complex environments.
Revised: We leave it to future work to assess the performance of this approach in more complex environments.

Original: D Variational Optimization
Revised: D Variational Optimization

Original: With VO we model the parameters of the policy with a Gaussian distribution w ∼ N ( w | µ, σ ) to form a diﬀerentiable upper bound that we can minimise w.r.t µ and σ .
Revised: With VO we model the parameters of the policy with a Gaussian distribution w ∼ N ( w | µ, σ ) to form a diﬀerentiable upper bound that we can minimise w.r.t µ and σ .

Original: [Equation]
Revised: [Equation]

Original: where E ( w ) is our memory model planning objective from equation 11.
Revised: where E ( w ) is our memory model planning objective from equation 11.

Original: After optimising this bound with respect to µ, σ , we take the ﬁnal µ as our value for w .
Revised: After optimising this bound with respect to µ, σ , we take the ﬁnal µ as our value for w .

Original: We compute gradients with the usual log-derivative trick, using J Monte Carlo samples to approximate the resulting expectation as follows.
Revised: We compute gradients with the usual log-derivative trick, using J Monte Carlo samples to approximate the resulting expectation as follows.

Original: In practice, we ﬁnd that learning σ does not improve performance, hence we ﬁx it to a value of 0 .
Revised: In practice, we ﬁnd that learning σ does not improve performance, hence we ﬁx it to a value of 0 .

Original: 2 throughout and minimize w.r.t µ ,
Revised: 2 throughout and minimize w.r.t µ ,

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

Original: where w j refers to the j th Monte Carlo sample of the policy parameters - see section Policy training in algorithm 1 for further details.
Revised: where w j refers to the j th Monte Carlo sample of the policy parameters - see section Policy training in algorithm 1 for further details.

Original: In practice, we also apply ﬁtness shaping Salimans et al.
Revised: In practice, we also apply ﬁtness shaping Salimans et al.

Original: (2017) to E ( w j ) in equation 18 to make the VO update invariant w.r.t.
Revised: (2017) to E ( w j ) in equation 18 to make the VO update invariant w.r.t.

Original: order-preserving transformations of the reward value.
Revised: order-preserving transformations of the reward value.

Original: E Experiment Details E. MNIST Game
Revised: E Experiment Details E. MNIST Game

Original: In this section we discuss the details of the architectures and hyper-parameter settings used for our toy MNIST game experiment from section 3.1.
Revised: In this section we discuss the details of the architectures and hyper-parameter settings used for our toy MNIST game experiment from section 3.1.

Original: Model Architecture
Revised: Model Architecture

Original: The recognition network is composed of a two-layer convolution neural network, followed by a two-layer feed-forward network.
Revised: The recognition network is composed of a two-layer convolution neural network, followed by a two-layer feed-forward network.

Original: The two convolution layers have 10 and 20 ﬁlters respectively with stride 2 and kernel size 5.
Revised: The two convolution layers have 10 and 20 ﬁlters respectively with stride 2 and kernel size 5.

Original: We use max pooling after convolutional layer.
Revised: We use max pooling after convolutional layer.

Original: We use ReLU activation functions for both the convolutional and feed-forward layers.
Revised: We use ReLU activation functions for both the convolutional and feed-forward layers.

Original: The forward network has 100 hidden units in each layer and the output size is equal to the number of solipsistic states (in this case 10).
Revised: The forward network has 100 hidden units in each layer and the output size is equal to the number of solipsistic states (in this case 10).

Original: We use a softmax function to create the probabilities of the categorical distribution for determining the solipsistic state assignments.
Revised: We use a softmax function to create the probabilities of the categorical distribution for determining the solipsistic state assignments.

Original: Model Training For our model objective equation 4 we set the hyper parameters as λ s = 1 , λ r = 2 and m = 5 .
Revised: Model Training For our model objective equation 4 we set the hyper parameters as λ s = 1 , λ r = 2 and m = 5 .

Original: We use ADAM (Kingma and Ba, 2015) as the optimizer with learning rate 10 − 4 to train the model for N model = 3000 iterations with batch size B = 64 .
Revised: We use ADAM (Kingma and Ba, 2015) as the optimizer with learning rate 10 − 4 to train the model for N model = 3000 iterations with batch size B = 64 .

Original: Solipsistic Contrast For computational eﬃciency, we approximate the contrast term in equation 3 using the following Monte Carlo approximation
Revised: Solipsistic Contrast For computational eﬃciency, we approximate the contrast term in equation 3 using the following Monte Carlo approximation

Original: [Equation]
Revised: [Equation]

Original: where s j ∈ M ( s t +1 ) \ s t +1 and M ( s t +1 ) is the mini-batch set that s t
Revised: where s j ∈ M ( s t +1 ) \ s t +1 and M ( s t +1 ) is the mini-batch set that s t

Original: +1 belongs to.
Revised: +1 belongs to.

Original: So the set
Revised: So the set

Original: M ( s t +1 ) \ s t +1 has size B − 1 .
Revised: M ( s t +1 ) \ s t +1 has size B − 1 .

Original: E. Gym Control
Revised: E. Gym Control

Original: In this section we discuss the details of the architectures and hyper-parameter settings used for our Cartpole experiment from section 3.2.
Revised: In this section we discuss the details of the architectures and hyper-parameter settings used for our Cartpole experiment from section 3.2.

Original: For model and policy training we follow algorithm 1.
Revised: For model and policy training we follow algorithm 1.

Original: Model Architecture
Revised: Model Architecture

Original: The recognition network is a four-layer convolution neural network with Batch Norm (Ioﬀe and Szegedy, 2015) and ReLU activation functions followed by a feed-forward layer with output size equal to the size of the solipsistic representation dim ( S )
Revised: The recognition network is a four-layer convolution neural network with Batch Norm (Ioﬀe and Szegedy, 2015) and ReLU activation functions followed by a feed-forward layer with output size equal to the size of the solipsistic representation dim ( S )

Original: = 16 .
Revised: = 16 .

Original: We set the kernel size to 3 for the ﬁrst three convolutional layers and 5 for the last convolutional layer.
Revised: We set the kernel size to 3 for the ﬁrst three convolutional layers and 5 for the last convolutional layer.

Original: We set the channel size to 8 and stride to 1 for the ﬁrst convolutional layer and 16 with stride 2 for the other three convolutional layers.
Revised: We set the channel size to 8 and stride to 1 for the ﬁrst convolutional layer and 16 with stride 2 for the other three convolutional layers.

Original: The policy network is a two layers feed forward neural network with 50 hidden units in each layer, which maps from the solipsistic state to a sigmoid function that parameterizes the probability of a Bernoulli distribution.
Revised: The policy network is a two layers feed forward neural network with 50 hidden units in each layer, which maps from the solipsistic state to a sigmoid function that parameterizes the probability of a Bernoulli distribution.

Original: The RNN we used is a single layer Gated Recurrent Unit (GRU) (Chung et al., 2014) with memory in the ﬁrst time-step initialized as the ﬁrst solipsistic state s 1 .
Revised: The RNN we used is a single layer Gated Recurrent Unit (GRU) (Chung et al., 2014) with memory in the ﬁrst time-step initialized as the ﬁrst solipsistic state s 1 .

Original: In each recurrent step, the GRU cell takes one action as input and outputs the prediction of the solipsistic state for the next time step.
Revised: In each recurrent step, the GRU cell takes one action as input and outputs the prediction of the solipsistic state for the next time step.

Original: Therefore, the size of the hidden memory of GRU is equal to the size of the solipsistic state.
Revised: Therefore, the size of the hidden memory of GRU is equal to the size of the solipsistic state.

Original: Model Training We use ADAM with learning rate 10 − 4 , on batches of size B = 10 of sampled environment trajectories of length T model = 50 , with λ s = 1 , λ r = 2 and m = 5 .
Revised: Model Training We use ADAM with learning rate 10 − 4 , on batches of size B = 10 of sampled environment trajectories of length T model = 50 , with λ s = 1 , λ r = 2 and m = 5 .

Original: Policy Training We use VO (section D), with J = 50 parameter perturbations, under model prediction roll-outs of length T policy = 200 and we train for N policy = 50 iterations, using ADAM with learning rate 10 − 2 .
Revised: Policy Training We use VO (section D), with J = 50 parameter perturbations, under model prediction roll-outs of length T policy = 200 and we train for N policy = 50 iterations, using ADAM with learning rate 10 − 2 .

Original: We found that re-initializing the policy parameters w after each model update helped overall performance.
Revised: We found that re-initializing the policy parameters w after each model update helped overall performance.

Original: Exploration Strategy At each iteration i of algorithm 1 during the Trajectory collection phase we collect C = 5 trajectories and add them to our memory for subsequent model training.
Revised: Exploration Strategy At each iteration i of algorithm 1 during the Trajectory collection phase we collect C = 5 trajectories and add them to our memory for subsequent model training.

Original: We use an exploration strategy (see line 9 of algorithm 1) during trajectory collection as follows: for c = 1 we follow the latest policy p w i ( a t | s t ) with re-planning (section B) to collect a full trajectory.
Revised: We use an exploration strategy (see line 9 of algorithm 1) during trajectory collection as follows: for c = 1 we follow the latest policy p w i ( a t | s t ) with re-planning (section B) to collect a full trajectory.

Original: Then, for c = 2 , . .
Revised: Then, for c = 2 , . .

Original: , 5 , we instead follow an (cid:15) -policy, where we take the action sampled from p w i ( a t | s t ) with probability (cid:15) and we take a random action with probability 1 − (cid:15) , where (cid:15)
Revised: , 5 , we instead follow an (cid:15) -policy, where we take the action sampled from p w i ( a t | s t ) with probability (cid:15) and we take a random action with probability 1 − (cid:15) , where (cid:15)

Original: = 0 .
Revised: = 0 .

Original: Solipsistic Contrast We use x mt to denote the state at time t of the m th trajectory.
Revised: Solipsistic Contrast We use x mt to denote the state at time t of the m th trajectory.

Original: We approximate the contrast term in equation 3 using the following Monte Carlo approximation
Revised: We approximate the contrast term in equation 3 using the following Monte Carlo approximation

Original: [Equation]
Revised: [Equation]

Original: where the n th ( n (cid:54)
Revised: where the n th ( n (cid:54)

Original: = m ) trajectory is sampled from memory.
Revised: = m ) trajectory is sampled from memory.

Original: We found this one-sample Monte
Revised: We found this one-sample Monte

Original: Carlo approximation works well in practice.
Revised: Carlo approximation works well in practice.

Original: E.2. Training Statistics
Revised: E.2. Training Statistics

Original: In Figure 9 we illustrate the training statistics from a single run of algorithm 1.
Revised: In Figure 9 we illustrate the training statistics from a single run of algorithm 1.

Original: This is representative of the typical behavior that appears for each run of this procedure from diﬀerent parameter initializations.
Revised: This is representative of the typical behavior that appears for each run of this procedure from diﬀerent parameter initializations.

Original: We plot the solipsistic consistency (equation 2), solipsistic contrast (equation 3) and reward consistency (section 2) in Figure 9a, Figure 9b and Figure 9c respectively.
Revised: We plot the solipsistic consistency (equation 2), solipsistic contrast (equation 3) and reward consistency (section 2) in Figure 9a, Figure 9b and Figure 9c respectively.

Original: In Figure 9d, we plot the average of the predicted rewards received during each iteration of policy training, within our overall learning procedure.
Revised: In Figure 9d, we plot the average of the predicted rewards received during each iteration of policy training, within our overall learning procedure.

Original: More speciﬁcally, we plot 1 J (cid:80)
Revised: More speciﬁcally, we plot 1 J (cid:80)

Original: Jj =1 E ( w j ) , where the j samples are used during our VO update step (section D)
Revised: Jj =1 E ( w j ) , where the j samples are used during our VO update step (section D)

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: E. Baseline Model-Free Methods E.3. PPO
Revised: E. Baseline Model-Free Methods E.3. PPO

Original: We implemented the PPO method (Schulman et al., 2017) as a baseline to evaluate our methods.
Revised: We implemented the PPO method (Schulman et al., 2017) as a baseline to evaluate our methods.

Original: We run PPO using two diﬀerent clip ratios 0.1 and 0.2.
Revised: We run PPO using two diﬀerent clip ratios 0.1 and 0.2.

Original: The policy and the value networks share similar architecture to our recognition convolutional neural network described in section E.2 with the only diﬀerence being the dimensionality of the output layers.
Revised: The policy and the value networks share similar architecture to our recognition convolutional neural network described in section E.2 with the only diﬀerence being the dimensionality of the output layers.

Original: For the Cartpole experiment, the output size of the policy network is 2 and we use a softmax activation function to parameterize the probabilities of choosing from two actions.
Revised: For the Cartpole experiment, the output size of the policy network is 2 and we use a softmax activation function to parameterize the probabilities of choosing from two actions.

Original: The output of the value network is a linear layer with output dimension 1.
Revised: The output of the value network is a linear layer with output dimension 1.

Original: For each training episode, we sample 5 trajectories from the true environment with a maximum length of 200.
Revised: For each training episode, we sample 5 trajectories from the true environment with a maximum length of 200.

Original: We then train the model for 10 epochs using ADAM where we take the hyper-parameters provided by the OpenAI Baselines implementation Dhariwal et al.
Revised: We then train the model for 10 epochs using ADAM where we take the hyper-parameters provided by the OpenAI Baselines implementation Dhariwal et al.

Original: (2017), where learning is 3 − 4 , γ = 0 .
Revised: (2017), where learning is 3 − 4 , γ = 0 .

Original: 99 , λ = 0 .
Revised: 99 , λ = 0 .

Original: 95 , the weight of the value term is 0 .
Revised: 95 , the weight of the value term is 0 .

Original: 5 and the weight of the entropy term is 0 . 01 .
Revised: 5 and the weight of the entropy term is 0 . 01 .

Original: We report results for two diﬀerent clip ratios of 0.1 and 0.2.
Revised: We report results for two diﬀerent clip ratios of 0.1 and 0.2.

Original: E.3. DQN
Revised: E.3. DQN

Original: The other model free baseline we compared to is DQN (Mnih et al., 2015).
Revised: The other model free baseline we compared to is DQN (Mnih et al., 2015).

Original: Like for the case of PPO, we keep the architecture of the Q-network and the target network similar to our recognition network, except for the ﬁnal output layer.
Revised: Like for the case of PPO, we keep the architecture of the Q-network and the target network similar to our recognition network, except for the ﬁnal output layer.

Original: In our Cartpole experiment, the
Revised: In our Cartpole experiment, the

Original: Q-network outputs the Q-value for the two possible actions given the states.
Revised: Q-network outputs the Q-value for the two possible actions given the states.

Original: The action with the larger Q-value is chosen during control.
Revised: The action with the larger Q-value is chosen during control.

Original: We train DQN for 2500 episodes (as shown in Figure 10).
Revised: We train DQN for 2500 episodes (as shown in Figure 10).

Original: For each episode, we sample one trajectory with the maximum length of 200.
Revised: For each episode, we sample one trajectory with the maximum length of 200.

Original: After a limited grid search of hyper-parmeters, we ﬁnd that the hyper-parameters from the
Revised: After a limited grid search of hyper-parmeters, we ﬁnd that the hyper-parameters from the

Original: PyTorch DQN tutorial (Paszke, 2020; Paszke et al., 2019).
Revised: PyTorch DQN tutorial (Paszke, 2020; Paszke et al., 2019).

Original: E. Activation Maps of PPO
Revised: E. Activation Maps of PPO

Original: In this section, we visualize the activations of the ﬁrst convolution layer of the policy network and the value network for PPO.
Revised: In this section, we visualize the activations of the ﬁrst convolution layer of the policy network and the value network for PPO.

Original: We plot all 8 ﬁlters’ activations and use a sigmoid function to create the following grey-scale images.
Revised: We plot all 8 ﬁlters’ activations and use a sigmoid function to create the following grey-scale images.

Original: Each column represents a ﬁlter.
Revised: Each column represents a ﬁlter.

Original: Each row represents the state of a trajectory as the cart is pushed to the right with monotonically increasing velocity (as in Figure 7).
Revised: Each row represents the state of a trajectory as the cart is pushed to the right with monotonically increasing velocity (as in Figure 7).

Original: In both Figure 11 and Figure 12, none of the activation maps have a change in activation that can clearly be interpreted as relating to the increase in velocity (unlike the case of the solipsistic recognition model illustrated Figure 7); only the positional information of the cart and pole is obviously captured.
Revised: In both Figure 11 and Figure 12, none of the activation maps have a change in activation that can clearly be interpreted as relating to the increase in velocity (unlike the case of the solipsistic recognition model illustrated Figure 7); only the positional information of the cart and pole is obviously captured.