Learning Robust Dynamics through Variational Sparse Gating
Authors: Arnav Kumar Jain
Shiva Kanth Sujit
Shruti Joshi
Vincent Michalski
Danijar Hafner
Samira Ebrahimi Kahou
Venue: NeurIPS 2022 Accept (2022)
Keywords: Deep Reinforcement Learning
Model Based Reinforcement Learning
World Models

Abstract: Learning world models from their sensory inputs enables agents to plan for actions by imagining their future outcomes. World models have previously been shown to improve sample-efficiency in simulated environments with few objects, but have not yet been applied successfully to environments with many objects. In environments with many objects, often only a small number of them are moving or interacting at the same time. In this paper, we investigate integrating this inductive bias of sparse interactions into the latent dynamics of world models trained from pixels. First, we introduce Variational Sparse Gating (VSG), a latent dynamics model that updates its feature dimensions sparsely through stochastic binary gates. Moreover, we propose a simplified architecture Simple Variational Sparse Gating (SVSG) that removes the deterministic pathway of previous models, resulting in a fully stochastic transition function that leverages the VSG mechanism. We evaluate the two model architectures in the BringBackShapes (BBS) environment that features a large number of moving objects and partial observability, demonstrating clear improvements over prior models.

Learning Robust Dynamies through Variational Sparse Gating
Arnay Kumar Jain 1.2,,, Shivakanth Sujitt,,,hhuti Shruti Joshi Vincent Michalski!.. 1.2 Danijar Hafner*. Sammia Ebrahimi-BR 2aahimiiKKhon23..
Abstract
Learning world models from their sensory inputs enables agents to plan for actions by imagining their future outcomes. World models have previously been shown to improve sample-efficiency in simulated environments with few objects, but have not yet been applied successfully to environments with many objects. In environments with many objects, often only small number of them are moving or interacting at the same time. In this paper, we investigate integrating this inductive bias of sparse interactions into the latent dynamies of world models trained from pixels. First, we introduce Variational Sparse Gating (VSG), a latent dynamics model that updates its feature dimensions sparsely through stochastic binary gates. Moreover, we propose simplified architecture Simple Variational Sparse Gating (SVSG) that removes the deterministic pathway of previous models, resulting in a fully stochastic transition function that leverages the VSG mechanism. We evaluate the two model architectures in the BringBackShapes (BBS) environment that features large number of moving objects and partial observability, demonstrating clear improvements over prior models.
Introduction
Latent dynamics models are models that generate agent's future states in the compact latent space without feeding the high-dimensional observations back to the model. They have shown promising results on various tasks like video prediction (Karl et al., 2016; Kalman, 1960; Krishnan et al., 2015), model-based Reinforcement Learning (RL) (Hafner et al., 2020; 2021: 2019; Ha and Schmidhuber, 2018), and robotics (Watter et al., 2015). Generating sequences in the compact latent space reduces the accumulating errors leading more accurate long-term predictions. Additionally, having lower dimensionality leads to a lower memory footprint. Solving tasks in model-based RL involves learning world model (Ha and Schmidhuber, 2018) that can predict outcomes of actions, followed by using them to derive behaviors (Sutton, 1991). Motivated by these benefits, the recently proposed DreamerV1 (Hafner et al., 2020) and DreamerV2 (Hafner al., 2021) agents achieved state-of-the-art results on wide range of visual control tasks.
Many complex tasks require reliable long-term prediction of dynamics. This is true especially in partially observable environments where only subspace is visible the agent, and it is usually required to accurately retain information over multiple time steps to solve the task. The Dreamer agents (Hafner et al., 2020; 2021) employ an Recurrent State-Space Model (RSSM) (Hafner et al., 2019) comprising of a Recurrent Neural Network (RNN). Training RNNs for long sequences is chal- lenging as they suffer from optimization problems like vanishing gradients (Hochreiter, 1991; Bengio et al., 1994). Different ways of applying sparse updates in RNNs have been investigated (Campos et al., 2017; Neil al., 2016; Goyal al., 2019), enabling subset of state dimensions to constant
36th Conference on Neural Information Processing Systems (NeurIPS 2022).
Universitéé de Montréal, "Mila- Quebee Al Institute, "École de technologie supérieure, University of Toronto, "Google Brain, "CIFAR. "Correspondence to Arny Kumar Jain <arnav- <arnav-kumar kumar.jain @ mila.quebeec>.

---

