Title: VARIATIONAL RECURRENT MODELS FOR SOLVING PARTIALLY OBSERVABLE CONTROL TASKS

Abstract: 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. In this study, we propose an RL algorithm for solving PO tasks. 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. 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. Our experiments show that the proposed algorithm achieved better data efficiency 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 .

1. INTRODUCTION:

Model-free deep reinforcement learning (RL) algorithms have been developed to solve difficult control and decision-making tasks by self-exploration (Sutton & Barto, 1998; Mnih et al., 2015; Silver et al., 2016) . While various kinds of fully observable environments have been well investigated, recently, partially observable (PO) environments (Hafner et al., 2018; Igl et al., 2018; Lee et al., 2019; 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.

There are many types of PO tasks; however, those that can be solved by taking the history of observations into account are more common. 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) . While humans are good at solving these tasks by extracting crucial information from the past observations, deep RL agents often have difficulty acquiring satisfactory policy and achieving good data efficiency, compared to those in fully observable tasks (Hafner et al., 2018; Lee et al., 2019) .

For solving such PO tasks, several categories of methods have been proposed. One simple, straightforward solution is to include a history of raw observations in the current "observation" (McCallum, 1993; Lee et al., 2019) . 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.

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. 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. 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 efficiently, since RNNs are relatively more difficult to train (Pascanu et al., 2013) than feedforward neural networks.

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; Ha & Schmidhuber, 2018; Lee et al., 2019) . The belief state is an agent-estimated variable encoding underlying states of the environment that determines state-transitions and rewards. Perfectly-estimated belief states can thus be taken as "observations" of an RL agent that contains complete information for solving the task. Therefore, solving a PO task is segregated into a representation learning problem and a fully observable RL problem. Since fully observable RL problems have been well explored by the RL community, the critical challenge here is how to estimate the belief state.

In this study, we developed a variational recurrent model (VRM) that models sequential observations and rewards using a latent stochastic variable. The VRM is an extension of the variational recurrent neural network (VRNN) model (Chung et al., 2015) that takes actions into account. Our approach falls into the third category by taking the internal states of the VRM together with raw observations as the belief state. We then propose an algorithm to solve PO tasks by training the VRM and a feed-forward RL controller network, respectively. The algorithm can be applied in an end-to-end manner, without fine tuning of a hyperparameters.

We then experimentally evaluated the proposed algorithm in various PO versions of robotic control tasks. The agents showed substantial policy improvement in all tasks, and in some tasks the algorithm performed essentially as in fully observable cases. 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.

2. RELATED WORK:

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; Kaiser et al., 2019) or for planning of future state-transitions (Watter et al., 2015; Hafner et al., 2018; Ke et al., 2019) . This usually requires a well-designed and finely tuned model so that its predictions are accurate and robust. In our case, we do not use VRMs for dreaming and planning, but for auto-encoding state-transitions. Actually, PO tasks can be solved without requiring VRMs to predict accurately (see Appendix E). This distinguishes our algorithm from typical model-based RL methods.

The work our method most closely resembles is known as stochastic latent actor-critic (SLAC, Lee et al. (2019) ), in which a latent variable model is trained and uses the latent state as the belief state for the critic. 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. 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. 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. Furthermore, SLAC did not include raw observations in the input of the critic, which may complicate training the critic before the model converges.

3. BACKGROUND:



3.1. PARTIALLY OBSERVABLE MARKOV DECISION PROCESSES:

The scope of problems we study can be formulated into a framework known as partially observable Markov decision processes (POMDP) (Kaelbling et al., 1998) . 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.

A POMDP is usually defined 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 × A → p(S) is the state-transition probability function that determines the distribution of the next state given current state and action. The reward function R : S × A → R decides the reward during a state-transition, which can also be probabilistic. Moreover, X is a set of observations, and observations are determined by the observation probability function O : S × A → p(X). By defining a POMDP, the goal is to maximize expected discounted future rewards t γ t r t by learning a good strategy to select actions (policy function).

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. 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 .

3.2. VARIATIONAL RECURRENT NEURAL NETWORKS:

