Original: V ARIATIONAL R ECURRENT M ODELS FOR S OLVING P ARTIALLY O BSERVABLE C ONTROL T ASKS
Revised: V ARIATIONAL R ECURRENT M ODELS FOR S OLVING P ARTIALLY O BSERVABLE C ONTROL T ASKS

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: In partially observable (PO) environments, deep reinforcement learning (RL) agents often suffer from unsatisfactory performance, since two problems need to be tackled together: how to extract information from the raw observations to solve the task, and how to improve the policy.
Revised: In partially observable (PO) environments, deep reinforcement learning (RL) agents often suffer from unsatisfactory performance, since two problems need to be tackled together: how to extract information from the raw observations to solve the task, and how to improve the policy.

Original: In this study, we propose an RL algorithm for solving PO tasks.
Revised: In this study, we propose an RL algorithm for solving PO tasks.

Original: Our method comprises two parts: a variational recurrent model (VRM) for modeling the environment and an RL controller that has access to both the environment and the VRM.
Revised: Our method comprises two parts: a variational recurrent model (VRM) for modeling the environment, and an RL controller that has access to both the environment and the VRM.

Original: The proposed algorithm was tested in two types of PO robotic control tasks, those in which either coordinates or velocities were not observable and those that require long-term memorization.
Revised: The proposed algorithm was tested in two types of PO robotic control tasks, those in which either coordinates or velocities were not observable and those that require long-term memorization.

Original: Our experiments show that the proposed algorithm achieved better data efﬁciency and/or learned more optimal policy than other alternative approaches in tasks in which unobserved states cannot be inferred from raw observations in a simple manner 
Revised: Our experiments show that the proposed algorithm achieved better data efﬁciency and/or learned more optimal policy than other alternative approaches in tasks in which unobserved states cannot be inferred from raw observations in a simple manner 1 .

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Model-free deep reinforcement learning (RL) algorithms have been developed to solve difﬁcult control and decision-making tasks by self-exploration (Sutton & Barto, 1998; Mnih et al., 2015;
Revised: Model-free deep reinforcement learning (RL) algorithms have been developed to solve difﬁcult control and decision-making tasks by self-exploration (Sutton & Barto, 1998; Mnih et al., 2015;

Original: Silver et al., 2016).
Revised: Silver et al., 2016).

Original: While various kinds of fully observable environments have been well investigated, recently, partially observable (PO) environments (Hafner et al., 2018;
Revised: While various kinds of fully observable environments have been well investigated, recently, partially observable (PO) environments (Hafner et al., 2018;

Original: Igl et al., 2018; Lee et al., 2019;
Revised: Igl et al., 2018; Lee et al., 2019;

Original: Jaderberg et al., 2019) have commanded greater attention, since real-world applications often need to tackle incomplete information and a non-trivial solution is highly desirable.
Revised: Jaderberg et al., 2019) have commanded greater attention, since real-world applications often need to tackle incomplete information and a non-trivial solution is highly desirable.

Original: There are many types of PO tasks; however, those that can be solved by taking the history of observations into account are more common.
Revised: There are many types of PO tasks; however, those that can be solved by taking the history of observations into account are more common.

Original: These tasks are often encountered in real life, such as videos games that require memorization of previous events (Kapturowski et al., 2018; Jaderberg et al., 2019) and robotic control using real-time images as input (Hafner et al., 2018; Lee et al., 2019).
Revised: These tasks are often encountered in real life, such as videos games that require memorization of previous events (Kapturowski et al., 2018; Jaderberg et al., 2019) and robotic control using real-time images as input (Hafner et al., 2018; Lee et al., 2019).

Original: While humans are good at solving these tasks by extracting crucial information from the past observations, deep RL agents often have difﬁculty acquiring satisfactory policy and achieving good data efﬁciency, compared to those in fully observable tasks (Hafner et al., 2018; Lee et al., 2019).
Revised: While humans are good at solving these tasks by extracting crucial information from the past observations, deep RL agents often have difﬁculty acquiring satisfactory policy and achieving good data efﬁciency, compared to those in fully observable tasks (Hafner et al., 2018; Lee et al., 2019).

Original: For solving such PO tasks, several categories of methods have been proposed.
Revised: For solving such PO tasks, several categories of methods have been proposed.

Original: One simple, straightforward solution is to include a history of raw observations in the current “observation” (McCallum, 1993; Lee et al., 2019).
Revised: One simple, straightforward solution is to include a history of raw observations in the current “observation” (McCallum, 1993; Lee et al., 2019).

Original: Unfortunately, this method can be impractical when decision-making requires a long-term memory because dimension of observation become unacceptably large if a long history is included.
Revised: Unfortunately, this method can be impractical when decision-making requires a long-term memory because dimension of observation become unacceptably large if a long history is included.

Original: Another category is based on model-free RL methods with recurrent neural networks (RNN) as function approximators (Schmidhuber, 1990; 1991; Igl et al., 2018; Kapturowski et al., 2018; Jaderberg et al., 2019), which is usually more tractable to implement.
Revised: Another category is based on model-free RL methods with recurrent neural networks (RNN) as function approximators (Schmidhuber, 1990; 1991; Igl et al., 2018; Kapturowski et al., 2018; Jaderberg et al., 2019), which is usually more tractable to implement.

Original: In this case, RNNs need to tackle two problems simultaneously (Lee et al., 2019): learning representation (encoded by hidden states of the RNN) of the underlying states of the environment from the state-transition data, and learning to maximize returns using the learned representation.
Revised: In this case, RNNs need to tackle two problems simultaneously (Lee et al., 2019): learning representation (encoded by hidden states of the RNN) of the underlying states of the environment from the state-transition data, and learning to maximize returns using the learned representation.

Original: As most RL algorithms use a bootstrapping strategy to learn the expected return and to improve the policy (Sutton & Barto, 1998), it is challenging to train the RNN stably and efﬁciently, since RNNs are relatively more difﬁcult to train (Pascanu et al., 2013) than feedforward neural networks.
Revised: As most RL algorithms use a bootstrapping strategy to learn the expected return and to improve the policy (Sutton & Barto, 1998), it is challenging to train the RNN stably and efﬁciently, since RNNs are relatively more difﬁcult to train (Pascanu et al., 2013) than feedforward neural networks.