during the update. A sparse update prior can also be motivated by the fact that in the real world, many factors of variation are constant over extended periods of time. For instance, several objects in physical simulation may be stationary until some force acts upon them. Additionally, this is useful in the partially observable setting where the agent observes constrained viewpoint and has to keep track of objects that are not visible for many time steps. In this work, we introduce Variational Sparse Gating (VSG), a stochastic gating mechanism that sparsely updates the latent states at each step.
Recurrent State-Space Model (RSSM) (Hafner et al., 2019) was introduced in PLaNet where the model state was composed of two paths, an image representation path and recurrent path. Dream- erVI (Hafner et al., 2020) and DreamerV2 (Hafner et al., 2021) utilized them to achieve state- of-the-art results in continuous and discrete control tasks (Hafner et al., 2019). While the image representation path which is stochastic accounts for multiple possible future states, the recurrent path is deterministic to retain information over multiple time steps to facilitate gradient-based optimiza- tion. (Hafner et al., 2019) showed that both components were important for solving tasks, where the stochastic part was more important to account for partial observability of the initial states. By leveraging the proposed gating mechanism ( (Variational Sparse Gating (VSG)), we demonstrate that a purely stochastic model with a single component can achieve competitive results, and call it Simple Variational Sparse Gating (SVSG). To the best of our knowledge, this is the first work that shows that purely stochastic models achieve competitive performance on continuous control tasks when compared to leading agents.
Existing benchmarks (Bellemare et al., 2013; Chevalier-Boisvert et al., 2018; Tassa et al., 2018) for RL do not test the capability of agents in both partial observability and stochasticity. The Atari (Bellemare et al., 2013) benchmark comprises of 55 games but most of the games are deterministic and a lot of compute is required to train on them. Some tasks in the Atari and Minigrid benchmarks are partially- observable but either lack stochasticity or are hard exploration tasks. Also, these benchmarks do not allow for controlling the factors of variation. We developed a new partially-observabll and stochastic environment, called BringBackShapes (BBS), where the task is to push objects to a predefined goal area. Solving tasks in BBS require agents to remember states of previously observed objects and avoid noisy distractor objects. Furthermore, VSG and SVSG outperformed leading model-based and freel-free baselines. We also present studies with varying partial- observabbbby and stochasticity to demonstrate that the proposed agents have better memory for tracking observed objects and are more robust to increasing levels of noise. Lastly, the proposed methods were also evaluated on existing benchmarks DeepMind Control (DMC) (Tassa et al., 201 18,, DMC with Natural Background (Zhang et al., 2021; Nguyen et al., 2021b), and Atari (Bellemare et al., 2013). On the existing benchmarks, the proposed method performed better on tasks with changing viewpoints and sparse rewards.
Our key contributions summarized as follows:
Variational Sparse Gating: We introduce Variational Sparse Gating (VSG), where the recurrent states are sparsely updated through a stochastic gating mechanism. A comprehensive empirical evaluation shows that VSG outperforms baselines on tasks requiring long-term memory.
Simple Variational Sparse Gating: We also propose Simple Variational Sparse Gating (SVSG) which has purely stochastic state, and achieves competitive results on continuous control tasks when compared with agents that also use deterministic component.
BringBackShapes: We developed the BringBack (BBS) environment to evaluate agents on partially-observabbe and stochastic settings where these variations can be controlled. Our experiments show that the proposed agents are more robust to such variations.
2 Variational Sparse Gating
Reinforcement Learning: The visual control task can be formulated as a Partially Observable Markov Decision Process (POMDP) with discrete time steps E [1:T]. The agent seleets action a, ~ p(a, loge, azee) to interact with the environment and receives the next observation and scalar reward oz, rt ~ p(o,, rz (oct, r'ce), respectively, each time step. The goal is to learn a policy that maximizes the expected discounted sum of rewards E, (Σtn 7'r)), where y is the discount factor.
Agent: Agent is composed of a world model and a policy (Fig. 1). World models (Sec. 2.1) encode a sequence of observations and actions into latent representations. The agents behavior (Appendix B) derived to maximize expected returns on the trajectories generated from the learned world model. While training, the world model is learned with collected experience, the policy is improved on

---

A
82
03
C12
VSG
VSG
VSG
VSG
VSG
VSG
$2
X1
(a)
(b)
Figure 1: (a) World Model: The VSG block takes the previous model state 8t-1 and action d;-1; and outputs the updated model state at next step 8t, which is further used to reconstruct image ft and reward Pr. (b) Policy: Comprises of actor to select optimal action â and critic to predict value é beyond the planning horizon. The world model unrolled using the prior model state 8, which does not contain information about image Tz.
trajectories unrolled using the world model and new episodes are collected by deploying the policy in the environment. An initial set of episodes are collected using a random policy. As training progresses, new episodes are collected using the latest policy to further improve the world model.
2.1 World Model
World Models (Ha and Schmidhuber, 2018) learn to mimic the environment using the collected experience and facilitate deriving behaviours in the abstract latent space. Given an abstract state the world and action, the model applies the learned transition dynamics to predict the resulting next state and reward. RSSM (Hafner et al., 2019) was introduced in PlaNet, where the model state was composed of two paths. The recurrent path consists of an RNN (See Figure 2 [a]), and is motivated with reliable long-tem information preservation, while the image representation path samples from a learned distribution account for multiple possible futures (Babaeizadeh al., 2017). In this work, we introduce Variational Sparse Gating (VSG), where the recurrent path selectively updates a subset of the latent states each step using stochastic gating network. Sparse updates enable the agent have long-term memory and learn robust representations to solve complex tasks.
The world model comprises of an image encoder, a VSG model, and predictors for image, discount and reward. The image encoder generates representations o, for the observation 2, using Convolutional Neural Networks (CNNs). The VSG model comprises of a recurrent model equipped with the stochastic gating mechanism get the recurrent state h,, and is used to compute two stochastic image representation states. The posterior representation state 2ς is obtained using the representation model and contains information about the current observation Zt. The prior state & is obtained from the transition predictor without observing the current observation 2,. This is useful while planning as sequences are generated in compact latent state, and the output from the transition predictor is utilized. This also results in lower memory footprint and enables predictions of thousands of trajectories in parallel on a single GPU. The representation states are sampled from known distribution with learned parameters like Gaussian (Hafner et al., 2020) or Categorical (Hafner al., 2021). The concatenation of outputs from the recurrent and image representation models gives the compact model state (s; = [/v, =x)). The posterior model state is further used to reconstruct the original image ét, predict the reward ft, and discount factor ât. The discount factor helps to predict the probability that an episode will end. The components of the world model are as follows:
Model Components:
Recurrent model:
h4 (1 14- a-11)
Representation model:
(z 11,
Transition predictor:
&. ~ p (2t
ht)
(1)
Image predictor:
f. p po (ix
h4, 21)
Reward predictor:
Pt ~ po (f,
h4s zt)
Discount predictor:
9, ~ (fe
h 2x)
3

---

Ianh
10
(h)
Figure 2: Architectures of (a) Recurrent State-Space Model (RSSM). (b) Variational Sparse Gating (VSG), and (c) Simple Variational Sparse Gating (SVSG), respectively. o and tanh denote the sigmoid and tanh non-linear activations, respectively. W. and b. are the corresponding weights and biases. ~, @ and @ denote sampling, vector concatenation, and element-wise multiplication, respectively. MM computes 2; = uyă; + (1 - uz - where = h, is used for RSSM and VSG, and xy = 8t is used for SVSG. B denotes Bernoulli distribution. fp and fq denote the prior and posterior distributions with learned parameters, respectively (See Appendix I for more details).
The representation model outputs posterior image representation state 2t conditioned the image encoding zt and recurrent state h4- The transition predictor provides the prior image representation state &. The image encoding o, is obtained by passing the image zt through CNN (LeCun et al., 1989) and Multi-layer Perceptron (MLP) layers. In VSG, propose to modify the Gated Recurrent Unit (GRU) used in RSSM to sparsely update the recurrent state at each step. The model state St. which passed through several layers of MLP to predict the discount and reward, and transposed CNN layers are used to reconstruct the image. The Exponential Linear Unit (ELU) activation is used for training all the components of the world model (Clevert et al., 2015).
Neural Networks:
concatenation of recurrent and image representation states
Sparse Gating: In light of training RNNs to capture long-term dependencies, different ways of applying sparse updates have been investigated (Campos al., 2017: Neil et al., 2016; Goyal et al., 2019), enabling subset of state dimensions to be constant during the update. They were found to alleviate the vanishing gradient problem by effectively reducing the number of sequential operations (Campos et al., 2017). Discrete gates may also improve long-term memory by avoiding the gradual change of state values introduced by repeated multiplication with continuous gate values in standard recurrent architectures. Previous works on sparsely updating hidden states (Campos et al., 2017; Neil et al., 2016) use separate layer applied over the outputs of RNN, and do not modify the RNN in itself. However, in this work, we modify the update gate in GRU (Cho et al., 2014) to take binary values by sampling from a Bernoulli distribution (Fig. [b] shows the architecture).
The input & to the recurrent model contains information about the action and obtained by concate- nating the previous image representation state 24-1 and action followed by passing them through MLP layer. Similar GRU (Cho al., 2014), there is reset and update gate. The reset gate d decides the extent of information flow from the previous recurrent state and inputs, and the update gate u, tells which parts of the recurrent state will be updated. The update gate takes only binary values, selecting whether the value will be updated or copied from previous time step. Binary values are obtained by sampling from a Bernoulli distribution where the probability of sampling obtained using the previous recurrent state h and input Straight- estimators (Bengio al., 2013) were used for propagating gradients backwards for training. The update equations are:
Vt (W." [ ir] bw)
ūe (W [ 1, bw)
h = tanh(u, (W?(hu-1s bc))
#: Bernoull((ū.)
h 4 h4 u+) h4-1;
where @ denotes element- wise multiplication, and tanh are the sigmoid and hyperbolic tangent activation function, W. and b. denotes the weights and biases, respectively. To control the sparsity of updates, we have used KL divergence between probability of sampling the update gate ū, and fixed prior probability k, where e is a tunable hyperparameter.
Loss function: The predictors for image and reward produces Gaussian distributions with unit variance, whereas the discount predictor predicts Bernoulli likelihood. The image representation

---

states are sampled from Gaussian (Hafner et al.. 2020) or a Categorical (Hafner et al.. 2021) distribution which are trained maximize the likelihood of targets. In addition, there is KL Divergence term between the prior and posterior distributions and similar to DreamerV2 (Hafner et al., 2021), we have also used KI. balancing with a factor of 0.8. We have also added sparsity loss to regularize the number of updates in hidden state at each step. All the components of the world model are optimized jointly using the loss function given by:
L( (()) = Egglsirr (i:r.
Ei 1==
In pop The zx) - In --llpppprr he; =x)
image log loss
reward log loss
In po hes 2)) +B KL go (z (ee xx) || po(2s he)
+ KL Ut
discount log loss
KI. loss
sparsity loss
where 3 and a are the scale for KL losses of the latent codes and the sparse update gates, respectively,
3 Simple Variational Sparse Gating
Stochastic State-Space Model (SSM) were proposed in PLaNet (Hafner et al., 2019), where it was discussed that it is not trivial to achieve competitive results without the deterministic recurrent path. Having deterministic component was motivated to allow the transition model to retain information for multiple time steps as the stochastic component induces variance (Hafner et al., 2019). this work, we show that having purely stochastic component achieves comparable performance with D2eamerV2 while significantly outperforming SSMs (refer to Appendix for more details). We introduce simplified version of VSG, called Simple Variational Sparse Gating (SVSG) where the world model has a model state with single path to preserve information over multiple steps and also account for partial observability in future states (Fig. 2 [c] presents the SVSG architecture).
Model Components: In SVSG, there no recurrent model and the posterior state 8, is obtained using the representation model by conditioning on the previous state $t-1. input image and the action ar Similar to VSG, there is a transition predictor that returns the prior state 6, which does not use the current image observation to imagine trajectories in the latent space. Both the modules sparsely update the model state at each step using the stochastic gating mechanism proposed in VSG. We have used Gaussian distribution for the stochastic state with learnable mean vector and learnable diagonal covariance matrix. Similar to VSG, the posterior state is used to reconstruct the image, and predict the reward and discount factor. The components of world model in SVSG are:
Representation model:
St go(st
St- T1: 0t)
Transition predictor:
st-1, 0t) 3. (5
Image predictor:
f. l0
st))
Reward predictor:
P lo 81)
Discount predictor:
st))
The representation model qa and transition predictor P are modified to output the posterior and prior & states, respectively. The reset gate U, and the update gate ü, is calculated using the previous state 8:-1 and input i, which has the information about the action ag. The candidate state ä, at each step obtained using input is, reset gate i and previous state 8,-1. Similar to VSG, the update gate ue is sampled from Bernoulli distribution to sparsely update the latent states at each step, given by:
2 = o oW [s ie] + bb)
ū = a T [s--1, dr] + bw) 5; tanh(v, O (W [st= bo))
Ut ~ Bernoulli(ū.),
where @ denotes the element-wist multiplication, and tanh are the sigmoid and hyperbolic tangent activation functions, and W, and b, denote the weights and biases, respectively.
The candidate state 8, is feeded through MLP layers to get the prior and posterior distributions. The image encoding F, was used to get posterior distribution, whereas the prior distribution was predicted without it. The prior & and posterior candidate states sampled from these distributions, where
5

---

5200
MV1
5o0e
80
15
30
25
Score
#0KE 50%
0.5
1.5
8.0
055
1.5
255
5.0
0.5
1.5
2.5
055 13
25
5p0r
5.0
95
2.5
2.0
25
8.0
05
1.5 Stops
2.0
25
0.0
0.5
1.5
5.0
0.5
1.5
2.5
Stops
365
Shops
36b
Stops
365
Dreaineti2
5190
VSG
Figure 3: a) (Top Left) Full arena of the Bhingees (BBS) where the gray region around agent shows the partial view received by it. The circular agent is located in the center of partial view and of teal blue color. The task is to push objects the green goal region on the right side of arena. b) (Bottom left) Scores obtained on BBS with Basic size and no distractors at IM and 2.5M steps. (Right) Performance (results
5 seeds reported) at different sizes of arena and number of distractor objects (#Dist). VSG and SVSG outperforms DreamerV significantly in most scenarios.
the update gate sparsely modifies the previous latent state and outputs the prior & and posterior st model states at each step, respectively. The update equations are given by:
2. fp((t)
2, ~ (5,
(6)
4e o 2 + (1 111-14 0 81-1
8t = 4 o 21 + (1 - -4)) 8t-1,
where Ip and fq denotes functions that output a distribution with learnable parameters for prior and posterior. respectively. For SVSG, Categorical latents (Hafner et al., 2021) were not performing well on our tasks. We attribute this to the fact that samples from a categorical distribution are binary vectors and it is difficult to accurately reconstruct with such sparse latent representations. Lastly, we observed that sparse gating mechanism introduced in VSG was important for convergence of SVSG.
We have used the same loss function as described in Sec. 2 and policy is similar to used in VSG (described in Appendix B). For training the SVSG model, replace the KL loss term between prior and posterior distributions in Eq. 3 with a masked KL loss that penalizes the state dimensions that were updated in the corresponding time step, i.e. those for which the corresponding element in de is equal to 1. We found this to be necessary, since the original. unmasked KL loss did not yield good performance, presumably due its effect on state dimensions that were not updated.
Loss function:
Experiments
4.1 BringBackShapes
Environment: this work, we developed the BringBackShapes (BBS) environment test the ability of agents to solve tasks in partially-observall task is to push the objects within the arena into a pre-specified goal area. Moreover, rewards are sparse and is +1 for successfully pushing an object into the goal. At each time step, the agent only receives an obfuscated view of the arena centered around current position. This requires agents to efficiently explore the arena to find new objects well as remember states of previously observed objects. The objects can collide with each other and the walls, which further requires the agent to account for these events while updating its state. Stochasticity was introduced in the environment by
and stochastic scenarios (see Fig. 3 [a]). The

---

DV2
SVSG VSG
Figure 4: Learned behaviors of DreamerV2 (DV2), VSG and SVSG agents on BringBackShapes on
Basic size and with distractor objects at different steps. DreamerV2 fails capture that distrac- tors (whitish cones) are noisy objects and tries to push them towards the goal, whereas VSG and SVSG learn to avoid the noisy objects and carefully maneuvers the right objects towards the goal.
First-Visit Time
Episode
Objects Not
Visited Objects
First- Visit Time
Episode
Objects Not
Visited Objects
Length
Visited (%)
Not Scored (%)
Length
Visited (%)
Not Scored (%)
Basic, Distrrctors=))
Medium, #Distractor=0))
DV2
500.25
2503 55
4.6
14.67
1800,79
2994.91 91
50.2
33.68
VSG
276.20
1881.0
0.08
1.38
1360.32
2953.21
32.6
24,98
SVSG
365.08
2196.37
2.00
5.15
1375.99
2964.85
30,00
38.53
Basic, #Distractos=22
Medium, #Distractors=2
DV2
593.53
2908.73
6.40
35.90
1669.42
2998.51
40.80
45.48
VSG
330.74
2482.40
0.4
9.67
1051.29
2944.34
16.4
32.12
SVSG
313.95
2292.00
0.6
7.61
1484.09
2988..7
79
31,80
51.42
Table I: Average values of the first-visit time, episode length, % of objects not visited, and % of objects visited but not scored within an episode for trained agents on Basic and Medium environments, and with 0 and 2 distractor objects respectively. Metrics were calculated for 50 episodes for 5 seeds. VSG and SVSG significantly outperformed DreamerV2 on most statistics.
using random distractor objects. The distractor objects follow Brownian motion in any direction and can't be pushed into the goal area. Additionally, they add noise to the reward signal as they might push objects into the goal, causing reward that not only partially related to the agent's 's behavior. such instances might not be visible to the agent, making the task even more challenging. Refer to Appendix A for further description of the environment.
Due to
Experimental Setup:
The BringBackShapes (BBS) environment returns high dimensional images of shape 64 x 64 x as observation. Action is a 2-dimensional continuous vector with acceleration and direction as components. Episodes last for 3000 environment steps and action repeat (Mnih et al., 2016) of 4 was used. Baseline agents include DreamerV2 (Hafner et al., 2021) and DrQ-v2 (Yarats al., 2022). In Appendix C, mention the hyperparameters for the proposed methods- VSG and SVSG. The model was implemented using Tensorflow Probabability (Dillon et al., 2017) and trained on single NVIDIA V100 GPU with 16GB memory. Training time for DreamerV2, VSG and SVSG methods on the BBS environment for 2.5M environment steps are around 12, 11 and 10.5 hours, respectively. Lastly, results are reported across 5 seeds."
Quantitative Results: Fig. 3 [b] compares the proposed methods VSG and SVSG with the leading RL agents- DreamerV2 (Hafner et al., 2021) and DrQ-V2 (Yarats et al., 2022). The score indicates how many objects average were scored within episode. As discussed in Section SVSG with Gaussian latents only. Upon evaluation at 2.5M timesteps, DreamerV2 achieves compet- itive scores when compared to the proposed methods. Whereas 1M steps, DreamerV2 does not perform as well as VSG, which has mean score of 4.9. This shows that learning with sparsity priors helps improve the sample efficiency. Furthermore, performance of SVSG is better than DreamerV2 but similar to VSG, demonstrating that a purely stochastic model can achieve similar performance.
trained
Varying Partiil-Obseevaablliyy and Stochasticy: We also study the effect of partial observability and stochasticity. For partial-observabiiiity we increased the size of the arena while reducing the portion visible to the agent. We consider 4 configurations of the arena- Basic, Small, Medium and Large. For stochasticity, we increase the number of distractor objects using values 0, 1, 2, and Fig. 3 [c] presents the plots of models trained at 2.5M steps at different sizes of the arena and number
'Code is available https: wtt.auu...coomraaaj..i.....
7

---

Cartpole Swingup
Cheebah Run
Finger Spin
Pensulum Sutngup
CA
0.6
0.2
1.0
02
1.2
1.0
0.6
1.
Quadruped Run
Quadruped Walk
Reacher Hard
Walker Run
10000
2000
50.00
1002
430
600
400-
600
230
220
202
0.2
166
1+c
1nd
165
DheamenV2 Disamatt
SLAC 12cb stapot
DVOV2
Figure 5: Comparison of VSG and SVSG with DreamerV1 (Hafner et al., 2020) and Dream- erV2 (Hafner et al., 2021) on the DeepMind Control Suite. VSG converges faster on many tasks as demonstrated by evaluation curves. Even with a single stochastic path, SVSG achieves performance competitive to the models that use a combination of multiple paths.
of distractor objects. It can be observed that increasing the size of arena makes it harder to score objects. VSG was found to outperform DreamerV2 across all arena sizes. However, SVSG did not perform well on larger arena sizes. Adding noisy distractor objects led drop final performance of all models. But VSG and SVSG still outperformed Dreamer indicating that the sparsity prior helps in ignoring the noisy objects in the arena while solving the task.
Ablation Studies: In this work, we also report statistics to describe the behavior of learned agents. First-visit time is the number of episode steps taken to visit an object (when object is completely visible in the agent's view) and is calculated by averaging the first-visit time of each object in the arena. Lower first-visit time indicates that an agent is able to quickly discover all the objects in the arena. Another metric is Episode Length which denotes the number of steps taken by the agent to complete the task. The maximum of these scores was set to 3000. We also report the percentage of objects that were visited within episode which represents the ability of agents explore all parts of the arena to find novel objects. Lastly, we report the percentage of visited objects that were not scored which indicates us that the agent might not be remembering positions of objects seen previously. Thus, agents might have to explore the arena again to find them leading to an increase in time taken to finish the task. Table 1 presents the results on different settings of the environment and it can be observed that VSG and SVSG outperform Dreamer V2 significantly.
Qualitative Results: We also observed the maneuvers taken by the agents push the objects to the goal. The D2eamerV2 agent was able to recognize objects and go behind them to push, but did not follow smooth trajectory and was spending more time around an object to push it in the goal. However, our methods showed smoother trajectories and were more efficient at pushing objects successfully in the goal area. Additionally, our methods learned avoid noisy distractor objects whereas the DreamerV2 agent was colliding with them and was trying to push them to the goal (see Fig. 4 and supplementary material for more videos).
Effect of Sparse Gating: We conducted an experiment where the learned world model was given the first 15 frames and 5 different rollouts were generated in the latent space for the next 35 frames. The sequence of actions is kept fixed across rollouts. The aim was to observe if the sparse gating mechanism is helping the model to retain information for longer time steps and the imagined trajectories are consistent. It was observed in Appendix that learned world model in VSG and SVSG remembers the color and location of objects, and is also cognizant about the goal location and walls. Furthermore, unrolled trajectories from the world model of DreamerV2 showed distortion in the shapes, and in some instances modifies the color of the objects.
4.2 DeepMind Control Suite
Experimental Sctup: The proposed method is evaluated on few tasks from DeepMind Control Suite (Tassa et al., 2018). Observations for the agents are high dimensional images of shape 64 x4x66644 x3, actions range between to 12 dimensions, and episodes last for 1000 steps. An action repeat (Mnih et al., 2016) of 2 was used. The model was implemented using Tensorflow Probabability (Dillon
8

---

et al., 2017) and trained on single NVIDIA V100 GPU with 16GB memory in less than 6 hours. The agents were trained for 1M environment steps. Baselines include DreamerV1 (Hafner et al.. 2020), DreamerV2 (Hafner et al., 2021), DrQ-v2 (Yarats et al., 2022). D4PG (Barth-Maron et al., 2018), and A3C (Mnih et al., 2016). Except A3C, all baselines learn policies from high dimensional pixel inputs. DreamerV2 was trained using the implementation provided by the authors. For other baselines, the metrics provided by the respective authors were used for comparison. Lastly, returns averaged across 5 seeds were reported.
Figure 5 presents comparison of VSG and SVSG with the baseline agents (Refer Ap- pendix D for comparison on more tasks). It can be observed that on most tasks, VSG performs comparable to better than DreamerV2 (Hafner et al., 2021). Notably, VSG significantly outper- forms V2 on Quadruped and Finger-Spin tasks. Furthermore, SVSG with purely stochastic component has similar performance DreamerV2, outperforming on Finger Spin and Quadruped tasks and performing worse on Cartpole Swingup and Cheetah Run tasks. In addition, SVSG signif- icantly outperforms DreamerV on many tasks which also used Gaussian latents and RSSM with multiple paths. Lastly, we also present the importance of sparsity loss in VSG (See Appendix I.3), and of KL Masking and Sparsity loss in SVSG (see Appendix E.2).
Results:
A Ablation Studies: In BBS, we added noise in the environment by having distractor objects. We also experimented with other forms of noise where natural videos are used in the background for DeepMind Control tasks (Zhang et al., 2021; Nguyen et al., 2021b). Since reconstruction-frr model based RL (Deng et al., 2021; ; Nguyen et al., 2021a) perform better than reconstruction based agents (Hafner et al., 2021) such scenarios, we updated the RSSM block in DreamerPro (Deng et al., 2021) with VSG and call it VSGPro, and trained it on DMC with natural background (Refer to Appendix F). We also experimented with VSG in discrete control tasks from the Atari benchmark in Appendix G, where VSG performed better on task with changing viewpoints.
5 Related Work
Latent Dynamics Models:
Latent dynamics models (Bourlard and Morgan, 2012; Kalman, 1960; Bengio et al., 1999) operate directly sequences predicted in the latent space rather than autoregres- sively feeding back the generated frames back to the model. Recent advancements in deep learning have allowed learning expressive latent dynamics models using stochastic backpropagation (Kingma and Welling, 2013: Chung et al., 2015; Krishnan et al., 2015; Karl et al., 2016). Recurrent State- Space Model (RSSM) (Hafner et al., 2019) comprises of stochastic and deterministic components. VideoFlow (Kumar et al., 2019) predicted future values of the latent state by normalizing flows for robotic object interactions. Hierarchical latent models for video prediction were proposed in CWVAEs (Saxena et al., 2021) with levels ticking at different intervals in time. (Franceschi et al., 2020) and (Donà et al., 2021) disentangled dynamic and static factors where 2-5 initial observations was used to estimate the static component.
Deep Reinforcement I Learning (DRL) methods fall into one of two categories: 1) Model-Based - where an explicit model of the environment and its dynamics are learned (Ha and Schmidhuber, 2018; Hafner et al., 2019; 2020; 2021; Zhang et al., 2019; Kaiser et al., 2019), and 2) Model-Free - where a policy is learned directly from the raw observations (Srinivas et al., 2020; Kostrikov et al., 2020b; Lillicrap al., 2015; Yarats et al., 2022; Schwarzer et al., 2021; Mondal et al., 2022). Deep Deterministic Policy Gradient (DDPG) (Lillicrap et al., 2015) combined actor-critic with insights from DQNs (Mnih et al., 2015) to learn agents for continuous action spaces. TD3 (Fujimoto et al., 2018) builds upon the DDPG algorithm and addresses the problem of overestimation bias in the value function. CURL (Srinivas et al., 2020) uses contrastive losses to learn discriminative representations. DrQ (Kostrikov et al., 2020a) and DrQ-v2 (Yarats et al., 2022) employed data augmentation techniques and do not use auxiliary losses or pre-training. RSSM was introduced in PlaNet (Hafner et al., 2019) and was employed for online planning in the latent space. DreamerV1 1 (Hafner al., 2020) and DreamerV2 (Hafner et al., 2021) achieved state-of-the-art results on DMC (Tassa et al., 2018) and Atari (Bellemare et al., 2013), respectively. SimPLe (Kaiser al., 2019) trains a PPO (Schulman et al., 2017) agent on the learned video generation model in pixel space. SOLAR (Zhang et al., 2019) solved robotics tasks via guided policy search.
RL for Visual Control:
Sparsity in RNN: Neural networks have widely adopted sparsity to reduce the memory footprint of weights and activations (LeCun et al., 1990; Chen et al., 2015; Han et al., 2015). Several works have explored sparsity in RNNs. Campos et al. (2017) introduced a mechanism in RNNs that learns