To model general state-transitions that can be stochastic and complicated, we employ a modified version of the VRNN (Chung et al., 2015) . 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. 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 .

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

EQUATION

where f s are parameterized mappings, such as feed-forward neural networks, and d t is the state variable of the RNN, which is recurrently updated by

EQUATION

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

EQUATION

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

EQUATION

where p and q are parameterized PDFs of z t by the generative model and the inference model, respectively. In a POMDP, a VRNN can be used to model the environment and to represent underlying states in its state variable d t . Thus an RL agent can benefit from a well-learned VRNN model since d t provides additional information about the environment beyond the current raw observation x t .

3.3. SOFT ACTOR CRITIC:

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) . A SAC agent learns to maximize reinforcement returns as well as entropy of its policy, so as to obtain more rewards while keeping actions sufficiently stochastic. A typical SAC implementation can be described as follows. 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. Also, an entropy coefficient factor (also known as the temperature parameter), denoted by α, is learned to control the degree of stochasticity of the policy. The parameters are learned by simultaneously minimizing the following loss functions.

EQUATION

EQUATION

EQUATION

EQUATION

where B is the replay buffer from which s t is sampled, and H tar is the target entropy. To compute the gradient of J π (η) (Equation. 7), the reparameterization trick (Kingma & Welling, 2013) is used on action, indicated by a η (s t ). Reparameterization of action is not required in minimizing J(α) (Equation. 8) since log π η (a|s t ) does not depends on α.

SAC was originally developed for fully observable environments; thus, the raw observation at the current step x t was used as network input. 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.

4. METHODS:



4.1. VARIATIONAL RECURRENT STATE-TRANSITION MODELS:

An overall diagram of the proposed algorithm is summarized in Fig. 1 (a), while a more detailed computational graph is plotted in Fig. 2 . 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. 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. Thus, we have the inference model (Fig. 1(c )), denoted by φ, as

EQUATION

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

EQUATION

)

For building recurrent connections, the choice of RNN types is not limited. In our study, the longshort term memory (LSTM) (Hochreiter & Schmidhuber, 1997) As in training a VRNN, the VRM is trained by maximizing an evidence lower bound (Fig. 1(c ))

ELBO = t E q φ [log p θ (x t |z 1:t , x 1:t-1 )] -D KL [q φ (z t |z 1:t-1 , x1:t , ā1:t )||p θ (z t |z 1:t-1 , x1:t-1 , ā1:t )]} .

(11) In practice, the first 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

EQUATION

For computation efficiency in experience replay, we train a VRM by sampling minibatchs of truncated sequences of fixed length, instead of whole episodes. Details are found in Appendix A.1.

Since training of a VRM is segregated from training of the RL controllers, there are several strategies for conducting them in parallel. For the RL controller, we adopted a smooth update strategy as in Haarnoja et al. (2018a) , i.e., performing one time of experience replay every n steps. To train the VRM, one can also conduct smooth update. However, in that case, RL suffers from instability of the representation of underlying states in the VRM before it converges. Also, stochasticity of RNN state variables d can be meaninglessly high at early stage of training, which may create problems in RL. 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. 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.

To resolve this conflict, we propose using two VRMs, which we call the first-impression model and the keep-learning model, respectively. As the names suggest, we pre-train the first-impression model and stop updating it when RL controllers and the keep-learning model start smooth updates. Then we take state variables from both VRMs, together with raw observations, as input for the RL controller. We found that this method yields better overall performance than using a single VRM (Appendix C). (Igl et al., 2018; Hafner et al., 2018) , in robotic control tasks with third-person image of the robot as observation2 .

Note that in our algorithm, we apply pre-training of the first-impression model. For a fair comparison, we also perform pre-training for the alternative algorithm with the same epochs. For SAC-MLP and SAC-LSTM, pre-training is conducted on RL networks; while for SLAC, its model is pre-trained.

5.1. PARTIALLY OBSERVABLE CLASSIC CONTROL TASKS:

The Pendlum and CartPole (Barto et al., 1983) tasks are the classic control tasks for evaluating RL algorithms (Fig. 3 , Left). 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). 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.

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. 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. 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.