Original: The third category considers learning a model of the environment and estimating a belief state , extracted from a sequence of state-transitions (Kaelbling et al., 1998;
Revised: The third category considers learning a model of the environment and estimating a belief state , extracted from a sequence of state-transitions (Kaelbling et al., 1998;

Original: Ha & Schmidhuber, 2018;
Revised: Ha & Schmidhuber, 2018;

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

Original: The belief state is an agent-estimated variable encoding underlying states of the environment that determines state-transitions and rewards.
Revised: The belief state is an agent-estimated variable encoding underlying states of the environment that determines state-transitions and rewards.

Original: Perfectly-estimated belief states can thus be taken as “observations” of an RL agent that contains complete information for solving the task.
Revised: Perfectly-estimated belief states can thus be taken as “observations” of an RL agent that contains complete information for solving the task.

Original: Therefore, solving a PO task is segregated into a representation learning problem and a fully observable RL problem.
Revised: Therefore, solving a PO task is segregated into a representation learning problem and a fully observable RL problem.

Original: Since fully observable RL problems have been well explored by the RL community, the critical challenge here is how to estimate the belief state.
Revised: Since fully observable RL problems have been well explored by the RL community, the critical challenge here is how to estimate the belief state.

Original: In this study, we developed a variational recurrent model (VRM) that models sequential observations and rewards using a latent stochastic variable.
Revised: In this study, we developed a variational recurrent model (VRM) that models sequential observations and rewards using a latent stochastic variable.

Original: The VRM is an extension of the variational recurrent neural network (VRNN) model (Chung et al., 2015) that takes actions into account.
Revised: The VRM is an extension of the variational recurrent neural network (VRNN) model (Chung et al., 2015) that takes actions into account.

Original: Our approach falls into the third category by taking the internal states of the VRM together with raw observations as the belief state.
Revised: Our approach falls into the third category by taking the internal states of the VRM together with raw observations as the belief state.

Original: We then propose an algorithm to solve PO tasks by training the VRM and a feed-forward RL controller network, respectively.
Revised: We then propose an algorithm to solve PO tasks by training the VRM and a feed-forward RL controller network, respectively.

Original: The algorithm can be applied in an end-to-end manner, without ﬁne tuning of a hyperparameters.
Revised: The algorithm can be applied in an end-to-end manner, without ﬁne tuning of a hyperparameters.

Original: We then experimentally evaluated the proposed algorithm in various PO versions of robotic control tasks.
Revised: We then experimentally evaluated the proposed algorithm in various PO versions of robotic control tasks.

Original: The agents showed substantial policy improvement in all tasks, and in some tasks the algorithm performed essentially as in fully observable cases.
Revised: The agents showed substantial policy improvement in all tasks, and in some tasks the algorithm performed essentially as in fully observable cases.

Original: In particular, our algorithm demonstrates greater performance compared to alternative approaches in environments where only velocity information is observable or in which long-term memorization is needed.
Revised: In particular, our algorithm demonstrates greater performance compared to alternative approaches in environments where only velocity information is observable or in which long-term memorization is needed.

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

Original: Typical model-based RL approaches utilize learned models for dreaming, i.e. generating statetransition data for training the agent (Deisenroth & Rasmussen, 2011; Ha & Schmidhuber, 2018;
Revised: Typical model-based RL approaches utilize learned models for dreaming, i.e. generating statetransition data for training the agent (Deisenroth & Rasmussen, 2011; Ha & Schmidhuber, 2018;

Original: Kaiser et al., 2019) or for planning of future state-transitions (Watter et al., 2015; Hafner et al., 2018; Ke et al., 2019).
Revised: Kaiser et al., 2019) or for planning of future state-transitions (Watter et al., 2015; Hafner et al., 2018; Ke et al., 2019).

Original: This usually requires a well-designed and ﬁnely tuned model so that its predictions are accurate and robust.
Revised: This usually requires a well-designed and ﬁnely tuned model so that its predictions are accurate and robust.

Original: In our case, we do not use VRMs for dreaming and planning, but for auto-encoding state-transitions.
Revised: In our case, we do not use VRMs for dreaming and planning, but for auto-encoding state-transitions.

Original: Actually, PO tasks can be solved without requiring VRMs to predict accurately (see Appendix E).
Revised: Actually, PO tasks can be solved without requiring VRMs to predict accurately (see Appendix E).

Original: This distinguishes our algorithm from typical model-based RL methods.
Revised: This distinguishes our algorithm from typical model-based RL methods.

Original: The work our method most closely resembles is known as stochastic latent actor-critic (SLAC, Lee et al.
Revised: The work our method most closely resembles is known as stochastic latent actor-critic (SLAC, Lee et al.

Original: (2019)), in which a latent variable model is trained and uses the latent state as the belief state for the critic.
Revised: (2019)), in which a latent variable model is trained and uses the latent state as the belief state for the critic.

Original: SLAC showed promising results using pixels-based robotic control tasks, in which velocity information needs to be inferred from third-person images of the robot.
Revised: SLAC showed promising results using pixels-based robotic control tasks, in which velocity information needs to be inferred from third-person images of the robot.

Original: Here we consider more general PO environments in which the reward may depend on a long history of inputs, e.g., in a snooker game one has to remember which ball was potted previously.
Revised: Here we consider more general PO environments in which the reward may depend on a long history of inputs, e.g., in a snooker game one has to remember which ball was potted previously.

Original: The actor network of SLAC did not take advantage of the latent variable, but instead used some steps of raw observations as input, which creates problems in achieving long-term memorization of reward-related state-transitions.
Revised: The actor network of SLAC did not take advantage of the latent variable, but instead used some steps of raw observations as input, which creates problems in achieving long-term memorization of reward-related state-transitions.

Original: Furthermore, SLAC did not include raw observations in the input of the critic, which may complicate training the critic before the model converges.
Revised: Furthermore, SLAC did not include raw observations in the input of the critic, which may complicate training the critic before the model converges.

Original:  B ACKGROUND 3. P ARTIALLY O BSERVABLE M ARKOV D ECISION P ROCESSES
Revised:  B ACKGROUND 3. P ARTIALLY O BSERVABLE M ARKOV D ECISION P ROCESSES

Original: The scope of problems we study can be formulated into a framework known as partially observable
Revised: The scope of problems we study can be formulated into a framework known as partially observable

Original: Markov decision processes (POMDP) (Kaelbling et al., 1998).
Revised: Markov decision processes (POMDP) (Kaelbling et al., 1998).

Original: POMDPs are used to describe decision or control problems in which a part of underlying states of the environment, which determine state-transitions and rewards, cannot be directly observed by an agent.
Revised: POMDPs are used to describe decision or control problems in which a part of underlying states of the environment, which determine state-transitions and rewards, cannot be directly observed by an agent.

Original: A POMDP is usually deﬁned as a 7-tuple ( S , A , T, R, X , O, γ ) , in which S is a set of states, A is a set of actions, and T : S ×
Revised: A POMDP is usually deﬁned as a 7-tuple ( S , A , T, R, X , O, γ ) , in which S is a set of states, A is a set of actions, and T : S ×

Original: A → p ( S ) is the state-transition probability function that determines the distribution of the next state given current state and action.
Revised: A → p ( S ) is the state-transition probability function that determines the distribution of the next state given current state and action.

Original: The reward function R : S ×
Revised: The reward function R : S ×

Original: A → R decides the reward during a state-transition, which can also be probabilistic.
Revised: A → R decides the reward during a state-transition, which can also be probabilistic.

Original: Moreover, X is a set of observations, and observations are determined by the observation probability function O : S ×
Revised: Moreover, X is a set of observations, and observations are determined by the observation probability function O : S ×

Original: A → p ( X ) .
Revised: A → p ( X ) .

Original: By deﬁning a POMDP, the goal is to maximize expected discounted future rewards (cid:80) t [Table] t [Table]
Revised: By deﬁning a POMDP, the goal is to maximize expected discounted future rewards (cid:80) t γ t r t by learning a good strategy to select actions (policy function).

Original: t by learning a good strategy to select actions (policy function).

Original: Our algorithm was designed for general POMDP problems by learning the representation of underlying states s t ∈ S via modeling observation-transitions and reward functions.
Revised: Our algorithm was designed for general POMDP problems by learning the representation of underlying states s t ∈ S via modeling observation-transitions and reward functions.

Original: However, it is expected to work in PO tasks in which s t or p ( s t ) can be (at least partially) estimated from the history of observations x 1: t .
Revised: However, it is expected to work in PO tasks in which s t or p ( s t ) can be (at least partially) estimated from the history of observations x 1: t .

Original: 3. V ARIATIONAL R ECURRENT N EURAL N ETWORKS
Revised: 3. V ARIATIONAL R ECURRENT N EURAL N ETWORKS

Original: To model general state-transitions that can be stochastic and complicated [Equation] we employ a modiﬁed version of the VRNN (Chung et al., 2015).
Revised: To model general state-transitions that can be stochastic and complicated, we employ a modiﬁed version of the VRNN (Chung et al., 2015).

Original: The VRNN was developed as a recurrent version of the variational auto-encoder (VAE, Kingma & Welling (2013)), composed of a variational generation model and a variational inference model.
Revised: The VRNN was developed as a recurrent version of the variational auto-encoder (VAE, Kingma & Welling (2013)), composed of a variational generation model and a variational inference model.

Original: It is a recurrent latent variable model that can learn to encode and predict complicated sequential observations x t with a stochastic latent variable z t .
Revised: It is a recurrent latent variable model that can learn to encode and predict complicated sequential observations x t with a stochastic latent variable z t .

Original: The generation model predicts future observations given the its internal states,
Revised: The generation model predicts future observations given the its internal states,

Original: [Equation]
Revised: [Equation]

Original: where f s are parameterized mappings [Equation] such as feed-forward neural networks
Revised: where f s are parameterized mappings, such as feed-forward neural networks

Original: [Equation]
Revised: [Equation]

Original: and d t is the state variable of the RNN [Equation] which is recurrently updated by
Revised: and d t is the state variable of the RNN, which is recurrently updated by

Original: [Equation]
Revised: [Equation]

Original: The inference model approximates the latent variable z t given x t and d t .
Revised: The inference model approximates the latent variable z t given x t and d t .

Original: [Equation]
Revised: [Equation]

Original: For sequential data that contain T time steps, learning is conducted by maximizing the evidence lower bound ELBO , like that in a VEA (Kingma & Welling, 2013), where
Revised: For sequential data that contain T time steps, learning is conducted by maximizing the evidence lower bound ELBO , like that in a VEA (Kingma & Welling, 2013) [Equation] where

Original: [Equation]
Revised: [Equation]

Original: where p and q are parameterized PDFs of z t by the generative model and the inference model, respectively.
Revised: where p and q are parameterized PDFs of z t by the generative model and the inference model, respectively.

Original: In a POMDP, a VRNN can be used to model the environment and to represent underlying states in its state variable d t .
Revised: In a POMDP, a VRNN can be used to model the environment and to represent underlying states in its state variable d t .

Original: Thus an RL agent can beneﬁt from a well-learned VRNN model since d t provides additional information about the environment beyond the current raw observation x t .
Revised: Thus an RL agent can beneﬁt from a well-learned VRNN model since d t provides additional information about the environment beyond the current raw observation x t .

Original: 3. S OFT A CTOR C RITIC
Revised: 3. S OFT A CTOR C RITIC

Original: Soft actor-critic (SAC) is a state-of-the-art model-free RL that uses experience replay for dynamic programming, which been tested on various robotic control tasks and that shows promising performance (Haarnoja et al., 2018a;b).
Revised: Soft actor-critic (SAC) is a state-of-the-art model-free RL that uses experience replay for dynamic programming, which been tested on various robotic control tasks and that shows promising performance (Haarnoja et al., 2018a;b).

Original: A SAC agent learns to maximize reinforcement returns as well as entropy of its policy, so as to obtain more rewards while keeping actions sufﬁciently stochastic. 
Revised: A SAC agent learns to maximize reinforcement returns as well as entropy of its policy, so as to obtain more rewards while keeping actions sufﬁciently stochastic. [Figure]

Original: A typical SAC implementation can be described as follows.
Revised: A typical SAC implementation can be described as follows.

Original: The state value function V ( s ) , the state-action value function Q ( s , a ) and the policy function π ( a | s ) are parameterized by neural networks, indicated by ψ, λ, η , respectively.
Revised: The state value function V ( s ) , the state-action value function Q ( s , a ) and the policy function π ( a | s ) are parameterized by neural networks, indicated by ψ, λ, η , respectively.

Original: Also, an entropy coefﬁcient factor (also known as the temperature parameter), denoted by α , is learned to control the degree of stochasticity of the policy.
Revised: Also, an entropy coefﬁcient factor (also known as the temperature parameter), denoted by α , is learned to control the degree of stochasticity of the policy.

Original: The parameters are learned by simultaneously minimizing the following loss functions.
Revised: The parameters are learned by simultaneously minimizing the following loss functions.

Original: [Equation]
Revised: [Equation]

Original: [Figure]

Original: where B is the replay buffer from which s t is sampled, and H tar is the target entropy.
Revised: where B is the replay buffer from which s t is sampled, and H tar is the target entropy.

Original: To compute the gradient of J π ( η ) (Equation. 7), the reparameterization trick (Kingma & Welling, 2013) is used on action, indicated by a η ( s t ) .
Revised: To compute the gradient of J π ( η ) (Equation. 7), the reparameterization trick (Kingma & Welling, 2013) is used on action, indicated by a η ( s t ) .

Original: Reparameterization of action is not required in minimizing J ( α ) (Equation.
Revised: Reparameterization of action is not required in minimizing J ( α ) (Equation.

Original: 8) since log π η ( a | s t ) does not depends on α .
Revised: 8) since log π η ( a | s t ) does not depends on α .

Original: SAC was originally developed for fully observable environments; thus, the raw observation at the current step x t was used as network input.
Revised: SAC was originally developed for fully observable environments; thus, the raw observation at the current step x t was used as network input.

Original: In this work, we apply SAC in PO tasks by including the state variable d t of the VRNN in the input of function approximators of both the actor and the critic.
Revised: In this work, we apply SAC in PO tasks by including the state variable d t of the VRNN in the input of function approximators of both the actor and the critic.

Original:  M ETHODS 4. V ARIATIONAL R ECURRENT S TATE -T RANSITION M ODELS
Revised:  M ETHODS 4. V ARIATIONAL R ECURRENT S TATE -T RANSITION M ODELS

Original: An overall diagram of the proposed algorithm is summarized in Fig.
Revised: An overall diagram of the proposed algorithm is summarized in Fig.

Original: 1(a), while a more detailed computational graph is plotted in Fig.
Revised: 1(a), while a more detailed computational graph is plotted in Fig.

Original: We extend the original VRNN model (Chung et al., 2015) to the proposed VRM model by adding action feedback, i.e., actions taken by the agent are used in the inference model and the generative model.
Revised: We extend the original VRNN model (Chung et al., 2015) to the proposed VRM model by adding action feedback, i.e., actions taken by the agent are used in the inference model and the generative model.

Original: Also, since we are modeling state-transition and reward functions, we include the reward r t − 1 in the current raw observation x t for convenience.
Revised: Also, since we are modeling state-transition and reward functions, we include the reward r t − 1 in the current raw observation x t for convenience.

Original: Thus, we have the inference model (Fig. 1(c)), denoted by φ , as
Revised: Thus, we have the inference model (Fig. 1(c)), denoted by φ , as

Original: [Equation]

Original: , where (cid:2) µ

Original: [Equation]
Revised: [Equation]

Original: The generative model (Fig. 1(b)), denoted by θ here, is
Revised: The generative model (Fig. 1(b)), denoted by θ here, is

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: For building recurrent connections, the choice of RNN types is not limited.
Revised: For building recurrent connections, the choice of RNN types is not limited.

Original: In our study, the longshort term memory (LSTM) (Hochreiter & Schmidhuber, 1997) is used since it works well in general cases.
Revised: In our study, the longshort term memory (LSTM) (Hochreiter & Schmidhuber, 1997) is used since it works well in general cases.

Original: So we have d t = LSTM
Revised: So we have d t = LSTM

Original: ( d t − 1 ; z t , x t ) .
Revised: ( d t − 1 ; z t , x t ) .

Original: As in training a VRNN, the VRM is trained by maximizing an evidence lower bound (Fig.
Revised: [Equation]

Original: 1(c))
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: In practice, the ﬁrst term E q φ [log p θ ( x t | z 1: t , x 1: t − 1 )] can be obtained by unrolling the RNN using the inference model (Fig. 1(c)) with sampled sequences of x t . Since q φ and p θ are parameterized Gaussian distributions, the KL-divergence term can be analytically expressed as
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: Deterministic node
Revised: Deterministic node

Original: Stochastic node
Revised: Stochastic node

Original: Generative model
Revised: Generative model

Original: Inference model
Revised: Inference model

Original: RL controllernetwork
Revised: RL controllernetwork

Original: Interactingwith the environment
Revised: Interactingwith the environment

Original: Error back- propagation
Revised: Error back- propagation

Original: RL controller
Revised: RL controller

Original: Excecution 
Revised: Excecution phase

Original: [Figure]

Original: Learning phase
Revised: Learning phase

Original: For computation efﬁciency in experience replay, we train a  VRM  by sampling minibatchs of truncated sequences of ﬁxed length, instead of whole episodes.
Revised: As in training a VRNN, the VRM is trained by maximizing an evidence lower bound (Fig. 1(c)) [Equation] [Equation]

Revised: In practice, the ﬁrst term E q φ [log p θ ( x t | z 1: t , x 1: t − 1 )] can be obtained by unrolling the RNN using the inference model (Fig. 1(c)) with sampled sequences of x t .

Revised: Since q φ and p θ are parameterized

Revised: Gaussian distributions, the KL-divergence term can be analytically expressed as

Revised: [Equation]

Revised: [Equation]

Revised: For computation efﬁciency in experience replay, we train a VRM by sampling minibatchs of truncated sequences of ﬁxed length, instead of whole episodes.

Original: Details are found in Appendix A.1.
Revised: Details are found in Appendix A.1.

Original: Since training of a VRM is segregated from training of the RL controllers, there are several strategies for conducting them in parallel.
Revised: Since training of a VRM is segregated from training of the RL controllers, there are several strategies for conducting them in parallel.

Original: For the RL controller, we adopted a smooth update strategy as in
Revised: For the RL controller, we adopted a smooth update strategy as in

Original: Haarnoja et al.
Revised: Haarnoja et al.

Original: (2018a), i.e., performing one time of experience replay every n steps.
Revised: (2018a), i.e., performing one time of experience replay every n steps.

Original: To train the
Revised: To train the

Original: VRM, one can also conduct smooth update.
Revised: VRM, one can also conduct smooth update.

Original: However, in that case, RL suffers from instability of the representation of underlying states in the VRM before it converges.
Revised: However, in that case, RL suffers from instability of the representation of underlying states in the VRM before it converges.

Original: Also, stochasticity of RNN state variables d can be meaninglessly high at early stage of training, which may create problems in RL.
Revised: Also, stochasticity of RNN state variables d can be meaninglessly high at early stage of training, which may create problems in RL.

Original: Another strategy is to pre-train the VRM for abundant epochs only before RL starts, which unfortunately, can fail if novel observations from the environment appear after some degree of policy improvement.
Revised: Another strategy is to pre-train the VRM for abundant epochs only before RL starts, which unfortunately, can fail if novel observations from the environment appear after some degree of policy improvement.

Original: Moreover, if pre-training and smooth update are both applied to the VRM, RL may suffer from a large representation shift of the belief state.
Revised: Moreover, if pre-training and smooth update are both applied to the VRM, RL may suffer from a large representation shift of the belief state.

Original: To resolve this conﬂict, we propose using two VRMs, which we call the ﬁrst-impression model and the keep-learning model , respectively.
Revised: To resolve this conﬂict, we propose using two VRMs, which we call the ﬁrst-impression model and the keep-learning model , respectively.

Original: As the names suggest, we pre-train the ﬁrst-impression model and stop updating it when RL controllers and the keep-learning model start smooth updates.
Revised: As the names suggest, we pre-train the ﬁrst-impression model and stop updating it when RL controllers and the keep-learning model start smooth updates.

Original: Then we take state variables from both VRMs, together with raw observations, as input for the RL controller.
Revised: Then we take state variables from both VRMs, together with raw observations, as input for the RL controller.

Original: We found that this method yields better overall performance than using a single VRM (Appendix C).
Revised: We found that this method yields better overall performance than using a single VRM (Appendix C).

Revised: Initialize the ﬁrst-impression VRM M f and the keep-learning VRM M k , the RL controller C , and the replay buffer D , global step t

Original: 4. R EINFORCEMENT L EARNING C ONTROLLERS
Revised: 4. R EINFORCEMENT L EARNING C ONTROLLERS

Original: As shown in Fig. 1(a), we use multi-layer perceptrons (MLP) as function approximators for V , Q , respectively.
Revised: As shown in Fig. 1(a), we use multi-layer perceptrons (MLP) as function approximators for V , Q , respectively.

Original: Inputs for the Q t network are ( x t , d t , a t ) , and V t is mapped from ( x t , d t ) .
Revised: Inputs for the Q t network are ( x t , d t , a t ) , and V t is mapped from ( x t , d t ) .

Original: Following
Revised: Following

Original: Haarnoja et al.
Revised: Haarnoja et al.

Original: (2018a)
Revised: (2018a)

Original: , we use two Q networks λ 1 and λ 2 and compute Q = min
Revised: , we use two Q networks λ 1 and λ 2 and compute Q = min

Original: ( Q λ 1 , Q λ 2 ) in
Revised: ( Q λ 1 , Q λ 2 ) in

Original: Eq.
Revised: Eq.

Original: 5 and 7 for better performance and stability.
Revised: 5 and 7 for better performance and stability.

Original: Furthermore, we also used a target value network for computing V in Eq.
Revised: Furthermore, we also used a target value network for computing V in Eq.

Original: 6 as in Haarnoja et al.
Revised: 6 as in Haarnoja et al.

Original: (2018a).
Revised: (2018a).

Original: The policy function π η follows a parameterized
Revised: The policy function π η follows a parameterized

Original: Gaussian distribution N ( µ η ( d t , x t ) , diag ( σ η ( d t , x t ))) where µ η and σ η are also MLPs.
Revised: Gaussian distribution N ( µ η ( d t , x t ) , diag ( σ η ( d t , x t ))) where µ η and σ η are also MLPs.

Original: Initialize the ﬁrst-impression VRM M f and the keep-learning VRM M k , the RL controller C , and the replay buffer D , global step t

Original: Compute if then

Original: For N if In the execution phase (Fig. 1(b)), observation and reward x t = ( X t , r t − 1 ) are received as VRM inputs to compute internal states d t using inference models.
Revised: In the execution phase (Fig. 1(b)), observation and reward x t = ( X t , r t − 1 ) are received as VRM inputs to compute internal states d t using inference models.

Original: Then, the agent selects an action, sampled from π η ( a t | d t , x t ) , to interact with the environment.
Revised: Then, the agent selects an action, sampled from π η ( a t | d t , x t ) , to interact with the environment.

Original: To train RL networks, we ﬁrst sample sequences of steps from the replay buffer as minibatches; thus, d t can be computed by the inference models using recorded observations ¯ x t and actions ¯ a t (See
Revised: To train RL networks, we ﬁrst sample sequences of steps from the replay buffer as minibatches; thus, d t can be computed by the inference models using recorded observations ¯ x t and actions ¯ a t (See

Original: Appendix A.1.2).
Revised: Appendix A.1.2).

Original: Then RL networks are updated by minimizing the loss functions with gradient descent.
Revised: Then RL networks are updated by minimizing the loss functions with gradient descent.

Original: Gradients stop at d t so that training of RL networks does not involve updating VRMs.
Revised: Gradients stop at d t so that training of RL networks does not involve updating VRMs.

Original:  R ESULTS
Revised:  R ESULTS

Original: To empirically evaluate our algorithm, we performed experiments in a range of (partially observable) continuous control tasks and compared it to the following alternative algorithms.
Revised: To empirically evaluate our algorithm, we performed experiments in a range of (partially observable) continuous control tasks and compared it to the following alternative algorithms.

Original: The overall procedure is summarized in Algorithm 1.
Revised: The overall procedure is summarized in Algorithm 1.

Original: For the

Original: RL controllers

Original: , we

Original:  adopted hyperparameters from the original SAC implementation (Haarnoja et al., 2018b).
Revised: For the RL controllers, we adopted hyperparameters from the original SAC implementation (Haarnoja et al., 2018b).

Original: Both the keep-learning and ﬁrst- impression
Revised: Both the keep-learning and ﬁrst-impression

Original: VRMs were trained using learning rate 0.0008.
Revised: VRMs were trained using learning rate 0.0008.

Original: We pre-trained the ﬁrst-impression VRM for 5, epochs, and updated the keep-learning VRM every 5 steps.
Revised: We pre-trained the ﬁrst-impression VRM for 5, epochs, and updated the keep-learning VRM every 5 steps.

Original: Batches of size 4, each containing a sequence of 64 steps, were used for training both the VRMs and the RL controllers.
Revised: Batches of size 4, each containing a sequence of 64 steps, were used for training both the VRMs and the RL controllers.

Original: All tasks used the same hyperparameters (Appendix A.1).
Revised: All tasks used the same hyperparameters (Appendix A.1).

Original: • SAC-MLP : The vanilla soft actor-critic implementation (Haarnoja et al., 2018a;b), in which each function is approximated by a 2-layer MLP taking raw observations as input.
Revised: • SAC-MLP : The vanilla soft actor-critic implementation (Haarnoja et al., 2018a;b), in which each function is approximated by a 2-layer MLP taking raw observations as input.

Original: • SAC-LSTM :
Revised: • SAC-LSTM :

Original: Soft actor-critic with recurrent networks as function approximators, where raw observations are processed through an LSTM layer followed by 2 layers of MLPs.
Revised: Soft actor-critic with recurrent networks as function approximators, where raw observations are processed through an LSTM layer followed by 2 layers of MLPs.

Original: This allows the agent to make decisions based on the whole history of raw observations.
Revised: This allows the agent to make decisions based on the whole history of raw observations.

Original: In this case, the network has to conduct representation learning and dynamic programming collectively.
Revised: In this case, the network has to conduct representation learning and dynamic programming collectively.

Original: Our algorithm is compared with SAC-LSTM to demonstrate the effect of separating representation learning from dynamic programming. 
Revised: Our algorithm is compared with SAC-LSTM to demonstrate the effect of separating representation learning from dynamic programming. [Figure]

Original: • SLAC :
Revised: • SLAC :

Original: The stochastic latent actor-critic algorithm introduced in Lee et al.
Revised: The stochastic latent actor-critic algorithm introduced in Lee et al.

Original: (2019), which is a state-of-the-art RL algorithm for solving POMDP tasks.
Revised: (2019), which is a state-of-the-art RL algorithm for solving POMDP tasks.

Original: It was shown that SLAC outperformed other model-based and model-free algorithms, such as (Igl et al., 2018; Hafner et al., 2018), in robotic control tasks with third-person image of the robot as observation 1 .[Figure]
Revised: It was shown that SLAC outperformed other model-based and model-free algorithms, such as (Igl et al., 2018; Hafner et al., 2018), in robotic control tasks with third-person image of the robot as observation 2 .

Original: Note that in our algorithm, we apply pre-training of the ﬁrst-impression model.
Revised: Note that in our algorithm, we apply pre-training of the ﬁrst-impression model.

Original: For a fair comparison, we also perform pre-training for the alternative algorithm with the same epochs.
Revised: For a fair comparison, we also perform pre-training for the alternative algorithm with the same epochs.

Original: For SAC-MLP and
Revised: For SAC-MLP and

Original: SAC-LSTM, pre-training is conducted on RL networks; while for SLAC, its model is pre-trained.
Revised: SAC-LSTM, pre-training is conducted on RL networks; while for SLAC, its model is pre-trained.

Original: 5. P ARTIALLY O BSERVABLE C LASSIC C ONTROL T ASKS
Revised: 5. P ARTIALLY O BSERVABLE C LASSIC C ONTROL T ASKS

Original: The Pendlum and CartPole (Barto et al., 1983) tasks are the classic control tasks for evaluating RL algorithms (Fig. 3, Left).
Revised: The Pendlum and CartPole (Barto et al., 1983) tasks are the classic control tasks for evaluating RL algorithms (Fig. 3, Left).

Original: The CartPole task requires learning of a policy that prevents the pole from falling down and keeps the cart from running away by applying a (1-dimensional) force to the cart, in which observable information is the coordinate of the cart, the angle of the pole, and their derivatives w.r.t time (i.e., velocities).
Revised: The CartPole task requires learning of a policy that prevents the pole from falling down and keeps the cart from running away by applying a (1-dimensional) force to the cart, in which observable information is the coordinate of the cart, the angle of the pole, and their derivatives w.r.t time (i.e., velocities).

Original: For the Pendulum task, the agent needs to learn a policy to swing an inverse-pendulum up and to maintain it at the highest position in order to obtain more rewards.
Revised: For the Pendulum task, the agent needs to learn a policy to swing an inverse-pendulum up and to maintain it at the highest position in order to obtain more rewards.

Original: We are interested in classic control tasks because they are relatively easy to solve when fully observable, and thus the PO cases can highlight the representation learning problem.
Revised: We are interested in classic control tasks because they are relatively easy to solve when fully observable, and thus the PO cases can highlight the representation learning problem.

Original: Experiments were performed in these two tasks, as well as their PO versions, in which either velocities cannot be observed or only velocities can be observed.
Revised: Experiments were performed in these two tasks, as well as their PO versions, in which either velocities cannot be observed or only velocities can be observed.

Original: The latter case is meaningful in real-life applications because an agent may not be able to perceive its own position, but can estimate its speed.
Revised: The latter case is meaningful in real-life applications because an agent may not be able to perceive its own position, but can estimate its speed.

Original: As expected, SAC-MLP failed to solve the PO tasks (Fig. 3).
Revised: As expected, SAC-MLP failed to solve the PO tasks (Fig. 3).

Original: While our algorithm succeeded in learning to solve all these tasks, SAC-LSTM showed poorer performance in some of them.
Revised: While our algorithm succeeded in learning to solve all these tasks, SAC-LSTM showed poorer performance in some of them.

Original: In particular, in the pendulum task with only angular velocity observable, SAC-LSTM may suffer from the periodicity of the angle.
Revised: In particular, in the pendulum task with only angular velocity observable, SAC-LSTM may suffer from the periodicity of the angle.

Original: SLAC performed well in the CartPole tasks, but showed less satisfactory sample efﬁciency in the Pendulum tasks.
Revised: SLAC performed well in the CartPole tasks, but showed less satisfactory sample efﬁciency in the Pendulum tasks.

Original: 5. P ARTIALLY O BSERVABLE R OBOTIC C ONTROL T ASKS
Revised: 5. P ARTIALLY O BSERVABLE R OBOTIC C ONTROL T ASKS

Original: To examine performance of the proposed algorithm in more challenging control tasks with higher degrees of freedom (DOF), we also evaluated performance of the proposed algorithm in the OpenAI
Revised: To examine performance of the proposed algorithm in more challenging control tasks with higher degrees of freedom (DOF), we also evaluated performance of the proposed algorithm in the OpenAI

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

Original: The Roboschool environments include a number  of continuous robotic control tasks, such as teaching a multiple-joint robot to walk as fast as possible without falling down (Fig. 4, Left).
Revised: The Roboschool environments include a number [Figure] of continuous robotic control tasks, such as teaching a multiple-joint robot to walk as fast as possible without falling down (Fig. 4, Left).

Original: The original Roboschool environments are nearly fully observable since observations include the robot’s coordinates and (trigonometric functions of) joint angles, as well as (angular and coordinate) velocities.
Revised: The original Roboschool environments are nearly fully observable since observations include the robot’s coordinates and (trigonometric functions of) joint angles, as well as (angular and coordinate) velocities.

Original: As in the PO classic control tasks, we also performed experiments in the PO versions of the Roboschool environments.
Revised: As in the PO classic control tasks, we also performed experiments in the PO versions of the Roboschool environments.

Original: Using our algorithm, experimental results (Fig. 4) demonstrated substantial policy improvement in all PO tasks (visualization of the trained agents is in Appendix D).
Revised: Using our algorithm, experimental results (Fig. 4) demonstrated substantial policy improvement in all PO tasks (visualization of the trained agents is in Appendix D).

Original: In some PO cases, the agents achieved comparable performance to that in fully observable cases.
Revised: In some PO cases, the agents achieved comparable performance to that in fully observable cases.

Original: For tasks with unobserved[Figure] velocities, our algorithm performed similarly to SAC-LSTM.
Revised: For tasks with unobserved velocities, our algorithm performed similarly to SAC-LSTM.

Original: This is because velocities can be simply estimated by one-step differences in robot coordinates and joint angles, which eases representation learning.
Revised: This is because velocities can be simply estimated by one-step differences in robot coordinates and joint angles, which eases representation learning.

Original: However, in environments where only velocities can be observed, our algorithm signiﬁcantly outperformed SAC-LSTM, presumably because SAC-LSTM is less efﬁcient at encoding underlying states from velocity observations.
Revised: However, in environments where only velocities can be observed, our algorithm signiﬁcantly outperformed SAC-LSTM, presumably because SAC-LSTM is less efﬁcient at encoding underlying states from velocity observations.

Original: Also, we found that learning of a SLAC agent was unstable, i.e., it sometimes could acquire a near-optimal policy, but often its policy converged to a poor one.
Revised: Also, we found that learning of a SLAC agent was unstable, i.e., it sometimes could acquire a near-optimal policy, but often its policy converged to a poor one.

Original: Thus, average performance of SLAC was less promising than ours in most of the PO robotic control tasks.
Revised: Thus, average performance of SLAC was less promising than ours in most of the PO robotic control tasks.

Original: 5. L ONG -T ERM M EMORIZATION T ASKS
Revised: 5. L ONG -T ERM M EMORIZATION T ASKS

Original: Another common type of PO task requires long-term memorization of past events.
Revised: Another common type of PO task requires long-term memorization of past events.

Original: To solve these tasks, an agent needs to learn to extract and to remember critical information from the whole history of raw observations.
Revised: To solve these tasks, an agent needs to learn to extract and to remember critical information from the whole history of raw observations.

Original: Therefore, we also examined our algorithm and other alternatives in a long-term memorization task known as the sequential target reaching task (Han et al., 2019), in which a robot agent needs to reach 3 different targets in a certain sequence (Fig. 5, Left).
Revised: Therefore, we also examined our algorithm and other alternatives in a long-term memorization task known as the sequential target reaching task (Han et al., 2019), in which a robot agent needs to reach 3 different targets in a certain sequence (Fig. 5, Left).

Original: The robot can control its two wheels to move or turn, and will get one-step small, medium, and large rewards when it reaches the ﬁrst, second, and third targets, respectively, in the correct sequence.
Revised: The robot can control its two wheels to move or turn, and will get one-step small, medium, and large rewards when it reaches the ﬁrst, second, and third targets, respectively, in the correct sequence.

Original: The robot senses distances and angles from the 3 targets, but does not receive any signal indicating which target to reach.
Revised: The robot senses distances and angles from the 3 targets, but does not receive any signal indicating which target to reach.

Original: In each episode, the robot’s initial position and those of the three targets are randomly initialized.
Revised: In each episode, the robot’s initial position and those of the three targets are randomly initialized.

Original: In order to obtain rewards, the agent needs to infer the current correct target using historical observations.
Revised: In order to obtain rewards, the agent needs to infer the current correct target using historical observations.

Original: [Figure] We found that agents using our algorithm achieved almost 100% success rate (reaching 3 targets in the correct sequence within maximum steps).
Revised: We found that agents using our algorithm achieved almost 100% success rate (reaching 3 targets in the correct sequence within maximum steps).

Original: SAC-LSTM also achieved similar success rate after convergence, but spent more training steps learning to encode underlying goal-related information  from sequential observations.
Revised: SAC-LSTM also achieved similar success rate after convergence, but spent more training steps learning to encode underlying goal-related information [Figure] from sequential observations.

Original: Also, SLAC struggled hard to solve this task since its actor only received a limited steps of observations, making it difﬁcult to infer the correct target.
Revised: Also, SLAC struggled hard to solve this task since its actor only received a limited steps of observations, making it difﬁcult to infer the correct target.

Original: 5. C ONVERGENCE OF THE KEEP - LEARNING VRM
Revised: 5. C ONVERGENCE OF THE KEEP - LEARNING VRM

Original: One of the most concerned problems of our algorithm is that input of the RL controllers can experience representation change, because the keep-learning model is not guaranteed to converge if novel observation appears due to improved policy (e.g. for a hopper robot, “in-the-air” state can only happen after it learns to hop).
Revised: One of the most concerned problems of our algorithm is that input of the RL controllers can experience representation change, because the keep-learning model is not guaranteed to converge if novel observation appears due to improved policy (e.g. for a hopper robot, “in-the-air” state can only happen after it learns to hop).

Original: To empirically investigate how convergence of the keep-learning VRM affect policy improvement, we plot the loss functions (negative ELBOs) of the the keep-learning
Revised: To empirically investigate how convergence of the keep-learning VRM affect policy improvement, we plot the loss functions (negative ELBOs) of the the keep-learning

Original: VRM for 3 example tasks (Fig. 6).
Revised: VRM for 3 example tasks (Fig. 6).

Original: For a simpler task (CartPole), the policy was already near optimal before the VRM fully converged.
Revised: For a simpler task (CartPole), the policy was already near optimal before the VRM fully converged.

Original: We also saw that the policy was gradually improved after the VRM mostly converged (RoboschoolAnt - no velocities), and that the policy and the VRM were being improved in parallel (RoboschoolAnt - velocities only).
Revised: We also saw that the policy was gradually improved after the VRM mostly converged (RoboschoolAnt - no velocities), and that the policy and the VRM were being improved in parallel (RoboschoolAnt - velocities only).

Original: The results suggested that policy could be improved with sufﬁcient sample efﬁciency even the keep-learning VRM did not converge.
Revised: The results suggested that policy could be improved with sufﬁcient sample efﬁciency even the keep-learning VRM did not converge.

Original: This can be explained by that the RL controller also extract information from the ﬁrst-impression model and the raw observations, which did not experience representation change during RL.
Revised: This can be explained by that the RL controller also extract information from the ﬁrst-impression model and the raw observations, which did not experience representation change during RL.

Original: Indeed, our ablation study showed performance degradation in many tasks without the ﬁrst-impression VRM (Appendix C).[Figure] r [Figure]
Revised: Indeed, our ablation study showed performance degradation in many tasks without the ﬁrst-impression VRM (Appendix C).

Original: r

Original: [Figure]

Original: t u r n

Original: [Figure]
Revised: [Figure]

Original: a v e r age r e t u r n

Original: [Figure]

Original: [Figure]

Original: nega t i v e E

Original: L B O

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: In this paper, we proposed a variational recurrent model for learning to represent underlying states of
Revised: In this paper, we proposed a variational recurrent model for learning to represent underlying states of

Original: PO environments and the corresponding algorithm for solving POMDPs.
Revised: PO environments and the corresponding algorithm for solving POMDPs.

Original: Our experimental results demonstrate effectiveness of the proposed algorithm in tasks in which underlying states cannot be simply inferred using a short sequence of observations.
Revised: Our experimental results demonstrate effectiveness of the proposed algorithm in tasks in which underlying states cannot be simply inferred using a short sequence of observations.

Original: Our work can be considered an attempt to understand how RL beneﬁts from stochastic Bayesian inference of state-transitions, which actually happens in the brain (Funamizu et al., 2016), but has been considered less often in RL studies.
Revised: Our work can be considered an attempt to understand how RL beneﬁts from stochastic Bayesian inference of state-transitions, which actually happens in the brain (Funamizu et al., 2016), but has been considered less often in RL studies.

Original: We used stochastic models in this work which we actually found perform better than deterministic ones, even through the environments we used are deterministic (Appendix C).
Revised: We used stochastic models in this work which we actually found perform better than deterministic ones, even through the environments we used are deterministic (Appendix C).

Original: The VRNN can be replaced with other alternatives (Bayer & Osendorfer, 2014; Goyal et al., 2017) to potentially improve performance, although developing model architecture is beyond the scope of the current study.
Revised: The VRNN can be replaced with other alternatives (Bayer & Osendorfer, 2014; Goyal et al., 2017) to potentially improve performance, although developing model architecture is beyond the scope of the current study.

Original: Moreover, a recent study (Ahmadi & Tani, 2019) showed a novel way of inference using back-propagation of prediction errors, which may also beneﬁt our future studies.
Revised: Moreover, a recent study (Ahmadi & Tani, 2019) showed a novel way of inference using back-propagation of prediction errors, which may also beneﬁt our future studies.

Original: Many researchers think that there are two distinct systems for model-based and model-free RL in the brain (Gl¨ascher et al., 2010; Lee et al., 2014) and a number of studies investigated how and when the brain switches between them (Smittenaar et al., 2013; Lee et al., 2014).
Revised: Many researchers think that there are two distinct systems for model-based and model-free RL in the brain (Gl¨ascher et al., 2010; Lee et al., 2014) and a number of studies investigated how and when the brain switches between them (Smittenaar et al., 2013; Lee et al., 2014).

Original: However, Stachenfeld et al.
Revised: However, Stachenfeld et al.

Original: (2017) suggested that the hippocampus can learn a successor representation of the environment that beneﬁts both model-free and model-based RL, contrary to the aforementioned conventional view.
Revised: (2017) suggested that the hippocampus can learn a successor representation of the environment that beneﬁts both model-free and model-based RL, contrary to the aforementioned conventional view.

Original: We further propose another possibility, that a model is learned, but not used for planning or dreaming.
Revised: We further propose another possibility, that a model is learned, but not used for planning or dreaming.

Original: This blurs the distinction between model-based and model-free RL.
Revised: This blurs the distinction between model-based and model-free RL.

Original: A I MPLEMENTATION D ETAILS
Revised: A CKNOWLEDGEMENT

Revised: This work was supported by Okinawa Institute of Science and Technology Graduate University funding, and was also partially supported by a Grant-in-Aid for Scientiﬁc Research on Innovative

Revised: Areas: Elucidation of the Mathematical Basis and Neural Mechanisms of Multi-layer Representation

Revised: Learning 16H06563.

Revised: We would like to thank the lab members in the Cognitive Neurorobotics

Revised: Research Unit and the Neural Computation Unit of Okinawa Institute of Science and Technology.

Revised: In particular, we would like to thank Ahmadreza Ahmadi for his help during model development.

Revised: We also would like to thank Steven Aird for assisting improving the manuscript.

Revised: A I MPLEMENTATION D ETAILS

Original: In this section we describe the details of implementing our algorithm as well as the alternative ones.
Revised: In this section we describe the details of implementing our algorithm as well as the alternative ones.

Original: Summaries of hyperparameters can be found in Table 1 and 2.
Revised: Summaries of hyperparameters can be found in Table 1 and 2.

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: N ETWORK ARCHITECTURES
Revised: N ETWORK ARCHITECTURES

Original: The ﬁrst-impression model and the keep-learning model adopted the same architecture.
Revised: The ﬁrst-impression model and the keep-learning model adopted the same architecture.

Original: Size of d and z is 256 and 64, respectively.
Revised: Size of d and z is 256 and 64, respectively.

Original: We used one-hidden-layer fully-connected networks with 128 hidden neurons for the inference models (cid:104) µ φ,t , σ 2 φ,t (cid:105) =
Revised: We used one-hidden-layer fully-connected networks with 128 hidden neurons for the inference models (cid:104) µ φ,t , σ 2 φ,t (cid:105) =

Original: φ ( x t , d t − 1 , a t − 1 ) , as well as for (cid:104) µ θ,t , σ 2 θ,t (cid:105) =
Revised: φ ( x t , d t − 1 , a t − 1 ) , as well as for (cid:104) µ θ,t , σ 2 θ,t (cid:105) =

Original: θ prior ( d t − 1 , a t − 1 ) in the generative models.
Revised: θ prior ( d t − 1 , a t − 1 ) in the generative models.

Original: For the decoder (cid:2) µ x,t , σ 2 x,t (cid:3) = θ decoder ( z t , d t − 1 ) in the generative models, we used 2-layers MLPs with 128 neurons in each layer.
Revised: For the decoder (cid:2) µ x,t , σ 2 x,t (cid:3) = θ decoder ( z t , d t − 1 ) in the generative models, we used 2-layers MLPs with 128 neurons in each layer.

Original: The input processing layer f x is also an one-layer MLP with size-128.
Revised: The input processing layer f x is also an one-layer MLP with size-128.

Original: For all the Gaussian variables, output functions for mean are linear and output functions for variance are softplus.
Revised: For all the Gaussian variables, output functions for mean are linear and output functions for variance are softplus.

Original: Other activation functions of the
Revised: Other activation functions of the

Original: VRMs are tanh.
Revised: VRMs are tanh.

Original: The RL controllers are the same as those in SAC-MLP (Section A.2.1) except that network inputs are raw observations together with the RNN states from the ﬁrst-impression model and the keep-learning model.
Revised: The RL controllers are the same as those in SAC-MLP (Section A.2.1) except that network inputs are raw observations together with the RNN states from the ﬁrst-impression model and the keep-learning model.

Original: A.1. I NITIAL STATES OF THE VRM S
Revised: A.1. I NITIAL STATES OF THE VRM S

Original: To train the VRMs, one can use a number of entire episodes as a mini-batch, using zero initial states, as in Heess et al.
Revised: To train the VRMs, one can use a number of entire episodes as a mini-batch, using zero initial states, as in Heess et al.

Original: However, when tackling with long episodes (e.g. there can be 1,000 steps in each episode in the robotic control tasks we used) or even inﬁnite-horizon problems, the computation consumption will be huge in back-propagation through time (BPTT).
Revised: However, when tackling with long episodes (e.g. there can be 1,000 steps in each episode in the robotic control tasks we used) or even inﬁnite-horizon problems, the computation consumption will be huge in back-propagation through time (BPTT).

Original: For better computation efﬁciency, we used 4 length-64 sequences for training the RNNs, and applied the burn-in method for providing the initial states (Kapturowski et al., 2018), or more speciﬁcally, unrolling the
Revised: For better computation efﬁciency, we used 4 length-64 sequences for training the RNNs, and applied the burn-in method for providing the initial states (Kapturowski et al., 2018), or more speciﬁcally, unrolling the

Original: RNNs using a portion of the replay sequence (burn-in period, up to 64 steps in our case) from zero initial states.
Revised: RNNs using a portion of the replay sequence (burn-in period, up to 64 steps in our case) from zero initial states.

Original: We assume that proper initial states can be obtained in this way.
Revised: We assume that proper initial states can be obtained in this way.

Original: This is crucial for the tasks that require long-term memorization, and is helpful to reduce bias introduces by incorrect initial states in general cases.
Revised: This is crucial for the tasks that require long-term memorization, and is helpful to reduce bias introduces by incorrect initial states in general cases.

Original: A. A LTERNATIVE ALGORITHMS A.2. SAC-MLP
Revised: A. A LTERNATIVE ALGORITHMS A.2. SAC-MLP

Original: We followed the original implementation of SAC in (Haarnoja et al., 2018a) including hyperparameters.
Revised: We followed the original implementation of SAC in (Haarnoja et al., 2018a) including hyperparameters.

Original: However, we also applied automatic learning of the entropy coefﬁcient α
Revised: However, we also applied automatic learning of the entropy coefﬁcient α

Original: (inverse of the the reward scale in Haarnoja et al.
Revised: (inverse of the the reward scale in Haarnoja et al.

Original: (2018a))
Revised: (2018a))

Original: as introduced by the authors in Haarnoja et al.
Revised: as introduced by the authors in Haarnoja et al.

Original: (2018b) to avoid tuning the reward scale for each task.
Revised: (2018b) to avoid tuning the reward scale for each task.

Original: A.2. SAC-LSTM
Revised: A.2. SAC-LSTM

Original: To apply recurrency to SAC’s function approximators, we added an LSTM network with sizereceiving raw observations as input.
Revised: To apply recurrency to SAC’s function approximators, we added an LSTM network with sizereceiving raw observations as input.

Original: The function approximators of actor and critic were the same as those in SAC except receiving the LSTM’s output as input.
Revised: The function approximators of actor and critic were the same as those in SAC except receiving the LSTM’s output as input.

Original: The gradients can pass through the
Revised: The gradients can pass through the

Original: LSTM so that the training of the LSTM and MLPs were synchronized.
Revised: LSTM so that the training of the LSTM and MLPs were synchronized.

Original: The training the network also followed Section A.1.2.
Revised: The training the network also followed Section A.1.2.

Original: A.2. SLAC
Revised: A.2. SLAC

Original: We mostly followed the implementation of SLAC explained in the authors’ paper (Lee et al., 2019).
Revised: We mostly followed the implementation of SLAC explained in the authors’ paper (Lee et al., 2019).

Original: One modiﬁcation is that since their work was using pixels as observations, convolutional neural networks (CNN) and transposed CNNs were chosen for input feature extracting and output decoding layers; in our case, we replaced the CNN and transposed CNNs by 2-layers MLPs with 256 units in each layer.
Revised: One modiﬁcation is that since their work was using pixels as observations, convolutional neural networks (CNN) and transposed CNNs were chosen for input feature extracting and output decoding layers; in our case, we replaced the CNN and transposed CNNs by 2-layers MLPs with 256 units in each layer.

Original: In addition, the authors set the output variance σ 2 y,t for each image pixel as 0.1.
Revised: In addition, the authors set the output variance σ 2 y,t for each image pixel as 0.1.

Original: However, σ 2 y,t = 0 . 1 can be too large for joint states/velocities as observations.
Revised: However, σ 2 y,t = 0 . 1 can be too large for joint states/velocities as observations.

Original: We found that it will lead to better performance by setting σ y,t as trainable parameters (as that in our algorithm).
Revised: We found that it will lead to better performance by setting σ y,t as trainable parameters (as that in our algorithm).

Original: We also used a 2-layer MLP with 256 units for approximating σ y ( x t , d t − 1 ) .
Revised: We also used a 2-layer MLP with 256 units for approximating σ y ( x t , d t − 1 ) .

Original: To avoid network weights being divergent, all the activation functions of the model were tanh except those for outputs.
Revised: To avoid network weights being divergent, all the activation functions of the model were tanh except those for outputs.

Original: B E NVIRONMENTS
Revised: B E NVIRONMENTS

Original: For the robotic control tasks and the Pendulum task, we used environments (and modiﬁed them for PO versions) from OpenAI Gym (Brockman et al., 2016).
Revised: For the robotic control tasks and the Pendulum task, we used environments (and modiﬁed them for PO versions) from OpenAI Gym (Brockman et al., 2016).

Original: The CartPole environment with a continuous action space was from Danforth (2018), and the codes for the sequential target reaching tasks were provided by the authors (Han et al., 2019).
Revised: The CartPole environment with a continuous action space was from Danforth (2018), and the codes for the sequential target reaching tasks were provided by the authors (Han et al., 2019).

Original: In the no-velocities cases, velocity information was removed from raw observations; while in the velocities-only cases, only velocity information was retained in raw observations.
Revised: In the no-velocities cases, velocity information was removed from raw observations; while in the velocities-only cases, only velocity information was retained in raw observations.

Original: We summarize key information of each environment in Table 3.
Revised: We summarize key information of each environment in Table 3.

Original: The performance curves were obtained in evaluation phases in which agents used same policy but did not update networks or record state-transition data.
Revised: The performance curves were obtained in evaluation phases in which agents used same policy but did not update networks or record state-transition data.

Original: Each experiment was repeated using 5 different random seeds.
Revised: Each experiment was repeated using 5 different random seeds.

Original: C A BLATION S TUDY
Revised: C A BLATION S TUDY

Original: This section demonstrated a ablation study in which we compared the performance of the proposed algorithm to the same but with some modiﬁcation:
Revised: This section demonstrated a ablation study in which we compared the performance of the proposed algorithm to the same but with some modiﬁcation:

Original: • With a single VRM .
Revised: • With a single VRM .

Original: In this case, we used only one VRM and applied both pre-training and smooth update to it.
Revised: In this case, we used only one VRM and applied both pre-training and smooth update to it.

Original: • Only ﬁrst-impression model .
Revised: • Only ﬁrst-impression model .

Original: In this case, only the ﬁrst-impression model was used and pre-trained.
Revised: In this case, only the ﬁrst-impression model was used and pre-trained.

Original: [Table]
Revised: [Table]

Original: • Only keep-learning model .
Revised: • Only keep-learning model .

Original: In this case, only the keep-learning model was used and smooth-update was applied.
Revised: In this case, only the keep-learning model was used and smooth-update was applied.

Original: • Deterministic model .
Revised: • Deterministic model .

Original: In this case, the ﬁrst-imporession model and the keep-learning model were deterministic RNNs which learned to model the state-transitions by minimizing mean-square error between prediction and observations instead of ELBO .
Revised: In this case, the ﬁrst-imporession model and the keep-learning model were deterministic RNNs which learned to model the state-transitions by minimizing mean-square error between prediction and observations instead of ELBO .

Original: The network architecture was mostly the same as the VRM expect that the inference model and the generative model were merged into a deterministic one.
Revised: The network architecture was mostly the same as the VRM expect that the inference model and the generative model were merged into a deterministic one.

Original: The learning curves are shown in Fig.
Revised: The learning curves are shown in Fig.

Original: It can be seen that the proposed algorithm consistently performed similar as or better than the modiﬁed ones.
Revised: It can be seen that the proposed algorithm consistently performed similar as or better than the modiﬁed ones.

Original: D V ISUALIZATION OF T RAINED A GENTS
Revised: D V ISUALIZATION OF T RAINED A GENTS

Original: Here we show actual movements of the trained robots in the PO robotic control tasks (Fig.
Revised: Here we show actual movements of the trained robots in the PO robotic control tasks (Fig.

Original: It can be seen that the robots succeeded in learning to hop or walk, although their policy may be sub-optimal.
Revised: It can be seen that the robots succeeded in learning to hop or walk, although their policy may be sub-optimal.

Original: E M ODEL A CCURACY
Revised: E M ODEL A CCURACY

Original: As we discussed in Section 2, our algorithm relies mostly on encoding capacity of models, but does not require models to make accurate prediction of future observations.
Revised: As we discussed in Section 2, our algorithm relies mostly on encoding capacity of models, but does not require models to make accurate prediction of future observations.

Original: Fig.
Revised: Fig.

Original: 9 shows open-loop (using the inference model to compute the latent variable z ) and close-loop (purely using the generative model) prediction of raw observation by the keep-learning models of randomly selected trained agents.
Revised: 9 shows open-loop (using the inference model to compute the latent variable z ) and close-loop (purely using the generative model) prediction of raw observation by the keep-learning models of randomly selected trained agents.

Original: Here we showcase “RoboschoolHopper - velocities only” and “Pendulum - no velocities” because in these tasks our algorithm achieved similar performance to those in fully-observable versions (Fig. 4), although the prediction accuracy of the models was imperfect.
Revised: Here we showcase “RoboschoolHopper - velocities only” and “Pendulum - no velocities” because in these tasks our algorithm achieved similar performance to those in fully-observable versions (Fig. 4), although the prediction accuracy of the models was imperfect.

Original: F S ENSITIVITY TO H YPERPARAMETERS OF THE VRM S
Revised: F S ENSITIVITY TO H YPERPARAMETERS OF THE VRM S

Original: To empirically show how choice of hyperparameters of the VRMs affect RL performance, we conducted experiments using hyperparameters different from those used in the main study.
Revised: To empirically show how choice of hyperparameters of the VRMs affect RL performance, we conducted experiments using hyperparameters different from those used in the main study.

Original: More speciﬁcally, the learning rate for both VRMs was randomly selected from { 0.0004, 0.0006, 0.0008, 0.001 } and the sequence length was randomly selected from { 16, 32, 64 } (the batch size was 256 / ( sequence length ) to ensure that the total number of samples in a batch was 256 which matched with the alternative approaches).
Revised: More speciﬁcally, the learning rate for both VRMs was randomly selected from { 0.0004, 0.0006, 0.0008, 0.001 } and the sequence length was randomly selected from { 16, 32, 64 } (the batch size was 256 / ( sequence length ) to ensure that the total number of samples in a batch was 256 which matched with the alternative approaches).

Original: The other hyperparameters were unchanged.
Revised: The other hyperparameters were unchanged.

Original: The results can be checked in Fig 10 for all the environments we used.
Revised: The results can be checked in Fig 10 for all the environments we used.

Original: The overall performance did not signiﬁcantly change using different, random hyperparameters of the VRMs, although we could observe signiﬁcant performance improvement (e.g. RoboshoolWalker2d) or degradation (e.g.
Revised: The overall performance did not signiﬁcantly change using different, random hyperparameters of the VRMs, although we could observe signiﬁcant performance improvement (e.g. RoboshoolWalker2d) or degradation (e.g.

Original: RoboshoolHopper - velocities only) in a few tasks using different haperparameters.
Revised: RoboshoolHopper - velocities only) in a few tasks using different haperparameters.

Original: Therefore, the representation learning part (VRMs) of our algorithm does not suffer from high sensitivity to hyperparameters.
Revised: Therefore, the representation learning part (VRMs) of our algorithm does not suffer from high sensitivity to hyperparameters.

Original: This can be explained by the fact that we do not use a bootstrapping (e.g. the estimation of targets of value functions depends on the estimation of value functions) (Sutton & Barto, 1998) update rule to train the VRMs.
Revised: This can be explained by the fact that we do not use a bootstrapping (e.g. the estimation of targets of value functions depends on the estimation of value functions) (Sutton & Barto, 1998) update rule to train the VRMs.

Original: G S CALABILITY
Revised: G S CALABILITY

Original: Table 4 showed scalability of our algorithm and the alternative ones.
Revised: Table 4 showed scalability of our algorithm and the alternative ones.

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

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