---

to skip state updates, effectively reducing the number of sequential operations on the latent state, thereby alleviating the problem of vanishing gradients in training on long sequences. Goyal al. (2019) presented Recurrent Independent Mechanism (RIM), an architecture that consists of separate recurrent modules which are sparsely updated using a learned attention mechanism. contrast to RIM, the number of updated state variables in VSG algorithm is not fixed.
Discussion
In this work, we introduce VSG and SVSG, two latent dynamics models leveraging sparse state updates. The sparse update prior was found to facilitate more efficient behaviors in tasks requiring long-horizon planning. Furthermore, SVSG is purely stochastic model with single component in the model state. We show that VSG and SVSG can outperform leading agents on the proposed Shanpess task, a challenging partiall--obserabbe and stochastic environment. BBS allows for controlling different factors variation like stochasticity and partial-observabiiity Experiments conducted various variations in BBS demonstrate that the proposed agents are robust to noise the environment and can better retain information of seen objects. Some limitations and potential research directions for future research are as follows:
the current implementation of VSG, the latent space does not exhibit disentanglement which could be an interesting direction for future research. Gating mechanisms in VSG can also be combined with other recurrent architectures like RIM (Goyal et al., 2019).
and stochasticity. More controllable factors like the nature of entities (shape, size, color of objects), underlying physics (mass, friction, elasticity), or procedural background generation can be introduced to further study generalization capabilities of RL agents.
In this work, BBS was explored with only 2 factors of variation:
SVSG being a purely stochastic model can further be used to estimate state uncertainty by marginalizing over multiple samples paths to efficiently explore in an unknown environment.
Evaluation on first-person view 3D games like tasks in DMLab (Beattie et al., 2016) would be interesting. Furthermore, a 3D version of the BBS environment with the viewpoints changing with rotation of agent and the underlying physics will make the task more challenging.
We have used small latent dimensions and it would be interesting to train such models with larger architectures and more complex tasks. Scaling the current architecture would also require optimizing the implementation to make them computationally feasible.
Categorical latents outperformed Gaussian latents the stochastic states of RSSM (Hafner et al., 2021), especially for discrete control tasks. However, SVSG was not found to work well with Categorical latents and we believe that sampled sparse states are hard optimize.
Model-based RL for visual control is still in early stages. However, a major challenge with deploying such models in the real world is safety especially during exploration. This would require an accurate world model that allows learning policies with stringent safety constraints that avoid mistakes when deployed in the real world. Such algorithms will rely on models that are robust when transferred from simulation to the real world.
Acknowledgements
The authors would like to thank David Meger, Lucas Lehnert and Ankesh Anand for their valuable feedback and discussions. The text also benefited from discussions with Abhinav Agarwalla, Rupali Bhati and Vineet Jain. The authors are also grateful to CIFAR for funding and the Digital Research Alliance Canada for computing resources.
References
Mohammad Babaeizadeh, Chelsea Finn, Dumitru Erhan, Roy H Campbell, and Sergey Levine. Stochastic variational video prediction. arXiv preprint arXiv:1710.11252 2017.
Gabriel Barth-Maron, Matthew W Hoffman, David Budden, Will Dabney, Dan Horgan, Alistair Muldal, Nicolas Heess, and Timothy Lillicrap. Distributed distributional deterministie policy gradients. arXiv preprint arXiv:1804.08677 2018.
10