As expected, SAC-MLP failed to solve the PO tasks (Fig. 3 ). While our algorithm succeeded in learning to solve all these tasks, SAC-LSTM showed poorer performance in some of them. In particular, in the pendulum task with only angular velocity observable, SAC-LSTM may suffer from the periodicity of the angle. SLAC performed well in the CartPole tasks, but showed less satisfactory sample efficiency in the Pendulum tasks.

5.2. PARTIALLY OBSERVABLE ROBOTIC CONTROL TASKS:

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 Roboschool environments (Brockman et al., 2016) . The Roboschool environments include a number from sequential observations. Also, SLAC struggled hard to solve this task since its actor only received a limited steps of observations, making it difficult to infer the correct target.

5.4. CONVERGENCE OF THE KEEP-LEARNING VRM:

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). 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 VRM for 3 example tasks (Fig. 6 ). For a simpler task (CartPole), the policy was already near optimal before the VRM fully converged. 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).

The results suggested that policy could be improved with sufficient sample efficiency even the keep-learning VRM did not converge. This can be explained by that the RL controller also extract information from the first-impression model and the raw observations, which did not experience representation change during RL. Indeed, our ablation study showed performance degradation in many tasks without the first-impression VRM (Appendix C). 

RoboschoolAnt -no veolocities:



6. DISCUSSION:

In this paper, we proposed a variational recurrent model for learning to represent underlying states of PO environments and the corresponding algorithm for solving POMDPs. 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. Our work can be considered an attempt to understand how RL benefits 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.

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). 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. Moreover, a recent study (Ahmadi & Tani, 2019) showed a novel way of inference using back-propagation of prediction errors, which may also benefit our future studies.

Many researchers think that there are two distinct systems for model-based and model-free RL in the brain (Gläscher 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 ). However, Stachenfeld et al. (2017) suggested that the hippocampus can learn a successor representation of the environment that benefits both model-free and model-based RL, contrary to the aforementioned conventional view. We further propose another possibility, that a model is learned, but not used for planning or dreaming. This blurs the distinction between model-based and model-free RL. The learning curves are shown in Fig. 7 . It can be seen that the proposed algorithm consistently performed similar as or better than the modified ones.

D VISUALIZATION OF TRAINED AGENTS:

Here we show actual movements of the trained robots in the PO robotic control tasks (Fig. 8 ). It can be seen that the robots succeeded in learning to hop or walk, although their policy may be sub-optimal.

E MODEL ACCURACY:

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. Fig. 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.

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.

F SENSITIVITY TO HYPERPARAMETERS OF THE VRMS:

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. More specifically, 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). The other hyperparameters were unchanged.

The results can be checked in Fig 10 for all the environments we used. The overall performance did not significantly change using different, random hyperparameters of the VRMs, although we could observe significant performance improvement (e.g. RoboshoolWalker2d) or degradation (e.g. RoboshoolHopper -velocities only) in a few tasks using different haperparameters. Therefore, the representation learning part (VRMs) of our algorithm does not suffer from high sensitivity to hyperparameters. 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. 

G SCALABILITY:



:

SLAC was developed for pixel observations. To compare it with our algorithm, we made some modifications of its implementation (see Appendix A.2.3). Nonetheless, we expect the comparison can demonstrate the effect of the key differences as aforementioned in Section 2.

Figures and Tables:

Figure 1: Diagrams of the proposed algorithm. (a) Overview. (b, c) The generative model and the inference model of a VRM.

Figure 2: Computation graph of the proposed algorithm. (a) The RL controller. (b) The execution phase. (c) The learning phase of a VRM. a: action; z: latent variable; d: RNN state variable; x: raw observation (including reward); Q: state-action value function; V : state value function. A bar on a variable means that it is the actual value from the replay buffer or the environment. Each stochastic variable follows a parameterized diagonal Gaussian distribution.

Figure 3: Learning curves of the classic control tasks. Shaded areas indicate S.E.M..

Figure 5: Learning curves of the sequential target reaching task.

Figure 6: Example tasks showing relationship between average return of the agent and negative ELBO (loss function, dashed) of the keep-learning VRM.

