{"title": "Variational Temporal Abstraction", "book": "Advances in Neural Information Processing Systems", "page_first": 11570, "page_last": 11579, "abstract": "We introduce a variational approach to learning and inference of temporally hierarchical structure and representation for sequential data. We propose the Variational Temporal Abstraction (VTA), a hierarchical recurrent state space model that can infer the latent temporal structure and thus perform the stochastic state transition hierarchically. We also propose to apply this model to implement the jumpy imagination ability in imagination-augmented agent-learning in order to improve the efficiency of the imagination. In experiments, we demonstrate that our proposed method can model 2D and 3D visual sequence datasets with interpretable temporal structure discovery and that its application to jumpy imagination enables more efficient agent-learning in a 3D navigation task.", "full_text": "Variational Temporal Abstraction\n\nTaesup Kim1,3,\u2020, Sungjin Ahn2\u2217, Yoshua Bengio1\u2217\n\n1Mila, Universit\u00e9 de Montr\u00e9al, 2Rutgers University, 3Kakao Brain\n\nAbstract\n\nWe introduce a variational approach to learning and inference of temporally hierar-\nchical structure and representation for sequential data. We propose the Variational\nTemporal Abstraction (VTA), a hierarchical recurrent state space model that can\ninfer the latent temporal structure and thus perform the stochastic state transition\nhierarchically. We also propose to apply this model to implement the jumpy imag-\nination ability in imagination-augmented agent-learning in order to improve the\nef\ufb01ciency of the imagination. In experiments, we demonstrate that our proposed\nmethod can model 2D and 3D visual sequence datasets with interpretable temporal\nstructure discovery and that its application to jumpy imagination enables more\nef\ufb01cient agent-learning in a 3D navigation task.\n\n1\n\nIntroduction\n\nDiscovering temporally hierarchical structure and representation in sequential data is the key to many\nproblems in machine learning. In particular, for an intelligent agent exploring an environment, it\nis critical to learn such spatio-temporal structure hierarchically because it can, for instance, enable\nef\ufb01cient option-learning and jumpy future imagination, abilities critical to resolving the sample\nef\ufb01ciency problem (Hamrick, 2019). Without such temporal abstraction, imagination would easily\nbecome inef\ufb01cient; imagine a person planning one-hour driving from her of\ufb01ce to home with future\nimagination at the scale of every second. It is also biologically evidenced that future imagination is\nthe very fundamental function of the human brain (Mullally & Maguire, 2014; Buckner, 2010) which\nis believed to be implemented via hierarchical coding of the grid cells (Wei et al., 2015).\nThere have been approaches to learn such hierarchical structure in sequences such as the HM-\nRNN (Chung et al., 2016). However, as a deterministic model, it has the main limitation that it\ncannot capture the stochastic nature prevailing in the data. In particular, this is a critical limitation\nto imagination-augmented agents because exploring various possible futures according to the un-\ncertainty is what makes the imagination meaningful in many cases. There have been also many\nprobabilistic sequence models that can deal with such stochastic nature in the sequential data (Chung\net al., 2015; Krishnan et al., 2017; Fraccaro et al., 2017). However, unlike HMRNN, these models\ncannot automatically discover the temporal structure in the data.\nIn this paper, we propose the Hierarchical Recurrent State Space Model (HRSSM) that combines the\nadvantages of both worlds: it can discover the latent temporal structure (e.g., subsequences) while also\nmodeling its stochastic state transitions hierarchically. For its learning and inference, we introduce a\nvariational approximate inference approach to deal with the intractability of the true posterior. We also\npropose to apply the HRSSM to implement ef\ufb01cient jumpy imagination for imagination-augmented\nagents. We note that the proposed HRSSM is a generic generative sequence model that is not tied to\nthe speci\ufb01c application to the imagination-augmented agent but can be applied to any sequential data.\nIn experiments, on 2D bouncing balls and 3D maze exploration, we show that the proposed model\n\n\u2217Equal advising, \u2020work also done while visiting Rutgers University.\n\nCorrespondence to taesup.kim@umontreal.ca and sungjin.ahn@rutgers.edu\n\n33rd Conference on Neural Information Processing Systems (NeurIPS 2019), Vancouver, Canada.\n\n\fcan model sequential data with interpretable temporal abstraction discovery. Then, we show that the\nmodel can be applied to improve the ef\ufb01ciency of imagination-augmented agent-learning.\nThe main contributions of the paper are:\n\n1. We propose the Hierarchical Recurrent State Space Model (HRSSM) that is the \ufb01rst stochas-\n\ntic sequence model that discovers the temporal abstraction structure.\n\n2. We propose the application of HRSSM to imagination-augmented agent so that it can\n\nperform ef\ufb01cient jumpy future imagination.\n\n3. In experiments, we showcase the temporal structure discovery and the bene\ufb01t of HRSSM\n\nfor agent learning.\n\n2 Proposed Model\n\n2.1 Hierarchical Recurrent State Space Models\n\n1:li\n\nhas length li such that T = (cid:80)T\n\nIn our model, we assume that a sequence X = x1:T = (x1, . . . , xT ) has a latent structure of temporal\nabstraction that can partition the sequence into N non-overlapping subsequences X = (X1, . . . , XN ).\ni=1 li and L = {li}. Unlike previous\nA subsequence Xi = xi\nworks (Serban et al., 2017), we treat the number of subsequences N and the lengths of subsequences\nL as discrete latent variables rather than given parameters. This makes our model discover the\nunderlying temporal structure adaptively and stochastically.\nWe also assume that a subsequence Xi is generated from a temporal abstraction zi and an observation\nxt has observation abstraction st. The temporal abstraction and observation abstraction have\na hierarchical structure in such a way that all observations in Xi are governed by the temporal\nabstraction zi in addition to the local observation abstraction st. As a temporal model, the two\nabstractions take temporal transitions. The transition of temporal abstraction occurs only at the\nsubsequence scale while the observation transition is performed at every time step. This generative\nprocess can then be written as follows:\n\np(X, S, L, Z, N ) = p(N )\n\np(Xi, Si|zi, li)p(li|zi)p(zi|z<i)\n\nwhere S = {si\n\nj} and Z = {zi} and the subsequence joint distribution p(Xi, Si|zi, li) is:\n\np(Xi, Si|zi, li) =\n\np(xi\n\nj|si\n\nj)p(si\n\nj|si\n\n<j, zi).\n\nN(cid:89)\nli(cid:89)\n\ni=1\n\n(1)\n\n(2)\n\nj=1\n\nWe note that it is also possible to use the traditional Markovian state space model in Eqn. (1) and\nEqn. (2) which has some desirable properties such as modularity and interpretability as well as\nhaving a closed-form solution for a limited class of models like the linear Gaussian model. However,\nit is known that this Markovian model has dif\ufb01culties in practice in capturing complex long-term\ndependencies (Auger-M\u00e9th\u00e9 et al., 2016). Thus, in our model, we take the recurrent state space\nmodel (RSSM) approach (Zheng et al., 2017; Buesing et al., 2018; Hafner et al., 2018b) which\nresolves this problem by adding a deterministic RNN path that can effectively encode the complex\nnonlinear long-term dependencies in the past, i.e., z<i and si\n<j in our model. Speci\ufb01cally, the\ntransition is performed by the following updates: ci = fz-rnn(zi\u22121, ci\u22121), zi \u223c p(zi|ci) for zi, and\nhi\nj = fs-rnn(si\n\nj) for si\nj.\n\nj|hi\n\nj\u22121||zi, hi\n\nj \u223c p(si\n2.2 Binary Subsequence Indicator\n\nj\u22121), si\n\nAlthough the above modeling intuitively explains the actual generation process, the discrete latent\nrandom variables N and {li}\u2014whose realization is an integer\u2014raise dif\ufb01culties in learning and\ninference. To alleviate this problem, we reformulate the model by replacing the integer latent variables\nby a sequence of binary random variables M = m1:T , called the boundary indicator. As the name\nimplies, the role of this binary variable is to indicate whether a new subsequence should start at the\nnext time step or not. In other words, it speci\ufb01es the end of a subsequence. This is a similar operation\nto the FLUSH operation in the HMRNN model (Chung et al., 2016). With the binary indicators, the\ngenerative process can be rewritten as follows:\n\np(X, Z, S, M ) =\n\np(xt|st)p(mt|st)p(st|s<t, zt, mt\u22121)p(zt|z<t, mt\u22121)\n\nT(cid:89)\n\nt=1\n\n2\n\n\fFigure 1: Sequence generative procedure (recurrent deterministic paths are excluded). Left: The\nmodel with the boundary indicators M = {0, 1, 0, 0}. Right: The corresponding generative proce-\ndure with a temporal structure derived from the boundary indicators M\n\nIn this representation of the generative process, we can remove the subsequence hierarchy and make\nboth transitions perform at every time step. Although this seemingly looks different to our original\ngeneration process, the control of the binary indicator\u2014selecting either COPY or UPDATE\u2014can\nmake this equivalent to the original generation process, which we explain later in more detail. In\nFigure 1, we provide an illustration on how the binary indicators induce an equivalent structure\nrepresented by the discrete random variables N and L.\nThis reformulation has the following advantages. First, we do not need to treat the two different\ntypes of discrete random variables N and L separately but instead can unify them by using only\none type of random variables M. Second, we do not need to deal with the variable range of N and\nL because each time step has \ufb01nite states {0, 1} while N and L depend on T that can be changed\nacross sequences. Lastly, the decision can be made adaptively while observing the progress of the\nsubsequence, instead of making a decision governing the whole subsequence.\n\n2.3 Prior on Temporal Structure\nWe model the binary indicator p(mt|st) by a Bernoulli distribution parameterized by \u03c3(fm-mlp(st))\nwith a multi-layer perceptron (MLP) fm-mlp and a sigmoid function \u03c3. In addition, it is convenient to\nexplicitly express our prior knowledge or constraint on the temporal structure using the boundary\ndistribution. For instance, it is convenient to specify the maximum number of subsequences Nmax or\nthe longest subsequence lengths lmax when we do not want too many or too long subsequences. To\nimplement, at each time step t, we can compute the number of subsequences discovered so far by\nusing a counter n(m<t) as well as the length of current subsequence with another counter l(m<t).\nBased on this, we can design the boundary distribution with our prior knowledge as follows:\n\n\uf8f1\uf8f2\uf8f30\n\nif n (m<t) \u2265 Nmax,\nelseif l (m<t) \u2265 lmax,\n\n1\n\u03c3 (fm-mlp (st))\n\notherwise.\n\np (mt = 1|st) =\n\n2.4 Hierarchical Transitions\n\nThe transition of temporal abstraction should occur only a subsequence is completed. This timing is\nindicated by the boundary indicator. Speci\ufb01cally, the transition of temporal abstraction is implemented\nas follows:\n\n(cid:26)\u03b4(zt = zt\u22121)\n\n\u02dcp(zt|ct)\n\nif mt\u22121 = 0 (COPY),\notherwise (UPDATE)\n\nwhere ct is the following RNN encoding of all previous temporal abstractions z<t (and m<t):\n\np (zt|z<t, mt\u22121) =\n\n(cid:26)ct\u22121\n\nct =\n\nif mt\u22121 = 0 (COPY),\nfz-rnn (zt\u22121, ct\u22121) otherwise (UPDATE).\n\nSpeci\ufb01cally, having mt\u22121 = 0 indicates that the time step t is still in the same subsequence as\nthe previous time step t \u2212 1 and thus the temporal abstraction should not be updated but copied.\nOtherwise, it indicates that the time step t \u2212 1 was the end of the previous subsequence and thus the\n\n3\n\n\f(a) Temporal abstraction\n\n(b) Observation abstraction\n\nFigure 2: State transitions: inference and generation with a given hierarchical temporal structure\nbased on the boundary indicators M.\n\ntemporal abstraction should be updated. This transition is implemented as a Gaussian distribution\nN (zt|\u00b5z(ct), \u03c3z(ct)) where both \u00b5z and \u03c3z are implemented with MLPs.\nAt test time, we can use this transition of temporal abstraction without the COPY mode, i.e., every\ntransition is UPDATE. This implements the jumpy future imagination which do not require to rollout\nat every raw time step and thus is computationally ef\ufb01cient.\nThe observation transition is similar to the transition of temporal abstraction except that we want\nto implement the fact that given the temporal abstraction zi, a subsequence is independent of other\nsubsequences. The observation transition is implemented as follows:\np(st|s<t, zt, mt\u22121) = \u02dcp(st|ht) where ht =\nHere, ht is computed by using an RNN fs-rnn to update (UPDATE), and a MLP fs-mlp to initialize\n(INIT). The concatenation is denoted by (cid:107). Note that if the subsequence is \ufb01nished, i.e., mt\u22121 = 1,\nwe sample a new observational abstraction st without conditioning on ht. That is, the underlying\nRNN is initialized.\n\n(cid:26)fs-rnn (st\u22121(cid:107)zt, ht\u22121)\n\nif mt\u22121 = 0 (UPDATE),\notherwise (INIT).\n\nfs-mlp (zt)\n\n3 Learning and Inference\n\nAs the true posterior is intractable, we apply variational approximation which gives the following\nevidence lower bound (ELBO) objective:\n\nlog p (X) \u2265(cid:88)\n\n(cid:90)\n\nq\u03c6 (Z, S, M|X) log\n\np\u03b8 (X, Z, S, M )\nq\u03c6 (Z, S, M|X)\n\ndZdS\n\nM\n\nZ,S\n\nThis is optimized w.r.t. \u03b8 and \u03c6 using the reparameterization trick (Kingma & Welling, 2014). In\nparticular, we use the Gumbel-softmax (Jang et al., 2017; Maddison et al., 2017) with straight-through\nestimators (Bengio et al., 2013) for the discrete variables M. For the approximate posterior, we use\nthe following factorization:\n\nq\u03c6(Z, S, M|X) = q(M|X)q(Z|M, X)q(S|Z, M, X).\n\nThat is, by sequence decomposition q(M|X), we \ufb01rst infer the boundary indicators independent of\nZ and S. Then, given the discovered boundary structure, we infer the two abstractions via the state\ninference q(Z|M, X) and q(S|Z, M, X).\nSequence Decomposition. Inferring the subsequence structure is important because the other state\ninference can be decomposed into independent subsequences. This sequence decomposition is\nimplemented by the following decomposition:\n\nq(M|X) =\n\nq(mt|X) =\n\nBern (mt|\u03c3(\u03d5(X))) ,\n\nT(cid:89)\n\nT(cid:89)\n\nt=1\n\nt=1\n\n4\n\n\fwhere \u03d5 is a convolutional neural network (CNN) applying convolutions over the temporal axis to\nextract dependencies between neighboring observations. This enables to sample all indicators M\nindependently and simultaneously. Empirically, we found this CNN-based architecture working\nbetter than an RNN-based architecture.\nState Inference. State inference is also performed hierarchically. The temporal abstraction predictor\nt=1 q(zt|M, X) does inference by encoding subsequences determined by M and\nX. To use the same temporal abstraction across the time steps of a subsequence, the distribution\nq(zt|M, X) is also conditioned on the boundary indicator mt\u22121:\n\nq(Z|M, X) =(cid:81)T\n\n(cid:26)\u03b4(zt = zt\u22121)\n\nt\u22121, \u03c8bwd\n\nt\n\nt\n\nt\u22121, \u03c8bwd\n\nt\n\n\u02dcq(zt|\u03c8fwd\n\nif mt\u22121 = 0 (COPY),\n) otherwise (UPDATE).\n\nq(zt|M, X) =\nWe use the distribution \u02dcq(zt|\u03c8fwd\n) to update the state zt. It is conditioned on all previous\nt\u22121 extracted from a forward RNN \u03c8fwd (X).\nobservations x<t and this is represented by a feature \u03c8fwd\nThe other is a feature \u03c8bwd\nrepresenting the current step\u2019s subsequence that is extracted from a\nbackward (masked) RNN \u03c8bwd (X, M ). In particular, this RNN depends on M, which is used as a\nmasking variable, to ensure independence between subsequences.\n\nThe observation abstraction predictor q(S|Z, M, X) =(cid:81)T\nobservational abstraction st is sampled from the distribution q(st|zt, M, X) = \u02dcq(cid:0)st|zt, \u03c6fwd\n\nt=1 q(st|zt, M, X) is factorized and each\n\n(cid:1). The\n\nfeature \u03c6fwd\nsequence X and resets hidden states when a new subsequence starts.\n\nis extracted from a forward (masked) RNN \u03c6fwd(X, M ) that encodes the observation\n\nt\n\nt\n\n4 Related Works\n\nThe most similar work with our model is the HMRNN (Chung et al., 2016). While it is similar in\nthe sense that both models discover the hierarchical temporal structure, HMRNN is a deterministic\nmodel and thus has a severe limitation to use for an imagination module. In the switching state-space\nmodel (Ghahramani & Hinton, 2000), the upper layer is a Hidden Markov Model (HMM) and the\nbehavior of the lower layer is modulated by the discrete state of the upper layer, and thus gives\nhierarchical temporal structure. Linderman et al. (2016) proposed a new class of switching state-space\nmodels that discovers the dynamical units and also explains the switching behavior depending on\nobservations or continuous latent states. The authors used inference based on message-passing. The\nhidden semi-Markov models (Yu, 2010; Dai et al., 2016) perform similar segmentation with discrete\nstates. However, unlike our model, there is no states for temporal abstraction. Kipf et al. (2018)\nproposed soft-segmentation of sequence for compositional imitation learning.\nThe variational recurrent neural networks (VRNN) (Chung et al., 2015) is a latent variable RNN\nbut uses auto-regressive state transition taking inputs from the observation. Thus, this can be\ncomputationally expensive to use as an imagination module. Also, the error can accumulate more\nseverely in the high dimensional rollout. To resolve this problem, Krishnan et al. (2017) and Buesing\net al. (2018) proposes to combine the traditional Markovian State Space Models with deep neural\nnetworks. Zheng et al. (2017) and Hafner et al. (2018a) proposed to use an RNN path to encode the\npast making non-Markovian state-space models which can alleviate the limitation of the traditional\nSSMs. Serban et al. (2017) proposed a hierarchical version of VRNN called Variational Hierarchical\nRecurrent Encoder-Decoder (VHRED) which results in a similar model as ours. However, it is a\nsigni\ufb01cant difference that our model learns the segment while VHRED uses a given structure. A\nclosely related work is TDVAE (Gregor et al., 2019). TDVAE is trained on pairs of temporally\nseparated time points. Jayaraman et al. (2019) and Neitz et al. (2018) proposed models that predict\nthe future frames that, unlike our approach, have the lowest uncertainty. The resulting models predict\na small number of easily predictable \u201cbottleneck\u201d frames through which any possible prediction\nmust pass. Pertsch et al. (2019) proposed to predict the keyframes with their temporal offsets using\nstochastic prediction and deterministically interpolate the intermediate frames.\n\n5 Experiments\n\nWe demonstrate our model on visual sequence datasets to show (1) how sequence data is decomposed\ninto perceptually plausible subsequences without any supervision, (2) how jumpy future prediction is\n\n5\n\n\fFigure 3: Left: Previously observed (context) data Xctx. Right: Each \ufb01rst row is the input observation\nsequence X and the second row is the corresponding reconstruction. The sequence decomposer\nq(M|X) predicts the starting frames of subsequences and it is indicated by arrows (red squared\nframes). Subsequences are newly de\ufb01ned when a ball hits the wall by changing the color.\n\ndone with temporal abstraction and (3) how this jumpy future prediction can improve the planning as\nan imagination module in a navigation problem. Moreover, we test conditional generation p(X|Xctx)\nwhere Xctx = x\u2212(Tctx\u22121):0 is the context observation of length Tctx. With the context, we preset the\nstate transition of the temporal abstraction by deterministically initializing c0 = fctx (Xctx) with\nfctx implemented by a forward RNN. The code of the implementation of our model is available at\nhttps://github.com/taesupkim/vta.\n5.1 Bouncing Balls\n\nWe generated a synthetic 2D visual sequence dataset called bouncing balls. The dataset is composed\nof two colored balls that are designed to bounce in hitting the walls of a square box. Each ball is\nindependently characterized with certain rules: (1) The color of each ball is randomly changed when\nit hits a wall and (2) the velocity (2D vector) is also slightly changed at every time steps with a\nsmall amount of noise. We trained a model to learn 1D state representations and all observation\ndata xt \u2208 R32\u00d732\u00d73 are encoded and decoded by convolutional neural networks. During training,\nthe length of observation sequence data X is set to T = 20 and the context length is Tctx = 5.\nHyper-parameters related to sequence decomposition are set as Nmax = 5 and lmax = 10.\nOur results in Figure 3 show that the sequence decomposer q(M|X) predicts reasonable subsequences\nby setting a new subsequence when the color of balls is changed or the ball is bounced. At the\nbeginning of training, the sequence decomposer is unstable with having large entropy and tends to\nde\ufb01ne subsequences with a small number of frames. It then began to learn to increase the length of\nsubsequences and this is controlled by annealing the temperature \u03c4 of Gumbel-softmax towards small\nvalues from 1.0 to 0.1. However, without our proposed prior on temporal structure, the sequence\ndecomposer fails to properly decompose sequences and our proposed model consequently converges\ninto RSSM.\n5.2 Navigation in 3D Maze\n\nAnother sequence dataset is generated from the 3D maze environment by an agent that navigates\nthe maze. Each observation data xt \u2208 R32\u00d732\u00d73 is de\ufb01ned as a partially observed view observed\nby the agent. The maze consists of hallways with colored walls and is de\ufb01ned on a 26 \u00d7 18 grid\nmap as shown in Figure 5. The agent is set to navigate around this environment and the viewpoint\nof the agent is constantly jittered with some noise. We set some constraints on the agent\u2019s action\n(forward, left-turn, right-turn) that the agent is not allowed to turn around when it is located on\nthe hallway. However, it can turn around when it arrives nearby intersections between hallways.\nDue to these constraints, the agent without a policy can randomly navigate the maze environment\nand collect meaningful data. To train an environment model, we collected 1M steps (frames)\nfrom the randomly navigating agent and used it to train both RSSM and our proposed HRSSM. For\nHRSSM, we used the same training setting as bouncing balls but different Nmax = 5 and lmax = 8\nfor the sequence decomposition. The corresponding learning curves are shown in Figure 4 that both\nreached a similar ELBO. This suggests that our model does not lose the reconstruction performance\n\n6\n\n\fFigure 4: The learning curve of RSSM and HRSSM: ELBO, KL-divergence and reconstruction loss.\n\nFigure 5: Left: Bird\u2019s-eye view (map) of the 3D maze with generated navigation paths. White\ndotted lines indicate the given context path Xctx and the corresponding frames are depicted below\nthe map. Solid lines are the generated paths (blue: top, red: bottom) conditioned on the same\ncontext. Circles are the starting points of subsequences where the temporal abstract transitions\nexactly take place. Right: The generated sequence data is shown with its temporal structure. Both\ngenerations are conditioned on the same context but different input actions as indicated. Frame\nsamples on each bottom row are generated with the temporal abstract transition \u02dcp (zt(cid:48)|ct(cid:48)) with\nct(cid:48) = fz-rnn (zt(cid:48)\u22121, ct(cid:48)\u22121) and this shows how the jumpy future prediction is done. Other samples on\ntop rows, which are not necessarily required for future prediction with our proposed HRSSM, are\ngenerated from the observation abstraction transition \u02dcp (st|ht) with ht = fs-rnn (st\u22121(cid:107)zt, ht\u22121). The\nboundaries between subsequences are determined by p (mt|st).\n\nwhile discovering the hierarchical structure. We trained state transitions to be action-conditioned\nand therefore this allows to perform action-controlled imagination. For HRSSM, only the temporal\nabstraction state transition is action-conditioned as we aim to execute the imagination only with\nthe jumpy future prediction. The overall sequence generation procedure is described in Figure 5.\nThe temporal structure of the generated sequence shows how the jumpy future prediction works and\nwhere the transitions of temporal abstraction occur. We see that our model learns to set each hallway\nas a subsequence and consequently to perform jumpy transitions between hallways without repeating\nor skipping a hallway. In Figure 6, a set of jumpy predicted sequences from the same context Xctx\nand different input actions are shown and this can be interpreted as imaginations the agent can use for\nplanning.\nGoal-Oriented Navigation We further use the trained model as an imagination module by aug-\nmenting it to an agent to perform the goal-oriented navigation. In this experiment, the task is to\nnavigate to a randomly selected goal position within the given life steps. The goal position in the\ngrid map is not provided to the agent, but a 3 \u00d7 3-cropped image around the goal position is given.\nTo reach the goal fast, the agent is augmented with the imagination model and allowed to execute\na rollout over a number of imagination trajectories (i.e., a sequence of temporal abstractions) by\nvarying the input actions. Afterward, it decides the best trajectory that helps to reach the goal faster.\nTo \ufb01nd the best trajectory, we use a simple strategy: a cosine-similarity based matching between\nall imagined state representations in imaginations and the feature of the goal image. The feature\nextractor for the goal image is jointly trained with the model. 2 This way, at every time step we let\n2During training, the 3\u00d73 window (image) around the agent position is always given as additional observation\ndata and we trained feature extractor by maximizing the cosine-similarity between the extracted feature and the\ncorresponding time step state representation.\n\n7\n\n\fFigure 6: Jumpy future prediction conditioned on the same context Xctx and different input actions\n\nFigure 7: Fully generate seqences conditioned on the same context Xctx and same input actions:\ngenerated paths are equal but the viewpoint and the lengths of subsequences are varied (red squared\nframes are jumpy future predictions).\n\nthe agent choose the \ufb01rst action resulting in the best trajectory. This approach can be considered as a\nsimple variant of the Monte Carlo Tree Search (MCTS) and the detailed overall procedure can be\nfound in Appendix. Each episode is de\ufb01ned by randomly initializing the agent position and the goal\nposition. The agent is allowed maximum 100 steps to reach the goal and the \ufb01nal reward is de\ufb01ned\nas the number of remaining steps when the agent reaches the goal or consumes all life-steps. The\nperformance highly depends on the accuracy and the computationally ef\ufb01ciency of the model and we\ntherefore compare between RSSM and HRSSM with varying the length of imagined trajectories. We\nmeasure the performance by randomly generated 5000 episodes and show how each setting performs\nacross the episodes by plotting the reward distribution in Figure 8. It is shown that the HRSSM signif-\nicantly improves the performance compared to the RSSM by having the same computational budget.\nHRSSM showed consistent performance\nover different lengths of imagined trajecto-\nries and most episodes were solved within\n50 steps. We believe that this is because\nHRSSM is able to abstract multiple time\nsteps within a single state transition and\nthis enables to reduce the computational\ncost for imaginations. The results also\nshow that \ufb01nding the best trajectory be-\ncomes dif\ufb01cult as the imagination length\ngets larger, i.e., the number of possible\nimagination trajectories increases. This\nsuggests that imaginations with temporal\nabstraction can bene\ufb01t both the accuracy\nand the computationally ef\ufb01ciency in effec-\ntive ways.\n\nFigure 8: Goal-oriented navigation with different\nlengths of imagined trajectories.\n\n6 Conclusion\n\nIn this paper, we introduce the Variational Temporal Abstraction (VTA), a generic generative temporal\nmodel that can discover the hierarchical temporal structure and its stochastic hierarchical state transi-\ntions. We also propose to use this temporal abstraction for temporally-extended future imagination\nin imagination-augmented agent-learning. Experiment results shows that in general sequential data\nmodeling, the proposed model discovers plausible latent temporal structures and perform hierarchical\nstochastic state transitions. Also, in connection to the model-based imagination-augmented agent for\na 3D navigation task, we demonstrate the potential of the proposed model in improving the ef\ufb01ciency\nof agent-learning.\n\n8\n\n\fAcknowledgments\n\nWe would like to acknowledge Kakao Brain cloud team for providing computing resources used in this\nwork. TK would also like to thank colleagues at Mila, Kakao Brain, and Rutgers Machine Learning\nGroup. SA is grateful to Kakao Brain, the Center for Super Intelligence (CSI), and Element AI for\ntheir support. Mila (TK and YB) would also like to thank NSERC, CIFAR, Google, Samsung, Nuance,\nIBM, Canada Research Chairs, Canada Graduate Scholarship Program, and Compute Canada.\n\nReferences\nAuger-M\u00e9th\u00e9, M., Field, C., Albertsen, C. M., Derocher, A. E., Lewis, M. A., Jonsen, I. D., and\nFlemming, J. M. State-space models\u2019 dirty little secrets: even simple linear gaussian models can\nhave estimation problems. Scienti\ufb01c reports, 6:26677, 2016.\n\nBengio, Y., L\u00e9onard, N., and Courville, A. C. Estimating or propagating gradients through stochastic\nneurons for conditional computation. CoRR, abs/1308.3432, 2013. URL http://arxiv.org/\nabs/1308.3432.\n\nBuckner, R. L. The role of the hippocampus in prediction and imagination. Annual review of\n\npsychology, 61:27\u201348, 2010.\n\nBuesing, L., Weber, T., Racaniere, S., Eslami, S., Rezende, D., Reichert, D. P., Viola, F., Besse, F.,\nGregor, K., Hassabis, D., et al. Learning and querying fast generative models for reinforcement\nlearning. arXiv preprint arXiv:1802.03006, 2018.\n\nChung, J., Kastner, K., Dinh, L., Goel, K., Courville, A. C., and Bengio, Y. A recurrent latent variable\nmodel for sequential data. In Advances in neural information processing systems, pp. 2980\u20132988,\n2015.\n\nChung, J., Ahn, S., and Bengio, Y. Hierarchical multiscale recurrent neural networks. arXiv preprint\n\narXiv:1609.01704, 2016.\n\nDai, H., Dai, B., Zhang, Y.-M., Li, S., and Song, L. Recurrent hidden semi-markov model. 2016.\n\nFraccaro, M., Kamronn, S., Paquet, U., and Winther, O. A disentangled recognition and nonlinear\ndynamics model for unsupervised learning. In Advances in Neural Information Processing Systems,\npp. 3601\u20133610, 2017.\n\nGhahramani, Z. and Hinton, G. E. Variational learning for switching state-space models. Neural\n\ncomputation, 12(4):831\u2013864, 2000.\n\nGregor, K., Papamakarios, G., Besse, F., Buesing, L., and Weber, T. Temporal difference variational\n\nauto-encoder. In International Conference on Learning Representations, 2019.\n\nHafner, D., Lillicrap, T., Fischer, I., Villegas, R., Ha, D., Lee, H., and Davidson, J. Learning latent\n\ndynamics for planning from pixels. arXiv preprint arXiv:1811.04551, 2018a.\n\nHafner, D., Lillicrap, T. P., Fischer, I., Villegas, R., Ha, D., Lee, H., and Davidson, J. Learning latent\n\ndynamics for planning from pixels. CoRR, abs/1811.04551, 2018b.\n\nHamrick, J. B. Analogues of mental simulation and imagination in deep learning. Current Opinion\nin Behavioral Sciences, 29:8 \u2013 16, 2019. ISSN 2352-1546. SI: 29: Arti\ufb01cial Intelligence (2019).\n\nJang, E., Gu, S., and Poole, B. Categorical reparameterization with gumbel-softmax. In International\n\nConference on Learning Representations, 2017.\n\nJayaraman, D., Ebert, F., Efros, A., and Levine, S. Time-agnostic prediction: Predicting predictable\nvideo frames. In International Conference on Learning Representations, 2019. URL https:\n//openreview.net/forum?id=SyzVb3CcFX.\n\nKingma, D. P. and Welling, M. Auto-encoding variational bayes. In 2nd International Conference on\nLearning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track\nProceedings, 2014.\n\n9\n\n\fKipf, T., Li, Y., Dai, H., Zambaldi, V. F., Grefenstette, E., Kohli, P., and Battaglia, P. Compositional\nimitation learning: Explaining and executing one task at a time. CoRR, abs/1812.01483, 2018.\nURL http://arxiv.org/abs/1812.01483.\n\nKrishnan, R. G., Shalit, U., and Sontag, D. Structured inference networks for nonlinear state space\n\nmodels. In Thirty-First AAAI Conference on Arti\ufb01cial Intelligence, 2017.\n\nLinderman, S. W., Miller, A. C., Adams, R. P., Blei, D. M., Paninski, L., and Johnson, M. J. Recurrent\n\nswitching linear dynamical systems. arXiv preprint arXiv:1610.08466, 2016.\n\nMaddison, C. J., Mnih, A., and Teh, Y. W. The concrete distribution: A continuous relaxation of\n\ndiscrete random variables. In International Conference on Learning Representations, 2017.\n\nMullally, S. L. and Maguire, E. A. Memory, imagination, and predicting the future: a common brain\n\nmechanism? The Neuroscientist, 20(3):220\u2013234, 2014.\n\nNeitz, A., Parascandolo, G., Bauer, S., and Sch\u00f6lkopf, B. Adaptive skip intervals: Temporal\nabstraction for recurrent dynamical models. In Bengio, S., Wallach, H., Larochelle, H., Grauman,\nK., Cesa-Bianchi, N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems\n31, pp. 9816\u20139826. 2018.\n\nPertsch, K., Rybkin, O., Yang, J., Derpanis, K. G., Lim, J., Daniilidis, K., and Jaegle, A. Keyin:\nDiscovering subgoal structure with keyframe-based video prediction. CoRR, abs/1904.05869,\n2019. URL http://arxiv.org/abs/1904.05869.\n\nSerban, I. V., Sordoni, A., Lowe, R., Charlin, L., Pineau, J., Courville, A., and Bengio, Y. A\nhierarchical latent variable encoder-decoder model for generating dialogues. In Thirty-First AAAI\nConference on Arti\ufb01cial Intelligence, 2017.\n\nWei, X.-X., Prentice, J., and Balasubramanian, V. A principle of economy predicts the functional\n\narchitecture of grid cells. Elife, 4:e08362, 2015.\n\nYu, S.-Z. Hidden semi-markov models. Artif. Intell., 174(2):215\u2013243, February 2010.\n\n0004-3702. doi: 10.1016/j.artint.2009.11.011.\n\nISSN\n\nZheng, X., Zaheer, M., Ahmed, A., Wang, Y., Xing, E. P., and Smola, A. J. State space LSTM models\n\nwith particle MCMC inference. arXiv preprint arXiv:1711.11179, 2017.\n\n10\n\n\f", "award": [], "sourceid": 6171, "authors": [{"given_name": "Taesup", "family_name": "Kim", "institution": "Mila / Kakao Brain"}, {"given_name": "Sungjin", "family_name": "Ahn", "institution": "Rutgers University"}, {"given_name": "Yoshua", "family_name": "Bengio", "institution": "Mila - University of Montreal"}]}