---

Charles Beattie, Joel Z Leibo, Denis Teplyashin, Tom Ward, Mareus Wainwright, Heinrich Küttler, Andrew Lefraneq, Simon Green, Víctor Valdés, Amir Sadik, al. Deepmind lab. arXiv preprint urXiv:/612.0388 2016.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Yoshua Bengio, Patrice Simard, and Paolo Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE transactions on neural networks, 5(2):157-166, 1994.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint 1308.. 2013.
Yoshua Bengio et al. Markovian models for sequential data. Neural computing surveys, 2(199): 129-162, 1999.
Herve A Bourlard and Nelson Morgan. Connectionist speech recognition: a hybrid approach, volume 247. Springer Science & Business Media, 2012.
Víctor Campos, Brendan Jou, Xavier Giró-i Nieto, Jordi Torres, and Shih-Fu Chang. Skip rnn: Learning to skip state updates recurrent neural networks. arXiv preprint 1.X.:v/708.0066 2017.
Wenlin Chen, James Wilson, Stephen Tyree, Kilian Weinberger, and Yixin Chen. Compressing neural networks with the hashing trick. In International conference on machine learning, pages 2285-2294. PMLR, 2015.
Maxime Chevalier-Boisvert,, Lucas Willems, and Suman Pal. Minimalistic gridworld environment for openai gym, 2018.
Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint urXiv:1406. 1078, 2014.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probabilistic dynamics models. Advances in neural information processing systems, 31, 2018.
Junyoung Chung, Kyle Kastner, Laurent Dinh, Kratarth Goel, Aaron Courville, and Yoshua Bengio. A recurrent latent variable model for sequential data. In Advances in neural information processing systems, pages 2980-2988, 2015.
Djork-Arné Clevert Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint wrXiv:1511.0720) 2015.
Fei Deng, Ingook Jang, and Sungjin Ahn. Dreamerpro: Reconstruction-free model-based reinforce- ment learning with prototypical representations. arXiv preprint 10Xiv:2110.14565, 2021.
Joshua Dillon, Ian Langmore, Dustin Tran, Fugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matt Hoffman, and Rif A Saurous. Tensorflow distributions. arXiv preprint arXiv: 2017.
Jérémie Donå, Jvan-Yves Franceschi, Sylvain Lamprier, and Patrick Gallinari. (PDE)-drivel spatiotemporal disentanglement. In International Conference on Learning Representations, 2021.
Lean-Yves Franceschi, Edouard Delasalles, Mickaël Chen, Sylvain Lamprier, and Patrick Gallinari. Stochastic latent residual video prediction. In International Conference on Machine Learning, pages 3233-3246. PMLR, 2020.
Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. arXiv preprint arXiv:1802.0904 1802.89077.7... 2018
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. arXiv preprint arXiv:1909. 10099 2019.
11