Information of environments we used. Only keep-learning model. In this case, only the keep-learning model was used and smooth-update was applied.• Deterministic model. In this case, the first-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. 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.

Table content: <table><tr><td>Name</td><td colspan="3">dim(X) DOF Maximum steps</td></tr><tr><td>Pendulum</td><td>3</td><td>1</td><td>200</td></tr><tr><td>Pendulum (velocities only)</td><td>1</td><td>1</td><td>200</td></tr><tr><td>Pendulum (no velocities)</td><td>2</td><td>1</td><td>200</td></tr><tr><td>CartPole</td><td>4</td><td>1</td><td>1,000</td></tr><tr><td>CartPole (velocities only)</td><td>2</td><td>1</td><td>1,000</td></tr><tr><td>CartPole (no velocities)</td><td>2</td><td>1</td><td>1,000</td></tr><tr><td>RoboschoolHopper</td><td>15</td><td>3</td><td>1,000</td></tr><tr><td>RoboschoolHopper (velocities only)</td><td>6</td><td>3</td><td>1,000</td></tr><tr><td>RoboschoolHopper (no velocities)</td><td>9</td><td>3</td><td>1,000</td></tr><tr><td>RoboschoolWalker2d</td><td>22</td><td>6</td><td>1,000</td></tr><tr><td colspan="2">RoboschoolWalker2d (velocities only) 9</td><td>6</td><td>1,000</td></tr><tr><td>RoboschoolWalker2d (no velocities)</td><td>13</td><td>6</td><td>1,000</td></tr><tr><td>RoboschoolAnt</td><td>28</td><td>8</td><td>1,000</td></tr><tr><td>RoboschoolAnt (velocities only)</td><td>11</td><td>8</td><td>1,000</td></tr><tr><td>RoboschoolAnt (no velocities)</td><td>17</td><td>8</td><td>1,000</td></tr><tr><td>Sequential goal reaching task</td><td>12</td><td>2</td><td>128</td></tr><tr><td>•</td><td/><td/><td/></tr></table>

showed scalability of our algorithm and the alternative ones.

Table content: <table><tr><td>Algorithm</td><td colspan="2">wall-clock time (100,000 steps) # parameters</td></tr><tr><td>Ours</td><td>8 hr</td><td>2.8M</td></tr><tr><td>SAC-MLP</td><td>1 hr</td><td>0.4M</td></tr><tr><td>SAC-LSTM</td><td>12 hr</td><td>1.1M</td></tr><tr><td>SLAC</td><td>5 hr</td><td>2.8M</td></tr></table>

Wall-clock time and number of parameters of our algorithm and the alternative ones. The working environment was a desktop computer using Intel i7-6850K CPU and the task is "Velocitiesonly RoboschoolHopper". The wall-clock time include training the first-impression VRM or pretrainings.

Table content: <table/>

Peer review (iclr2020):
The paper proposes that for partially observable reinforcement learning tasks it might be simpler to decompose the problem in two parts: a recurrent world model and a feedforward agent, as opposed to using just a recurrent agent. Intuitively the decomposition makes sense, though it*s not very clear to me that the problems encountered training a recurrent agent should be dramatically simpler when training a recurrent world model instead. For the world model the paper proposes using a variational recurrent state-transition model, which is essentially a VRNN which also conditions on the actions. It is not clear to me that this is easier to learn than a recurrent agent because the probability distribution of the data used to train the modified VRNN is highly dependent on the current policy in the same way that the non-iid training data makes it tricky for the RNN policy models to converge. While the experimental results seem to show the new algorithm outperforming the existing ones over a large set of random seeds the paper does not specify how the fixed hyperparameters for all models were chosen, leaving open the possibility that different hyperparameter settings would have shown reversals in the experimental results. It*s also not clear that the complexity of the alternate models was adequately accounted for (specially since two VRNNs are required to match the performance of a single RNN agent). That said the incompleteness of the experimental results is my only reservation against accepting this paper. ---- The author response somewhat addressed my major concern, which was the lack of hyperparameter tuning in the baselines, as well as clarifying some of the other questions I had. Given that my score has been updated.