---

David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 1803 2018.
Danijar Hafner, Timothy Lillicrap, lan Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th Interaational Conference on Muchine Learning, volume 97 of Proceedings of Machine Learning Research, pages 2565-2565. PMLR, 09-15 Jun 2019. URL https: https////rooeeggeeeee mlr wwwppaeesssseeee htt///afner1999 hmm.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2020.
URL httpp://opentei //ppnee//pee wwt.//eeeeeeeeeeeeeeeeeee net//orummidd-S11111
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference Learning Representations, 2021. URL
Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint 15Xiv:1510.00049 2015.
Sepp Hochreiter. Untersuchungen zu dynamischen neuronalen netzen. Diploma, Technische Univer- sität München, 91(1), 1991.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for atari. arXiv preprint arXiv: w:xiv:/933.00377 2019.
Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. Journal of basic Engineering. 82(1):35 45, 1960.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick van der Smagt. Deep variational bayes filters: Unsupervised learning of state space models from raw data. arXiv preprint 16Xi:16005.06432. 2016.
Diederik Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible Ixl convolutions. arXiv preprint urXiv:1807.03033 2018.
Diederik P Kingma and Max Welling. 11X1v:1312.6114, 2013.
Auto-encoding variational bayes.
arXiv preprint
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint 13X4,,2004..3649, 2020a.
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXhv:2004.13649 2020b.
Rahul G Krishnan, Uri Shalit, and David Sontag. arXiv:1511.05121, 2015.
Deep kalman filters.
arXiv preprint
Manoj Kumar, Mohammad Babaeizadeh, Dumitru Erhan, Chelsea Finn, Sergey Levine, Laurent Dinh, and Durk Kingma. Videoflow: A flow-based generative model for video. arXiv preprint arXiv:1903.01434 2(5), 2019.
Yann LeCun, Bernhard Boser, John S Denker, Donnie Henderson, Richard Howard, Wayne Hubbard, and Lawrence D Jackel. Backpropagation applied to handwritten zip code recognition. Neural computation, 1(4):541-551 1989.
Yann LeCun, John S Denker, and Sara A Solla. Optimal brain damage. In Advances in neural information processing systems, pages 398-605, 1990.
Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint
2015.
12

---

Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, al. Human-level control through deep reinforcement learning. Nature, 518(7540):529, 2015.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. International Conference on Machine Learning, pages 1928-1937, 2016.
Arnab Kumar Mondal, Vineet Jain, Kaleem Siddiqi, and Siamak Ravanbakhsh. EqR: Equiv- ariant representations for data-efficient reinforcement learning. In Kamalika Chaudhuri, Ste- fanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceed- ings of the 39th International Conference on Machine Learning, volume 162 of Proceed- ings of Machine Learning Research, pages https: httpss////cccccg.....om mlr. press matrsss/v162mmmddaa22222 html.
PMLR, 17-23 Jul 2022. URL
Daniel Neil, Michael Pfeiffer, and Shih-Chii Liu. Phased Istm: Accelerating recurrent network training for long or event- sequences. arXiv preprint 1rxiv:1610.09511,, 2016.
Tung D Nguyen, Rui Shu, Tuan Pham, Hung Bui, and Stefano Ermon. Temporal predictive coding for model-based planning in latent space. In International Conference on Machine Learning, pages 8130-8139 PMLR, 2021a.
Tung D Nguyen, Rui Shu, Tuan Pham, Hung Bui, and Stefano Ermon. Temporal predictive coding for model-based planning in latent space. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages -3399.19. PMLR, 18 24 Jul 2021b. URL btttpss//proceeeiii mlr. har.press/vee hatssssvvv9/////uu hwtp///ggunnnn..nmmen//// html.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In Eric P. Xing and Tony Jebara, editors, Proceedings of the 31st International Conference on Machine Learning, volume 32 of Proceedings of Machine Learning Research, pages 1278- 1286, Bejing, China. 24 Jun 2014. PMLR. URL attps://pree //pocc/e/sss.csnngs.... mlr.
wttps//v2/rreeeeee htmml.. html.
Reuven Y Rubinstein. Optimization of computer simulation models with rare events. European Journal of Operational Research, 99(1)889-112, 1997.
Vaibhav Saxena, Jimmy Ba, and Danijar Hafner. Clockwork variational autoencoders. In A. Beygelz- imer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL https:: wttp://wwwweeeoooomimeeeeeeeee ftt
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: (700::0070077777 2017.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with peef-prediccive representations. In International Conference on Learning Representations, 2021. URL https: h //ppnreeiee......e///// wwtp:/.ee//oooooeoeeeeeeee/eeeo formmff uCQfPZwWaUu.
Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. arXiv preprint arXiv:2004.044 041166 2020.
Richard S Sutton. Dyna, an integrated architecture for learning, planning, and reacting, ACM SIGART Bulletin, 264):160-163, 1991.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv:1801.00690. 2018
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: locally linear latent dynamics model for control from raw images. In Advances neural information processing systems, pages 2746-2754, 2015.
13

---

Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. International Conference on Learning Representations, 2022. URL. htt.ttt///peprrrraaaii
-yy-eye
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning
invariant representations for reinforcement learning without reconstruction. International Conference on Leurning Representations, 2021. URL httpp://openreeiee..eee. wwtp://e///oommmeee../
-2FCwDRRREu.
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine. Solar: deep structured representations for model-based reinforcement learning. In International Conference on Machine Learning, 2019.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section
(c) Did you discuss any potential negative societal impacts your work? [Yes] Described in the section with limitations.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (cither in the supplemental material or as a URL)? [Yes] The code for the model and dataset is in the supplementury material. There is a README file with the instructions to run them.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] We have specified the hyperparameters in Appendix C.
(c) Did you report error bars (e.g., with respect to the random seed after running exper- iments multiple times)? [Yes] All the results and plots presented in Section 4 and Appendix were obtained after training on multiple seeds ranging from 3-5.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] We have mentioned about the GPUs and time taken to run on a single seed the implementation details for each environment.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [No] We used the implementation of DreamerV2, DreamerPro and DrQ-v2 provided by the authors with MIT license. For comparison with DBC, we thank the authors for sharing the evaluation logs. URL? [Yes]
(c) Did you include any assets either in the supplemental material The code for the new environment (BBS) is included the supplementary material.
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [Yes]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
14

---

(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
15


Reviews:
[{"review_id": "rKBOu2AgXl2", "reviewer": "Reviewer_Dxye", "summary": "This paper proposes two new latent dynamics models (VSG and SVSG), which aims to incorporate a certain inductive bias: in some real world problems, the changes are very sparse, so the latent can be updated sparsely to help long term memory. They test the proposed algorithm on a new environment BBS, which shows that VSG and SVSG outperforms existing methods, and on DMC benchmark, which shows that VSG and SVSG ad comparable with existing methods.", "questions": "See above", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- This is a very interesting paper to incorporate the inductive bias to existing latent dynamics model RSSM and SSM, and the results show that the proposed methods seem to be effective on specific problems.\n- It is also every interesting to see that after the modification, the SVSG can outperform SSM and comparable with RSSM/VSG.\n\nWeakness:\n\nI agree that this is a very interesting work, but I think it will be better if the analysis can be more sufficient. For instance, by *randomly* selecting some latents $u_t$ to be zero in GRU, the VSG can achieve better performance on long-memory cases, which is good. However, this may mismatch the motivation “the real world changes are sparse” since the selection is purely stochastic. And this is even more significant for SVSG. I think it is important to understand why the proposed methods may work.\nCan I think of this as a regularization trick, which ask for the encoder to encode the information duplicately for each dimension, since each of them might be zeroed? Is it possible for you to show an experiment that shows the effects of the randomness? For instance, given a sequence, and several different samples from Bernoulli distributions, to see if the next state can decode the same image?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "dH8uFfsPY16", "reviewer": "Reviewer_TPk3", "summary": "This paper proposes two recurrent models for model-based RL: Variational Sparse Gating (VSG), improved from DreamerV2 with binary gate, and Simple Variational Sparse Gating (SVSG) which has purely stochastic states. It also introduces a new partially-observable environment, BringBackShapes(BBS), to evaluate agents under POMDP setting. The experimental results show that the proposed VSG outperforms DreamerV2 in continuous control tasks including BBS and DeepMind Control Suite.", "questions": "1. Setting $u_t$ as binary values seems make the model not recurrent anymore. The model parameters $W_v, b_c, W_c$ and $b_c$ only get update when $u_t = 1$. \n2. I am afraid that $u_t=0$ would be a big issue when planning in world model (as Figure 1(b)). If $ u_t=0 $, then $ h_t = h_{t-1} $ due to Equation (2). \nThe transition predictor $ p_{\\phi}(z_t|h_t) $ that predicts prior state $ \\hat{z}_t $ \ndepends only on $ h_t $. \n\n  In the case of $h_t = h_{t-1} $, the predicted state $ \\hat{ z }_t $ \n\n  would be similar to $ \\hat{z}_{t-1} $ at time step $t-1$. \nConsequently, the actor conditioned on $ \\hat{z}_t $ output the similar policy as time step $t-1$. \n\n  If the distributions of $ p_{\\phi}(z_t|h_{t-1}) $ and the policy have low variance, making $u_{k \\geq t}=0$ almost surely, then the whole planning in world model takes input of the same hidden state $ h_{k \\geq t} = h_{t-1} $ .\n\n3. Does SVSG have the same losses as VSG ?\n4. Why VSG and SVSG are called \"variational\"? It seems that there is nothing related to variational inference.\n5. What setting of BBS is used in Figure 3b), i.e., the number of distractors and the size of arena?\n6. In Table 1, what does first-visit time mean? Is it the time when visiting the first object, or the first time when visiting the last object (all objects)? \n7. The experiments and related work lack of an important baseline method, SLAC[1], for continuous control tasks. SLAC solves POMDP by stochastic sequential models and also defines stochastic latent states.\n8. DrQV2 is a strong baseline algorithm in continuous pixel control. Showing training curves of DrQV2 in figures would be good comparisons. \n9. DeepMind Control Suite is POMDP but many algorithms assuming fully-observability solve DMC by stacking frames. It would be nice if you evaluate the methods on other public environments that are \"truly\" POMDPs, e.g., first-person view 3D games.\n\nOther comments:\n1. In line 85, $o_t$ is defined as state but in line 115 $s_t=[h_t,z_t]$ is also defined as state. I suppose the latter one should be correct for the definition of POMDP, since $o_t$ is the image encoding as in line 120. \n1. SLAC is listed in the legend of Figure 5, but not shown in any curves or main text. \n\nReference\n\n[1] Lee et al. Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable Model. NeurIPS 2020\n", "limitations": "The effect of sparse gating needs to be further discussed and verified. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n1. VSG achieves impressive experimental results comparing to DreamerV2. \n1. A new environment with various difficulty levels is designed for POMDP. \n\n## Weakness\n1. The significance of the proposed sparse gating needs to be clarified. In VSG, the sparse gating is the only improvement from DreamerV2, while the rest part including architecture and losses are the same as DreamerV2. The authors should discuss more about why sparse gating can improve performance.\nFor example, list some hypotheses to the reasons and verify them in experiments.\nBesides, additional ablation studies about sparse gating may also help to demonstrate its effect. \nOtherwise, without such discussion and verification, the technical contribution of VSG is considered as minor.\n2. The motivation of SVSG is unclear. Previously Stochastic State-Space Model (SSM) that has stochastic state has been empirically proved to have worse results than RSSM that has deterministic state. In the experiments of this paper, SVSG also has worse performance than VSG. The authors should clarify the motivation of proposing a new model with stochastic state. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "N8N2ilfZDbZ", "reviewer": "Reviewer_4uKx", "summary": "This paper studies modelling the latent dynamics model by exploiting the idea that different elements of the latent space might exhibit varying degree of changes. The authors propose variational sparse gating (VSG), which is a modification of update gate of the GRU network (Cho, et al. 2014) in the recurrent state space model (RSSM; Hafner, et al. 2019) with a stochastic gating mechanism. The authors also propose Simple VSG (SVSG), which follows the similar stochastic gating updates as VSG, but contains only a purely stochastic latent state. VSG and SVSG are empirically evaluated on DeepMind Control Suite and a novel environment (BringBackShapes) to test their ability of solving stochastic tasks with partial-observability. ", "questions": "The major concerns are listed in the previous section, some minor questions are listed below.\n\n- What is the relative computational complexity between VSG and SVSG?\n- Is the sparse reward setting necessary in BBS? A complicated (stochastic, non-Markovian) reward structure might be less trivial for standard agents without sparse gating (e.g., Dreamer v2) to acquire, according to the statements in the main text;\n- The Private Eye evaluation (Figure 10) and the accompanying discussion is not so convincing given the large variance in the evaluation curves;", "limitations": "The authors have provided thoughtful discussion on the limitation and potential future areas for investigation of the VSG model, such as combination with other recurrent architectures, increasing the complexity of BBS environment, VSG dealing with categorical latents, etc.\n\nI see no potential negative societal impact.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros: \n- The paper is well-written;\n- The proposed new environment, BringBackShapes, is designed thoughtfully and sufficiently complicated to assess the ability of agents for solving tasks with stochastic transition and reward structures;\n\nCons:\n- The motivation for SVSG is not clear, having a purely stochastic latent state does not show performance improvement in general, and no evidence/intuition of why this could be preferred over VSG (e.g., model complexity, stronger generalisability, etc.);\n- There is inconsistency in the performance of VSG/SVSG with respect to varying level of stochasticity and partial-observability, e.g., one would expect monotonic increasing size of performance gap between VSG/SVSG and Dreamer V2 in the BBS environment, but none is observed (at least from the reported results in Figure 3c). Moreover, in the Atari evaluations (Figure 10), tasks such as Ms Pacman and Seaquest clearly exhibits strong stochasticity, but VSG fail to outperform Dreamer V2. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "FZ1jjAqamsi", "reviewer": "Reviewer_ANCN", "summary": "In summary, this paper provides a modification to the Gated Recurrent Unit (GRU) that *sparsely* updates the recurrent state by samplying from a Bernoulli distribution for the *update* term. Thus, rather than being a continuous value in GRU, this is now from $\\{0, 1\\}$. This is the Variational Sparse Gate (VSG). \n\nThe Simple Variational Sparse Gate (SVSG) further modifies VSG to be made purely stochastic (supposedly - see Questions), and used to evaluate claims from a previous paper that a deterministic recurrent path is required fro competitive results.\n\nExperiments were carried out on a novel environment, and evironments from the DeepMind Control Suite. The novel environment specifically emphasizes the key aspects of this work, namely partial observability, and remembering past states. ", "questions": "* Figure 2: in the caption, the description of $M$ does not match how it is used in Figure 2(c). \n* line 54 and line162-164: mentions difference between a *purely stochastic* component, and one that contains some deterministic parts. I'm confused by what this means. From Figure 2, it looks like all units contain both deterministic and stochastic calculations. I'm probably missing something, could you say what that could be?\n* line 263: Can you really say that, generally speaking, SVSG is comparable to DreamerV2? SVSG seems to perform worse in the Cartpole and Cheetah environments. \n\n\n*General*\n* What exactly is the role of the transition predictor?\n* What is th meaning of *variational* in the name? \n* How exactly do the components in Equations 1, 2, 4 fit together (in their separate cases)? \n* How many seeds were used in Figure 3? Was it 3 (inferred from the checklist). If so, why is it different to the 5 seeds used in other figures?\n", "limitations": "I think limitations and future work are discussed well. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "*Strengths*\n* Well written for the most part. Explanations and arguments are generally clear. \n* Contributions are clearly stated, and the paper is structured will to argue these.\n* A satisfying range of experiments was carried out, at least in the simpler case. The ablation studies were nice to see. \n* Experiments were also carried out on more complex environments too. \n\n\n*Weaknesses*\n* Interpretations of the notation (Equation 1, 2, 4) could have been written more clearly. \n* Lacks a baseline description of how GRU works (at least intuitively). This could help readers understand why the proposed changes are meaningful, and significant. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept